continue work on amd_uop_matmul

write out kernel 3 in uops (#11352 )
* write out kernel 3 in uops * matmul is correct * gemm passes spec * bugfix to match speed * cleanups
2025-07-23 18:49:19 -07:00 · 2025-07-23 17:32:38 -07:00 · 2025-07-23 17:40:55 -04:00 · 2025-07-23 13:25:55 -07:00 · 2025-07-23 11:26:05 -07:00 · 2025-07-23 20:30:18 +03:00
172 changed files with 4651 additions and 2555 deletions
@@ -70,8 +70,8 @@ jobs:
      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - name: Run Tensor Core GEMM (float)
      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
@@ -10,6 +10,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
+    timeout-minutes: 100

    steps:
    - name: Checkout Code
@@ -12,6 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
+    timeout-minutes: 240

    steps:
    - name: Checkout Code
@@ -239,8 +239,8 @@ jobs:
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated AMD MFMA tensor cores
@@ -252,8 +252,8 @@ jobs:
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated CUDA tensor cores
@@ -326,16 +326,21 @@ jobs:
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
+        python3 -m ruff check extra/onnx.py extra/onnx_parser.py
        python3 -m ruff check examples/mlperf/ --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
    - name: Run mypy
-      run: python -m mypy --strict-equality --lineprecision-report . && cat lineprecision.txt
+      run: |
+        python -m mypy --strict-equality --lineprecision-report .
+        cat lineprecision.txt
+        python -m mypy --strict-equality extra/onnx_parser.py
+        python -m mypy --strict-equality extra/onnx.py

  unittest:
    name: Unit Tests
    runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15

    steps:
    - name: Checkout Code
@@ -368,8 +373,8 @@ jobs:
        PYTHONPATH=. python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 14600 lines
-      run: MAX_LINE_COUNT=14600 python sz.py
+    - name: Repo line count < 15500 lines
+      run: MAX_LINE_COUNT=15500 python sz.py

  fuzzing:
    name: Fuzzing
@@ -422,6 +427,29 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  testgendataset:
+    name: 'GPU Generate Kernel Dataset'
+    runs-on: ubuntu-22.04
+    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: gen-dataset
+          deps: testing_minimal
+          opencl: 'true'
+      - name: Generate Dataset
+        run: PYTHONPATH="." extra/optimization/generate_dataset.sh
+      - name: Upload artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: sops.gz
+          path: /tmp/sops.gz
+
  testopenpilot:
    name: 'openpilot Compile Tests'
    runs-on: ubuntu-22.04
@@ -440,11 +468,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2137 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      # - name: Test openpilot simple_plan vision model correctness (float32)
+      #   run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot compile4
@@ -475,15 +505,10 @@ jobs:
        run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX Runner (CPU)
        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
-      - name: Test ONNX Runner (WEBGPU)
-        run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
      - name: Test Additional ONNX Ops (CPU)
        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
-      - name: Run REMOTE=1 Test
-        run: |
-          REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -507,10 +532,6 @@ jobs:
          opencl: 'true'
      - name: Test ONNX (GPU)
        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Run REMOTE=1 Test
-        run: |
-          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
-          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Test Optimization Helpers
        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
@@ -524,6 +545,20 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  testllm:
+    name: Test LLM
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: apps_llm
+      - name: Test 1B LLM
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
+
  testmodels:
    name: Models (llvm+cpu+gpu)
    runs-on: ubuntu-22.04
@@ -563,7 +598,7 @@ jobs:
      with:
        key: dsp-minimal
        deps: testing_minimal
-        pydeps: "onnx==1.17.0 onnxruntime pillow"
+        pydeps: "onnx==1.18.0 onnxruntime pillow"
        llvm: "true"
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
@@ -651,7 +686,9 @@ jobs:
        if: matrix.backend=='amdllvm'
        run: python test/test_amd_llvm.py
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py test/external/external_test_am.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py --durations=20
+      - name: Run pytest (amd)
+        run: python -m pytest test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
@@ -804,7 +841,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: osx-webgpu
-        deps: testing_minimal
+        deps: testing
        webgpu: 'true'
    - name: Test infinity math in WGSL
      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
@@ -827,6 +864,8 @@ jobs:
    #    pip install $GITHUB_WORKSPACE
    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
    #    node test_viz.js
+    - name: Test ONNX Runner (WEBGPU)
+      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py

  osxremote:
    name: MacOS (remote metal)
@@ -853,12 +892,11 @@ jobs:
          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

  amdremote:
-    name: Linux (remote amd)
+    name: Linux (remote)
    runs-on: ubuntu-22.04
    timeout-minutes: 20
    env:
      REMOTE: 1
-      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
      PYTHONPATH: ${{ github.workspace }}
    steps:
      - name: Checkout Code
@@ -866,38 +904,58 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: linux-remote-amd
+          key: linux-remote
          deps: testing_minimal
          amd: 'true'
          llvm: 'true'
+          opencl: 'true'
      - name: Start remote server
        run: |
          start_server() {
            systemd-run --user \
              --unit="$1" \
-              --setenv=REMOTEDEV=AMD \
+              --setenv=REMOTEDEV="$2" \
              --setenv=MOCKGPU=1 \
              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$2" \
+              --setenv=PORT="$3" \
              --working-directory="$(pwd)" \
              python tinygrad/runtime/ops_remote.py
          }

-          start_server "remote-server-1" 6667
-          start_server "remote-server-2" 6668
+          start_server "remote-server-amd-1" "AMD" 6667
+          start_server "remote-server-amd-2" "AMD" 6668
+          start_server "remote-server-gpu" "GPU" 7667
+          start_server "remote-server-cpu" "CPU" 8667
      - name: Check Device.DEFAULT and print some source
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test
+      - name: Run REMOTE=1 Test (AMD)
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
        run: |
          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
+      - name: Run REMOTE=1 Test (GPU)
+        env:
+          HOST: 127.0.0.1:7667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
+      - name: Run REMOTE=1 Test (CPU)
+        env:
+          HOST: 127.0.0.1:8667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
      - name: Show remote server logs
        if: always()
        run: |
-          journalctl --user -u remote-server-1 --no-pager
-          journalctl --user -u remote-server-2 --no-pager
+          journalctl --user -u remote-server-amd-1 --no-pager
+          journalctl --user -u remote-server-amd-2 --no-pager
+          journalctl --user -u remote-server-gpu --no-pager
+          journalctl --user -u remote-server-cpu --no-pager

  osxtests:
    strategy:
@@ -938,7 +996,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu]
+        backend: [llvm, cpu, webgpu]

    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
@@ -951,11 +1009,12 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: windows-minimal
+          key: windows-${{ matrix.backend }}-minimal
          deps: testing_unit
+          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
      - name: Run unit tests
        if: matrix.backend=='llvm'
        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py
@@ -39,9 +39,8 @@ Try a matmul. See how, despite the style, it is fused into one kernel with the p

 ```sh
 DEBUG=3 python3 -c "from tinygrad import Tensor;
-N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
-c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2);
-print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
+N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N);
+(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
 ```

 And we can change `DEBUG` to `4` to see the generated code.
@@ -149,6 +149,7 @@ generate_nv() {
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
    --clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
    -o $BASE/nv_gpu.py
  fixup $BASE/nv_gpu.py
@@ -166,6 +167,7 @@ generate_nv() {
  sed -n '1i\
 nv_status_codes = {}
 /^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.nv_gpu"

  clang2py -k cdefstum \
    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
@@ -180,6 +182,7 @@ nv_status_codes = {}
    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.h \
    extra/nv_gpu_driver/g_rpc-message-header.h \
    extra/nv_gpu_driver/gsp_static_config.h \
    extra/nv_gpu_driver/vbios.h \
@@ -187,7 +190,7 @@ nv_status_codes = {}
    -o $BASE/nv/nv.py

  fixup $BASE/nv/nv.py
-  python3 -c "import tinygrad.runtime.autogen.nv_gpu"
+  python3 -c "import tinygrad.runtime.autogen.nv.nv"
 }

 generate_amd() {
@@ -441,7 +444,7 @@ generate_libusb() {
    -o $BASE/libusb.py

  fixup $BASE/libusb.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
@@ -7,28 +7,30 @@

 print("******** first, the runtime ***********")

-from tinygrad.runtime.ops_cpu import ClangJITCompiler, MallocAllocator, CPUProgram
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
+
+cpu = CPUDevice()

 # allocate some buffers
-out = MallocAllocator.alloc(4)
-a = MallocAllocator.alloc(4)
-b = MallocAllocator.alloc(4)
+out = cpu.allocator.alloc(4)
+a = cpu.allocator.alloc(4)
+b = cpu.allocator.alloc(4)

 # load in some values (little endian)
-MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
-MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))

 # compile a program to a binary
 lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")

 # create a runtime for the program
-fxn = CPUProgram("add", lib)
+fxn = cpu.runtime("add", lib)

 # run the program
 fxn(out, a, b)

 # check the data out
-print(val := MallocAllocator._as_buffer(out).cast("I").tolist()[0])
+print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
 assert val == 5


@@ -46,7 +48,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 out = Buffer(DEVICE, 1, dtypes.int32).allocate()
 a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-# NOTE: a._buf is the same as the return from MallocAllocator.alloc
+# NOTE: a._buf is the same as the return from cpu.allocator.alloc

 # describe the computation
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
@@ -78,7 +80,7 @@ print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.kernelize.kernelize import get_kernelize_map
+from tinygrad.schedule.kernelize import get_kernelize_map

 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -52,7 +52,7 @@ Signals are device-dependent structures used for synchronization and timing in H
 The following Python code demonstrates the usage of signals:

 ```python
-signal = your_device.signal_t()
+signal = your_device.new_signal(value=0)

 HWQueue().timestamp(signal) \
         .signal(signal, value_to_fire) \
@@ -6,11 +6,11 @@ Directories are listed in order of how they are processed.

 ---

-## tinygrad/kernelize
+## tinygrad/schedule

 Group UOps into kernels.

-::: tinygrad.kernelize.kernelize.get_kernelize_map
+::: tinygrad.schedule.kernelize.get_kernelize_map
    options:
        members: false
        show_labels: false
@@ -12,7 +12,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
 | [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
 | [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
+| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |

 ## Interoperability

@@ -26,5 +26,6 @@
 ::: tinygrad.Tensor.transpose
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
+::: tinygrad.Tensor.diag
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange
@@ -1,11 +1,10 @@
 import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
-  onnx_model = onnx_load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)
  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
  return run_onnx_jit, run_onnx.graph_inputs

@@ -10,6 +10,7 @@ import tensorflow as tf
 import tf2onnx
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
+from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model

 def get_uncompiled_model2(dataset_size=32, output_size=4):
@@ -25,7 +26,7 @@ class TinyOnnx:
  def __init__(self, keras_model):
    input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-    self.run_onnx = OnnxRunner(onnx_model)
+    self.run_onnx = OnnxRunner(Tensor(onnx_model.SerializeToString(), device="PYTHON"))

  def forward(self, x):
    return self.run_onnx({"x": x}, debug=False)['predictions']
@@ -47,8 +48,8 @@ def compile_onnx_model(onnx_model):
  cprog.append("void initialize(float *weights) {")
  weights = bytes()
  for name,cl in bufs_to_save.items():
-    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl._buf)*4});")
-    weights += bytes(cl._buf)
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
+    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
  cprog.append("}")

  # write the weights to disk
@@ -9,6 +9,7 @@ from tinygrad.device import Compiled
 from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
+from tinygrad.engine.realize import get_program

 def get_sched_resnet():
  mdl = ResNet50()
@@ -108,7 +109,7 @@ if __name__ == "__main__":
    choices = []
    for lin, nm in lins:
      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=lin.to_program()).estimates.ops
+      ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
      choices.append((tm, gflops, lin, prg, nm))

@@ -121,7 +122,7 @@ if __name__ == "__main__":
    if getenv("SRC"):
      print(si.ast)
      print(lin.applied_opts)
-      print(lin.to_program().src)
+      print(get_program(lin.get_optimized_ast(), lin.opts).src)
    total_tm += tm
    running_gflops += gflops * tm
    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
@@ -5,29 +5,26 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

-from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
-from tinygrad.tensor import _from_np_dtype
 from tinygrad.engine.realize import CompiledRunner

 import onnx
-from onnx.helper import tensor_dtype_to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"

 def compile(onnx_file):
-  onnx_model = onnx_load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)
  print("loaded model")

-  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
+  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
  # Float inputs and outputs to tinyjits for openpilot are always float32
-  input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
+  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

@@ -57,11 +54,11 @@ def compile(onnx_file):
      gated_read_image_count += ei.prg.p.src.count("?read_image")
  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
-    assert kernel_count <= allowed_kernel_count, f"too many kernels! {kernel_count=}, {allowed_kernel_count=}"
+    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
  if (allowed_read_image:=getenv("ALLOWED_READ_IMAGE", -1)) != -1:
    assert read_image_count == allowed_read_image, f"different read_image! {read_image_count=}, {allowed_read_image=}"
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
-    assert gated_read_image_count <= allowed_gated_read_image, f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

  with open(OUTPUT, "wb") as f:
    pickle.dump(run_onnx_jit, f)
@@ -1,8 +1,8 @@
-import sys, onnx
+import sys
 from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.kernelize.kernelize import get_kernelize_map
+from tinygrad.schedule.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule

@@ -12,10 +12,8 @@ OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/comm
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"

 if __name__ == "__main__":
-  fn = fetch(OPENPILOT_MODEL)
  onnx_file = fetch(OPENPILOT_MODEL)
-  onnx_model = onnx.load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)

  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
@@ -321,7 +321,7 @@ if __name__ == "__main__":
        log_spec = prep_audio(total.reshape(1, -1), model.batch_size, truncate=True)
        encoded_audio = model.encoder.encode(Tensor(log_spec))
      # pass the previously inferred tokens as 'prefix' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
-      out = model.decoder(Tensor([lst]), 0, encoded_audio, streaming=True).realize()
+      out = model.decoder(Tensor([lst]), 0, encoded_audio).realize()
      idx = int(out[0,-1].argmax().numpy().item())
      lst.append(idx)
      dec = enc.decode(lst)
@@ -2,13 +2,12 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs

 os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():
  model = YOLO("yolov8n-seg.pt")
  model.export(format="onnx", imgsz=[480,640])
-onnx_model = onnx_load(open("yolov8n-seg.onnx", "rb"))
-run_onnx = OnnxRunner(onnx_model)
+run_onnx = OnnxRunner("yolov8n-seg.onnx")
 run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)
@@ -8,8 +8,6 @@ from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA

-AM_VERSION = 0xA0000005
-
 def bold(s): return f"\033[1m{s}\033[0m"

 def trim(s:str, length:int) -> str:
@@ -73,21 +71,12 @@ class AMSMI(AMDev):
    self._run_discovery()
    self._build_regs()

-    if self.reg("regSCRATCH_REG7").read() != AM_VERSION:
+    if self.reg("regSCRATCH_REG7").read() != AMDev.Version:
      raise Exception(f"Unsupported AM version: {self.reg('regSCRATCH_REG7').read():x}")

-    self.is_booting, self.smi_dev = True, True
+    self.is_booting = True
+    self.init_sw(smi_dev=True)
    self.partial_boot = True # do not init anything
-    self.mm = AMMemoryManager(self, self.vram_size)
-
-    # Initialize IP blocks
-    self.soc:AM_SOC = AM_SOC(self)
-    self.gmc:AM_GMC = AM_GMC(self)
-    self.ih:AM_IH = AM_IH(self)
-    self.psp:AM_PSP = AM_PSP(self)
-    self.smu:AM_SMU = AM_SMU(self)
-
-    for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]: ip.init_sw()

  def read_pci_state(self):
    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
@@ -136,7 +125,7 @@ class SMICtx:
        if d.pci_state == "D0": d._init_from_d0()
        os.system('clear')

-      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AM_VERSION:
+      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AMDev.Version:
        self.devs.remove(d)
        self.opened_pcidevs.remove(d.pcibus)
        os.system('clear')
@@ -295,8 +284,8 @@ if __name__ == "__main__":
          while True:
            try: pid = subprocess.check_output(['sudo', 'lsof', '-t', dev]).decode('utf-8').split('\n')[0]
            except subprocess.CalledProcessError: break
-            if stopped_pids[pid] > 0: time.sleep(0.5)
-            if stopped_pids[pid] == 10:
+            if stopped_pids[pid] > 0: time.sleep(0.1)
+            if stopped_pids[pid] == 64:
              print(f"{dev[8:-5]}: can't stop process {pid}, exitting")
              exit(1)

@@ -15,8 +15,8 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
        u.vin = tuple(n if x == o else x for x in u.vin)
    # pointer indexing
    if u.uop in {UOps.LOAD, UOps.STORE} and u.vin[0].dtype.itemsize > 1:
-      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_before=uops.uops.index(u))
-      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_before=uops.uops.index(u))
+      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_at=uops.uops.index(u))
+      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_at=uops.uops.index(u))
      u.vin = (u.vin[0], ptr) + u.vin[2:]
  #uops.print()

@@ -4,7 +4,7 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, to_mv
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
@@ -68,7 +68,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in

  if not wasm:
    for name,cl in bufs_to_save.items():
-      weight = ''.join(["\\x%02X"%x for x in bytes(cl._buf)])
+      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
@@ -1,96 +1,39 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
+# sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
-import numpy as np
 from dataclasses import replace
-from tinygrad import Tensor, Device, Context
+from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.helpers import getenv
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem
-from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
-
-# TODO: on METAL for `DEBUG=4 python3 extra/gemm/amd_matmul.py`
-#  * fix load grouping (like float4). idk why it's not working, need new devectorizer (this is a Monday project)
-#  * DONE - remove extra barrier
-#  * DONE (moved Ops.ADD) - fix load order to be in order (the +0 one is last!)
-#  * explore async (fast) global load -> local store
-#  * why is TC=3 broken for 4096x4096?
-#  * write syntactic sugar for these local additions + use it in tensor core kernel.py
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 N = 4096
-LN = 16
 run_count = 5

-from tinygrad.shape.shapetracker import ShapeTracker, View
-def transform_load(ctx:tuple[Kernel, set[UOp]], x:UOp):
-  if x.src[0].op is not Ops.DEFINE_GLOBAL: return None
-  if x in ctx[1]: return None
-  print(ctx[0].colored_shape())
-  ctx[1].add(x)
-  input_st: ShapeTracker = x.src[1].arg
-  #strides = input_st.real_strides()
-  #strides = (0,0)+strides[2:]
-  if input_st.real_strides()[2] == 0:
-    perm = (0,1,5,3,4,2)
-    strides = (0,0,LN*4,4,0,0,1,0)
-  elif input_st.real_strides()[3] == 0:
-    perm = (0,1,2,5,4,3)
-    strides = (0,0,LN*4,4,0,0,0,1)
-  else:
-    return None
-  if len(input_st.shape) == 8:
-    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1,4,4), strides),))
-    perm = perm + (6,7)
-  else:
-    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
-  #local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
-  load_st = local_st.permute(perm)
-  input_st = input_st.permute(perm)
-  lcl = UOp(Ops.DEFINE_LOCAL, x.dtype.ptr(local_st.real_size(), local=True), (), f"temp{x.src[0].arg}")
-  global_load = x.replace(src=(x.src[0], input_st.to_uop()))
-  ret = UOp(Ops.STORE, src=(lcl, local_st.to_uop(), global_load))
-  return UOp(Ops.LOAD, x.dtype, src=(lcl, load_st.to_uop(), ret))
-
-local_loads_pm = PatternMatcher([
-  (UPat(Ops.LOAD, name="x"), transform_load),
-])
-
-def ast_transform(k, ast):
-  #return ast
-  ast = graph_rewrite(ast, local_loads_pm, ctx=(k, set()))
-  #ast = ast.replace(arg=replace(ast.arg, upcasted=0))
-  print(ast)
-  return ast
-
 if __name__ == "__main__":
-  rng = np.random.default_rng()
-  a = Tensor(na:=rng.random((4096, 4096), dtype=np.float32)).realize()
-  b = Tensor(nb:=rng.random((4096, 4096), dtype=np.float32)).realize()
-  c = a @ b
-  si = c.schedule()[-1]
-  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
-  #opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16),
-  #        Opt(op=OptOps.LOCAL, axis=0, arg=8),
-  #        Opt(op=OptOps.UPCAST, axis=2, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=1, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=0, arg=2)]
-  #opts = [Opt(op=OptOps.UPCAST, axis=1, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=0, arg=4),
-  #        Opt(op=OptOps.LOCAL, axis=1, arg=8),
-  #        Opt(op=OptOps.LOCAL, axis=0, arg=4)]
-  opts = [Opt(op=OptOps.UNROLL, axis=0, arg=LN),
-          #Opt(op=OptOps.UPCAST, axis=0, arg=4),
-          #Opt(op=OptOps.UPCAST, axis=1, arg=4),
-          Opt(op=OptOps.LOCAL, axis=1, arg=LN),
-          Opt(op=OptOps.LOCAL, axis=0, arg=LN)]
-  k.apply_opts(opts)
-  prg = k.to_program(ast_transform=ast_transform)
-  if getenv("FAST", 1) and Device.DEFAULT == "AMD":
-    #src = (pathlib.Path(__file__).parent / "fp32_sgemm_amd" / "src" / "kernel8_batched_gmem.s").read_text()
-    src = (pathlib.Path(__file__).parent / "kernel8_batched_gmem.s").read_text()
-    prg = replace(prg, src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
-  print(prg.global_size, prg.local_size)
-  ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)
+  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
+  prg = get_program(ast, Device.default.renderer)
+
+  if getenv("ASM") == 1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
+    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  elif getenv("ASM") == -1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
+    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+  else:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
+    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  runner = CompiledRunner(prgfast)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  c = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  ei = ExecItem(runner, [a.uop.buffer, b.uop.buffer, c.uop.buffer])
  with Context(DEBUG=2):
    for _ in range(run_count): ei.run(wait=True)
-  nc = c.numpy()
-  np.testing.assert_allclose(na@nb, nc, rtol=1e-5)
+  print(f"custom  {(c-tc).square().mean().item()}")
@@ -0,0 +1,142 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 256
+extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 64;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM];
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    __syncthreads();
+
+    // We populate the Shared Memory with Ks row and columns
+    for (int i = 0; i < nbReadsB; i++) {
+      int index_x = BN * blockIdx.x + rBIdx;
+      int index_y = rBIdy + i * strideReadB + kId;
+      Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+    }
+
+    for (int i = 0; i < nbReadsA; i++) {
+      int index_x = rAIdx + kId;
+      int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+      As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+    }
+
+    __syncthreads();
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}
@@ -0,0 +1,172 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 128
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel5_lds_optim(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size . 4x4
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 128;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM+4]; // 4 padding to avoid bank conflicts
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  // initial copy into shared memory
+  for (int i = 0; i < nbReadsB; i++) {
+    int index_x = BN * blockIdx.x + rBIdx;
+    int index_y = rBIdy + i * strideReadB;
+    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+  }
+  for (int i = 0; i < nbReadsA; i++) {
+    int index_x = rAIdx;
+    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+  }
+
+  __syncthreads();
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    float regA[nbReadsA];
+    float regB[nbReadsB];
+    if (kId < N - BK) {
+      // We populate the Shared Memory with Ks row and columns
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        regB[i] = b[N * index_y + index_x];
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        regA[i] = a[N * index_y + index_x];
+      }
+    }
+
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    if (kId < N - BK) {
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        Bs[index_y % BK][index_x % BN] = regB[i]; // row
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        As[(index_x % BK)][(index_y % BM)] = regA[i];
+      }
+      __syncthreads();
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}
@@ -1,6 +1,5 @@
 	.text
 	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
-	;.amdhsa_code_object_version 5
 	.protected	kernel                  ; -- Begin function kernel
 	.globl	kernel
 	.p2align	8
@@ -9,7 +8,7 @@ kernel:                                 ; @kernel
 ; %bb.0:                                ; %.preheader193

 	;; Init code for matrix A and B buffer Loads  - START
-	s_load_b128 s[20:23], s[0:1], 0x8 ; Matrix A and B
+	s_load_b128 s[20:23], s[0:1], 0x0 ; Matrix A and B
 	s_waitcnt lgkmcnt(0)

 	; Matrix B offsets:
@@ -76,14 +75,12 @@ kernel:                                 ; @kernel


 	s_clause 0x1
-	;s_load_b128 s[4:7], s[0:1], 0x18 ; N, alpha, beta, ???
-	s_load_b128 s[8:11], s[0:1], 0x8 ; Matrix A and B
-
-	s_mov_b32 s4, 4096 ; hardcode 4096
-	s_mov_b32 s5, 0x3f800000 ; alpha
-	s_mov_b32 s6, 0    ; beta
-	s_mov_b32 s7, 0
-
+	; s_load_b128 s[4:7], s[0:1], 0x18
+	; N=4096, alpha=1.0, beta=0.0
+	s_mov_b32 s4, 4096
+	s_mov_b32 s5, 0x3F800000
+	s_mov_b32 s6, 0
+	s_load_b128 s[8:11], s[0:1], 0x0
 	s_lshl_b32 s2, s14, 7
 	v_lshrrev_b32_e32 v4, 3, v0
 	v_or_b32_e32 v1, s2, v0
@@ -93,7 +90,7 @@ kernel:                                 ; @kernel
 	v_or_b32_e32 v22, s3, v4
 	v_ashrrev_i32_e32 v2, 31, v1
 	s_lshr_b32 s12, s12, 25
-	s_load_b64 s[0:1], s[0:1], 0  ; Matrix C
+	s_load_b64 s[0:1], s[0:1], 0x10
 	v_lshlrev_b32_e32 v135, 2, v118
 	s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_3) | instid1(VALU_DEP_3)
 	v_lshlrev_b64 v[5:6], 2, v[1:2]
@@ -463,7 +460,7 @@ kernel:                                 ; @kernel

 	v_mov_b32_e32 v5, 0
 	v_mov_b32_e32 v3, 0
-	s_add_i32 s7, s4, -1
+	s_add_i32 s7, s4, -8
 	s_add_u32 s8, s8, 32
 	s_addc_u32 s9, s9, 0
 	s_mov_b32 s12, 0
@@ -2398,18 +2395,9 @@ amdhsa.kernels:
        .offset:         16
        .size:           8
        .value_kind:     global_buffer
-      - .offset:         24
-        .size:           4
-        .value_kind:     by_value
-      - .offset:         28
-        .size:           4
-        .value_kind:     by_value
-      - .offset:         32
-        .size:           4
-        .value_kind:     by_value
    .group_segment_fixed_size: 8320
    .kernarg_segment_align: 8
-    .kernarg_segment_size: 36
+    .kernarg_segment_size: 24
    .language:       OpenCL C
    .language_version:
      - 2
@@ -0,0 +1,148 @@
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.dtype import AddrSpace
+
+N = 4096
+run_count = 5
+
+def hand_spec_kernel3():
+  # block tile size
+  BN = 128
+  BM = 128
+  # number of row/column we read per batch
+  BK = 8
+
+  # thread tile size 4x4
+  TN = 4
+  TM = 4
+
+  BLOCK_SIZE = 128
+  nbWaves = BLOCK_SIZE // 32
+  # wave tile size
+  WN = 128
+  WM = BN * BM // nbWaves // WN
+
+  nbWaveX = BN // WN
+  nbWaveY = BM // WM
+
+  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
+  waveIndex = threadIdx_x // 32
+  waveIdx = waveIndex % nbWaveX
+  waveIdy = waveIndex // nbWaveX
+  indexInWave = threadIdx_x % 32
+
+  nbThreadXPerWave = 8
+  nbThreadYPerWave = 4
+
+  idxInWave = indexInWave % nbThreadXPerWave
+  idyInWave = indexInWave // nbThreadXPerWave
+
+  nbIterWaveN = WN // (nbThreadXPerWave * TN)
+  nbIterWaveM = WM // (nbThreadYPerWave * TM)
+
+  SUBWN = WN // nbIterWaveN
+  SUBWM = WM // nbIterWaveM
+
+  # Thread mapping to read BKxBN block from A
+  rAIdx = threadIdx_x % BK
+  rAIdy = threadIdx_x // BK
+  # Thread mapping to read BNxBK block from B
+  rBIdx = threadIdx_x % BN
+  rBIdy = threadIdx_x // BN
+
+  strideReadB = BLOCK_SIZE // BN
+  strideReadA = BLOCK_SIZE // BK
+  nbReadsB = BN * BK // BLOCK_SIZE
+  nbReadsA = BM * BK // BLOCK_SIZE
+
+  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
+  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
+
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+
+  junk = UOp.const(dtypes.float, 0)
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), src=(junk,), arg=0)
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), src=(junk,), arg=1)
+
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*(BM+4), AddrSpace.LOCAL), arg=0) # 4 padding to avoid bank conflicts
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
+
+  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), src=(junk,), arg=2)
+
+  kId_range = UOp.range(dtypes.int, N//BK, 0)
+  kId = kId_range*BK
+
+  # load from globals into locals
+  i = UOp.range(dtypes.int, nbReadsB, 1)
+  index_x = BN * blockIdx_x + rBIdx
+  index_y = rBIdy + i * strideReadB + kId
+  Bs_store = Bs[(index_y % BK)*BN + index_x%BN].store(b[N * index_y + index_x].load(), i)
+
+  i = UOp.range(dtypes.int, nbReadsA, 2)
+  index_x = rAIdx + kId
+  index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+  As_store = As[(index_x % BK)*(BM+4) + index_y%BM].store(a[N * index_y + index_x].load(), i)
+
+  barrier = UOp(Ops.BARRIER, src=(As_store, Bs_store))
+
+  k = UOp.range(dtypes.int, BK, 3)
+
+  # load from locals into registers
+  iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
+  i = UOp.range(dtypes.int, TN, 5)
+  index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
+  B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
+
+  iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
+  i = UOp.range(dtypes.int, TM, 7)
+  index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
+  A_col_store = A_col[iterWave*TM + i].store(As[k*(BM+4) + index].load(barrier), iterWave, i)
+
+  # do the GEMM math
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 9)
+  yt = UOp.range(dtypes.int, TM, 10)
+  xt = UOp.range(dtypes.int, TN, 11)
+  x = iterWaveN * TN + xt
+  y = iterWaveM * TM + yt
+  c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+  sink = c_regs_idx.store(c_regs_idx.load() + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                          iterWaveM, iterWaveN, yt, xt, k, kId_range)
+
+  # store c_regs into c
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 12)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 13)
+  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
+  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
+  yt = UOp.range(dtypes.int, TM, 14)
+  xt = UOp.range(dtypes.int, TN, 15)
+  indexC = N * (yOut + yt) + xOut + xt
+  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink), iterWaveM, iterWaveN, yt, xt)
+
+  return sink.sink(arg=KernelInfo(name="tinygemm"))
+
+if __name__ == "__main__":
+  hprg = hand_spec_kernel3()
+  prg = get_program(hprg, Device.default.renderer)
+  print(prg.src)
+  hrunner = CompiledRunner(prg)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  hc = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  ei = ExecItem(hrunner, [a.uop.buffer, b.uop.buffer, hc.uop.buffer])
+  with Context(DEBUG=2):
+    for _ in range(run_count): ei.run(wait=True)
+  err = (hc-tc).square().mean().item()
+  print(f"hrunner {err}")
+  if err > 1e-06: raise RuntimeError("matmul is wrong!")
@@ -2,6 +2,7 @@ import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
+from tinygrad.engine.realize import get_program

 # for copied uops
 from tinygrad.opt.kernel import Kernel, KernelOptError
@@ -55,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
  k = Kernel(ast)
  k.apply_opts(opts)
-  p = k.to_program()
+  p = get_program(k.get_optimized_ast(), k.opts)
  return CUDAProgram(device, p.function_name, compiler.compile(p.src))

 if __name__ == "__main__":
@@ -1,7 +1,7 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from dataclasses import replace

 N = 4096
@@ -29,7 +29,7 @@ if __name__ == "__main__":
      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
  k.apply_opts(opts)
-  prg = k.to_program()
+  prg = get_program(k.get_optimized_ast(), k.opts)
  new_src = prg.src
  # can mod source here
  prg = replace(prg, src=new_src)
@@ -40,7 +40,7 @@ sched = C.schedule()
 from tinygrad.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
-lin.linearize()
+lin.to_program()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)
 print(src)
@@ -5,7 +5,7 @@ start = time.perf_counter()
 # *** ioctl lib ***
 libc = ctypes.CDLL(ctypes.util.find_library("c"))
 # platform.processor calls `uname -p` which can return `unknown` on some systems
-processor = os.getenv("IOCTL_PROCESSOR") or platform.processor()
+processor = os.getenv("IOCTL_PROCESSOR") or platform.processor() or platform.machine()
 IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]

 def get_struct(argp, stype):
@@ -1,6 +1,6 @@
 import onnx, yaml, tempfile, time, collections, pprint, argparse, json
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs

@@ -13,8 +13,7 @@ def get_config(root_path: Path):
  return ret

 def run_huggingface_validate(onnx_model_path, config, rtol, atol):
-  onnx_model = onnx_load(onnx_model_path)
-  onnx_runner = OnnxRunner(onnx_model)
+  onnx_runner = OnnxRunner(onnx_model_path)
  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)

@@ -46,7 +45,7 @@ def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
  for model_id, (root_path, relative_path) in models.items():
    print(f"examining {model_id}")
    model_path = root_path / relative_path
-    onnx_runner = OnnxRunner(onnx.load(model_path))
+    onnx_runner = OnnxRunner(model_path)
    for node in onnx_runner.graph_nodes:
      op_counter[node.op] += 1
      if node.op not in supported_ops:
@@ -7,6 +7,7 @@ from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache
 from tinygrad.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
 from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.engine.realize import get_program

 class MCTSNode:
  def __init__(self, kernel:Kernel, parent=None):
@@ -110,7 +111,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
      seen_asts[opt_ast.key] = node

      # lowering (50% of the time)
-      p = node.kernel.to_program(name_override="test")
+      p = get_program(node.kernel.get_optimized_ast(name_override="test"), node.kernel.opts)

      # rollout
      tm1 = time.perf_counter()
@@ -27,6 +27,20 @@
 #include "gpu/vbios/bios_types.h"
 #define FALCON_APPLICATION_INTERFACE_ENTRY_ID_DMEMMAPPER     (0x4)

+typedef struct
+{
+    NvU8 version;
+    NvU8 headerSize;
+    NvU8 entrySize;
+    NvU8 entryCount;
+} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_HEADER_V1;
+
+typedef struct
+{
+    NvU32 id;
+    NvU32 dmemOffset;
+} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_ENTRY_V1;
+
 typedef struct
 {
    NvU32 signature;
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys
+
+def scan_devs_based_on_lock(prefix:str) -> list[str]:
+  devs = []
+  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
+    dev_id = dev[8:-5]
+    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
+  if not to_unload:
+    print("NVIDIA kernel modules are not loaded")
+  else:
+    print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  if not os.path.exists("/sys/module/nvidia"):
+    print("Inserting nvidia kernel module")
+    subprocess.run(["nvidia-smi"], check=True)
+  else: print("Nvidia kernel module already loaded")
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock("nv")
+  dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  for dev in devs:
+    if dev.startswith(dev_to_reset):
+      print(f"Resetting device {dev}")
+      _do_reset_device(dev)
+      time.sleep(0.2)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  subparsers = parser.add_subparsers(required=True, dest="cmd")
+
+  parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
+  parser_insmod.set_defaults(func=cmd_insert_module)
+
+  parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
+  parser_rmmod.set_defaults(func=cmd_remove_module)
+
+  parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
+  parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
+  parser_reset.set_defaults(func=cmd_reset_devices)
+
+  args = parser.parse_args()
+  if args.cmd is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)
@@ -1,10 +1,12 @@
+# mypy: disable-error-code="misc, list-item, assignment, attr-defined, operator, index, arg-type"
 from types import SimpleNamespace
-from typing import Any, Sequence, cast, Literal, Callable
-import dataclasses, functools, io, math, types, warnings, sys
+from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
+import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
-from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort
+from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
 from tinygrad.device import is_dtype_supported, Device
+from extra.onnx_parser import onnx_load

 # https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
 data_types: dict[int, DType] = {
@@ -24,8 +26,7 @@ attribute_types: dict[int, Callable] = {
 }

 # ***** protobuf parsing ******
-from onnx import AttributeProto, ModelProto, TensorProto, TypeProto, helper
-import numpy as np
+from onnx import AttributeProto, TensorProto, TypeProto

 def has_field(onnx_type: TypeProto|SimpleNamespace, field):
  if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
@@ -48,30 +49,19 @@ def attribute_parse(onnx_attribute: AttributeProto):

 def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
  if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
-  dtype, shape = dtype_parse(onnx_tensor.data_type, "buffer parse"), tuple(onnx_tensor.dims)
-  data = None
-  if len(onnx_tensor.float_data): data = onnx_tensor.float_data
-  elif len(onnx_tensor.int32_data): data = onnx_tensor.int32_data
-  elif len(onnx_tensor.int64_data): data = onnx_tensor.int64_data
-  elif len(onnx_tensor.double_data): data = onnx_tensor.double_data
-  elif len(onnx_tensor.uint64_data): data = onnx_tensor.uint64_data
-  if isinstance(data, Tensor):
-    if len(data) == 1: return Tensor(data.tolist()[0], dtype=dtype).reshape(shape)
-    return data.cast(dtype).reshape(shape).to(Device.DEFAULT)
-  if has_field(onnx_tensor, "raw_data"):
-    raw_data = onnx_tensor.raw_data
-    if not isinstance(raw_data, Tensor): raw_data = Tensor(raw_data)
-    if not is_dtype_supported(data_types[onnx_tensor.data_type]):
-      np_buffer = np.frombuffer(raw_data.data().tobytes(),
-                                dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
-      if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
-      return Tensor(np_buffer, dtype=dtype)
-    ret = raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
-    if shape == ():
-      if ret.dtype is dtypes.float16 and sys.version_info < (3, 12): ret = ret.cast(dtypes.float32)
-      ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
-    return ret
-  return Tensor(None)
+  to_dtype, true_dtype = dtype_parse(onnx_tensor.data_type, "buffer parse"), data_types[onnx_tensor.data_type]
+  shape = tuple(onnx_tensor.dims)
+  keys = ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', "raw_data"]
+  data = next((val for k in keys if (val := getattr(onnx_tensor, k)) is not None), None)
+  if data is None: raise RuntimeError("empty buffer")
+  if not isinstance(data, Tensor): return Tensor(data, dtype=to_dtype).reshape(shape)
+  assert data.dtype is dtypes.uint8, data.dtype
+  data = data.bitcast(true_dtype).reshape(shape)
+  data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
+  if shape == ():
+    if data.dtype is dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
+    return Tensor(data.item(), dtype=to_dtype).reshape(shape)
+  return data

 def type_parse(onnx_type: TypeProto):
  elem_type = onnx_type
@@ -94,10 +84,24 @@ class OnnxValue:
  is_optional: bool
  is_sequence: bool

+class Domain(enum.Enum):
+  ONNX = "ai.onnx"
+  ONNX_ML = "ai.onnx.ml"
+  AI_ONNX_TRAINING = "ai.onnx.training"
+  AI_ONNX_PREVIEW_TRAINING = "ai.onnx.preview.training"
+  MICROSOFT_CONTRIB_OPS = "com.microsoft"
+  @classmethod
+  def from_onnx(cls, domain: str | None) -> "Domain": return cls.ONNX if domain is None or domain == "" else cls(domain)
+
+class OpSetId(NamedTuple):
+  domain: Domain
+  version: int
+
@dataclasses.dataclass(frozen=True)
 class OnnxNode:
  num: int
  op: str
+  opset_id: OpSetId
  inputs: tuple[str, ...]
  outputs: tuple[str, ...]
  opts: dict[str, Any]
@@ -132,17 +136,27 @@ def to_python_const(t:Any, op:str, idx:int) -> list[ConstType]|ConstType|bytes:
 debug = int(getenv("DEBUGONNX", "0"))
 limit = int(getenv("ONNXLIMIT", "-1"))
 class OnnxRunner:
-  def __init__(self, model: ModelProto|SimpleNamespace):
-    # parse model protobuf
-    self.is_training = any(n.domain in {"ai.onnx.training", "ai.onnx.preview.training"} for n in model.graph.node)
+  """
+  `OnnxRunner` executes an ONNX model using Tinygrad.
+
+  Args:
+    model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
+  """
+  def __init__(self, model_path: Tensor | str | pathlib.Path):
+    model = onnx_load(model_path)
+    self.is_training = any(n.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in model.graph.node)
    self.old_training = Tensor.training
    Tensor.training = True if self.is_training else False
    self.graph_values = {"": None, **{x.name:buffer_parse(x) for x in model.graph.initializer}}
    self.graph_inputs = {x.name:type_parse(x.type) for x in model.graph.input if x.name not in self.graph_values}
    self.graph_outputs = tuple(x.name for x in model.graph.output)
-    self.graph_nodes = tuple(OnnxNode(num, n.op_type, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute})
-                       for num,n in enumerate(model.graph.node))
-    self.opset_version = model.opset_import[0].version
+    opset_imports = {Domain.from_onnx(getattr(x, "domain", "")):x.version for x in model.opset_import}
+    self.graph_nodes = []
+    for num, n in enumerate(model.graph.node):
+      domain = Domain.from_onnx(n.domain)
+      opset_id = OpSetId(domain, opset_imports.get(domain, 1))
+      self.graph_nodes.append(OnnxNode(num, n.op_type, opset_id, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute}))
+    self.graph_nodes = tuple(self.graph_nodes)
    self.variable_dims: dict[str, int] = {}

    self.onnx_ops = onnx_ops
@@ -155,7 +169,7 @@ class OnnxRunner:
      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
      if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
      return sequence
-    dtype = _from_np_dtype(value.dtype) if str(type(value)) == "<class 'numpy.ndarray'>" else spec.dtype
+    dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) else spec.dtype
    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
    for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
@@ -164,20 +178,25 @@ class OnnxRunner:
      if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
    return tensor

-  def _dispatch_op(self, op, inps, opts):
-    if op in self.onnx_ops:
-      fxn = self.onnx_ops[op]
-      if isinstance(fxn, dict):
-        for k in sorted(fxn.keys()):
-          if k <= self.opset_version:
-            real_fxn = fxn[k]
-      else: real_fxn = fxn
-      return real_fxn(*inps, **opts)
-    raise NotImplementedError(f"{op=} not supported")
+  def _select_op(self, op:str, required_opset:OpSetId) -> types.FunctionType:
+    if op not in self.onnx_ops: raise NotImplementedError(f"{op=} is not supported")
+    # return default implementation if no opset_id is specified
+    if isinstance(impl := self.onnx_ops[op], types.FunctionType): return impl
+    # match domain and select implementation with latest compatible version
+    eligible_ops = {impl_opset.version:impl_fxn for impl_opset,impl_fxn in impl.items()
+                    if impl_opset.domain == required_opset.domain and impl_opset.version <= required_opset.version}
+    if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
+    return eligible_ops[max(eligible_ops.keys())]

  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}

+  def to(self, device:str|None):
+    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
+    self.graph_nodes = tuple(OnnxNode(n.num, n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
+                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
+    return self
+
  def __call__(self, inputs:dict[str, Any], debug=debug):
    for name, input_spec in self.graph_inputs.items():
      if name not in inputs: raise RuntimeError(f"Please provide input data for {name}")
@@ -193,7 +212,7 @@ class OnnxRunner:

      if debug >= 1: print(f"{node.num}: op '{node.op}' opt {opts}")
      if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
-      ret = self._dispatch_op(node.op, inps, opts)
+      ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
      ret = ret if isinstance(ret, tuple) else (ret,)
      if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{x} - {o!r}" for x,o in zip(node.outputs, ret)))

@@ -208,8 +227,10 @@ class OnnxRunner:
 ####################
 ##### ONNX OPS #####
 ####################
-def get_onnx_ops():
+def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
  # ***** helper functions *****
+  def _resolve_const(x: Sequence[ConstType]|ConstType): return x if isinstance(x, get_args(ConstType)) else get_single_element(x)
+
  def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)

  # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
@@ -280,7 +301,8 @@ def get_onnx_ops():
    if value_string is not None or value_strings is not None and sparse_value is not None:
      raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value')

-  def Range(start:float|int, limit:float|int, delta:float|int): return Tensor.arange(start=start, stop=limit, step=delta)
+  def Range(start:float|int|list[float|int], limit:float|int|list[float|int], delta:float|int|list[float|int]):
+    return Tensor.arange(start=_resolve_const(start), stop=_resolve_const(limit), step=_resolve_const(delta))

  def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
    try: import PIL.Image
@@ -307,13 +329,13 @@ def get_onnx_ops():

  # ***** Unary Ops (math) *****
  def Not(x:Tensor): return x.logical_not()
-  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)
+  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002
  def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))

  # ***** Unary Ops (activation) *****
-  def Softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
-  def Softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
-  Softmax = {1:Softmax_1, 13:Softmax_13}
+  def softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
+  def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
+  Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
  def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
  def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
  def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
@@ -450,7 +472,7 @@ def get_onnx_ops():
                        zip(strides, input_shape, output_padding, kernel_shape, dilations, output_shape)], auto_pad)
    if pads is None: # we generate pads
      output_shape = output_shape or [X.shape[i+2] * strides[i] for i in range(len(strides))]
-      pads = [strides[i]*(input_shape[i]-1) + output_padding[i] + ((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
+      pads = [strides[i]*(input_shape[i]-1)+output_padding[i]+((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
      pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape) * 2
    pads = _onnx_pads_to_tiny_pads(pads)
    return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads, output_padding=output_padding)
@@ -468,14 +490,16 @@ def get_onnx_ops():

  def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)

-  def CumSum(X:Tensor, axis:int|list, exclusive:int=0, reverse:int=0):
-    axis = X._resolve_dim(axis[0] if isinstance(axis, list) else axis)
+  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
+    axis = X._resolve_dim(_resolve_const(axis))
    if reverse: X = X.flip(axis)
    if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
                        .shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
    return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)

-  def Trilu(x:Tensor, k:int=0, upper:int=1): return x.triu(k) if upper else x.tril(k)
+  def Trilu(x:Tensor, k:int|list[int]=0, upper:int=1):
+    k_ = _resolve_const(k)
+    return x.triu(k_) if upper else x.tril(k_)

  def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
            axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
@@ -536,8 +560,8 @@ def get_onnx_ops():
    return X.permute(*argsort(perm)) if perm else X
  def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated

-  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):
-    val, idx = X.topk(K if isinstance(K, int) else K[0], axis, largest, sorted)
+  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002
+    val, idx = X.topk(_resolve_const(K), axis, largest, sorted)
    return val, idx.cast(dtypes.int64)

  # ***** Neural Network Ops *****
@@ -599,9 +623,9 @@ def get_onnx_ops():
  def MeanVarianceNormalization(x:Tensor, axis:list[int]=[0,2,3]):
    return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)

-  def OneHot(indices:Tensor, depth:float|int|list, values:Tensor, axis:int=-1):
+  def OneHot(indices:Tensor, depth:float|int|list[int|float], values:Tensor, axis:int=-1):
    # Scalar or Rank 1 tensor containing exactly one element
-    depth = int(depth[0] if isinstance(depth, list) else depth)
+    depth = int(_resolve_const(depth))
    indices = indices.int()
    indices = (indices < 0).where(indices+depth, indices)
    return indices.unsqueeze(axis)._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
@@ -612,7 +636,8 @@ def get_onnx_ops():
    return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)

  # Reimplemented here because you need legacy RNG for passing ONNX tests.
-  def Dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
+  def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
+    import numpy as np
    if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
    if seed is not None:
      rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
@@ -621,8 +646,8 @@ def get_onnx_ops():
    mask = rand >= ratio
    return data * mask / (1.0 - ratio), mask
  # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
-  def Dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return Dropout_7(data, ratio, training_mode=not is_test)
-  Dropout = {6:Dropout_6, 7:Dropout_7}
+  def dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return dropout_7(data, ratio, training_mode=not is_test)
+  Dropout = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 7):dropout_7}

  def LRN(x:Tensor, size:int, alpha:float=1e-4, beta:float=0.75, bias:float=1.0):
    pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
@@ -637,16 +662,17 @@ def get_onnx_ops():
  def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
    N, _, *spatial_dims = size
    def generate_grid(steps):
-      return Tensor.linspace(-1, 1, steps, device=theta.device) if align_corners else Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
+      if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
+      return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
    grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
    base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
    base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
    return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)

-  def Attention(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None, attention_bias:Tensor|None=None,
-                past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0, num_heads:int|None=None,
-                past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None, rotary_embedding_dim:int|None=None,
-                scale:float|None=None, unidirectional:int=0):
+  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
+                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
+                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
+                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
    assert not do_rotary and not attention_bias, "TODO"
    if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
    qkv = x.linear(weights, bias)
@@ -687,6 +713,86 @@ def get_onnx_ops():
    output = output.transpose(1, 2).reshape(batch_size, seq_len, -1)
    return output, present

+  def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, past_key:Tensor|None=None, past_value:Tensor|None=None,
+                     is_causal:int=0, kv_num_heads:int|None=None, q_num_heads:int|None=None, qk_matmul_output_mode:int=0, scale:float|None=None,
+                     softcap:float=0.0, softmax_precision:int|None=None):
+    input_shape_len = Q.ndim
+    if input_shape_len == 3:
+      assert q_num_heads is not None and kv_num_heads is not None
+      Q = Q.reshape(Q.shape[0], q_num_heads, Q.shape[1], -1)
+      K = K.reshape(K.shape[0], kv_num_heads, K.shape[1], -1)
+      V = V.reshape(V.shape[0], kv_num_heads, V.shape[1], -1)
+
+    if past_key is not None: K = past_key.cat(K, dim=2)
+    if past_value is not None: V = past_value.cat(V, dim=2)
+    present_key, present_value = K, V
+
+    _q_heads, _kv_heads = q_num_heads or Q.shape[1], kv_num_heads or K.shape[1]
+    if _q_heads != _kv_heads:
+      K = K.repeat((1, _q_heads // _kv_heads, 1, 1))
+      V = V.repeat((1, _q_heads // _kv_heads, 1, 1))
+
+    effective_scale = scale if scale is not None else 1.0 / (Q.shape[-1] ** 0.5)
+    scores = (Q @ K.transpose(-1, -2)) * effective_scale
+    qk_matmul_return_val = scores
+
+    if is_causal:
+      causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0)
+      scores = scores.masked_fill(causal_mask.logical_not(), -float("inf"))
+
+    if attn_mask is not None:
+      mask_to_add = attn_mask.where(0, -float("inf")) if attn_mask.dtype == dtypes.bool else attn_mask
+      scores = scores + mask_to_add
+    if qk_matmul_output_mode == 1: qk_matmul_return_val = scores
+
+    if softcap > 0.0: scores = (scores / softcap).tanh() * softcap
+    if qk_matmul_output_mode == 2: qk_matmul_return_val = scores
+
+    if softmax_precision: scores = scores.cast({1: dtypes.float32, 10: dtypes.float16, 16: dtypes.bfloat16}[softmax_precision])
+    qk_softmax = scores.softmax(-1).cast(Q.dtype)
+    if qk_matmul_output_mode == 3: qk_matmul_return_val = qk_softmax
+
+    output = (qk_softmax @ V).cast(Q.dtype)
+    if input_shape_len == 3: output = output.permute(0, 2, 1, 3).reshape(Q.shape[0], Q.shape[2], -1)
+    return output, present_key, present_value, qk_matmul_return_val
+  Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}
+
+  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
+    norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
+    return X * norm * scale
+
+  def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
+                      rotary_embedding_dim:int=0):
+    original_input_shape = X.shape
+
+    if X.ndim == 4: X = X.permute(0, 2, 1, 3)
+    elif X.ndim == 3:
+      assert num_heads is not None, "num_heads must be provided for 3D input"
+      X = X.reshape(*X.shape[:-1], num_heads, X.shape[-1] // num_heads)
+
+    head_size = X.shape[-1]
+    rot_dim = rotary_embedding_dim or head_size
+    x_rotate, x_pass = X[..., :rot_dim], X[..., rot_dim:]
+
+    cos = cos_cache[position_ids] if position_ids is not None else cos_cache[:X.shape[1]]
+    sin = sin_cache[position_ids] if position_ids is not None else sin_cache[:X.shape[1]]
+    cos = cos[..., :rot_dim//2].unsqueeze(2)
+    sin = sin[..., :rot_dim//2].unsqueeze(2)
+
+    if interleaved:
+      x1, x2 = x_rotate[..., ::2], x_rotate[..., 1::2]
+      real = x1 * cos - x2 * sin
+      imag = x1 * sin + x2 * cos
+      x_rotated = Tensor.stack(real, imag, dim=-1).flatten(start_dim=-2)
+    else:
+      x1, x2 = x_rotate.chunk(2, dim=-1)
+      real = x1 * cos - x2 * sin
+      imag = x1 * sin + x2 * cos
+      x_rotated = real.cat(imag, dim=-1)
+
+    output = x_rotated.cat(x_pass, dim=-1)
+    return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
+
  # ***** Indexing Ops *****
  def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]

@@ -1,6 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx import OnnxValue
 import numpy as np
 import onnxruntime as ort
@@ -46,7 +46,7 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  return ret

 def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
-  run_onnx = OnnxRunner(onnx_load(onnx_file))
+  run_onnx = OnnxRunner(onnx_file)

  ort_options = ort.SessionOptions()
  ort_options.log_severity_level = 3
@@ -2,10 +2,9 @@

 import os, pathlib, struct
 from io import BufferedReader
-from typing import Tuple, Union
 from types import SimpleNamespace
 from tinygrad.nn.state import TensorIO
-from tinygrad.tensor import Tensor, dtypes
+from tinygrad.tensor import Tensor

 # Protobuf Wire Types
 WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
@@ -22,13 +21,13 @@ class AttributeType:

 class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702

-PB_INFOS = {
+PB_INFOS: dict[str, dict] = {
  "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
  "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
-  # TODO: support uint64 parsing (11: "uint64_data") and double parsing (10: "double_data")
  "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
    13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
-    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},
+    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES),
+    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
  "TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
  "TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
  "ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
@@ -37,16 +36,16 @@ PB_INFOS = {
    8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
  "GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
    1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
-    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [],
-                                                                  "int64_data": [], "double_data": [], "uint64_data": []})),
+    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
+                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
    11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
  "NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
    4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
    5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
  "AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
    2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
-    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [], "int64_data": [],
-                                                        "double_data": [], "uint64_data": []}))},
+    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
+                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
  "ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
  "TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
    9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
@@ -55,7 +54,7 @@ PB_INFOS = {
  "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
 }

-def onnx_load(fn: Union[Tensor, str, pathlib.Path], load_external_data: bool=True):
+def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
  parser = OnnxParser(fn, load_external_data)
  onnx_model = parser.parse()
  model = dict_to_namespace(onnx_model)
@@ -71,8 +70,8 @@ def dict_to_namespace(d):
  return d

 class OnnxParser:
-  def __init__(self, inp: Union[Tensor, str, pathlib.Path], load_external_data: bool=True):
-    self.file_path: Union[pathlib.Path, None] = None
+  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
+    self.file_path: pathlib.Path|None = None
    self.load_external_data = load_external_data
    if not isinstance(inp, Tensor):
      self.file_path = pathlib.Path(inp)
@@ -90,7 +89,6 @@ class OnnxParser:
        elif len(config) == 4: name, attr, repeated, parser_fn = config
        handler_fn = self.attr_func_dict[attr]
        def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
-        _wrapper_handler._debug_info = f"{fid}, {name} => {handler_fn}"
        res[fid] = _wrapper_handler
      self.registered_handles[pb_name] = res

@@ -131,16 +129,19 @@ class OnnxParser:
    if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
    return obj

-  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tuple[bytes, Tensor]:
+  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
    str_len = self.decode_varint(reader)
    if not use_tensor: return reader.read(str_len)
-    res = reader.raw._tensor[reader.tell():(reader.tell()+str_len)]
+    raw = reader.raw
+    assert isinstance(raw, TensorIO)
+    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
    reader.seek(str_len, os.SEEK_CUR)
    return res

  def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
    value = self._handle_delimited(reader)
+    assert isinstance(value, bytes)
    gen_result(obj, key_name, value.decode("utf-8"), repeated)

  def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
@@ -166,16 +167,17 @@ class OnnxParser:
    while reader.tell() < total_bytes_len + old_pos:
      val = self.decode_varint(reader) # need copy here because packed ints are varint
      values.append(val - 2**64 if val & (1 << 63) else val)
-    obj[key_name] = Tensor(values, dtype=dtypes.int64)
+    obj[key_name] = values

  def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
    value = self._handle_delimited(reader, use_tensor=True)
-    obj[key_name] = value.bitcast(dtypes.float32)
+    obj[key_name] = value

  def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
    if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
    value = self._handle_delimited(reader, use_tensor=True)
+    assert isinstance(value, Tensor)
    if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
    elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
    else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
@@ -194,7 +196,7 @@ class OnnxParser:
    if self.file_path is None:
      # get onnx file path from Tensor
      if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
-        self.file_path = self.tensor.device[5:]
+        self.file_path = pathlib.Path(self.tensor.device[5:])
        if not (ext_path := self.file_path.parent.joinpath(location)).exists():
          raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
      else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")
@@ -13,5 +13,5 @@ GPU=1 python3 -m pytest test/test_tiny.py
 extra/optimization/extract_dataset.py
 sort -u /tmp/ops > /tmp/sops
 ls -lh /tmp/ops /tmp/sops
-# gzip -k /tmp/sops
+gzip -k /tmp/sops
 # mv /tmp/sops.gz extra/datasets/
@@ -6,6 +6,7 @@ from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.helpers import getenv
+from tinygrad.engine.realize import get_program
 inf, nan = float('inf'), float('nan')
 UOps = Ops

@@ -115,7 +116,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_

  rawbufs = _ensure_buffer_alloc(rawbufs)
  var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
-  p = lin.to_program()
+  p = get_program(lin.get_optimized_ast(), lin.opts)
  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))

@@ -1,6 +1,6 @@
 import sys, pickle, decimal, json
-from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent
-from tinygrad.helpers import tqdm, temp
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
+from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent

 devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
 def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
@@ -4,7 +4,7 @@ from tinygrad import Device, Context, Tensor, GlobalCounters
 from tinygrad.device import Buffer
 from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
-from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
 from tinygrad.renderer import ProgramSpec
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.opt.heuristic import hand_coded_optimizations
@@ -58,7 +58,7 @@ if __name__ == "__main__":
            GlobalCounters.kernel_count -= 1

        if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = k.to_program()
+        p2 = get_program(k.get_optimized_ast(), k.opts)
        new_ei = replace(ei, prg=CompiledRunner(p2))
        new_ei.run()
        new_jit.append(new_ei)
@@ -26,7 +26,7 @@ extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp

 Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding

-If multiplle gpus are used you can select which one to export with `-d` like this:
+If multiple gpus are used you can select which one to export with `-d` like this:

 ```bash
 extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
@@ -352,6 +352,11 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
  unwrap(indices).assign(out_indices.cast(dtypes.int64))
  return wrap(out_values), wrap(out_indices)

+@torch.library.impl("aten::_linalg_svd", "privateuseone")
+def _linalg_svd(self, full_matrices=False):
+  U, S, Vh = unwrap(self).svd(full_matrices)
+  return wrap(U), wrap(S), wrap(Vh)
+
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
@@ -412,6 +417,7 @@ decomps = [
  #aten.lgamma,
  # this needs copy_strided
  #aten.lerp,
+  aten.norm,
 ]
 for k,v in get_decompositions(decomps).items():
  key = str(k._schema).split("(")[0]
@@ -473,6 +479,7 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
  "aten.fmax.out": lambda input,other: Tensor.where(input.isnan() & ~other.isnan(), other, Tensor.where(~input.isnan() & other.isnan(), input, Tensor.maximum(input, other))),
  "aten.fmin.out": lambda input,other: Tensor.where(input.isnan() & ~other.isnan(), other, Tensor.where(~input.isnan() & other.isnan(), input, Tensor.minimum(input, other))),
  "aten.amax.out": lambda self,dim=None: self.max(axis=dim),
+  "aten.amin.out": lambda self,dim=None: self.min(axis=dim),
  # TODO: this gets the shape wrong
  #"aten.arange.start_out": Tensor.arange,
  "aten.lerp.Scalar_out": Tensor.lerp,
@@ -103,6 +103,27 @@ class TestTorchBackend(unittest.TestCase):
    expected = np.array([[4.7, 12.9, 12.3], [16.9, 24.9, 23.6]], dtype=np.float32)
    np.testing.assert_equal(y3.cpu().numpy(), expected)

+
+  def test_amin(self):
+    x = torch.tensor([[[ 1.5,  2.3,  3.1,  4.7],
+                       [ 5.2,  6.8,  7.4,  12.9],
+                       [ 9.0, 12.3, 11.6, 10.1]],
+                      [[13.2, 16.9, 15.5, 14.1],
+                       [17.1, 24.9, 19.8, 20.2],
+                       [21.0, 22.3, 23.6, 18.4]]], device=device)
+
+    y1 = torch.amin(x)
+    expected = np.array([1.5], dtype=np.float32)
+    np.testing.assert_equal(y1.cpu().numpy(), expected)
+
+    y2 = torch.amin(x, dim=(1,2))
+    expected = np.array([1.5, 13.2], dtype=np.float32)
+    np.testing.assert_equal(y2.cpu().numpy(), expected)
+
+    y3 = torch.amin(x, dim=2)
+    expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
+    np.testing.assert_equal(y3.cpu().numpy(), expected)
+
  def test_isfinite(self):
    a = torch.ones(4, device=device)
    np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -25,9 +25,9 @@ setup(name='tinygrad',
      long_description=long_description,
      long_description_content_type='text/markdown',
      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
-                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize',
+                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv'],
+                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
      classifiers=[
        "Programming Language :: Python :: 3",
@@ -55,7 +55,7 @@ setup(name='tinygrad',
        ],
        'testing': testing_minimal + [
            "pillow",
-            "onnx==1.17.0",
+            "onnx==1.18.0",
            "onnx2torch",
            "onnxruntime",
            "opencv-python",
@@ -1,24 +1,21 @@
 import time, sys, hashlib
 from pathlib import Path
-from onnx.helper import tensor_dtype_to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange
-from tinygrad.tensor import _from_np_dtype
 import numpy as np
 from extra.bench_log import BenchEvent, WallTimeEvent

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"

 if __name__ == "__main__":
-  onnx_model = onnx_load(onnx_path := fetch(OPENPILOT_MODEL))
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))

  Tensor.manual_seed(100)
-  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
-  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in input_shapes.items()}
+  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
+  new_inputs_junk = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
  new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}

  # benchmark
@@ -1,4 +1,5 @@
 from tinygrad import Tensor, dtypes, GlobalCounters
+from tinygrad.engine.realize import get_program

 if __name__ == "__main__":
  t = Tensor.empty(81920, 4096, dtype=dtypes.half)
@@ -23,5 +24,5 @@ if __name__ == "__main__":
  #k.apply_opt(Opt(OptOps.GROUP, 1, 32))
  #k.apply_opt(Opt(OptOps.GROUP, 0, 32))
  from tinygrad.engine.realize import CompiledRunner, ExecItem
-  run = CompiledRunner(prg:=k.to_program())
+  run = CompiledRunner(prg:=get_program(k.get_optimized_ast(), k.opts))
  ExecItem(run, si.bufs).run()
@@ -1,7 +1,7 @@
 # ruff: noqa: E501
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.opt.search import bufs_from_lin
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -35,7 +35,7 @@ k = Kernel(ast)
 k.apply_opts(opts)
 bufs = bufs_from_lin(k)

-prg = CompiledRunner(k.to_program())
+prg = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))

 for i in range(10):
  speed = prg(bufs, var_vals={}, wait=True)
@@ -0,0 +1,35 @@
+# eval for tinygrad.apps.llm
+import pyarrow.parquet as pq
+from tinygrad.helpers import fetch, colored
+from tinygrad.apps.llm import Transformer, SimpleTokenizer, models
+from tinygrad import Tensor
+
+if __name__ == "__main__":
+  dat = fetch("https://huggingface.co/datasets/allenai/ai2_arc/resolve/main/ARC-Challenge/test-00000-of-00001.parquet")
+  table = pq.read_table(dat)
+
+  model, kv = Transformer.from_gguf(Tensor.from_url(models["1B"]), max_context=4096)
+
+  tok = SimpleTokenizer(kv["tokenizer.ggml.tokens"])
+  bos_id: int = kv['tokenizer.ggml.bos_token_id']
+  eos_id: int = kv['tokenizer.ggml.eos_token_id']
+
+  num_correct, num_answered = 0, 0
+  total_questions = len(table["question"])
+  for question, choices, answer in zip(table["question"], table["choices"], table["answerKey"]):
+    phrasing = f"Question: {question}\n\n" + \
+               '\n'.join([f"{k}) {v}" for k,v in zip(choices['label'], choices['text'])]) +\
+               "\n\nReply with the letter of the correct answer only."
+    try:
+      ids = [bos_id] + tok.role("user") + tok.encode(phrasing) + [eos_id] + tok.role("assistant") + tok.encode("Answer: ")
+    except RuntimeError:
+      # TODO: fix the tokenizer
+      pass
+    next_id = next(model.generate(ids))
+    correct, given = answer.as_py().strip(), tok.decode([next_id]).strip()
+    num_correct += correct == given
+    num_answered += 1
+    print(f"{num_answered:4d}/{total_questions:4d}  "+\
+          f"Correct Answer: {correct}  "+\
+          f"Given Answer: {colored(given, 'green' if correct==given else 'red')}  "+\
+          f"Percent: {num_correct*100.0/num_answered:.2f}%")
@@ -2,11 +2,11 @@ import csv, pathlib, time
 import numpy as np
 import torch
 torch.set_num_threads(1)
-from onnx.helper import tensor_dtype_to_np_dtype
 import onnxruntime as ort
 from onnx2torch import convert
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.helpers import OSX, DEBUG, fetch, getenv
+from tinygrad.dtype import _to_np_dtype
 from tinygrad import Tensor, Device, dtypes

 MODELS = {
@@ -50,20 +50,19 @@ def benchmark_model(m, devices, validate_outs=False):
  CSV = {"model": m}

  fn = fetch(MODELS[m])
-  onnx_model = onnx_load(fn)
-  output_names = [out.name for out in onnx_model.graph.output]
-  excluded = {inp.name for inp in onnx_model.graph.initializer}
-  input_shapes = {inp.name:tuple(x.dim_value if hasattr(x, "dim_value") and x.dim_value != 0 else 1 for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input if inp.name not in excluded}  # noqa: E501
-  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input if inp.name not in excluded}
-  np_inputs = {k:torch.randn(shp).numpy().astype(input_types[k]) for k,shp in input_shapes.items()}
+  runner = OnnxRunner(fn)
+  output_names = runner.graph_outputs
+  input_shapes = {name: tuple(s if isinstance(s, int) and s != 0 else 1 for s in spec.shape) for name, spec in runner.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in runner.graph_inputs.items()}
+  np_inputs = {k:torch.randn(shp).numpy().astype(_to_np_dtype(input_types[k])) for k,shp in input_shapes.items()}
  assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"

  # print input names
-  if DEBUG >= 2: print([inp.name for inp in onnx_model.graph.input if inp.name not in excluded])
+  if DEBUG >= 2: print(list(runner.graph_inputs))
  for device in devices:
    Device.DEFAULT = device
    inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
-    tinygrad_model = OnnxRunner(onnx_model)
+    tinygrad_model = runner.to(device)
    benchmark(m, f"tinygrad_{device.lower()}_jitless", lambda: {k:v.numpy() for k,v in tinygrad_model(inputs).items()})

    from tinygrad.engine.jit import TinyJit
@@ -107,12 +106,12 @@ def benchmark_model(m, devices, validate_outs=False):
      rtol, atol = 2e-3, 2e-3  # tolerance for fp16 models
      Device.DEFAULT = device
      # force half inputs to float for numerical stability when validating
-      # this will reply on automatic dtype promotion for converting half weights inside the graph
+      # this will rely on automatic dtype promotion for converting half weights inside the graph
      if m in half_models:
        inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
      else:
        inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
-      tinygrad_model = OnnxRunner(onnx_model)
+      tinygrad_model = runner.to(device)
      tinygrad_out = tinygrad_model(inputs)

      ort_sess = ort.InferenceSession(str(fn), ort_options, ["CPUExecutionProvider"])
@@ -21,12 +21,13 @@ class FakeAM:
  def __init__(self):
    self.is_booting, self.smi_dev = True, False
    self.pcidev = FakePCIDev()
-    self.vram_mv = memoryview(bytearray(4 << 30))
+    self.vram_size = (512 << 20)
+    self.vram_mv = memoryview(bytearray(self.vram_size))
    self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
    self.gmc = FakeGMC(self)
-    self.mm = AMMemoryManager(self, 4 << 30, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
-      pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, first_page_lv=am.AMDGPU_VM_PDB2,
-      va_base=AMMemoryManager.va_allocator.base)
+    self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
+      first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
+      palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
    self.is_booting = False
    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
  def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
@@ -55,6 +56,8 @@ def helper_read_entry_components(entry_val):
          "read": (entry_val >> 5) & 0x1, "write": (entry_val >> 6) & 0x1, "exec": (entry_val >> 4) & 0x1,
          "mtype": (entry_val >> 48) & 0x7, "T": (entry_val >> 51) & 0x1, "L": (entry_val >> 55) & 0x1, "F": (entry_val >> 56) & 0x1}

+def helper_va(va:int): return va + AMMemoryManager.va_allocator.base
+
 class TestAMPageTable(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -65,10 +68,9 @@ class TestAMPageTable(unittest.TestCase):

    for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                  (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      exteranl_va = va + AMMemoryManager.va_allocator.base
-      mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
+      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])

-      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
      results = list(ctx.next(sz))

      total_covered = 0
@@ -85,7 +87,7 @@ class TestAMPageTable(unittest.TestCase):
          assert pte['paddr'] == va + _offset + i * _pte_covers, f"Expected paddr {pte['paddr']:#x} to be {va + _offset + i * _pte_covers:#x}"
          assert pte['valid'] == 1

-      mm.unmap_range(va, sz)
+      mm.unmap_range(helper_va(va), sz)

      for tup in results:
        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
@@ -98,18 +100,16 @@ class TestAMPageTable(unittest.TestCase):
    mm0 = self.d[0].mm

    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
-      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
-      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
-      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
-      mm0.unmap_range(va2, sz2)
-      mm0.unmap_range(va1, sz1)
+      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
+      mm0.unmap_range(helper_va(va2), sz2)
+      mm0.unmap_range(helper_va(va1), sz1)

  def test_double_map(self):
    mm0 = self.d[0].mm

    for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
-      exteranl_va = va + AMMemoryManager.va_allocator.base
+      exteranl_va = helper_va(va)
      mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])

      with self.assertRaises(AssertionError):
@@ -143,36 +143,36 @@ class TestAMPageTable(unittest.TestCase):
    mm0 = self.d[0].mm

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(0x10000, 0x3000)
+      mm0.unmap_range(helper_va(0x10000), 0x3000)

-    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(0x10000, 0x3000)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(helper_va(0x10000), 0x3000)

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(0x10000, 0x3000)
+      mm0.unmap_range(helper_va(0x10000), 0x3000)

-    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(0x10000, 0x3000)
+    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(helper_va(0x10000), 0x3000)

    with self.assertRaises(AssertionError):
-      mm0.unmap_range(0x10000, 0x3000)
+      mm0.unmap_range(helper_va(0x10000), 0x3000)

  def test_free_pt(self):
    mm0 = self.d[0].mm

    # offset from start
    for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(0x1000000 + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(0x1000000 + off, (2 << 20) - off)
-      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(0x1000000, 2 << 20)
+      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(helper_va(0x1000000), 2 << 20)

    # offset from end
    for off in [0x1000, 0x20000]:
-      mm0.map_range(0x1000000, (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(0x1000000, (2 << 20) - off)
-      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(0x1000000, 2 << 20)
+      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
+      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(helper_va(0x1000000), 2 << 20)

  def test_frag_size(self):
    mm0 = self.d[0].mm
@@ -15,8 +15,8 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

 def helper_test_lin(lin: Kernel, opts, failed_platforms, validate_device, rtol=1e-2, atol=1e-2):
-  if any(b.dtype.base == dtypes.half for b in lin.membufs) and not is_dtype_supported(dtypes.half): return
-  if any(b.dtype.base == dtypes.bfloat16 for b in lin.membufs) and not is_dtype_supported(dtypes.bfloat16): return
+  if any(b.dtype.base == dtypes.half for b in lin.bufs) and not is_dtype_supported(dtypes.half): return
+  if any(b.dtype.base == dtypes.bfloat16 for b in lin.bufs) and not is_dtype_supported(dtypes.bfloat16): return

  try:
    lin.apply_opts(opts)
@@ -11,7 +11,7 @@ class TestHIPCompileSpeed(unittest.TestCase):
    a, b = Tensor([1,2,3,4,5]), Tensor([1,2,3,4,5])
    out = a + b
    lin = Kernel(create_schedule([out.uop])[-1].ast[0])
-    lin.linearize()
+    lin.to_program()

    reference = """
 #include <hip/hip_common.h>
@@ -3,8 +3,7 @@ from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
 from tinygrad.opt.search import Opt, OptOps
-from test.test_linearizer_failures import helper_test_lin
-from tinygrad.engine.realize import get_runner, CompiledRunner
+from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from test.external.fuzz_linearizer import get_fuzz_rawbufs

 from tinygrad.opt.kernel import Kernel
@@ -24,11 +23,6 @@ class TestNV(unittest.TestCase):
    TestNV.b.uop.buffer.allocate()
    TestNV.addr = struct.pack("QQ", TestNV.b.uop.buffer._buf.va_addr, TestNV.a.uop.buffer._buf.va_addr)

-  def test_oor_kernels(self):
-    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=Ops.MUL, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 256, 1, 512, 4, 16, 4, 16), strides=(0, 100352, 0, 196, 0, 14, 0, 1), offset=-15, mask=((0, 1), (0, 256), (0, 1), (0, 512), (0, 4), (1, 15), (0, 4), (1, 15)), contiguous=False), View(shape=(256, 1, 512, 7, 7, 512, 3, 3), strides=(2097152, 0, 0, 128, 2, 4096, 1088, 17), offset=0, mask=None, contiguous=False))))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(256, 1, 512, 7, 7, 512, 3, 3), strides=(25088, 0, 49, 7, 1, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=(dtypes.float, False)),), arg=((0, 3, 4), dtypes.float)),), arg=(dtypes.half, False)),), arg=MemBuffer(idx=0, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 512, 1, 1, 512, 3, 3), strides=(0, 0, 4608, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=True),)))) # noqa: E501
-    opts = [Opt(op=OptOps.TC, axis=6, arg=(-1, 2, 1)), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=3), Opt(op=OptOps.UPCAST, axis=1, arg=2)] # noqa: E501
-    helper_test_lin(Kernel(ast), opts=opts, failed_platforms=["NV"])
-
  def test_error_on_huge_dims(self):
    ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=Ops.MUL, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 0, 1), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.half, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 683), strides=(0, 0, 683, 1), offset=0, mask=None, contiguous=True),))))), arg=None),), arg=dtypes.float),), arg=(3,)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 1, 1024, 1), strides=(0, 0, 1, 0), offset=0, mask=None, contiguous=True),)))) # noqa: E501
    opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=2), Opt(op=OptOps.LOCAL, axis=0, arg=2)] # noqa: E501
@@ -36,7 +30,7 @@ class TestNV(unittest.TestCase):
      lin = Kernel(ast)
      lin.apply_opts(opts)
      rawbufs = get_fuzz_rawbufs(lin)
-      prg = CompiledRunner(lin.to_program())
+      prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
      prg(rawbufs, {}, wait=True)
    self.assertEqual(str(cm.exception), "This is a runtime error message")

@@ -1,4 +1,4 @@
-import tempfile, unittest
+import unittest
 from typing import Any, Tuple
 from onnx.backend.base import Backend, BackendRep
 import onnx.backend.test
@@ -6,12 +6,11 @@ import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.helpers import getenv, OSX
 from tinygrad.device import is_dtype_supported
+from tinygrad.frontend.onnx import OnnxRunner

 # pip3 install tabulate
 pytest_plugins = 'onnx.backend.test.report',

-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
-
 class TinygradModel(BackendRep):
  def __init__(self, run_onnx, input_names):
    super().__init__()
@@ -30,11 +29,8 @@ class TinygradBackend(Backend):
    input_initializer = [x.name for x in model.graph.initializer]
    net_feed_input = [x for x in input_all if x not in input_initializer]
    print("prepare", cls, device, net_feed_input)
-    with tempfile.NamedTemporaryFile(suffix='.onnx') as f:
-      onnx.save(model, f.name)
-      f.flush()
-      new_model = onnx_load(f.name)
-    run_onnx = OnnxRunner(new_model)
+    model = Tensor(model.SerializeToString(), device="PYTHON")
+    run_onnx = OnnxRunner(model)
    return TinygradModel(run_onnx, net_feed_input)

  @classmethod
@@ -44,9 +40,6 @@ class TinygradBackend(Backend):

 backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)

-# BUG: segfaults
-backend_test.exclude('test_MaxPool1d_stride_padding_dilation_cpu')
-
 # BUG: buggy onnx tests
 backend_test.exclude('test_adam_multiple_cpu')

@@ -94,6 +87,7 @@ backend_test.exclude('FLOAT8')
 backend_test.exclude('INT4')
 backend_test.exclude('UINT4')
 backend_test.exclude('BFLOAT16')  # not supported in numpy
+backend_test.exclude('FLOAT4E2M1')

 backend_test.exclude('test_dequantizelinear_int4_cpu')
 backend_test.exclude('test_dequantizelinear_uint4_cpu')
@@ -105,10 +99,12 @@ backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_quantizelinear_e5m2_cpu')
 backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_quantizelinear_e5m2_cpu')
+backend_test.exclude('test_quantizelinear_float4e2m1_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
 backend_test.exclude('test_dequantizelinear_e5m2_cpu')
+backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')

 # we don't support indexes
 backend_test.exclude('test_nonzero_*')
@@ -188,6 +184,11 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

+# regression from removing StrEnum in Domain
+backend_test.exclude('test_adam_cpu')
+backend_test.exclude('test_gradient_of_add_and_mul_cpu')
+backend_test.exclude('test_gradient_of_add_cpu')
+
 if Device.DEFAULT in ['GPU', 'METAL']:
  backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')
  backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu')
@@ -4,12 +4,17 @@

 from typing import Any
 import unittest, onnx, tempfile
-from tinygrad import dtypes
+from tinygrad import dtypes, Tensor
 from tinygrad.frontend.onnx import OnnxRunner
 import numpy as np
 from extra.onnx_helpers import validate
 from onnx.defs import ONNX_DOMAIN, AI_ONNX_PREVIEW_TRAINING_DOMAIN
 MICROSOFT_CONTRIB_OPS_DOMAIN = "com.microsoft"
+# TODO: remove this once ORT supports 1.18.0
+from onnx.helper import VERSION_TABLE
+VERSION_MAP = {row[0]: row[1:] for row in VERSION_TABLE}
+IR_VERSION, ai_onnx, ai_onnx_ml, ai_onnx_training = VERSION_MAP["1.17.0"]
+

 class TestOnnxOps(unittest.TestCase):
  DOMAIN = None
@@ -18,7 +23,14 @@ class TestOnnxOps(unittest.TestCase):
    onnx_outputs = [onnx.helper.make_empty_tensor_value_info(name) for name in outs]
    nodes = [onnx.helper.make_node(op, list(inps), list(outs), domain=self.DOMAIN, **opts)]
    graph = onnx.helper.make_graph(nodes, f"test_{op.lower()}", onnx_inputs, onnx_outputs)
-    model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}")
+    #model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}")
+    # TODO: remove this once ORT supports 1.18.0
+    opset_id = None
+    if type(self).__name__ == "TestMainOnnxOps": opset_id = ai_onnx
+    if type(self).__name__ == "TestTrainingOnnxOps": opset_id = ai_onnx_training
+    if type(self).__name__ == "TestContribOnnxOps": opset_id = 1
+    model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}", ir_version=IR_VERSION,
+                                    opset_imports=[onnx.helper.make_opsetid(self.DOMAIN, opset_id)])
    return model

  def helper_test_single_op(self, op:str, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str], rtol=1e-3, atol=1e-6):
@@ -88,7 +100,8 @@ class TestMainOnnxOps(TestOnnxOps):
    attributes = {"detect_negative":1, "detect_positive":1}
    outputs = ["y"]
    model = self.helper_build_model("IsInf", inputs, attributes, outputs)
-    outputs = OnnxRunner(model)(inputs)
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
+    outputs = runner(inputs)
    assert outputs["y"].dtype is dtypes.bool

  def test_quantize_linear(self):
@@ -203,7 +216,7 @@ class TestTrainingOnnxOps(TestOnnxOps):
  def _validate_training(self, op:str, onnx_fxn, inps:dict[str, np.ndarray], opts:dict[str, Any], outs:list[str]):
    model = self.helper_build_model(op, inps, opts, outs)
    if op == "Momentum": del opts['mode']
-    runner = OnnxRunner(model)
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
    tiny_out = runner(inps)
    onnx_out = onnx_fxn(**inps, **opts)
    for (nm, t_out), o_out in  zip(tiny_out.items(), onnx_out):
@@ -238,6 +251,7 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

+  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
  def test_adam_t_greater_than_zero(self):
    from onnx.backend.test.case.node.adam import apply_adam
    for t in [1, 3, 100]:
@@ -1,77 +1,139 @@
-import unittest, onnx, tempfile
-from tinygrad import dtypes
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+import unittest, onnx, tempfile, pathlib
+import numpy as np
+from tinygrad import dtypes, Tensor
+from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
 from extra.onnx import data_types
-from hypothesis import given, settings, strategies as st
-import numpy as np
+from tinygrad.frontend.onnx import OnnxRunner
+from hypothesis import given, strategies as st

-data_types.pop(16) # TODO: this is bf16, need to support double parsing first.
-device_supported_dtypes = [odt for odt, dtype in data_types.items() if is_dtype_supported(dtype)]
-device_unsupported_dtypes = [odt for odt, dtype in data_types.items() if not is_dtype_supported(dtype)]
+# copied from test_const_folding.py
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  schedule = t.schedule()
+  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+def build_onnx(nodes, from_disk:bool=True, **kwargs):
+  """Helper to build and return an OnnxRunner from ONNX nodes."""
+  graph = onnx.helper.make_graph(nodes, 'test', kwargs.get('inputs', []), kwargs.get('outputs', []), kwargs.get('initializers', []))
+  model = onnx.helper.make_model(graph)
+  if from_disk:
+    with tempfile.TemporaryDirectory() as tmpdir:
+      tmp_path = pathlib.Path(tmpdir)
+      model_path = tmp_path / "model.onnx"
+      onnx.save(model, model_path)
+      runner = OnnxRunner(model_path)
+  else:
+    # use the in-memory method
+    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
+  return runner
+
+class TestOnnxRunner(unittest.TestCase):
+  def _test_const_fold_unary_op(self, from_disk:bool):
+    runner = build_onnx(
+        nodes=[
+          onnx.helper.make_node('Expand', ['inp', 'shape'], ['expanded']),
+          onnx.helper.make_node('Exp', ['expanded'], ['output'])
+        ],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (5,))],
+        initializers=[
+          onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (), [1.0]),
+          onnx.helper.make_tensor('shape', onnx.TensorProto.INT64, (1,), [5])
+        ],
+        from_disk=from_disk)
+    output = runner({'inp': Tensor([1.0])})['output']
+    _check_ast_count(0, output)
+
+  def _test_const_fold_binary_op(self, from_disk:bool):
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Add', ['inp', 'const'], ['output'])],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
+        initializers=[
+          onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (4,), [1, 2, 3, 4]),
+          onnx.helper.make_tensor('const', onnx.TensorProto.FLOAT, (), [0])
+        ],
+        from_disk=from_disk)
+    output = runner({'inp': Tensor([1, 2, 3, 4])})['output']
+    _check_ast_count(0, output)
+
+  def test_const_fold_from_disk(self):
+    self._test_const_fold_unary_op(True)
+    self._test_const_fold_binary_op(True)
+
+  def test_const_fold_from_memory(self):
+    self._test_const_fold_unary_op(False)
+    # TODO: understand this and fix this, bitcast related
+    # self._test_const_fold_binary_op(False)
+
+  def test_external_data_loading(self):
+    weights = np.arange(4, dtype=np.float32)
+    tensor_with_data = onnx.helper.make_tensor('weights', onnx.TensorProto.FLOAT, weights.shape, weights.tobytes(), raw=True)
+    graph = onnx.helper.make_graph(
+        nodes=[onnx.helper.make_node('Add', ['inp', 'weights'], ['output'])],
+        name='test_external',
+        inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (1,))],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, weights.shape)],
+        initializer=[tensor_with_data]
+    )
+    model = onnx.helper.make_model(graph)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+      tmp_path = pathlib.Path(tmpdir)
+      model_path = tmp_path / "model.onnx"
+      onnx.save_model(model, model_path, save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=0, location="weights.onnx_data")
+      runner = OnnxRunner(model_path)
+      output = runner({'inp': Tensor([1])})['output']
+      np.testing.assert_equal(output.numpy(), weights + 1)
+
+all_dtypes = list(data_types.keys())
+device_supported_dtypes = {odt for odt, dtype in data_types.items() if is_dtype_supported(dtype)}

 class TestOnnxRunnerDtypes(unittest.TestCase):
-  def _test_input_spec_dtype(self, onnx_data_type, tinygrad_dtype):
-    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
-    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
-    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
-    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor])
-    model = onnx.helper.make_model(graph)
-    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
-    onnx.save(model, tmp.name)
-    tmp.flush()
-    model = onnx_load(tmp.name)
-    runner = OnnxRunner(model)
-    self.assertEqual(len(runner.graph_inputs), 1)
-    self.assertEqual(runner.graph_inputs['input'].dtype, tinygrad_dtype)
+  """
+  Internal tensors (initializers, attributes) fallback to default dtype if unsupported by device.
+  External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
+  """
+  def _get_expected_dtype(self, onnx_dtype: int, is_input: bool):
+    true_dtype = data_types[onnx_dtype]
+    # inputs always preserve their true dtype.
+    if is_input:
+      return true_dtype
+    # supported types are always themselves.
+    if onnx_dtype in device_supported_dtypes:
+      return true_dtype
+    # otherwise it's an unsupported dtype that's internal to the ONNX model, which should fallback to default.
+    return dtypes.default_int if dtypes.is_int(true_dtype) else dtypes.default_float

-  def _test_initializer_dtype(self, onnx_data_type, tinygrad_dtype):
-    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
-    initializer = onnx.helper.make_tensor('initializer', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
-    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
-    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
-    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
-    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor], [initializer])
-    model = onnx.helper.make_model(graph)
-    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
-    onnx.save(model, tmp.name)
-    tmp.flush()
-    model = onnx_load(tmp.name)
-    runner = OnnxRunner(model)
-    self.assertEqual(len(runner.graph_inputs), 1)
-    self.assertEqual(runner.graph_values['initializer'].dtype, tinygrad_dtype)
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_input_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, True)
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
+        inputs=[onnx.helper.make_tensor_value_info('input', onnx_dtype, ())],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, ())],
+        from_disk=False)
+    self.assertEqual(runner.graph_inputs['input'].dtype, expected_dtype)

-  def _test_node_attribute_dtype(self, onnx_data_type, tinygrad_dtype):
-    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
-    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, arr.shape)
-    value_tensor = onnx.helper.make_tensor('value', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
-    node = onnx.helper.make_node('Constant', inputs=[], outputs=['output'], value=value_tensor)
-    graph = onnx.helper.make_graph([node], 'attribute_test', [], [output_tensor])
-    model = onnx.helper.make_model(graph)
-    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
-    tmp.flush()
-    onnx.save(model, tmp.name)
-    model = onnx_load(tmp.name)
-    runner = OnnxRunner(model)
-    self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, tinygrad_dtype)
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_initializer_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, False)
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Identity', ['initializer'], ['output'])],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
+        initializers=[onnx.helper.make_tensor('initializer', onnx_dtype, (2,), [1, 2])],
+        from_disk=False)
+    self.assertEqual(runner.graph_values['initializer'].dtype, expected_dtype)

-  @settings(deadline=1000) # TODO investigate unreliable timing
-  @given(onnx_data_type=st.sampled_from(device_supported_dtypes))
-  def test_supported_dtype_spec(self, onnx_data_type):
-    tinygrad_dtype = data_types[onnx_data_type]
-    self._test_input_spec_dtype(onnx_data_type, tinygrad_dtype)
-    self._test_initializer_dtype(onnx_data_type, tinygrad_dtype)
-    self._test_node_attribute_dtype(onnx_data_type, tinygrad_dtype)
-
-  @unittest.skipUnless(device_unsupported_dtypes, "No unsupported dtypes for this device to test.")
-  @settings(deadline=1000) # TODO investigate unreliable timing
-  @given(onnx_data_type=st.sampled_from(device_unsupported_dtypes))
-  def test_unsupported_dtype_spec(self, onnx_data_type):
-    true_dtype = data_types[onnx_data_type]
-    default_dtype = dtypes.default_int if dtypes.is_int(true_dtype) else dtypes.default_float
-    self._test_input_spec_dtype(onnx_data_type, true_dtype)
-    self._test_initializer_dtype(onnx_data_type, default_dtype)
-    self._test_node_attribute_dtype(onnx_data_type, default_dtype)
+  @given(onnx_dtype=st.sampled_from(all_dtypes))
+  def test_node_attribute_dtype(self, onnx_dtype: int):
+    expected_dtype = self._get_expected_dtype(onnx_dtype, False)
+    value_tensor = onnx.helper.make_tensor('value', onnx_dtype, (2,), [1, 2])
+    runner = build_onnx(
+        nodes=[onnx.helper.make_node('Constant', [], ['output'], value=value_tensor)],
+        outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
+        from_disk=False)
+    self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, expected_dtype)

 if __name__ == '__main__':
  unittest.main()
@@ -0,0 +1,39 @@
+from transformers import AutoTokenizer
+from datasets import load_dataset
+from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.helpers import tqdm, getenv
+
+# use ALLOW_FAILED=-1 to go over the entire dataset without printing.
+if __name__ == "__main__":
+  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+  vocab_words = [ word for word, _ in sorted(base_tokenizer.get_vocab().items(), key=lambda t: t[1]) ]
+  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
+  simple_tokenizer = SimpleTokenizer(vocab_words)
+
+  color_codes = [ 91, 92, 94, 93, 95 ]
+  def color_tokens(tids): return "".join(f"\033[{color_codes[i%len(color_codes)]}m{inv_vocab[t]}" for i, t in enumerate(tids)) + "\033[0m"
+
+  ds = load_dataset("OpenAssistant/oasst1")
+  allow_failed = getenv("ALLOW_FAILED", 10)
+
+  fail_count, total = 0, 0
+
+  for idx, el in enumerate(tqdm(ds["train"])):
+    total += 1
+
+    try: simple_tokens = tuple(simple_tokenizer.encode(el["text"]))
+    except RuntimeError: simple_tokens = ()
+    base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False))
+
+    if simple_tokens != base_tokens:
+      fail_count += 1
+      allow_failed -= 1
+
+      if allow_failed >= 0:
+        print(f"tokens mismatch at index: {idx}.\n")
+
+        print("simple:  ", color_tokens(simple_tokens))
+        print("official:", color_tokens(base_tokens) + "\n")
+
+      if allow_failed == 0: break
+  print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
@@ -1,121 +0,0 @@
-# ruff: noqa: E501
-import unittest
-
-from tinygrad import Device
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.opt.kernel import Kernel
-
-class TestOpenpilotValidhack(unittest.TestCase):
-  def test_valid_removal(self):
-    Device.DEFAULT = "GPU"
-
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((64, 1024, 4)), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 64, 128, 1, 1, 8, 4, 1, 1, 1, 1), strides=(0, 4096, 32, 0, 0, 4, 1, 0, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.ADD, dtypes.float, arg=None, src=(
-          UOp(Ops.MAX, dtypes.float, arg=None, src=(
-            x5:=UOp(Ops.ADD, dtypes.float, arg=None, src=(
-              UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (7, 8, 9, 10)), src=(
-                UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                  UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                    UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                      UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((128, 768, 4)), arg=1, src=()),
-                      UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 1, 1, 1, 1, 3, 1, 4, 4, 130, 4, 258), strides=(0, 0, 0, 0, 0, 4, 0, 1, 0, 3072, 0, 12), offset=-3084, mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 3), (0, 1), (0, 4), (0, 4), (1, 129), (0, 4), (1, 257)), contiguous=False), View(shape=(1, 64, 128, 1, 1, 8, 4, 3, 4, 3, 3), strides=(0, 2064, 2, 0, 0, 0, 0, 2146560, 536640, 135192, 259), offset=0, mask=None, contiguous=False))), src=()),)),
-                    UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((8, 108, 4)), arg=2, src=()),
-                        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 64, 128, 1, 1, 8, 4, 3, 4, 3, 3), strides=(0, 0, 0, 0, 0, 432, 1, 48, 4, 144, 16), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
-              UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=3, src=()),
-                UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 64, 128, 1, 1, 8, 4, 1, 1, 1, 1), strides=(0, 0, 0, 0, 0, 4, 1, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-            x19:=UOp(Ops.CONST, dtypes.float, arg=0.0, src=(
-              x20:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 64, 128, 1, 1, 8, 4, 1, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
-          UOp(Ops.MUL, dtypes.float, arg=None, src=(
-            UOp(Ops.MAX, dtypes.float, arg=None, src=(
-              UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                UOp(Ops.CONST, dtypes.float, arg=1.0, src=(
-                  x20,)),
-                UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                  UOp(Ops.EXP2, dtypes.float, arg=None, src=(
-                    UOp(Ops.MUL, dtypes.float, arg=None, src=(
-                      x5,
-                      UOp(Ops.CONST, dtypes.float, arg=1.4426950408889634, src=(
-                        x20,)),)),)),
-                  x29:=UOp(Ops.CONST, dtypes.float, arg=-1.0, src=(
-                    x20,)),)),)),
-              x19,)),
-            x29,)),)),)),))
-
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None)]
-    kernel = Kernel(ast)
-
-    kernel.apply_opts(opts)
-
-    p = kernel.to_program()
-    print(p.src)
-
-  def test_const_idx(self):
-    Device.DEFAULT = "GPU"
-
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((10, 128, 4)), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 512, 1), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.CAST, dtypes.float, arg=None, src=(
-          UOp(Ops.ADD, dtypes.float, arg=None, src=(
-            UOp(Ops.ADD, dtypes.float, arg=None, src=(
-              UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((1, 128, 4)), arg=1, src=()),
-                UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=0, mask=((0, 1), (0, 1), (0, 512)), contiguous=False),)), src=()),)),
-              UOp(Ops.CAST, dtypes.float, arg=None, src=(
-                UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                  UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                    UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                      UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                        UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                          UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                            UOp(Ops.ADD, dtypes.float, arg=None, src=(
-                              UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                                x18:=UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=2, src=()),
-                                UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=48128, mask=((0, 1), (1, 2), (0, 512)), contiguous=False),)), src=()),)),
-                              UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                                x18,
-                                UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=45568, mask=((0, 1), (2, 3), (0, 512)), contiguous=False),)), src=()),)),)),
-                            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                              x18,
-                              UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=43008, mask=((0, 1), (3, 4), (0, 512)), contiguous=False),)), src=()),)),)),
-                          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                            x18,
-                            UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=40448, mask=((0, 1), (4, 5), (0, 512)), contiguous=False),)), src=()),)),)),
-                        UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                          x18,
-                          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=37888, mask=((0, 1), (5, 6), (0, 512)), contiguous=False),)), src=()),)),)),
-                      UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                        x18,
-                        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=35328, mask=((0, 1), (6, 7), (0, 512)), contiguous=False),)), src=()),)),)),
-                    UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                      x18,
-                      UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=32768, mask=((0, 1), (7, 8), (0, 512)), contiguous=False),)), src=()),)),)),
-                  UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-                    x18,
-                    UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=30208, mask=((0, 1), (8, 9), (0, 512)), contiguous=False),)), src=()),)),)),)),)),
-            UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-              UOp(Ops.DEFINE_GLOBAL, dtypes.imagef((1, 128, 4)), arg=3, src=()),
-              UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 10, 512), strides=(0, 0, 1), offset=0, mask=((0, 1), (9, 10), (0, 512)), contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None)]
-    kernel = Kernel(ast)
-
-    kernel.apply_opts(opts)
-
-    p = kernel.to_program()
-    # ((idx1<1)?read_imagef(data1, smp, (int2)(idx0,0)):(float4)(0.0f,0.0f,0.0f,0.0f))
-    print(p.src)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -2,7 +2,7 @@ import random
 from z3 import Int, Solver, sat
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher
-from tinygrad.codegen.devectorizer import fast_idiv
+from tinygrad.codegen.optional import fast_idiv
 random.seed(42)

 z3_renderer = PatternMatcher([
@@ -3,6 +3,7 @@ from typing import Any
 import numpy as np
 from collections import defaultdict
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, kern_str_to_lin
+from tinygrad.engine.realize import get_program

 # We need to insert ioctl before opening devices.
 if os.getenv("VALIDATE_HCQ", 0) != 0:
@@ -93,7 +94,7 @@ def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]:

  # TODO: images needs required_optimization
  try:
-    prg = CompiledRunner(lin.to_program())
+    prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
  except KeyboardInterrupt: raise
  except Exception:
    traceback.print_exc()
@@ -114,7 +115,7 @@ def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]:

 def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=None, rtol=1e-2, atol=1e-2):
  # TODO: for bfloat16 it compiles linearizer, but it does not run because numpy cannot generate bf16 buffer.
-  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.membufs)
+  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.bufs)

  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
  try:
@@ -206,7 +207,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")

        # stop if kernel uops repeat
-        try: tuops = tuplize_uops(test_lin.linearize().uops)
+        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.opts).uops)
        except KeyboardInterrupt: raise
        except BaseException as e:
          print(test_lin.ast)
@@ -1,12 +1,20 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
 from typing import Callable, Any
-from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
-from tinygrad.kernelize.kernelize import get_kernelize_map
-from tinygrad.renderer import Renderer, ProgramSpec
-from tinygrad.engine.realize import get_program
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+
+ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
+if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
+
+try:
+  from tinygrad.schedule.kernelize import get_kernelize_map
+  from tinygrad.renderer import Renderer, ProgramSpec
+  from tinygrad.engine.realize import get_program
+  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+except ImportError as e:
+  print(repr(e))
+  exit(int(ASSERT_DIFF))

 # *** process replay settings

@@ -20,12 +28,11 @@ early_stop = multiprocessing.Event()
 logging.basicConfig(level=logging.INFO, format="%(message)s")
 MAX_LINES = 500
 def trunc_log(x):
-  if len(lines:=repr(x).splitlines()) > MAX_LINES: lines = lines[:MAX_LINES]+[f"WARN: truncated string with {len(lines)} lines"]
+  if len(lines:=(x if isinstance(x, str) else repr(x)).splitlines()) > MAX_LINES:
+    lines = lines[:MAX_LINES]+[f"WARN: truncated string with {len(lines)} lines"]
  logging.info("\n".join(lines))

 # user config
-ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in os.getenv("PR_TITLE", flag))
-if not getenv("ASSERT_PROCESS_REPLAY", 1): ASSERT_DIFF = 0
 SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
 if REF == "master": SKIP_PROCESS_REPLAY = True
 class ProcessReplayWarning(Warning): pass
@@ -41,9 +48,15 @@ def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[
  return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
-  p2 = get_program(ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast, renderer)
-  def to_str(ret:ProgramSpec) -> str: return ret.src
-  return to_str(p2), to_str(p), (p.ast, renderer, p.applied_opts)
+  input_ast = ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast
+  p2 = get_program(input_ast, renderer)
+  def to_str(ret:ProgramSpec) -> str:
+    # PYTHON renderer pickles UOps, first unpickle and decode here
+    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
+    return ret.src
+  # properly color the name arg
+  ast_repr = codecs.decode(str(input_ast), "unicode_escape")
+  return to_str(p2), to_str(p), (ast_repr, renderer)

 replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program}

@@ -4,7 +4,7 @@ from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_amd import AMDDevice
 from contextlib import contextmanager
@@ -77,9 +77,9 @@ if __name__ == "__main__":
    with run_amd():
      amdlin = ast_str_to_lin(ast, opts=amddev.renderer)
      amdlin.apply_opts(hand_coded_optimizations(amdlin))
-      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.membufs)
+      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.bufs)

-      amd_prg = CompiledRunner(amdlin.to_program())
+      amd_prg = CompiledRunner(get_program(amdlin.get_optimized_ast(), amdlin.opts))
      amdbufs = bufs_from_lin(amdlin)
      test_amdbufs = get_fuzz_rawbufs(amdlin) if not has_bf16 else amdbufs
      if not has_bf16: contents = [buf.as_buffer() for buf in test_amdbufs]
@@ -89,7 +89,7 @@ if __name__ == "__main__":
      rdr.device = "AMD:1"
      amlin = ast_str_to_lin(ast, opts=amdev.renderer)
      amlin.apply_opts(hand_coded_optimizations(amlin))
-      am_prg = CompiledRunner(amlin.to_program())
+      am_prg = CompiledRunner(get_program(amlin.get_optimized_ast(), amlin.opts))
      ambufs = bufs_from_lin(amlin)
      test_ambufs = get_fuzz_rawbufs(amlin) if not has_bf16 else ambufs
      if not has_bf16:
@@ -100,7 +100,7 @@ if __name__ == "__main__":
      cpu_rdr.device = "CPU"
      cpulin = ast_str_to_lin(ast, opts=cpu_rdr)
      cpulin.apply_opts(hand_coded_optimizations(cpulin))
-      cpu_prg = CompiledRunner(cpulin.to_program())
+      cpu_prg = CompiledRunner(get_program(cpulin.get_optimized_ast(), cpulin.opts))
      cpubufs = bufs_from_lin(cpulin)
      test_cpubufs = get_fuzz_rawbufs(cpulin) if not has_bf16 else ambufs
      if not has_bf16:
@@ -4,7 +4,7 @@ from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.tensor import _to_np_dtype
 import numpy as np

@@ -23,9 +23,9 @@ if __name__ == "__main__":
    # cuda compile
    culin = ast_str_to_lin(ast, opts=cudev.renderer)
    culin.apply_opts(hand_coded_optimizations(culin))
-    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.membufs)
+    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.bufs)

-    cuda_prg = CompiledRunner(culin.to_program())
+    cuda_prg = CompiledRunner(get_program(culin.get_optimized_ast(), culin.opts))
    cubufs = bufs_from_lin(culin)
    test_cubufs = get_fuzz_rawbufs(culin) if not has_bf16 else cubufs

@@ -33,7 +33,7 @@ if __name__ == "__main__":
    rdr.device = "NV"
    nvlin = ast_str_to_lin(ast, opts=rdr)
    nvlin.apply_opts(hand_coded_optimizations(nvlin))
-    nv_prg = CompiledRunner(nvlin.to_program())
+    nv_prg = CompiledRunner(get_program(nvlin.get_optimized_ast(), nvlin.opts))
    nvbufs = bufs_from_lin(nvlin)
    test_nvbufs = get_fuzz_rawbufs(nvlin) if not has_bf16 else nvbufs
    if not has_bf16:
@@ -1,6 +1,6 @@
 import itertools
 from tinygrad import Device
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
@@ -25,7 +25,7 @@ if __name__ == "__main__":
    dev.compiler = CUDACompiler(dev.arch)
    lin = ast_str_to_lin(ast, opts=dev.renderer)
    lin.apply_opts(hand_coded_optimizations(lin))
-    cuda_prg = CompiledRunner(lin.to_program())
+    cuda_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))

    bufs = bufs_from_lin(lin)

@@ -33,7 +33,7 @@ if __name__ == "__main__":
    dev.compiler = PTXCompiler(dev.arch)
    lin = ast_str_to_lin(ast, opts=ptx)
    lin.apply_opts(hand_coded_optimizations(lin))
-    ptx_prg = CompiledRunner(lin.to_program())
+    ptx_prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))

    # warmup
    try:
@@ -7,7 +7,7 @@ try:
  import onnx
 except ModuleNotFoundError:
  raise unittest.SkipTest("onnx not installed, skipping onnx test")
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import CI, fetch, temp

@@ -25,7 +25,7 @@ np.random.seed(1337)

 class TestOnnxModel(unittest.TestCase):
  def test_benchmark_openpilot_model(self):
-    onnx_model = onnx_load(fetch(OPENPILOT_MODEL))
+    onnx_model = fetch(OPENPILOT_MODEL)
    run_onnx = OnnxRunner(onnx_model)
    def get_inputs():
      np_inputs = {
@@ -69,7 +69,7 @@ class TestOnnxModel(unittest.TestCase):
      ps.print_stats(30)

  def test_openpilot_model(self):
-    onnx_model = onnx_load(fetch(OPENPILOT_MODEL))
+    onnx_model = fetch(OPENPILOT_MODEL)
    run_onnx = OnnxRunner(onnx_model)
    print("got run_onnx")
    inputs = {
@@ -121,10 +121,9 @@ class TestOnnxModel(unittest.TestCase):
      input_name, input_new)

  def _test_model(self, fn, input_name, input_new, debug=False):
-    onnx_model = onnx_load(fn)
+    run_onnx = OnnxRunner(fn)
    print("onnx loaded")
    from test.models.test_efficientnet import chicken_img, car_img, preprocess, _LABELS
-    run_onnx = OnnxRunner(onnx_model)

    def run(img):
      inputs = {input_name: preprocess(img, new=input_new)}
@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.opt.kernel import Opt, OptOps, Kernel, KernelOptError
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops

@@ -17,7 +17,7 @@ class TestArange(unittest.TestCase):
    k = Kernel(sched[-1].ast)
    if opts is not None:
      for o in opts: k.apply_opt(o)
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    print(p.name)
    #print(p.src)
    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
@@ -96,6 +96,14 @@ class TestDType(unittest.TestCase):
     get_available_cast_dtypes(self.DTYPE)
    ))

+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
+  @unittest.skipIf(getenv("PTX"), "skip for now")
+  def test_uint_overflow(self):
+    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
+    v = dtypes.max(self.DTYPE)
+    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
+    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
+
  def test_dtypes_fields(self):
    fields = dtypes.fields()
    self.assertIn("float", fields)
@@ -94,42 +94,20 @@ class TestEmptyTensorEdgeCases(unittest.TestCase):
    out = Tensor([], dtype=dtypes.float32).masked_select(Tensor([], dtype=dtypes.bool))
    np.testing.assert_equal(out.numpy(), torch_out.numpy())

-class TestRollEdgeCases(unittest.TestCase):
-  # we don't need more of these
-
-  @unittest.expectedFailure
-  def test_roll_mismatched_dims(self):
-    with self.assertRaises(RuntimeError):
-      torch.roll(torch.arange(9).reshape(3, 3), 1, dims=(0, 1))
-    with self.assertRaises(RuntimeError):
-      Tensor.arange(9).reshape(3, 3).roll(1, dims=(0, 1))
-
-  @unittest.expectedFailure
-  def test_roll_extra_shift(self):
-    # tinygrad ignores extra shift values instead of raising
-    with self.assertRaises(RuntimeError):
-      torch.roll(torch.arange(10), (1, 2), dims=0)
-    with self.assertRaises(RuntimeError):
-      Tensor.arange(10).roll((1, 2), dims=0)
-
 class TestDropoutProbabilityEdgeCases(unittest.TestCase):
  # we don't need more of these

-  @unittest.expectedFailure
  def test_dropout_rate_one(self):
-    # out is full of NaNs it should be 0s
    with Tensor.train():
      out = Tensor.ones(100).dropout(1.0)
      np.testing.assert_allclose(out.numpy(), np.zeros(100))

-  @unittest.expectedFailure
  def test_dropout_invalid_prob(self):
-    # negative dropout probability should raise an error
    with self.assertRaises(ValueError):
      torch.nn.functional.dropout(torch.ones(10), -0.1, True)
-    with Tensor.train():
-      out = Tensor.ones(10).dropout(-0.1)
-      np.testing.assert_allclose(out.numpy(), np.ones(10))
+    with self.assertRaises(ValueError):
+      with Tensor.train():
+        Tensor.ones(10).dropout(-0.1)

 class TestInputValidation(unittest.TestCase):
  # we don't need more of these, input validation bugs are not very interesting, many are WONTFIX
@@ -1,10 +1,11 @@
-import unittest, ctypes, struct, os, random
+import unittest, ctypes, struct, os, random, numpy as np
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import getenv, CI, mv_address
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
-from tinygrad.engine.realize import get_runner, CompiledRunner
+from tinygrad.runtime.support.system import PCIIfaceBase
+from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad import Variable

@@ -67,20 +68,20 @@ class TestHCQ(unittest.TestCase):
      if queue_type is None: continue

      with self.subTest(name=str(queue_type)):
-        fake_signal = TestHCQ.d0.signal_t()
+        fake_signal = TestHCQ.d0.new_signal()
        fake_signal.value = 1
        queue_type().wait(fake_signal, 1) \
                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
        TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
        TestHCQ.d0.timeline_value += 1

-  @unittest.skipIf(MOCKGPU, "Can't handle async update on MOCKGPU for now")
+  @unittest.skipIf(MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"}, "Can't handle async update on MOCKGPU for now")
  def test_wait_late_set(self):
    for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
      if queue_type is None: continue

      with self.subTest(name=str(queue_type)):
-        fake_signal = TestHCQ.d0.signal_t()
+        fake_signal = TestHCQ.d0.new_signal()
        queue_type().wait(fake_signal, 1) \
                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)

@@ -100,7 +101,7 @@ class TestHCQ(unittest.TestCase):
        virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
        virt_signal = TestHCQ.d0.signal_t(base_buf=HCQBuffer(Variable("sig_addr", 0, 0xffffffffffffffff, dtypes.uint64), 16))

-        fake_signal = TestHCQ.d0.signal_t()
+        fake_signal = TestHCQ.d0.new_signal()
        q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)

        fake_signal.value = 0x30
@@ -136,6 +137,7 @@ class TestHCQ(unittest.TestCase):
    val = TestHCQ.a.uop.buffer.as_buffer().cast("f")[0]
    assert val == 200.0, f"got val {val}"

+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "No globals/locals on LLVM/CPU")
  def test_exec_update(self):
    sint_global = (Variable("sint_global", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.runner.p.global_size[1:])
    sint_local = (Variable("sint_local", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.runner.p.local_size[1:])
@@ -153,6 +155,7 @@ class TestHCQ(unittest.TestCase):
    val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
    assert val == 0.0, f"got val {val}, should not be updated"

+  @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "No globals/locals on LLVM/CPU")
  def test_exec_update_fuzz(self):
    virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
    virt_local = [Variable(f"local_{i}", 0, 0xffffffff, dtypes.uint32) for i in range(3)]
@@ -163,7 +166,7 @@ class TestHCQ(unittest.TestCase):
    k = Kernel(si.ast, opts=TestHCQ.d0.renderer)
    for i in range(3): k.apply_opt(Opt(op=OptOps.LOCAL, axis=0, arg=3))

-    runner = CompiledRunner(k.to_program())
+    runner = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))

    zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
    zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -292,7 +295,7 @@ class TestHCQ(unittest.TestCase):
      virt_signal = TestHCQ.d0.signal_t(base_buf=HCQBuffer(Variable("sig_addr", 0, 0xffffffffffffffff, dtypes.uint64), 16))

      with self.subTest(name=str(queue_type)):
-        fake_signal = TestHCQ.d0.signal_t()
+        fake_signal = TestHCQ.d0.new_signal()
        q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
        q.bind(TestHCQ.d0)

@@ -309,7 +312,7 @@ class TestHCQ(unittest.TestCase):
    try: d1 = Device[f"{Device.DEFAULT}:1"]
    except Exception: self.skipTest("no multidevice, test skipped")

-    TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0.signal_t(value=0), value=0xfff) \
+    TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0.new_signal(value=0), value=0xfff) \
                                .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)

    d1.hw_copy_queue_t().wait(sig, value=0xfff) \
@@ -323,7 +326,7 @@ class TestHCQ(unittest.TestCase):

  # Test profile api
  def test_speed_exec_time(self):
-    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
+    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
    TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \
                                   .exec(TestHCQ.runner._prg, TestHCQ.kernargs_ba_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size) \
                                   .timestamp(sig_en) \
@@ -335,7 +338,7 @@ class TestHCQ(unittest.TestCase):
    et = float(sig_en.timestamp - sig_st.timestamp)

    print(f"exec kernel time: {et:.2f} us")
-    assert 0.1 <= et <= (15000 if MOCKGPU else 100)
+    assert 0.1 <= et <= (15000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)

  def test_speed_copy_bandwidth(self):
    if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -345,7 +348,7 @@ class TestHCQ(unittest.TestCase):
    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
    b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()

-    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
+    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf.va_addr, b._buf.va_addr, SZ) \
                                .timestamp(sig_en) \
@@ -372,7 +375,7 @@ class TestHCQ(unittest.TestCase):
    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
    TestHCQ.d0.allocator.map(b._buf)

-    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
+    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf.va_addr, b._buf.va_addr, SZ) \
                                .timestamp(sig_en) \
@@ -530,10 +533,28 @@ class TestHCQ(unittest.TestCase):
    try: nv_dev = Device["NV"]
    except Exception: self.skipTest("no NV device, test skipped")

-    x = amd_dev.signal_t()
-    y = nv_dev.signal_t()
+    x = amd_dev.new_signal()
+    y = nv_dev.new_signal()
    assert type(x) is amd_dev.signal_t
    assert type(y) is nv_dev.signal_t

+  def test_multidevice_p2p(self):
+    try:
+      amd_dev = Device["AMD"]
+      if not issubclass(type(amd_dev.iface), PCIIfaceBase): self.skipTest("Not a pci dev")
+    except Exception: self.skipTest("no AMD device, test skipped")
+
+    try:
+      nv_dev = Device["NV"]
+      if not issubclass(type(nv_dev.iface), PCIIfaceBase): self.skipTest("Not a pci dev")
+    except Exception: self.skipTest("no NV device, test skipped")
+
+    def _check_copy(dev1, dev2):
+      buf1 = Tensor.randn(10, 10, device=dev1).realize()
+      buf2 = buf1.to(dev2).realize()
+      np.testing.assert_equal(buf1.numpy(), buf2.numpy(), "p2p failed")
+    _check_copy("AMD", "NV")
+    _check_copy("NV", "AMD")
+
 if __name__ == "__main__":
  unittest.main()
@@ -0,0 +1,66 @@
+import numpy as np
+import unittest
+from tinygrad import Tensor
+from typing import List
+import functools
+
+def orthogonality_helper(A:Tensor,tolerance=1.0e-5):
+  b_shape,m = A.shape[0:-2],A.shape[-2]  #outer dimension should be the dim along orthogonality
+  A_identity = (Tensor.eye(m).reshape((1,) * len(b_shape)+(m,m)).expand(b_shape+(m,m)))
+  np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)
+
+def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
+  reconstructed_tensor = functools.reduce(Tensor.matmul, A)
+  np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
+
+class TestLinAlg(unittest.TestCase):
+
+  def test_svd_general(self):
+    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
+    for size in sizes:
+      a = Tensor.randn(size).realize()
+      U,S,V = Tensor.svd(a)
+      b_shape,m,n = size[0:-2],size[-2],size[-1]
+      k = min(m,n)
+      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
+      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
+      orthogonality_helper(U)
+      orthogonality_helper(V)
+      reconstruction_helper([U,s_diag,V],a)
+
+  def test_svd_nonfull(self):
+    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
+    for size in sizes:
+      a = Tensor.randn(size).realize()
+      U,S,V = Tensor.svd(a,full_matrices=False)
+      b_shape,m,n = size[0:-2],size[-2],size[-1]
+      k = min(m,n)
+      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
+      #reduced U,V is only orthogonal along smaller dim
+      if (m < n): orthogonality_helper(U),orthogonality_helper(V)
+      else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
+      reconstruction_helper([U,s_diag,V],a)
+
+  @unittest.skip("very big. recommend wrapping with TinyJit around inner function")
+  def test_svd_large(self):
+    size = (1024,1024)
+    a = Tensor.randn(size).realize()
+    U,S,V = Tensor.svd(a)
+    b_shape,m,n = size[0:-2],size[-2],size[-1]
+    k = min(m,n)
+    s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
+    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
+    orthogonality_helper(U,tolerance=1.0e-3)
+    orthogonality_helper(V,tolerance=1.0e-3)
+    reconstruction_helper([U,s_diag,V],a,tolerance=1.0e-3)
+
+  def test_qr_general(self):
+    sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
+    for size in sizes:
+      a = Tensor.randn(size).realize()
+      Q,R = Tensor.qr(a)
+      orthogonality_helper(Q)
+      reconstruction_helper([Q,R],a)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,9 +1,8 @@
-from typing import Union
 import numpy as np
 import unittest
 from dataclasses import replace

-from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel
+from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
 from tinygrad.codegen.gpudims import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -13,9 +12,9 @@ from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
-from tinygrad.dtype import DType, dtypes
+from tinygrad.dtype import DType, dtypes, AddrSpace

-def helper_realized_ast(r:Union[Tensor, list[Tensor]]) -> tuple[UOp, list[Buffer]]:
+def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  if isinstance(r, Tensor): r = [r]
  s = Tensor.schedule(*r)
  run_schedule(s[:-1])  # run all kernels except the last one
@@ -33,8 +32,8 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
  realized_ast, bufs = helper_realized_ast(r)
  k = Kernel(realized_ast)
  k.apply_tensor_cores(use_tensor_cores, axis=axis, tc_select=tc_select, tc_opt=tc_opt)
-  prg = CompiledRunner(replace(k.to_program(), device=Device.DEFAULT))
-  if use_tensor_cores == 1: assert len([uop for uop in k.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
+  prg = CompiledRunner(replace(get_program(k.get_optimized_ast(), k.opts), device=Device.DEFAULT))
+  if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
  assert len([x for x in k.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
  prg.exec(bufs)
  if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
@@ -100,8 +99,9 @@ class TestLinearizer(unittest.TestCase):
    a_t = Tensor.full(st.shape, 2).contiguous().realize()
    b_t = Tensor.full(st.shape, 3).contiguous().realize()
    lin = helper_linearizer_ast(sink, [a_t, b_t], wanna_output=[a_t.numpy()+b_t.numpy(), a_t.numpy()*b_t.numpy()])[0]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops

-    stores = [u for u in lin.uops if u.op is Ops.STORE]
+    stores = [u for u in uops if u.op is Ops.STORE]
    mutable_bufs = dedup(flatten([[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_GLOBAL] for u in stores]))
    assert len(mutable_bufs) == len(stores) == 2
    self.assertSetEqual(set([u.arg for u in mutable_bufs]), set([0,1]))
@@ -148,76 +148,85 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum()
    lin = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])[0]
-    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now
    # RANGE -> LOAD -> RANGE -> ASSIGN
-    #assert any(x.op is Ops.LOAD for x in lin.uops[ranges[0]:ranges[1]])
+    #assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])

  def test_three_nested_range(self):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
    lin = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])[0]
-    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now
    # RANGE -> RANGE -> LOAD -> RANGE -> ASSIGN
    # NOTE: nothing should toposort between the first two ranges
    #assert ranges[0]+1 == ranges[1]
-    #assert any(x.op is Ops.LOAD for x in lin.uops[ranges[1]:ranges[2]])
+    #assert any(x.op is Ops.LOAD for x in uops[ranges[1]:ranges[2]])

  def test_two_nested_range_alt_indexing(self):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
    lin = helper_linearizer_opt(out, wanna_output=[24])[0]
-    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> ASSIGN
-    assert any(x.op in GroupOp.ALU for x in lin.uops[ranges[0]:ranges[1]])
-    assert not any(x.op is Ops.LOAD for x in lin.uops[ranges[0]:ranges[1]])
-    assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in lin.uops[ranges[1]:])
+    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
+    assert not any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
+    assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])

  def test_range_outer_op_before_phi(self):
    a = Tensor.randn(4, 1).realize()
    b = Tensor.randn(1, 1).realize()
    out = (a + b[0]).sum() + b[0]
    lin = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])[0]
-    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> ASSIGN
-    assert len([x for x in lin.uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
+    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1

  def test_range_outer_op_before_phi_nested_range(self):
    a = Tensor.randn(2, ).realize()
    b = Tensor.randn(1, 1).realize()
    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
    lin = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])[0]
-    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
+    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now
    #if getenv("PTX"):
    # LOAD -> RANGE -> CAST -> ALU -> ALU -> LOAD -> ALU -> RANGE -> ALU -> ASSIGN
-    #  assert lin.uops[ranges[0]-2].op is Ops.LOAD
+    #  assert uops[ranges[0]-2].op is Ops.LOAD
    #  assert ranges[1] == ranges[0]+6
-    #  assert [x.op for x in lin.uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
+    #  assert [x.op for x in uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
    # LOAD -> RANGE -> LOAD -> ALU -> RANGE -> ASSIGN
    #else:
-    #  assert lin.uops[ranges[0]-2].op is Ops.LOAD
+    #  assert uops[ranges[0]-2].op is Ops.LOAD
    #  assert ranges[1] == ranges[0]+3
-    #  assert [x.op for x in lin.uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
+    #  assert [x.op for x in uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]

+  @unittest.skip("fragile crap")
  def test_range_outer_op_after_phi(self):
    a = Tensor.randn(4, 1).realize()
    out = a.sum() * a.sum()
    lin = helper_linearizer_opt(out, wanna_output=[a.numpy().sum()*a.numpy().sum()])[0]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
    # RANGE -> LOAD -> ASSIGN -> ALU
-    end = max(i for i,u in enumerate(lin.uops) if u.op is Ops.ENDRANGE)
+    end = max(i for i,u in enumerate(uops) if u.op is Ops.ENDRANGE)
    # the INDEX can be first
-    assert lin.uops[end+1].op in GroupOp.ALU or lin.uops[end+2].op in GroupOp.ALU
+    assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU

+  @unittest.skip("fragile crap")
  def test_range_outer_op_after_phi_nested_range(self):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum() + a.reshape(2, 1).expand(2, 3).sum()
    lin = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3))).sum()*2])[0]
+    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
    # RANGE -> LOAD -> ASSIGN -> ALU
-    end = max(i for i,u in enumerate(lin.uops) if u.op is Ops.ENDRANGE)
+    end = max(i for i,u in enumerate(uops) if u.op is Ops.ENDRANGE)
    # the INDEX can be first
-    assert lin.uops[end+1].op in GroupOp.ALU or lin.uops[end+2].op in GroupOp.ALU
+    assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU

  def test_load_dedup(self):
    # for different leaves in the AST, the same loads may occur.
@@ -227,9 +236,9 @@ class TestLinearizer(unittest.TestCase):
    r = a[:-1] + a[1:]

    k = Kernel(r.schedule()[-1].ast)
-    k.upcast()
-    k.linearize()
-    num_loads = len([uop for uop in k.uops if uop.op is Ops.LOAD])
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"

@@ -240,9 +249,9 @@ class TestLinearizer(unittest.TestCase):
    r = a.expand([2]) + b.expand([2])

    k = Kernel(r.schedule()[-1].ast)
-    k.upcast()
-    k.linearize()
-    num_ops = len([uop for uop in k.uops if uop.op in GroupOp.ALU])
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -251,14 +260,14 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.conv2d(x,w,padding=1).relu()

    k = Kernel(r.schedule()[-1].ast)
-    k.upcast()
-    k.upcast()
-    k.linearize()
-    accs = [u for u in k.uops if u.op is Ops.DEFINE_REG]
-    stores = [u for u in k.uops if u.op is Ops.STORE]
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    accs = [u for u in uops if u.op is Ops.DEFINE_REG]
+    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
    assert len(stores) == 1
-    assert stores[0].src[-1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[1].dtype == dtypes.float.vec(4)

  # NOTE: can reenable, it does work. it just makes BEAM slow
  @unittest.expectedFailure
@@ -267,7 +276,7 @@ class TestLinearizer(unittest.TestCase):
    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
    k = Kernel(out.schedule()[-1].ast)
    k.apply_opt(Opt(OptOps.LOCAL, axis=0, arg=4))
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    self.assertEqual(len(prg.src.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -282,13 +291,13 @@ class TestLinearizer(unittest.TestCase):
    realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
    program = get_program(realized_ast, Device[Device.DEFAULT].renderer)

-    stores = [u for u in program.uops if u.op is Ops.STORE]
+    stores = [u for u in program.uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
-    assert stores[0].src[-1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[1].dtype == dtypes.float.vec(4)
    assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
    # the second store is to gds with no upcasts
-    assert stores[1].src[-1].dtype == dtypes.float
+    assert stores[1].src[1].dtype == dtypes.float
    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())

  def test_zero_fold(self):
@@ -296,9 +305,9 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.stack(a, b)

    k = Kernel(r.schedule()[-1].ast)
-    k.upcast()
-    k.linearize()
-    num_ops = len([uop for uop in k.uops if uop.op in GroupOp.ALU])
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

  def test_sum_acc_dtype(self):
@@ -310,7 +319,7 @@ class TestLinearizer(unittest.TestCase):
        realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
        program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
        local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
-        assert local[0].dtype == acc_dtype
+        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
@@ -318,7 +327,7 @@ class TestLinearizer(unittest.TestCase):
      realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple()))
      program = get_program(realized_ast, Device[Device.DEFAULT].renderer)
      local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
-      assert local[0].dtype == expected_dtype
+      self.assertEqual(local[0].dtype.base, expected_dtype)

    tests = (
      (dtypes.float16, None, dtypes.float),
@@ -356,7 +365,7 @@ class TestLinearizer(unittest.TestCase):
      realized_ast = sched[-1].ast
      kernel = Kernel(realized_ast)
      kernel.apply_tensor_cores(1, axis=0, tc_select=-1, tc_opt=2)
-      prg = kernel.to_program()
+      prg = get_program(kernel.get_optimized_ast(), kernel.opts)
      if Device.DEFAULT == "LLVM":
        assert "0x201000" in prg.src
      elif Device.DEFAULT == "AMD" and AMD_LLVM:
@@ -447,7 +456,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
    k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in k.uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.ASSIGN

@@ -458,7 +467,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
    k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in k.uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.ASSIGN
@@ -471,7 +480,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out).relu()
    k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4)]], apply_tc=True, atol=3e-2, rtol=1e-3)[-1]
-    for u in k.uops:
+    for u in get_program(k.get_optimized_ast(), k.opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.ASSIGN
@@ -482,13 +491,14 @@ class TestLinearizer(unittest.TestCase):
    r = (x@y).relu()
    k = helper_linearizer_opt(r, [[Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]])[-1]
    # the uops graph is RANGE -> DEFINE_ACC -> 4x ALU -> 4x ASSIGN -> ENDRANGE
-    for u in k.uops:
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    for u in uops:
      if u.op is Ops.ASSIGN:
        assert u.src[1].op in GroupOp.ALU
      # children of ASSIGN are placed after ENDRANGE
      if any(x.op is Ops.ASSIGN for x in u.src):
-        end_range = [i for i, x in enumerate(k.uops) if x.op is Ops.ENDRANGE][0]
-        assert end_range < k.uops.index(u)
+        end_range = [i for i, x in enumerate(uops) if x.op is Ops.ENDRANGE][0]
+        assert end_range < uops.index(u)

  def test_grouped_dims(self):
    def _assert_grouped_dims(prefix, dims, max_sizes, reverse_dims, expected_sizes, assert_same_length = True):
@@ -566,7 +576,8 @@ class TestLinearizer(unittest.TestCase):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
    k = helper_linearizer_opt(t+1)[0]
-    idxs = dedup([uop for uop in k.uops if uop.op is Ops.SPECIAL])
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg[0])
    assert idxs[0].arg == ('gidx0', 6), idxs[0].arg
    assert idxs[1].arg == ('gidx1', 5), idxs[1].arg
@@ -605,7 +616,7 @@ class TestLinearizer(unittest.TestCase):
  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      k = helper_linearizer_opt(t)[-1]
-      uops = list(k.linearize().uops)
+      uops = get_program(k.get_optimized_ast(), k.opts).uops
      # ignore kernel optimized IF statements for now
      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
        uops = uops[:uops.index(if_op)]
@@ -635,8 +646,9 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.randn(64,64), Tensor.randn(64,64)
    out = x.matmul(y)
    k = helper_linearizer_opt(out)[-1]
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
    # check that the float4 cast collapses
-    store_vals = [u.src[-1] for u in k.uops if u.op is Ops.STORE]
+    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
      assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE

@@ -659,7 +671,7 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    k = helper_linearizer_opt(out)[-1]
-    store_val = [u.src[-1] for u in k.uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -672,16 +684,17 @@ class TestLinearizer(unittest.TestCase):
            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
    k = helper_linearizer_opt(out, opts=[opt])[-1]
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    local_stores = [u for u in k.uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
-    global_stores = [u for u in k.uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
-    barrier = [u for u in k.uops if u.op is Ops.BARRIER][0]
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
+    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
+    barrier = [u for u in uops if u.op is Ops.BARRIER][0]
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
-      assert store.src[-1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
+      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
-    assert len([u for u in k.uops if u.op is Ops.IF and u.src[-1] == barrier]) == 1
+    assert len([u for u in uops if u.op is Ops.IF and u.src[-1] == barrier]) == 1

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -690,14 +703,15 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    k = helper_linearizer_opt(r)[-1]
-    stores = [u for u in k.uops if u.op is Ops.STORE]
+    uops = get_program(k.get_optimized_ast(), k.opts).uops
+    stores = [u for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
-    self.assertEqual(stores[0].src[-1].dtype, dtypes.float.vec(4))
+    self.assertEqual(stores[0].src[1].dtype, dtypes.float.vec(4))
    #assert stores[0].src[-1].op is not Ops.VECTORIZE

    # the global store doesn't change
-    assert stores[1].src[-1].dtype == dtypes.float
+    assert stores[1].src[1].dtype == dtypes.float

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -715,8 +729,8 @@ class TestLinearizer(unittest.TestCase):
        Opt(op=OptOps.LOCAL, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=3, arg=2)
    ]
    k = helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt])[-1]
-    out = [u for u in k.uops if u.op is Ops.STORE][0]
-    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype == dtypes.float.vec(4)
+    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
+    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -733,19 +747,19 @@ class TestLinearizer(unittest.TestCase):
            Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8),
            Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
    k = helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt])[-1]
-    out = [u for u in k.uops if u.op is Ops.STORE][0]
-    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype.count != 1
+    out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
+    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1

@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
 class TestFloat4(unittest.TestCase):
  @staticmethod
  def count_float4(uops: list[UOp], n=4):
    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.float.vec(n)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)]))
  @staticmethod
  def count_half4(uops: list[UOp]):
    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.half.vec(4)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)]))

  def test_float4_basic(self):
    a = Tensor.empty(2, 8).realize()
@@ -768,13 +782,11 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.shift_to(0, 4)  # float4 dimension
-    k.shift_to(0, 2, insert_before=k.shape_len-1)
-    k.upcast()
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=2))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) == (4, 2)
+    assert TestFloat4.count_float4(uops) == (4, 2)

  @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
  def test_float4_multidim_amx(self):
@@ -785,20 +797,17 @@ class TestFloat4(unittest.TestCase):

      s = c.schedule()[0]
      k = Kernel(s.ast)
-      k.shift_to(0, 4)
-      k.shift_to(0, shift, insert_before=k.shape_len-1)
-      k.upcast()
-      k.upcast()
-      k.linearize()
-      return k
+      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+      k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=shift))
+      return get_program(k.get_optimized_ast(), k.opts).uops

    sizes = [12, 8, 16]
    shifts = [3, 2, 4]
-    excepted_upcast_size = [4, 8, 16]
+    expected_upcast_size = [4, 8, 16]
    expected_output = [(6,3), (2,1), (2,1)]

    for i in range(len(sizes)):
-      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), excepted_upcast_size[i]) == expected_output[i]
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), expected_upcast_size[i]) == expected_output[i]

  def test_float4_unaligned_load(self):
    a = Tensor.empty(9).realize().shrink(((1, 9),))
@@ -821,13 +830,11 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.shift_to(len(k.full_unupcasted_shape)-1, 4)  # manual trigger float4 dim
-    k.upcast()
-    k.shift_to(len(k.full_unupcasted_shape)-1, 2, insert_before=k.shape_len-1)
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=4))
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=1, arg=2))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) == (0, 2)
+    assert TestFloat4.count_float4(uops) == (0, 2)

  @unittest.skipUnless(Device.DEFAULT in {"CPU", "LLVM"} and AMX, "Only CPU with AMX upcasts float up to size 16")
  def test_float4_multidim_unaligned_load_amx(self):
@@ -838,20 +845,17 @@ class TestFloat4(unittest.TestCase):

      s = c.schedule()[0]
      k = Kernel(s.ast)
-      k.shift_to(len(k.full_unupcasted_shape)-1, 4)  # manual trigger float4 dim
-      k.upcast()
-      k.shift_to(len(k.full_unupcasted_shape)-1, shift, insert_before=k.shape_len-1)
-      k.upcast()
-      k.linearize()
-      return k
+      k.shift_to(1, 4, AxisType.UPCAST)  # manual trigger float4 dim
+      k.shift_to(1, shift, AxisType.UPCAST, insert_at=k.shape_len-1)
+      return get_program(k.get_optimized_ast(), k.opts).uops

    sizes = [13, 9, 17]
    shifts = [3, 2, 4]
-    excepted_upcast_size = [4, 8, 16]
+    expected_upcast_size = [4, 8, 16]
    expected_output = [(0,3), (0,1), (0,1)]

    for i in range(len(sizes)):
-      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]).uops, excepted_upcast_size[i]) == expected_output[i]
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), expected_upcast_size[i]) == expected_output[i]

  def test_float4_sometimes_unaligned(self):
    a = Tensor.empty(1, 1, 8).realize()
@@ -862,10 +866,10 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) == (0, 0)
+    assert TestFloat4.count_float4(uops) == (0, 0)

  def test_float4_multidim_sometimes_unaligned(self):
    a = Tensor.empty(1, 1, 7).realize()
@@ -878,27 +882,11 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.upcast()
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=0))
+    k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=0))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) in {(0,1), (1,1)}
-
-  def test_float4_noncontiguous(self):
-    a = Tensor.empty(4, 2).realize()
-    b = Tensor.empty(4, 2).realize()
-    c = a + b
-
-    # we will upcast the top axis of sz 4. they should not be coalesced into float4,
-    # since the top axis is not contiguous.
-
-    s = c.schedule()[0]
-    k = Kernel(s.ast)
-    k.shift_to(0, 4, top=True)  # top axes are float4 axes
-    k.upcast()
-    k.linearize()
-
-    assert TestFloat4.count_float4(k.uops) == (0, 0)
+    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}

  def test_float4_expand(self):
    a = Tensor.empty(9).realize().shrink(((1, 9),))
@@ -910,11 +898,10 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.shift_to(0, 4)  # float4 axis
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) == (0, 1)
+    assert TestFloat4.count_float4(uops) == (0, 1)

  def test_float4_heterogeneous(self):
    a = Tensor.empty(8).realize()
@@ -925,11 +912,10 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    k.shift_to(0, 4)  # float4 axis
-    k.upcast()
-    k.linearize()
+    k.apply_opt(Opt(op=OptOps.UPCAST, axis=0, arg=4))
+    uops = get_program(k.get_optimized_ast(), k.opts).uops

-    assert TestFloat4.count_float4(k.uops) == (1, 1)
+    assert TestFloat4.count_float4(uops) == (1, 1)

  def test_half4_load_unrolled(self):
    # from llama 7B shard 4 gpus
@@ -1084,7 +1070,7 @@ class TestHandCodedOpts(unittest.TestCase):
    k = helper_linearizer_opt(c)[-1]

    assert k.group_for_reduces == 1
-    assert k.local_dims == 1
+    assert k.axis_types.count(AxisType.LOCAL) == 1
    assert k.upcasted == 1

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -1094,7 +1080,7 @@ def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
      for out in ast.src]
  return _helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)

-def helper_linearizer_opt(r:Union[Tensor, list[Tensor]], *args, **kwargs):
+def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
  realized_ast, real_bufs = helper_realized_ast(r)
  return _helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)

@@ -1114,7 +1100,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
  outbufs = [real_bufs[x.src[0].base.arg] for x in realized_ast.src]
  device = real_bufs[0].device

-  def get_prg(k:Kernel): return CompiledRunner(replace(k.to_program(), device=device))
+  def get_prg(k:Kernel): return CompiledRunner(replace(get_program(k.get_optimized_ast(), k.opts), device=device))

  def check_opt(opts, create_k, expected_color_size):
    k = create_k()
@@ -1371,9 +1357,9 @@ class TestKernelOpts(unittest.TestCase):
      [Opt(OptOps.PADTO, 2, 8)],
    ])
    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 0, 0), Opt(OptOps.PADTO, 2, 8)]])
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 0, 0), Opt(OptOps.PADTO, 1, 8)]])
    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 2, 8)]])
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 1, 8)]])
    with self.assertRaises(KernelOptError):
      helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])

@@ -1476,7 +1462,7 @@ class TestKernelOpts(unittest.TestCase):
    opts_shapes = [
      ([Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("red",32)]),
      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.GROUP, 0, 2)], [("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]),
-      # check to ensure local_dims are stable for full UNROLL of first_reduce
+      # check to ensure local_dims are stable for full UNROLL of the first reduce
      ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.UNROLL, 0, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
      ([Opt(OptOps.UNROLL, 0, 0),Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
      # check behavior for full UNROLL on an existing GROUP
@@ -10,6 +10,7 @@ from tinygrad.helpers import getenv
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.opt.search import Opt, OptOps
 from tinygrad.opt.kernel import Kernel
+from tinygrad.engine.realize import get_program

 class TestLinearizerDumb(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -37,12 +38,12 @@ class TestLinearizerDumb(unittest.TestCase):
    opts = [Opt(op=OptOps.TC, axis=2, arg=(-1, 2, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0)]
    k = Kernel(ast, opts=Device["METAL"].renderer)
    k.apply_opts(opts)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)
    Device[Device.DEFAULT].compiler.compile_cached(prg.src)
    gate_count = len([x for x in prg.src.splitlines() if "if" in x])
    assert gate_count == 1, f"must have only one gate {gate_count} != 1"
-    assert len([u for u in k.uops if u.op is Ops.IF]) == 1, "must have a single IF"
+    assert len([u for u in prg.uops if u.op is Ops.IF]) == 1, "must have a single IF"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
  def test_max_simplify_and_cancel(self):
@@ -76,7 +77,7 @@ class TestLinearizerDumb(unittest.TestCase):
    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)]
    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
    k.apply_opts(opts)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)
    assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"

@@ -93,9 +94,9 @@ class TestLinearizerDumb(unittest.TestCase):
    opts = [Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.PADTO, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=0)]
    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
    k.apply_opts(opts)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)
-    if_uops = [u for u in k.uops if u.op is Ops.IF]
+    if_uops = [u for u in prg.uops if u.op is Ops.IF]
    self.assertIn(len(if_uops), {1,2,3})
    conditions = if_uops[0].src[0].toposort()
    self.assertLessEqual(len(conditions), 9)
@@ -134,7 +135,7 @@ class TestLinearizerDumb(unittest.TestCase):
                  UOp(Ops.VIEW, dtypes.half.ptr(131072000), arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)), src=(
                    UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=()),)),)),)),)),)),)),)),))
    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)

  @unittest.expectedFailure
@@ -163,7 +164,7 @@ class TestLinearizerDumb(unittest.TestCase):
    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=0)]
    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
    k.apply_opts(opts)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)
    load_idxs = [x.src[1] for x in k.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
    assert load_idxs[0] < load_idxs[1], f"first loaded idx {load_idxs[0].arg} then {load_idxs[1].arg}!"
@@ -187,7 +188,7 @@ class TestLinearizerDumb(unittest.TestCase):
    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
    k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
    k.apply_opts(opts)
-    prg = k.to_program()
+    prg = get_program(k.get_optimized_ast(), k.opts)
    print(prg.src)
    store_idxs = [x.src[1] for x in k.uops if x.op is Ops.STORE]
    for i in range(len(store_idxs) - 1):
@@ -14,7 +14,6 @@ from tinygrad.shape.view import View
 def _test_overflow(ast, opts):
  lin = Kernel(ast)
  lin.apply_opts(opts)
-  lin.linearize()
  bufs = bufs_from_lin(lin)
  print(bufs)
  time_linearizer(lin, bufs)
@@ -13,29 +13,6 @@ from test.helpers import not_support_multi_device

@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
 class TestNN(unittest.TestCase):
-  def test_sparse_cat_cross_entropy(self):
-    # create in tinygrad
-    input_tensor = Tensor.randn(6, 5) # not square to test that mean scaling uses the correct dimension
-    target = Tensor([0, 0, 0, 1, 2, 3])  # torch doesn't support target=-1
-    torch_input = torch.tensor(input_tensor.numpy())
-    torch_target = torch.tensor(target.numpy(), dtype=torch.long)
-
-    for smoothing in [0.0, 0.1, 0.5, 1.0]:
-      for ignore_index in [-1, 0, 2]:
-        for reduction in ["none", "sum", "mean"]:
-          loss = input_tensor.sparse_categorical_crossentropy(target, label_smoothing=smoothing, ignore_index=ignore_index, reduction=reduction)
-          torch_loss = torch.nn.CrossEntropyLoss(reduction=reduction, label_smoothing=smoothing, ignore_index=ignore_index)(torch_input, torch_target)
-          np.testing.assert_allclose(loss.numpy(), torch_loss.detach().numpy(), atol=1e-5, rtol=1e-6)
-
-          # also test with a batch dimension (of size 1)
-          loss = input_tensor.unsqueeze(0).sparse_categorical_crossentropy(
-            target.unsqueeze(0), label_smoothing=smoothing, ignore_index=ignore_index, reduction=reduction
-          )
-          torch_loss = torch.nn.CrossEntropyLoss(reduction=reduction, label_smoothing=smoothing, ignore_index=ignore_index)(
-            torch_input.unsqueeze(0).permute(0,2,1), torch_target.unsqueeze(0)
-          )
-          np.testing.assert_allclose(loss.numpy(), torch_loss.detach().numpy(), atol=1e-5, rtol=1e-6)
-
  def test_batchnorm2d(self, training=False, threed=False, track_running_stats=True):
    with Tensor.train(training):
      szs = [4, 8, 16, 32]
@@ -1,4 +1,4 @@
-import time, math, unittest, functools, warnings
+import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
@@ -86,9 +86,11 @@ def prepare_test_op(low, high, shps, vals, forward_only=False):

 class TestOps(unittest.TestCase):

-  def helper_test_exception(self, shps, torch_fxn, tinygrad_fxn, expected, forward_only=False, exact=False, vals=None, low=-1.5, high=1.5):
+  def helper_test_exception(self, shps, torch_fxn, tinygrad_fxn=None, expected=None, forward_only=False, exact=False, vals=None, low=-1.5, high=1.5):
    if getenv("MOCKGPU") and Device.DEFAULT == "NV": self.skipTest('helper_test_exception fails in CI CUDA')
    ts, tst = prepare_test_op(low, high, shps, vals, forward_only)
+    if tinygrad_fxn is None:
+      tinygrad_fxn = torch_fxn
    with self.assertRaises(expected) as torch_cm:
      torch_fxn(*ts)
    with self.assertRaises(expected) as tinygrad_cm:
@@ -234,10 +236,10 @@ class TestOps(unittest.TestCase):
    helper_test_op([(3,3,3)], lambda x: x.unfold(1, 0, 8))
    helper_test_op([(3,3,3,3,3)], lambda x: x.unfold(-1, 2, 2))

-    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 9, 3), lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
-    self.helper_test_exception([(8,)], lambda x: x.unfold(1, 8, 3), lambda x: x.unfold(1, 8, 3), expected=IndexError)
-    self.helper_test_exception([(8,)], lambda x: x.unfold(0, -1, 3), lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
-    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 1, -1), lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
+    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
+    self.helper_test_exception([(8,)], lambda x: x.unfold(1, 8, 3), expected=IndexError)
+    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
+    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 1, -1), expected=RuntimeError)

  def test_meshgrid(self):
    x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.], requires_grad=True)
@@ -546,7 +548,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65), (45,65)], lambda x,y: x/y)
    helper_test_op([(), ()], lambda x,y: x/y)

-  @unittest.skipIf(AMD_LLVM, "AMD with LLVM backend generate rcp in FP division causes trunc/floor errors")
+  @unittest.skipIf(Device.DEFAULT == "AMD" and AMD_LLVM, "AMD with LLVM backend generate rcp in FP division causes trunc/floor errors")
  def test_div_rounding_mode(self):
    for denominator in [-10, -5, -3, -2, -1, 1, 2, 3, 5, 10]:
      # int numerator
@@ -574,8 +576,7 @@ class TestOps(unittest.TestCase):
        helper_test_op(None, lambda x,y: x.div(y, rounding_mode="trunc"), forward_only=True, vals=[[numerator], [denominator]])
        helper_test_op(None, lambda x,y: x.div(y, rounding_mode="floor"), forward_only=True, vals=[[numerator], [denominator]])

-    self.helper_test_exception(None, lambda x,y: x.div(y, rounding_mode="typo"), lambda x,y: x.div(y, rounding_mode="typo"), forward_only=True,
-                               vals=[[5], [0]], expected=RuntimeError)
+    self.helper_test_exception(None, lambda x,y: x.div(y, rounding_mode="typo"), forward_only=True, vals=[[5], [0]], expected=RuntimeError)

  def test_div_int(self):
    helper_test_op(None, lambda x,y: x/y, Tensor.div, forward_only=True, vals=[[5, 6, 7],[1, 2, 3]])
@@ -587,12 +588,6 @@ class TestOps(unittest.TestCase):
    if is_dtype_supported(dtypes.uint64):
      x = Tensor(2**64 - 1, dtype=dtypes.uint64).idiv(1)
      np.testing.assert_equal(x.numpy(), 2**64 - 1)
-    # 1 // 0 is device dependent, but it should not raise
-    Tensor([1]).idiv(1).realize()
-    if not CI:  # TODO: crashed in CI on some devices
-      # ... because if might be in a where branch that the output is well defined
-      t = Tensor([-1, 0, 1, 2])
-      np.testing.assert_equal((t > 0).where(1//t, t).numpy(), [-1, 0, 1, 0])

  def test_scalar_div(self):
    helper_test_op([(45,65)], lambda x: x/255)
@@ -741,7 +736,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
    helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)

-    self.helper_test_exception([(4), (4)], torch.bitwise_xor, Tensor.bitwise_xor, expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)

  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
@@ -758,7 +753,7 @@ class TestOps(unittest.TestCase):

    helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])

-    self.helper_test_exception([(4), (4)], torch.bitwise_and, Tensor.bitwise_and, expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)

  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
@@ -773,7 +768,7 @@ class TestOps(unittest.TestCase):
    ten0, ten1 = Tensor(data[0], dtype=dtypes.bool), Tensor(data[1], dtype=dtypes.bool)
    helper_test_op([], lambda: tor0|tor1, lambda: ten0|ten1, forward_only=True)

-    self.helper_test_exception([(4), (4)], torch.bitwise_or, Tensor.bitwise_or, expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_or(y), expected=RuntimeError)

  def test_bitwise_not(self):
    data = [[1,-8,1],[32,1,6]]
@@ -788,7 +783,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor.bitwise_not(), lambda: ten.bitwise_not(), forward_only=True)
    helper_test_op([], lambda: ~tor, lambda: ~ten, forward_only=True)

-    self.helper_test_exception([(4)], torch.bitwise_not, Tensor.bitwise_not, expected=RuntimeError)
+    self.helper_test_exception([(4)], lambda x: x.bitwise_not(), expected=RuntimeError)

  def test_lshift(self):
    data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -826,6 +821,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
      helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
  def test_cos(self):
    helper_test_op([(45,65)], lambda x: x.cos())
    helper_test_op([()], lambda x: x.cos())
@@ -833,6 +829,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
      helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
  def test_tan(self):
    # NOTE: backward has much higher diff with input close to pi/2 and -pi/2
    helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)
@@ -1108,6 +1105,12 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x.sort(stable=True, descending=True).indices.type(torch.int32),
                   lambda x: x.sort(descending=True)[1], forward_only=True, vals=[[0, 1] * 9])

+  def test_argsort(self):
+    for dim in [-1, 0, 1]:
+      for descending in [True, False]:
+        helper_test_op([(8,8,6)], lambda x: torch.argsort(x, dim=dim, descending=descending, stable=True).type(torch.int32),
+                                  lambda x: x.argsort(dim, descending), forward_only=True)
+
  def test_topk(self):
    helper_test_op([(10)], lambda x: x.topk(3).values, lambda x: x.topk(3)[0], forward_only=True)
    helper_test_op([(10)], lambda x: x.topk(3).indices.type(torch.int32), lambda x: x.topk(3)[1], forward_only=True)
@@ -1127,7 +1130,7 @@ class TestOps(unittest.TestCase):
    value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3, largest=False)
    np.testing.assert_equal(value.numpy(), [0, 0, 0])
    np.testing.assert_equal(indices.numpy(), [2, 4, 6])
-    self.helper_test_exception([(4)], lambda x: x.topk(5), lambda x: x.topk(5), expected=(RuntimeError, ValueError))
+    self.helper_test_exception([(4)], lambda x: x.topk(5), expected=(RuntimeError, ValueError))

  def test_einsum(self):
    # matrix transpose
@@ -1281,7 +1284,8 @@ class TestOps(unittest.TestCase):
                                                                         np.arange(64,128,dtype=np.float32).reshape(8,8)])
  def test_small_gemm_eye(self):
    helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
-  @unittest.skipIf(CI and Device.DEFAULT in ["NV", "LLVM", "GPU", "CUDA"] or IMAGE, "not supported on these in CI/IMAGE")
+  @unittest.skipIf(CI and Device.DEFAULT in ["NV", "LLVM", "GPU", "CUDA"] or IMAGE
+  or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
  def test_gemm_fp16(self):
    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3)
  def test_gemm(self):
@@ -1329,9 +1333,9 @@ class TestOps(unittest.TestCase):
    helper_test_op([()], lambda x: x.sum(0))
    helper_test_op([()], lambda x: x.sum(-1))
    helper_test_op([()], lambda x: x.sum(()))
-    self.helper_test_exception([(3,4,5,6)], lambda x: x.sum(5), lambda x: x.sum(5), expected=IndexError)
-    self.helper_test_exception([()], lambda x: x.sum(1), lambda x: x.sum(1), expected=IndexError)
-    self.helper_test_exception([()], lambda x: x.sum((1,)), lambda x: x.sum((1,)), expected=IndexError)
+    self.helper_test_exception([(3,4,5,6)], lambda x: x.sum(5), expected=IndexError)
+    self.helper_test_exception([()], lambda x: x.sum(1), expected=IndexError)
+    self.helper_test_exception([()], lambda x: x.sum((1,)), expected=IndexError)

  def test_sum_dtype_arg(self):
    helper_test_op([(45,3)], lambda x: x.sum(), lambda x: x.sum(dtype=dtypes.float32))
@@ -1846,9 +1850,9 @@ class TestOps(unittest.TestCase):
    helper_test_op([(3,4,5,6)], lambda x: x.permute((3,2,1,0)))
    helper_test_op([(3,4,5,6)], lambda x: x.permute((-2,-1,1,0)))
    helper_test_op([()], lambda x: x.permute(()))
-    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,2)), lambda x: x.permute((0,2)), expected=RuntimeError)
-    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,1,2,3,3,3)), lambda x: x.permute((0,1,2,3,3,3)), expected=RuntimeError)
-    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,0,1,2,3)), lambda x: x.permute((0,0,1,2,3)), expected=RuntimeError)
+    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,2)), expected=RuntimeError)
+    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,1,2,3,3,3)), expected=RuntimeError)
+    self.helper_test_exception([(3,4,5,6)], lambda x: x.permute((0,0,1,2,3)), expected=RuntimeError)

  def test_reshape(self):
    helper_test_op([(4,3,6,6)], lambda x: x.reshape((12,6,6)))
@@ -1859,8 +1863,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([(1,)], lambda x: x.reshape(()))
    helper_test_op([()], lambda x: x.reshape((1,)))
    helper_test_op([()], lambda x: x.reshape((1,1,1)))
-    self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,2)), lambda x: x.reshape((-1,-1,2)), expected=RuntimeError)
-    self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,-1,2)), lambda x: x.reshape((-1,-1,-1,2)), expected=RuntimeError)
+    self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,2)), expected=RuntimeError)
+    self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,-1,2)), expected=RuntimeError)

    with self.assertRaises(ValueError):
      x = Tensor.ones((4,3,6,6))
@@ -1885,16 +1889,16 @@ class TestOps(unittest.TestCase):
    helper_test_op([()], lambda x: x.flip(()))
    helper_test_op([(1,)], lambda x: x.flip(()))
    helper_test_op([(4,3,6,6)], lambda x: x.flip(()))
-    self.helper_test_exception([(3,4)], lambda x: x.flip((0,0)), lambda x: x.flip((0,0)), expected=RuntimeError)
-    self.helper_test_exception([(3,4)], lambda x: x.flip((1,1)), lambda x: x.flip((1,1)), expected=RuntimeError)
-    self.helper_test_exception([(3,4)], lambda x: x.flip((1,-1)), lambda x: x.flip((1,-1)), expected=RuntimeError)
+    self.helper_test_exception([(3,4)], lambda x: x.flip((0,0)), expected=RuntimeError)
+    self.helper_test_exception([(3,4)], lambda x: x.flip((1,1)), expected=RuntimeError)
+    self.helper_test_exception([(3,4)], lambda x: x.flip((1,-1)), expected=RuntimeError)

  def test_squeeze(self):
    helper_test_op([(1,3,6,6)], lambda x: x.squeeze(0))
    helper_test_op([(4,3,1,6)], lambda x: x.squeeze(1))
    helper_test_op([(4,3,6,6)], lambda x: x.squeeze(3))
-    self.helper_test_exception([(4,3,6,6)], lambda x: torch.squeeze(x, 50), lambda x: x.squeeze(dim=50), expected=IndexError)
-    self.helper_test_exception([(4,3,6,6)], lambda x: torch.squeeze(x, -50), lambda x: x.squeeze(dim=-50), expected=IndexError)
+    self.helper_test_exception([(4,3,6,6)], lambda x: x.squeeze(50), expected=IndexError)
+    self.helper_test_exception([(4,3,6,6)], lambda x: x.squeeze(50), expected=IndexError)
    helper_test_op([(4,3,6,1)], lambda x: x.squeeze(-1))
    helper_test_op([(4,3,6,6)], lambda x: x.squeeze())
    helper_test_op([(1,3,6,6)], lambda x: x.squeeze())
@@ -1902,9 +1906,9 @@ class TestOps(unittest.TestCase):
    helper_test_op([()], lambda x: x.squeeze(-1))
    helper_test_op([()], lambda x: x.squeeze(0))
    helper_test_op([()], lambda x: x.squeeze())
-    self.helper_test_exception([()], lambda x: torch.squeeze(x, 10), lambda x: x.squeeze(dim=10), expected=IndexError)
-    self.helper_test_exception([()], lambda x: torch.squeeze(x, 1), lambda x: x.squeeze(dim=1), expected=IndexError)
-    self.helper_test_exception([()], lambda x: torch.squeeze(x, -2), lambda x: x.squeeze(dim=-2), expected=IndexError)
+    self.helper_test_exception([()], lambda x: x.squeeze(10), expected=IndexError)
+    self.helper_test_exception([()], lambda x: x.squeeze(1), expected=IndexError)
+    self.helper_test_exception([()], lambda x: x.squeeze(-2), expected=IndexError)

  def test_unsqueeze(self):
    helper_test_op([(4,3,6,6)], lambda x: x.unsqueeze(0))
@@ -1927,20 +1931,31 @@ class TestOps(unittest.TestCase):
    helper_test_op([(4,3,6,6)], lambda x: x.unflatten(3, (3, 2)))
    helper_test_op([(4,3,6,6)], lambda x: x.unflatten(-1, (3, 2, 1)))

+  def test_diag(self):
+    helper_test_op([(5,)], lambda x: x.diag())
+
+  def test_diagonal(self):
+    helper_test_op([(5,5)], lambda x: x.diagonal())
+
  def test_roll(self):
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, 1, 0), lambda x: x.roll(1, 0))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, -1, 0), lambda x: x.roll(-1, 0))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, shifts=(2, 1), dims=(0, 1)), lambda x: x.roll(shifts=(2, 1), dims=(0, 1)))
-    helper_test_op([(2, 4, 6)], lambda x: torch.roll(x, 1, 0), lambda x: x.roll(1, 0))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, 1, -1), lambda x: x.roll(1, -1))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, -1, -1), lambda x: x.roll(-1, -1))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, 5, 0), lambda x: x.roll(5, 0))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, -5, 0), lambda x: x.roll(-5, 0))
-    helper_test_op([(2, 4, 6)], lambda x: torch.roll(x, shifts=(2, -3), dims=(0, 2)), lambda x: x.roll(shifts=(2, -3), dims=(0, 2)))
-    helper_test_op([(2, 4, 6)], lambda x: torch.roll(x, shifts=(1, 2, -1), dims=(0, 1, 2)), lambda x: x.roll(shifts=(1, 2, -1), dims=(0, 1, 2)))
-    helper_test_op([(2, 4)], lambda x: torch.roll(x, 0, 0), lambda x: x.roll(0, 0))
-    helper_test_op([(2, 4, 6)], lambda x: torch.roll(x, shifts=(0, 0), dims=(0, 1)), lambda x: x.roll(shifts=(0, 0), dims=(0, 1)))
-    helper_test_op([(2, 4, 6)], lambda x: torch.roll(x, shifts=(0, 2), dims=(0, 1)), lambda x: x.roll(shifts=(0, 2), dims=(0, 1)))
+    helper_test_op([(2, 4)], lambda x: x.roll(1))
+    helper_test_op([(2, 4)], lambda x: x.roll((1,)))
+    self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
+    helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))
+    helper_test_op([(2, 4)], lambda x: x.roll(shifts=(2, 1), dims=(0, 1)))
+    helper_test_op([(2, 4, 6)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 4)], lambda x: x.roll(1, -1))
+    helper_test_op([(2, 4)], lambda x: x.roll(-1, -1))
+    helper_test_op([(2, 4)], lambda x: x.roll(5, 0))
+    helper_test_op([(2, 4)], lambda x: x.roll(-5, 0))
+    helper_test_op([(2, 4, 6)], lambda x: x.roll(shifts=(2, -3), dims=(0, 2)))
+    helper_test_op([(2, 4, 6)], lambda x: x.roll(shifts=(1, 2, -1), dims=(0, 1, 2)))
+    helper_test_op([(2, 4)], lambda x: x.roll(0, 0))
+    helper_test_op([(2, 4, 6)], lambda x: x.roll(shifts=(0, 0), dims=(0, 1)))
+    helper_test_op([(2, 4, 6)], lambda x: x.roll(shifts=(0, 2), dims=(0, 1)))
+    self.helper_test_exception([(3, 3)], lambda x: x.roll(shifts=1, dims=(0, 1)), expected=RuntimeError)
+    self.helper_test_exception([(10,)], lambda x: x.roll(shifts=(1, 2), dims=0), expected=RuntimeError)

  def test_detach(self):
    helper_test_op([(4,3,6,6)], lambda x: x.detach(), forward_only=True)
@@ -2594,10 +2609,13 @@ class TestOps(unittest.TestCase):

  def test_stack(self):
    for dim in range(-1, 3):
-      helper_test_op([(45,65,3), (45,65,3), (45,65,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack(x, y, z, dim=dim))
+      helper_test_op([(5,6,3), (5,6,3), (5,6,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack(x, y, z, dim=dim))
+      helper_test_op([(5,6,3), (5,6,3), (5,6,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack((x, y, z), dim=dim))

    with self.assertRaises(IndexError):
      Tensor.stack(Tensor.randn(45, 65, 3), dim=77)
+    with self.assertRaises(ValueError):
+      Tensor.stack((Tensor([1, 2]), Tensor([3, 4])), Tensor([5, 6]))

    a = Tensor(3.14)
    np.testing.assert_allclose(Tensor.stack(a, a).numpy(), Tensor([3.14, 3.14]).numpy())
@@ -2636,7 +2654,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65)], lambda x: x.clip(3, 0))  # min > max
    helper_test_op([(45,65)], lambda x: x.clip(None, 0))
    helper_test_op([(45,65)], lambda x: x.clip(0, None))
-    self.helper_test_exception([(45,65)], lambda x: x.clip(None, None), lambda x: x.clip(None, None), RuntimeError)
+    self.helper_test_exception([(45,65)], lambda x: x.clip(None, None), expected=RuntimeError)

  def test_matvecmat(self):
    helper_test_op([(1,128), (128,128), (128,128)], lambda x,y,z: (x@y).relu()@z)
@@ -2897,6 +2915,17 @@ class TestOps(unittest.TestCase):
      lambda x,y,z,m: Tensor.scaled_dot_product_attention(x,y,z,is_causal=True,attn_mask=m),
      expected=RuntimeError)

+  def test_scaled_dot_product_attention_gqa(self):
+    helper_test_op([(32,32,16,64), (32,8,16,64), (32,8,16,64)],
+                   lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z,enable_gqa=True),
+                   lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,enable_gqa=True))
+
+  def test_scaled_dot_product_attention_gqa_errors(self):
+    self.helper_test_exception([(32,31,16,64), (32,8,16,64), (32,8,16,64)],
+      lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z),
+      lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,enable_gqa=True),
+      expected=(AssertionError, RuntimeError, ValueError))
+
  def test_binary_crossentropy(self):
    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),y.clip(0,1)),
                                       lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1)))
@@ -2944,6 +2973,39 @@ class TestOps(unittest.TestCase):
      helper_test_op([(32,10)], lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes), label_smoothing=ls),
                                lambda x: x.cross_entropy(Tensor(classes), label_smoothing=ls))

+  def test_sparse_categorical_crossentropy(self):
+    classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist()
+    helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss()(x, torch.tensor(classes)),
+                              lambda x: x.sparse_categorical_crossentropy(Tensor(classes)))
+
+    # combine args
+    helper_test_op([(12,10)],
+                   lambda x: torch.nn.CrossEntropyLoss(reduction="mean", ignore_index=classes[0], label_smoothing=0.3)(x, torch.tensor(classes)),
+                   lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction="mean", ignore_index=classes[0], label_smoothing=0.3))
+
+    # with batch. somehow this does not match torch
+    classes = np.random.randint(0, 10, (3,12), dtype=np.int32).tolist()
+    helper_test_op([(3,12,10)], lambda x: torch.nn.CrossEntropyLoss()(x.permute(0,2,1), torch.tensor(classes)),
+                                lambda x: x.sparse_categorical_crossentropy(Tensor(classes)))
+
+  def test_sparse_categorical_crossentropy_reductions(self):
+    for r in ("mean", "sum", "none"):
+      classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist()
+      helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(reduction=r)(x, torch.tensor(classes)),
+                                lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction=r))
+
+  def test_sparse_categorical_crossentropy_ignore_index(self):
+    classes = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
+    for i in (-1, 0, 3):
+      helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(ignore_index=i)(x, torch.tensor(classes)),
+                                lambda x: x.sparse_categorical_crossentropy(Tensor(classes), ignore_index=i))
+
+  def test_sparse_categorical_crossentropy_label_smoothing(self):
+    for s in (0.3, 0.9):
+      classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist()
+      helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(label_smoothing=s)(x, torch.tensor(classes)),
+                                lambda x: x.sparse_categorical_crossentropy(Tensor(classes), label_smoothing=s))
+
  def test_nll_loss(self):
    target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
    helper_test_op([(32,10)],
@@ -3022,6 +3084,20 @@ class TestOps(unittest.TestCase):
  def test_bitcast(self):
    helper_test_op([(3, 3)], lambda x: x.view(torch.int32), lambda x: x.bitcast(dtypes.int32), forward_only=True)

+  def test_svd(self):
+    # test for tiny backend. real svd tests are in test_linalg
+    A = torch.randn(5, 5)
+    U, S, Vh = torch.linalg.svd(A)
+    np.testing.assert_equal(U.shape, (5,5))
+    np.testing.assert_equal(Vh.shape, (5,5))
+    np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
+
+    A = torch.randn(5, 3)
+    U, S, Vh = torch.linalg.svd(A, full_matrices=False)
+    np.testing.assert_equal(U.shape, (5,3))
+    np.testing.assert_equal(Vh.shape, (3,3))
+    np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
+
@unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}")
 class TestOpsUint8(unittest.TestCase):
  def test_cast(self):
@@ -3,7 +3,7 @@ import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 class TestOptGemm(unittest.TestCase):
  @classmethod
@@ -19,7 +19,7 @@ class TestOptGemm(unittest.TestCase):
    si = get_single_element(t.schedule())
    k = Kernel(si.ast)
    k.apply_opts(opts)
-    run = CompiledRunner(k.to_program())
+    run = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
    ExecItem(run, si.bufs).run()
    test = si.bufs[0].numpy().reshape(self.res.shape)
    np.testing.assert_allclose(self.res, test, atol=1e-4)
@@ -0,0 +1,148 @@
+import unittest
+from tinygrad import Tensor, nn, Variable, UOp, dtypes
+
+# outerworld range should support three things
+# 1. full optimizer steps (test_model_bound_range)
+# 2. gradient accumulation (you want to end the range before running the optimizer)
+# 3. stacked linear layers
+
+class Model:
+  def __init__(self): self.w = nn.Linear(64, 8, bias=False)
+  def __call__(self, x:Tensor) -> Tensor: return self.w(x)
+
+def get_model_and_opt():
+  Tensor.manual_seed(1337)
+  m = Model()
+  opt = nn.optim.SGD(nn.state.get_parameters(m), lr=0.1, weight_decay=0)
+  return m, opt
+
+class TestOuterworldRange(unittest.TestCase):
+  STEPS = 5
+  BS = 20
+
+  @classmethod
+  def setUpClass(cls):
+    Tensor.manual_seed(1338)
+    # it learns to compute mean
+    cls.X = Tensor.randn(cls.STEPS, cls.BS, 64).contiguous().realize()
+    cls.Y = cls.X.reshape(cls.STEPS, cls.BS, 8, 8).mean(axis=-1).contiguous().realize()
+    cls.losses = cls._get_model_baseline()
+
+  def _compare(self, losses):
+    for i,(x,y) in enumerate(zip(self.losses, losses)):
+      self.assertAlmostEqual(x, y, places=5, msg=f"mismatch at {i} in {self.losses} vs {losses}")
+
+  @classmethod
+  @Tensor.train()
+  def _get_model_baseline(self):
+    m, opt = get_model_and_opt()
+    losses = []
+    for i in range(self.STEPS):
+      opt.zero_grad()
+      loss = (m(self.X[i]) - self.Y[i]).square().mean()
+      loss.backward()
+      loss.realize(*opt.schedule_step())
+      losses.append(loss.item())
+    return losses
+
+  @Tensor.train()
+  def test_model_grad_acc(self):
+    m, opt = get_model_and_opt()
+    losses = []
+    for i in range(self.STEPS):
+      opt.zero_grad()
+      sub_batch_size = self.BS//2
+      loss = 0
+      scaling_factor = self.BS//sub_batch_size
+      for j in range(0, self.BS, sub_batch_size):
+        sub_loss = (m(self.X[i][j:j+sub_batch_size]) - self.Y[i][j:j+sub_batch_size]).square().mean() / scaling_factor
+        sub_loss.backward()
+        loss += sub_loss
+      loss.realize(*opt.schedule_step())
+      losses.append(loss.item())
+    self._compare(losses)
+
+  @Tensor.train()
+  def test_model_variable(self):
+    m, opt = get_model_and_opt()
+    losses = []
+    vi = Variable('i', 0, self.STEPS-1)
+    for i in range(self.STEPS):
+      vib = vi.bind(i)
+      opt.zero_grad()
+      loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
+      loss.backward()
+      loss.realize(*opt.schedule_step())
+      losses.append(loss.item())
+    self._compare(losses)
+
+  @Tensor.train()
+  def test_model_scheduled(self):
+    m, opt = get_model_and_opt()
+    losses = []
+    for i in range(self.STEPS):
+      opt.zero_grad()
+      loss = (m(self.X[i]) - self.Y[i]).square().mean()
+      loss.backward()
+      opt.schedule_step()
+      losses.append(loss)
+    self._compare(Tensor.stack(*losses).tolist())
+
+  @Tensor.train()
+  def test_model_scheduled_setitem(self):
+    m, opt = get_model_and_opt()
+    losses = Tensor.empty(self.STEPS)
+    for i in range(self.STEPS):
+      opt.zero_grad()
+      loss = (m(self.X[i]) - self.Y[i]).square().mean()
+      loss.backward()
+      opt.schedule_step()
+      # TODO: this shouldn't realize
+      losses[i] = loss.requires_grad_(False)
+    self._compare(losses.tolist())
+
+  @unittest.expectedFailure
+  @Tensor.train()
+  def test_model_scheduled_variable(self):
+    m, opt = get_model_and_opt()
+    losses = []
+    vi = Variable('i', 0, self.STEPS-1)
+    for i in range(self.STEPS):
+      vib = vi.bind(i)
+      opt.zero_grad()
+      loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
+      loss.backward()
+      opt.schedule_step()
+      losses.append(loss)
+    self._compare(Tensor.stack(*losses).tolist())
+
+  @unittest.expectedFailure
+  @Tensor.train()
+  def test_model_scheduled_variable_setitem(self):
+    m, opt = get_model_and_opt()
+    losses = Tensor.empty(self.STEPS)
+    vi = Variable('i', 0, self.STEPS-1)
+    for i in range(self.STEPS):
+      vib = vi.bind(i)
+      opt.zero_grad()
+      loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
+      loss.backward()
+      opt.schedule_step()
+      losses[vib] = loss.requires_grad_(False)
+    self._compare(losses.tolist())
+
+  @unittest.expectedFailure
+  @Tensor.train()
+  def test_model_bound_range(self):
+    m, opt = get_model_and_opt()
+    # TODO: should ranges be unique so you don't have to pass in the -1?
+    rng = UOp.range(dtypes.int, self.STEPS, -1)
+    vib = Variable('i', 0, self.STEPS-1).bind(rng)
+    loss = (m(self.X[vib]) - self.Y[vib]).square().mean()
+    loss.backward()
+    losses = Tensor.empty(self.STEPS)
+    losses[vib] = loss
+    losses.realize(*opt.schedule_step())
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,7 +1,7 @@
-import unittest, struct, contextlib, statistics
+import unittest, struct, contextlib, statistics, time, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context
-from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileRangeEvent, ProfileDeviceEvent, ProfileGraphEvent
+from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events
+from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner

@@ -10,7 +10,11 @@ MOCKGPU = getenv("MOCKGPU")
@contextlib.contextmanager
 def helper_collect_profile(*devs):
  for dev in devs: dev.synchronize()
-  Compiled.profile_events = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]
+  saved = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]
+  Compiled.profile_events.clear()
+  for x in saved: Compiled.profile_events.append(x)
+
+  cpu_events.clear()

  profile_list = []
  with Context(PROFILE=1):
@@ -18,6 +22,7 @@ def helper_collect_profile(*devs):
    for dev in devs: dev.synchronize()
    for dev in devs: dev._at_profile_finalize()
    for x in Compiled.profile_events: profile_list.append(x)
+    profile_list.extend(cpu_events)

 def helper_profile_filter_device(profile, device:str):
  assert any(getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent) for x in profile), f"device {device} is not registred"
@@ -25,7 +30,10 @@ def helper_profile_filter_device(profile, device:str):
  assert len(dev_events) == 1, "only one device registration event is expected"
  return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]

-@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled) or Device.DEFAULT in {"METAL"}, "HCQ device required to run")
+# TODO: support in HCQCompiled
+is_cpu_hcq = Device.DEFAULT in {"CPU", "LLVM"}
+
+@unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")
 class TestProfiler(unittest.TestCase):
  @classmethod
  def setUpClass(self):
@@ -73,13 +81,15 @@ class TestProfiler(unittest.TestCase):
    evs = [x for x in profile if isinstance(x, ProfileRangeEvent)]

    assert len(evs) == 3, "3 kernel runs are expected"
-    assert evs[0].is_copy, "kernel should be copy"
-    assert evs[1].name == runner_name, "kernel name is not correct"
-    assert not evs[1].is_copy, "kernel should not be copy"
-    assert evs[2].is_copy, "kernel should be copy"
+    # NOTE: order of events does not matter, the tool is responsible for sorting them
+    copy_events = [e for e in evs if e.is_copy]
+    self.assertEqual(len(copy_events), 2)

-    for i in range(1, 3):
-      assert evs[i].st > evs[i-1].en, "timestamp not aranged"
+    prg_events = [e for e in evs if not e.is_copy]
+    assert prg_events[0].name == runner_name, "kernel name is not correct"
+
+    #for i in range(1, 3):
+    #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"

  def test_profile_multidev(self):
    d1 = Device[f"{Device.DEFAULT}:1"]
@@ -159,5 +169,45 @@ class TestProfiler(unittest.TestCase):
      assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5ms"
    print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))

+  def test_cpu_profile(self):
+    def test_fxn(err=False):
+      time.sleep(0.1)
+      if err: raise Exception()
+      time.sleep(0.1)
+
+    with helper_collect_profile(dev:=TestProfiler.d0) as profile:
+      with cpu_profile("test_1", dev.device):
+        test_fxn(err=False)
+      with self.assertRaises(Exception):
+        with cpu_profile("test_2", dev.device):
+          test_fxn(err=True)
+
+    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent)]
+    self.assertEqual(len(range_events), 2)
+    # record start/end time up to exit (error or success)
+    for e in range_events:
+      self.assertGreater(e.en, e.st)
+    e1, e2 = range_events
+    self.assertEqual([e1.name, e2.name], ["test_1", "test_2"])
+    # TODO: this is flaky
+    #self.assertLess(e1.st, e2.st)
+    #self.assertGreater(e1.en-e1.st, e2.en-e2.st)
+
+  @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
+  def test_graph(self):
+    from test.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs
+    device = TestProfiler.d0.device
+    bufs = [helper_alloc_rawbuffer(device, fill=True) for _ in range(5)]
+    graphs = [[helper_exec_op(device, bufs[0], [bufs[1], bufs[2]]), helper_exec_op(device, bufs[0], [bufs[3], bufs[4]]),]]
+    with helper_collect_profile(dev:=TestProfiler.d0) as profile:
+      helper_test_graphs(dev.graph, graphs, runs:=2)
+      # NOTE: explicitly trigger deletion of all graphs
+      graphs.clear()
+      gc.collect()
+    graphs = [e for e in profile if isinstance(e, ProfileGraphEvent)]
+    self.assertEqual(len(graphs), runs)
+    for ge in graphs:
+      self.assertEqual(len(ge.ents), len(graphs))
+
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -5,7 +5,7 @@ from dataclasses import replace
 from tinygrad import Tensor, Context, Device, dtypes
 from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
 from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item
+from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
 from tinygrad.opt.search import bufs_from_lin
 from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import

@@ -32,7 +32,9 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
    graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init])

  # Create and save the model
-  model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
+  #model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
+  # TODO remove this once ORT supports 1.18.0
+  model_def = helper.make_model(graph_def, producer_name="single_gemm_example", ir_version=10, opset_imports=[helper.make_opsetid("", 22)])
  onnx.save_model(model_def, model_path)
  return model_path

@@ -41,7 +43,7 @@ def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
  #opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128)] #, Opt(op=OptOps.UNROLL, axis=0, arg=4)]
  k.apply_opts(opts)
-  prg = k.to_program()
+  prg = get_program(k.get_optimized_ast(), k.opts)
  if replace_src is not None:
    old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
    prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
@@ -63,6 +65,7 @@ def get_quantized_model(sz):
                  extra_options={"ActivationSymmetric": False})
  return out_file

+@unittest.skip("this is broken")
@unittest.skipIf(Device.DEFAULT != "CPU", "only tests for CPU")
 class TestQuantizeOnnxCPU(unittest.TestCase):
  def test_quant_128(self, sz=128):
@@ -70,10 +73,9 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
      import onnx # noqa: F401 # pylint: disable=unused-import
    except ImportError:
      raise unittest.SkipTest()
-    from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+    from tinygrad.frontend.onnx import OnnxRunner
    out_file = get_quantized_model(sz)
-    onnx_model = onnx_load(out_file)
-    run_onnx = OnnxRunner(onnx_model)
+    run_onnx = OnnxRunner(out_file)
    inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
    with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1):
      sched = run_onnx({"input":inp})["output"].schedule()
@@ -297,7 +299,7 @@ class TestDSPCache(unittest.TestCase):
    with Context(DEVECTORIZE=0, QUANTIZE=1):
      k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
      k.apply_opts(opts)
-      prg = k.to_program()
+      prg = get_program(k.get_optimized_ast(), k.opts)
      #print(prg.src)

    new_src = """
@@ -306,7 +308,7 @@ typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128))
 typedef unsigned char unsigned_char8 __attribute__((aligned(8),vector_size(8)));
 typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
 typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
-__attribute__((noinline)) void r_196_24_8_32_4(unsigned char* restrict __attribute__((align_value(128))) data0, unsigned char* restrict __attribute__((align_value(128))) data1, signed char* restrict __attribute__((align_value(
+__attribute__((noinline)) void r_196_32_4_24_8(unsigned char* restrict __attribute__((align_value(128))) data0, unsigned char* restrict __attribute__((align_value(128))) data1, signed char* restrict __attribute__((align_value(
 128))) data2, int* restrict __attribute__((align_value(128))) data3) {
  int32 cast0 = (int32){0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
  int32 val0 = *((int32*)((data3+0)));
@@ -251,6 +251,9 @@ class TestRandomness(unittest.TestCase):
    self.assertTrue(normal_test(Tensor.randn))
    self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x)))

+  def test_randn_device(self):
+    self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU")
+
  @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
  @unittest.skipIf(Device.DEFAULT in ["HSA", "AMD"], "bfloat16 local buffer broken in HSA")
  def test_randn_finite(self, default_float):
@@ -1,8 +1,9 @@
 import numpy as np, unittest, string
 from hypothesis import given, strategies as st
-from tinygrad import Device, Tensor, TinyJit
+from tinygrad import Device, Tensor, TinyJit, dtypes
 from tinygrad.runtime.ops_remote import RemoteDevice, parse_hosts
-from tinygrad.helpers import LazySeq, all_same
+from tinygrad.runtime.graph.remote import RemoteGraph
+from tinygrad.helpers import LazySeq, all_same, Context

 def multihost_env(devices):
  def same_hosts(devices): return all_same([h for h,_ in devices])
@@ -15,10 +16,11 @@ class TestRemoteMultiHost(unittest.TestCase):
    b = a.to('REMOTE:6').contiguous().realize()
    np.testing.assert_equal(b.numpy(), np.arange(0, 16))

-  # NOTE: remote graph currently throws GraphException on host mismatch, this just checks that it is being handled, not that jit graph is being used
-  def test_multihost_matmul_jit(self):
+  @Context(JIT_BATCH_SIZE=2**32)
+  def test_multihost_matmul_jit_graph(self):
    @TinyJit
    def do(a:Tensor, b:Tensor): return (a @ b).contiguous().realize()
+
    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
    for _ in range(3):
      na, nb = np.random.rand(128, 128).astype(np.float32), np.random.rand(128, 128).astype(np.float32)
@@ -27,6 +29,38 @@ class TestRemoteMultiHost(unittest.TestCase):
      c = do(a, b)
      np.testing.assert_allclose(nc, c.numpy(), rtol=3e-2, atol=1e-4) # tolerances from extra/gemm/simple_matmul.py

+    # Verify that everything is in one big cross-host graph
+    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
+
+  @Context(JIT_BATCH_SIZE=2**32)
+  def test_multihost_aware_schedule(self):
+    @TinyJit
+    def do(*ts:Tensor):
+      acc = Tensor.zeros(1, dtype=dtypes.float32)
+      for t in ts: acc += t.sum()
+      return acc.realize()
+
+    def do_np(*ts:np.ndarray):
+      acc = np.zeros(1, np.float32)
+      for t in ts: acc += t.sum()
+      return acc
+
+    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
+    TS = 64
+    for _ in range(3):
+      inp_np = [np.random.rand(256).astype(np.float32) for _ in range(TS)]
+      inp = [Tensor(inp).shard(ds, 0).contiguous().realize() for inp in inp_np]
+      out_np = do_np(*inp_np)
+      out = do(*inp)
+      np.testing.assert_allclose(out_np, out.numpy(), rtol=3e-2, atol=1e-4)
+
+    # Verify that everything is in one big cross-host graph and that the scheduling is reasonable
+    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
+    # At the time of writing this: 2050 graph breaks without multihost aware scheduling, 14 with it. I've set fail threshold to 28 to not fail on
+    # unrelated scheduling changes. Maybe 2x is a bit too pessimistic, but remote should perform just fine as long as this is not like a half hundred
+    # or more here.
+    self.assertLess(len(do.captured._jit_cache[0].prg.template), 28, "Very bad scheduling! Many unnecesary graph breaks!")
+
 class TestParseHosts(unittest.TestCase):
  def assert_seq(self, result:LazySeq, host:str):
    self.assertIsInstance(result, LazySeq)
@@ -5,7 +5,7 @@
 import unittest
 import numpy as np
 import functools
-from typing import List, Optional, Union, cast
+from typing import cast
 from hypothesis import assume, given, strategies as strat

 from tinygrad import nn, dtypes, Device, Tensor
@@ -15,16 +15,16 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.kernelize.kernelize import merge_views, get_kernelize_map, Kernel
+from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule

 class KernelCountException(Exception): pass
-def check_schedule(t:Union[Tensor, List[Tensor], UOp], allowed:int, to_prerealize:Optional[List[Tensor]]=None, filter_sink=True):
+def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
  if to_prerealize:
    with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
  if isinstance(t, Tensor): sched = t.schedule()
-  elif isinstance(t, List) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
+  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
  else:
    assert isinstance(t, UOp), f"can't schedule {t}"
    sink = UOp.sink(t) if t.op is not Ops.SINK else t
@@ -1727,7 +1727,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]])

 class TestIndexing(unittest.TestCase):
-  def check_schedule(self, xt:Union[Tensor,List[Tensor]], cnt:int):
+  def check_schedule(self, xt:Tensor|list[Tensor], cnt:int):
    with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):
      lst = [xt] if isinstance(xt, Tensor) else xt
      s = Tensor.schedule(*lst)
@@ -176,5 +176,12 @@ class TestWithGrad(unittest.TestCase):
    with self.assertRaises(NotImplementedError):
      z[:3] = x

+class TestSetitemLoop(unittest.TestCase):
+  def test_arange(self):
+    N = 10
+    cmp = Tensor.empty(N)
+    for i in range(N): cmp[i] = i
+    self.assertListEqual(Tensor.arange(N).tolist(), cmp.tolist())
+
 if __name__ == '__main__':
  unittest.main()
@@ -512,6 +512,24 @@ class TestTinygrad(unittest.TestCase):
    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
                    shell=True, check=True)

+    if Device.DEFAULT != "CPU":
+      # setting multiple devices fail
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+      # setting device via DEV
+      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
  def test_no_attributeerror_after_apply_uop_exception(self):
    try:
      Tensor.arange(4).reshape(3,2)
@@ -1,6 +1,7 @@
 from typing import List
 import unittest, pytest
 from tinygrad import dtypes, Variable
+from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
 from tinygrad.uop.symbolic import sym
@@ -241,6 +242,22 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(out.op, Ops.CONST)
    self.assertEqual(out.arg, 0)

+  def test_const_bitcast(self):
+    bf = UOp(Ops.CONST, dtypes.float, arg=1.0)
+    out = UOp(Ops.BITCAST, dtypes.uint32, (bf,))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 1)
+    out = uops[-1]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.arg, 0x3F800000)
+
+  @unittest.expectedFailure
+  def test_const_shape_change_bitcast(self):
+    bf = UOp(Ops.CONST, dtypes.uint8, arg=0x3F)
+    out = UOp(Ops.BITCAST, dtypes.half, (bf,))
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 1)
+
  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
@@ -423,6 +440,7 @@ class TestUOpGraph(unittest.TestCase):
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

+  @unittest.skip("outdated")
  def test_in_out_of_bounds_access_gated_store(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
@@ -433,11 +451,12 @@ class TestUOpGraph(unittest.TestCase):
      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
      with self.assertRaises(RuntimeError): to_uops_list([st1])

+  @unittest.skip("outdated")
  def test_in_bounds_access_gated_local(self):
    with Context(IGNORE_OOB=0):
      # Define buffers
      gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
-      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, local=True), (), "temp0")
+      sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")

      # Define indices, valids and barrier
      gidx = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 416))
@@ -512,7 +531,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_fold_gated_load_local(self):
    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, local=True), (), "temp")
+    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 16))
    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx), UOp.load(glbl0.index(lidx), dtype=dtypes.int)))
    barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
@@ -529,12 +548,12 @@ class TestUOpGraph(unittest.TestCase):
    idx0 = UOp.const(dtypes.int, 0)
    idx1 = UOp.const(dtypes.int, 0)
    val = UOp.const(dtypes.int, 42)
-    st0 = UOp(Ops.STORE, dtypes.void, (glbl.index(idx0, UOp.const(dtypes.bool, False)), val))
-    st1 = UOp(Ops.STORE, dtypes.void, (glbl.index(idx1, UOp.const(dtypes.bool, True)), val))
+    st0 = glbl.index(idx0, UOp.const(dtypes.bool, False)).store(val)
+    st1 = glbl.index(idx0, UOp.const(dtypes.bool, True)).store(val)
    uops = to_uops_list([st0, st1])
    # only the second store happens
    self.assertEqual(len(uops), 5)
-    self.assertEqual(uops[-1], UOp.store(glbl.index(idx1), val))
+    self.assertEqual(uops[-1], glbl.index(idx1).store(val))

  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
@@ -710,7 +729,7 @@ class TestExpander(unittest.TestCase):
 class TestIFUOps(unittest.TestCase):
  def test_create_ifs(self):
    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, local=True), (), "smem")
+    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, addrspace=AddrSpace.LOCAL), (), "smem")
    valid = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 10))<5
    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 4))
    gate = valid&(lidx.ne(2))
@@ -729,7 +748,7 @@ class TestIFUOps(unittest.TestCase):

  def test_expand_ifs_one_gate(self):
    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, local=True), (), "smem")
+    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, addrspace=AddrSpace.LOCAL), (), "smem")
    valid = UOp(Ops.SPECIAL, dtypes.int, (), ("gidx0", 4))<1
    lidx = UOp(Ops.SPECIAL, dtypes.int, (), ("lidx0", 16))
    gate = valid&(lidx.ne(2))
@@ -5,12 +5,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View # noqa F401
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.helpers import CI, DEBUG, getenv, Timing
-from tinygrad.dtype import dtypes, DType
+from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
 from tinygrad.uop.spec import spec
 from tinygrad.renderer import ProgramSpec
-from tinygrad.kernelize.kernelize import fix_kernel_ops
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
@@ -304,7 +303,7 @@ class TestLocalAccess(unittest.TestCase):
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
  def test_local_basic(self):
    uops = []
-    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, local=True), (), 'smem')
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
@@ -314,7 +313,7 @@ class TestLocalAccess(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local access with packed data type")
  def test_local_packed(self):
    uops = []
-    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, local=True), (), 'smem')
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
    sres = uop(uops, Ops.LOAD, dtypes.uint8, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
@@ -326,7 +325,7 @@ class TestLocalAccess(unittest.TestCase):
    _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
    size = 16
    for dtype in _dtypes:
-      temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, local=True), (), 'smem')
+      temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, addrspace=AddrSpace.LOCAL), (), 'smem')
      uops = to_uops_list([temp], opts=Device[Device.DEFAULT].renderer)
      out = Device[Device.DEFAULT].renderer.render(uops)
      # half is supported in wgsl, so it doesn't have to be packed
@@ -337,7 +336,7 @@ class TestLocalAccess(unittest.TestCase):
  @unittest.skip("tinygrad doesn't support this behavior")
  def test_local_indirect(self):
    uops = []
-    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.int32.ptr(size=16, local=True), (), 'smem')
+    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.int32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st1 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), uop(uops, Ops.CONST, dtypes.int32, (), 2)))
    st2 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 2)), uop(uops, Ops.CONST, dtypes.int32, (), 42)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st1,st2))
@@ -475,50 +474,9 @@ class TestUOpStr(unittest.TestCase):
    sum_uop = Tensor.empty(32, 32).sum().uop
    assert str(eval(str(sum_uop))) == str(sum_uop)

-@unittest.skip("uop no longer has order like this")
-class TestIndexingOrdering(unittest.TestCase):
-  # NOTE: these tests skip type_verify since they add dtype to STORE
-  @unittest.expectedFailure
-  def test_simple_order(self):
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    st0 = UOp(Ops.STORE, dtypes.float.vec(4), (buf, UOp.const(dtypes.int, 0), UOp.const(dtypes.float.vec(4), 42)))
-    st1 = UOp(Ops.STORE, dtypes.float, (buf, UOp.const(dtypes.int, 4), UOp.const(dtypes.float, 10)))
-    uops = to_uops_list([st1, st0], skip_check=True)
-    stores = [st for st in uops if st.op is Ops.STORE]
-    assert stores[0].src[1] < stores[1].src[1], f"stored at idx {stores[1].src[1].arg} AFTER {stores[0].src[1].arg}"
-
-  @unittest.expectedFailure
-  def test_ordering_multi_output(self):
-    buf0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
-    st0_0 = UOp(Ops.STORE, dtypes.float.vec(4), (buf0, UOp.const(dtypes.int, 0), UOp.const(dtypes.float.vec(4), 42)))
-    st1_0 = UOp(Ops.STORE, dtypes.float, (buf0, UOp.const(dtypes.int, 4), UOp.const(dtypes.float, 10)))
-    st0_1 = UOp(Ops.STORE, dtypes.float.vec(4), (buf1, UOp.const(dtypes.int, 0), UOp.const(dtypes.float.vec(4), 42)))
-    st1_1 = UOp(Ops.STORE, dtypes.float, (buf1, UOp.const(dtypes.int, 4), UOp.const(dtypes.float, 10)))
-    uops = to_uops_list([st0_0, st1_0, st0_1, st1_1], skip_check=True)
-    stores = [st for st in uops if st.op is Ops.STORE]
-    print("\n".join(map(str, stores)))
-    # buf0 stores come first
-    self.assertEqual(stores[0].src[0].arg, stores[1].src[0].arg)
-    # buf1 stores come next
-    self.assertEqual(stores[2].src[0].arg, stores[3].src[0].arg)
-    # both stores are aligned based on idx
-    assert stores[0].src[1] < stores[1].src[1], f"stored at idx {stores[1].src[1].arg} AFTER {stores[0].src[1].arg}"
-    assert stores[2].src[1] < stores[3].src[1], f"stored at idx {stores[1].src[1].arg} AFTER {stores[0].src[1].arg}"
-
-  def test_simple_order_with_special(self):
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (), ('gidx0', 4))
-    st0 = UOp(Ops.STORE, dtypes.float.vec(4), (buf, gidx0+UOp.const(dtypes.int, 0), UOp.const(dtypes.float.vec(4), 42)))
-    st1 = UOp(Ops.STORE, dtypes.float, (buf, UOp.const(dtypes.int, 4), UOp.const(dtypes.float, 10)))
-    uops = full_rewrite(UOp.sink(st1, st0))
-    stores = [st for st in uops if st.op is Ops.STORE]
-    assert stores[0].src[1] < stores[1].src[1], f"stored at idx {stores[1].src[1].arg} AFTER {stores[0].src[1].arg}"
-
 class TestUPatHelpers(unittest.TestCase):
  def test_location(self):
    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
-    self.assertEqual(fix_kernel_ops.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "kernelize.py")
    self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
    test_upat = UPat(Ops.CONST, dtypes.bool)
    self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import getenv, GlobalCounters
-from tinygrad.engine.realize import lower_schedule_item, ProgramSpec
+from tinygrad.engine.realize import lower_schedule_item, ProgramSpec, get_program
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
@@ -173,7 +173,8 @@ class TestStatsOptimized(unittest.TestCase):
    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats

  def test_gemm(self):
-    p = Kernel(self.ast_gemm).to_program()
+    k = Kernel(self.ast_gemm)
+    p = get_program(k.get_optimized_ast(), k.opts)
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)

@@ -181,7 +182,7 @@ class TestStatsOptimized(unittest.TestCase):
    k = Kernel(self.ast_gemm)
    if not k.apply_tensor_cores(): self.skipTest("no tensor cores")
    k.apply_opt(Opt(OptOps.UNROLL, 0, 2))
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    print(p.src)
    self.check_gemm(p)

@@ -190,7 +191,7 @@ class TestStatsOptimized(unittest.TestCase):
  def test_gemm_one_upcasted(self):
    k = Kernel(self.ast_gemm)
    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)

@@ -199,7 +200,7 @@ class TestStatsOptimized(unittest.TestCase):
    k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
    k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
    k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)

@@ -212,7 +213,7 @@ class TestStatsOptimized(unittest.TestCase):
      k.apply_opt(Opt(OptOps.LOCAL, 1, 5))
    except KernelOptError:
      raise unittest.SkipTest("no locals")
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)

@@ -223,14 +224,14 @@ class TestStatsOptimized(unittest.TestCase):
    except KernelOptError:
      raise unittest.SkipTest("no locals")
    SZ = N*N*4
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    # NOTE: these are sort of wrong. they aren't honoring the IF statement
    self.check_gemm(p, extra_flops=SZ*4)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)

  def test_reduce(self):
    k = Kernel(self.ast_reduce)
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
    self.assertEqual(p.estimates.ops, N*N)
    self.assertEqual(p.estimates.mem, N*N*4 + 4)
@@ -241,7 +242,7 @@ class TestStatsOptimized(unittest.TestCase):
      k.apply_opt(Opt(OptOps.GROUP, 0, 50))
    except KernelOptError:
      raise unittest.SkipTest("no locals")
-    p = k.to_program()
+    p = get_program(k.get_optimized_ast(), k.opts)
    # NOTE: these are wrong, they don't respect the if statement
    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)

@@ -35,7 +35,7 @@ class TestBenchLog(unittest.TestCase):
      self.assertGreater(_events[event]["wall"][0], 0)
      self.assertGreater(_events[event]["wall"][1], 0)

-  @skipIf(CI and Device.DEFAULT == "CUDA", "ci cuda timing is not accurate")
+  @skipIf(CI, "ci timing is not accurate")
  def test_log_single_kernel_time(self):
    wall_times = []

@@ -1,4 +1,4 @@
-import pathlib, tempfile, unittest
+import os, pathlib, tempfile, unittest
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DType
@@ -410,5 +410,13 @@ class TestPathTensor(unittest.TestCase):
    self.assertEqual(t_cpu.device, "CPU")
    np.testing.assert_array_equal(t_cpu.numpy(), np.frombuffer(self.test_data, dtype=np.uint8))

+  def test_path_tensor_disk_device_bug(self):
+    test_file = pathlib.Path(self.temp_dir.name) / "disk_device_bug"
+    with open(test_file, "wb") as f: f.write(bytes(range(10)))
+    os.chmod(test_file, 0o000)
+    with self.assertRaises(PermissionError):
+      Tensor(pathlib.Path(test_file)).tolist()
+    os.chmod(test_file, 0o644)
+    assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
 if __name__ == "__main__":
  unittest.main()
@@ -1,8 +1,9 @@
 import unittest, math
 from tinygrad import dtypes
 from tinygrad.helpers import all_same
-from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
 from tinygrad.codegen import full_rewrite_to_sink
+from hypothesis import given, strategies as strat

 # Helper function to apply the graph rewrite
 def apply_rewrite(expr):
@@ -284,13 +285,17 @@ class TestSubstitute(unittest.TestCase):
    # the srcs are rewritten but we keep tag
    self.assertIs(ret, (b+4).replace(tag=1))

+matchers = strat.sampled_from([PatternMatcher, TrackedPatternMatcher])
+
 class TestRecurse(unittest.TestCase):
-  def test_no_inf_loop(self):
+  @given(matchers)
+  def test_no_inf_loop(self, PatternMatcher):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
    graph_rewrite(a, pm)

-  def test_no_inf_loop_bottom_up(self):
+  @given(matchers)
+  def test_no_inf_loop_bottom_up(self, PatternMatcher):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
    graph_rewrite(a, pm, bottom_up=True)
@@ -1,8 +1,8 @@
-import ctypes, gzip, unittest
+import ctypes, gzip, unittest, timeit
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
-from tinygrad.tensor import get_shape
+from tinygrad.tensor import Tensor, get_shape
 from tinygrad.shape.view import get_contraction, get_contraction_with_reduce
 import numpy as np

@@ -186,6 +186,38 @@ class TestMemoryview(unittest.TestCase):
    mv[0] = 2
    assert base[0] == 2

+  @unittest.skipIf(CI, "dangerous for CI, it allocates tons of memory")
+  def test_to_mv(self):
+    sizes = [
+      (16, "16 B"),
+      (64, "64 B"),
+      (256, "256 B"),
+      (1024, "1 KB"),
+      (4 * 1024, "4 KB"),
+      (16 * 1024, "16 KB"),
+      (64 * 1024, "64 KB"),
+      (256 * 1024, "256 KB"),
+      (1 * 1024 * 1024, "1 MB"),
+      (10 * 1024 * 1024, "10 MB"),
+      (200 * 1024 * 1024, "200 MB"),
+    ]
+
+    for sz, label in sizes:
+      buf = np.random.randint(0, 256, sz, dtype=np.uint8)
+      ptr = buf.ctypes.data
+
+      iters = 100_000
+      t_us = timeit.timeit(lambda: to_mv(ptr, sz), number=iters) * 1e6 / iters
+      print(f"Size {label:>9} | Time: {t_us:8.3f} µs")
+
+  def test_speed_from_mv_vs_mv_address(self):
+    x = memoryview(bytearray(1))
+
+    iters = 100000
+    fmv_us = timeit.timeit(lambda: from_mv(x), number=iters) * 1e6 / iters
+    mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
+    print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")
+
 class TestGetContraction(unittest.TestCase):
  def test_contraction_with_reduce(self):
    r = get_contraction((16, 1, 1, 1), (16, 1, 1))
@@ -376,5 +408,32 @@ class TestWordWrap(unittest.TestCase):
    st2 = word_wrap(st, wrap=wrap)
    self.assertEqual(len(st2.splitlines()), 2)

+  def test_wrap_explicit_newline(self):
+    wrap = 10
+    st = "\n".join(["x"*wrap, "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines()))
+
+    st = "\n".join(["x"*(wrap+1), "x"*wrap, "x"*wrap])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+1)
+
+    st = "\n".join(["x"*(wrap+1), "x"*(wrap+1), "x"*(wrap+1)])
+    st2 = word_wrap(st, wrap=wrap)
+    self.assertEqual(len(st2.splitlines()), len(st.splitlines())+3)
+
+class TestIsNumpyNdarray(unittest.TestCase):
+  def test_ndarray(self):
+    self.assertTrue(is_numpy_ndarray(np.array([1, 2, 3])))
+  def test_ndarray_tolist(self):
+    self.assertFalse(is_numpy_ndarray(np.array([1, 2, 3]).tolist()))
+  def test_list(self):
+    self.assertFalse(is_numpy_ndarray([1, 2, 3]))
+  def test_tensor(self):
+    self.assertFalse(is_numpy_ndarray(Tensor([1, 2, 3])))
+    self.assertFalse(is_numpy_ndarray(Tensor(np.array([1, 2, 3]))))
+  def test_tensor_numpy(self):
+    self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
+
 if __name__ == '__main__':
  unittest.main()
@@ -38,5 +38,14 @@ class TestKeccak(unittest.TestCase):
    out = Tensor(b"abc").keccak()
    self.assertEqual(bytes(out.tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))

+  def test_long(self):
+    data = b"\x00" * 4
+    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
+
+    data = b"\x00" * 4096
+    with self.assertRaises(RecursionError):
+      # TODO: fix
+      self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
+
 if __name__ == "__main__":
  unittest.main()
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Context, Device
 from tinygrad.engine.realize import get_program
-from tinygrad.renderer import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import KernelInfo

 class TestLinearizerRewrite(unittest.TestCase):
@@ -28,5 +28,19 @@ class TestLinearizerRewrite(unittest.TestCase):
      prg = get_program(ast, Device["CPU"].renderer)
      print(prg.src)

+  def test_kernel_info(self):
+    out = Tensor.arange(4, device="NULL")
+    si = out.schedule()[-1]
+
+    ast = si.ast.replace(arg=KernelInfo(opts_to_apply=()))
+    prg = get_program(ast, Device["CPU"].renderer)
+    assert prg.applied_opts == (), f"expected no opts, got {prg}"
+
+    prg = get_program(ast.replace(arg=None), Device["CPU"].renderer)
+    assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"
+
+    prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
+    self.assertEqual(prg.name, "custom")
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
-from tinygrad.kernelize.kernelize import sym, merge_views
+from tinygrad.schedule.kernelize import sym, merge_views

 class TestRewriteTrackedChildren(unittest.TestCase):
  @unittest.skip("track_children no longer supported")
@@ -15,8 +15,8 @@ class TestSearchUtil(unittest.TestCase):
  def test_bufs_from_lin(self):
    a = Tensor([1,2,3,4]).realize()
    si = (a+1).schedule()[0]
-    rawbufs = bufs_from_lin(lin:=Kernel(si.ast))
-    assert len(rawbufs) == len(lin.membufs) == 2
+    rawbufs = bufs_from_lin(Kernel(si.ast))
+    assert len(rawbufs) == 2
    assert all(r is not None for r in rawbufs)
    assert all(isinstance(r, Buffer) for r in rawbufs)
    assert all(r.size > 0 for r in rawbufs)
@@ -25,8 +25,8 @@ class TestSearchUtil(unittest.TestCase):
    a = Tensor.randn(4, 4).realize()
    b = a+a[0]
    si = b.schedule()[0]
-    rawbufs = bufs_from_lin(k:=Kernel(si.ast))
-    assert len(rawbufs) == len(k.membufs) == 2
+    rawbufs = bufs_from_lin(Kernel(si.ast))
+    assert len(rawbufs) == 2
    assert all(r is not None for r in rawbufs)
    assert all(isinstance(r, Buffer) for r in rawbufs)
    assert all(r.size > 0 for r in rawbufs)
@@ -39,10 +39,10 @@ class TestUOpSpec(unittest.TestCase):
    bufs = [UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), i) for i in range(6)]
    a = UOp(Ops.LOAD, dtype, (bufs[2].view(ShapeTracker.from_shape((32, 1))),))
    b = UOp(Ops.LOAD, dtype, (bufs[3].view(ShapeTracker.from_shape((32, 1))),))
-    st0 = UOp.store(bufs[0], ShapeTracker.from_shape((32, 1)).to_uop(), a+b)
+    st0 = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), a+b)
    a = UOp(Ops.LOAD, dtype, (bufs[4].view(ShapeTracker.from_shape((32, 32))),))
    b = UOp(Ops.LOAD, dtype, (bufs[5].view(ShapeTracker.from_shape((32, 32))),))
-    st1 = UOp.store(bufs[1], ShapeTracker.from_shape((32, 32)).to_uop(), a+b)
+    st1 = UOp.store(bufs[1].view(ShapeTracker.from_shape((32, 32))), a+b)
    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st0, st1)

  def test_no_implicit_broadcasting(self):
@@ -63,14 +63,14 @@ class TestUOpSpec(unittest.TestCase):
    bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
    a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),))
    r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
-    st = UOp.store(bufs[0], ShapeTracker.from_shape((32, 1)).to_uop(), r)
+    st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r)
    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)

  def test_reduce_add_store(self):
    bufs = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), i) for i in range(2)]
    a = UOp(Ops.LOAD, dtypes.float, (bufs[1].view(ShapeTracker.from_shape((32, 1))),))
    r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
-    st = UOp.store(bufs[0], ShapeTracker.from_shape((32, 1)).to_uop(), r+a)
+    st = UOp.store(bufs[0].view(ShapeTracker.from_shape((32, 1))), r+a)
    with self.assertRaises(InvalidASTException): helper_test_verify_ast(st)

  def test_buffer_uops_st(self):
@@ -85,7 +85,7 @@ class TestUOpSpec(unittest.TestCase):
    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    a = UOp(Ops.LOAD, dtypes.float, (buf.view(ShapeTracker.from_shape((32, 1))),))
    r = UOp(Ops.REDUCE_AXIS, dtypes.float, (a,), (Ops.ADD, (0,)))
-    st = UOp.store(buf, ShapeTracker.from_shape((32, 1)).to_uop(), r.view(r.st.expand((32, 1)))+a)
+    st = UOp.store(buf.view(ShapeTracker.from_shape((32, 1))), r.view(r.st.expand((32, 1)))+a)
    with self.assertRaisesRegex(InvalidASTException, "UOp verification failed"): helper_test_verify_ast(st)

  def test_const_view_always_valid(self):
@@ -14,7 +14,7 @@ def render(self) -> tuple[str, ConstType, ConstType]:
  # NOTE: we need STORE so the ALU op has children
  glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
  uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), self)).sink())
-  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]
+  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
  return rewritten_uop.render(simplify=False), rewritten_uop.vmin, rewritten_uop.vmax

 def uconst(val): return UOp.const(dtypes.int, val)
@@ -169,6 +169,12 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
    self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")

+  def test_div_mod_zero(self):
+    with self.assertRaises(ZeroDivisionError):
+      (Variable("a", 0, 7) // 0).simplify()
+    with self.assertRaises(ZeroDivisionError):
+      (Variable("a", 0, 7) % 0).simplify()
+
  def test_sum_div_remove(self):
    self.helper_test_variable(usum([Variable("a", 0, 7), Variable("b", 0, 3)]) // 20, 0, 0, "0")

@@ -636,7 +642,7 @@ class TestSymbolic(unittest.TestCase):
    # TODO: copied from render, render does not support cast
    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
    uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]
+    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]

    self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))

@@ -110,6 +110,18 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(x.vmin, 0)
    self.assertEqual(x.vmax, 10 >> 2)

+  def test_vmin_vmax_cast(self):
+    x = UOp.variable('x', -10, 10, dtypes.int)
+    x_float = x.cast(dtypes.float)
+    self.assertEqual(x_float.vmin, -10)
+    self.assertEqual(x_float.vmax, 10)
+    x_bool = x.cast(dtypes.bool)
+    self.assertEqual(x_bool.vmin, False)
+    self.assertEqual(x_bool.vmax, True)
+    x_uint = x.cast(dtypes.uint)
+    self.assertEqual(x_uint.vmin, dtypes.min(dtypes.uint))
+    self.assertEqual(x_uint.vmax, dtypes.max(dtypes.uint))
+
 class TestVminVmaxDivMod(unittest.TestCase):
  def test_vmin_vmax_division_positive(self):
    # vmin and vmax for division of a variable by a positive constant
@@ -185,6 +185,17 @@ class TestMergeViews(unittest.TestCase):
    self.assertIsNotNone(v)
    self.assertEqual(v, target)

+  def test_merge_views_variable(self):
+    from tinygrad import Variable
+    N = 100
+    start_pos = Variable("start_pos", 1, N-1)
+    v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
+    v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
+    target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
  def test_view_padded_area1(self):
    # test_multinomial
    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
--- a/Show More
+++ b/Show More