fine for ptx

that hack broke things
Revert "fix tests"
2025-08-05 18:26:06 -07:00 · 2025-08-05 18:24:58 -07:00 · 2025-08-05 18:06:27 -07:00 · 2025-08-05 17:50:14 -07:00 · 2025-08-05 17:45:35 -07:00 · 2025-08-05 17:07:12 -07:00
203 changed files with 14441 additions and 3393 deletions
@@ -112,7 +112,16 @@ runs:
        fi

    # ******************* apt *******************
+    - name: Setup apt
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      shell: bash
+      run: |
+        sudo chown -R $USER:$USER /var/cache/apt/archives

+        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
+        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
+    
    - name: Add OpenCL Repo
      if: inputs.opencl == 'true' && runner.os == 'Linux'
      shell: bash
@@ -135,14 +144,11 @@ runs:
        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

-    - name: apt-get update + install
+    - name: Compute Package List + Hash
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      id: apt-pkgs
      shell: bash
      run: |
-        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        sudo apt -qq update || true
-
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
@@ -153,7 +159,7 @@ runs:
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
        fi
        # **** CUDA ****
        if [[ "${{ inputs.cuda }}" == "true" ]]; then
@@ -168,14 +174,31 @@ runs:
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
+        
+        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
+        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

+    - name: Cache apt
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      uses: actions/cache@v4
+      with:
+        path: /var/cache/apt/archives/
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
+
+    - name: Run apt Update + Install
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      shell: bash
+      run: |
+        sudo apt -qq update || true
+        
        # ******** do install ********
-        if [[ -n "$pkgs" ]]; then
-          sudo apt-get -y --allow-unauthenticated --no-install-recommends install $pkgs
+        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
+          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi
-
+        
+        sudo chown -R $USER:$USER /var/cache/apt/archives/
+    
    # **** AMD ****
-
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
@@ -228,7 +251,7 @@ runs:
      shell: bash
      run: |
        cd ${{ github.workspace }}/gpuocelot/ocelot/build
-        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || ''}}lib/
+        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || '' }}lib/

    # **** WebGPU ****

@@ -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)
@@ -617,6 +617,10 @@ jobs:
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
    - name: openpilot dmonitoring compile3 0.9.7
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 Space Lab policy + vision
+      run: |
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -637,3 +641,127 @@ jobs:
          openpilot_0_9_7.txt
          openpilot_image_0_9_4.txt
          openpilot_image_0_9_7.txt
+
+  testreddriverbenchmark:
+    name: AM Benchmark
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 15
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: Test driver cold start time
+      run: time DEBUG=3 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test driver warm start time
+      run: time DEBUG=3 AMD=1 python3 test/test_tiny.py TestTiny.test_plus
+    # Fails on 9070
+    # - name: Test tensor cores
+    #   run: |
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (AMD)
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
+    - name: Test AMD=1
+      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
+    - name: Test DISK copy time
+      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf Bert training steps (1 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (AM Driver)
+        path: |
+          am_matmul_amd.txt
+          am_train_cifar_one_gpu.txt
+          am_train_resnet_one_gpu.txt
+          am_train_bert_one_gpu.txt
+    - name: Run process replay tests
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+
+  testgreendriverbenchmark:
+    name: NV Benchmark
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 15
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove nv modules
+      run: ./extra/hcq/hcq_smi.py nv rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py nv kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: Test driver start time
+      run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test tensor cores
+      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test DISK copy time
+      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+    - name: Test LLAMA-3
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf Bert training steps (1 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (NV Driver)
+        path: |
+          nv_llama3_beam.txt
+          nv_train_cifar_one_gpu.txt
+          nv_train_resnet_one_gpu.txt
+          nv_train_bert_one_gpu.txt
+    - name: Run process replay tests
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -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: 360

    steps:
    - name: Checkout Code
@@ -26,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -132,10 +132,13 @@ jobs:
      run: |
        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
        ./autogen_stubs.sh libc
        ./autogen_stubs.sh io_uring
+        ./autogen_stubs.sh ib
        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
+        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
    - name: Verify WebGPU autogen
      run: |
        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
@@ -239,8 +242,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 +255,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 +329,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 +376,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 +430,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 +471,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 +508,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 +535,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 +548,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 +601,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 +689,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 +844,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,38 +867,39 @@ 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)
-    runs-on: macos-15
-    timeout-minutes: 10
-    env:
-      REMOTE: 1
-      REMOTEDEV: METAL
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: macos-remote
-          deps: testing_minimal
-      - name: Check Device.DEFAULT and print some source
-        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 == 'METAL', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test
-        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_tensor_variable.py
+  #osxremote:
+  #  name: MacOS (remote metal)
+  #  runs-on: macos-15
+  #  timeout-minutes: 10
+  #  env:
+  #    REMOTE: 1
+  #    REMOTEDEV: METAL
+  #  steps:
+  #    - name: Checkout Code
+  #      uses: actions/checkout@v4
+  #    - name: Setup Environment
+  #      uses: ./.github/actions/setup-tinygrad
+  #      with:
+  #        key: macos-remote
+  #        deps: testing_minimal
+  #    - name: Check Device.DEFAULT and print some source
+  #      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 == 'METAL', Device.default.properties.real_device"
+  #        DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+  #    - name: Run REMOTE=1 Test
+  #      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_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 +907,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 +999,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu]
+        backend: [llvm, cpu, webgpu]

    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
@@ -951,11 +1012,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() {
@@ -237,6 +240,21 @@ generate_io_uring() {
  fixup $BASE/io_uring.py
 }

+generate_ib() {
+  clang2py -k cdefstum \
+    /usr/include/infiniband/verbs.h \
+    /usr/include/infiniband/verbs_api.h \
+    /usr/include/infiniband/ib_user_ioctl_verbs.h \
+    /usr/include/rdma/ib_user_verbs.h \
+    -o $BASE/ib.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
+  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
+
+  fixup $BASE/ib.py
+}
+
 generate_libc() {
  clang2py -k cdefstum \
    $(dpkg -L libc6-dev | grep sys/mman.h) \
@@ -441,7 +459,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"
@@ -462,6 +480,7 @@ elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
+elif [ "$1" == "ib" ]; then generate_ib
 elif [ "$1" == "libc" ]; then generate_libc
 elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl
@@ -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
@@ -126,7 +126,7 @@ print(t_log_grad.uop)
 """
 void E_(float* restrict data0, float* restrict data1) {
  float val0 = *(data1+0);
-  *(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
+  *(data0+0) = (1/val0);
 }
 """
 # the derivative is close to 1/3
@@ -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)
@@ -1,4 +1,6 @@
-import os, random, pickle, queue
+import functools
+import hashlib
+import os, random, pickle, queue, struct, math
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -6,6 +8,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
 import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.nn.state import TensorIO

 ### ResNet

@@ -510,6 +513,202 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

+# llama3
+
+class BinIdxDataset:
+  def __init__(self, base_path:Path):
+    self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
+    self.idx = TensorIO(self.idx_t)
+
+    # parse idx file
+    magic = self.idx.read(9)
+    assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
+    version, = struct.unpack("<Q", self.idx.read(8))
+    assert version == 1, "unsupported index version"
+    dtype_code, = struct.unpack("<B", self.idx.read(1))
+    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
+    self.count, = struct.unpack("<Q", self.idx.read(8))
+    doc_count, = struct.unpack("<Q", self.idx.read(8))
+
+    start = self.idx.tell()
+    end = start + self.count * dtypes.int32.itemsize
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
+
+    start = end
+    end = start + self.count * dtypes.int64.itemsize
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
+
+    start = end
+    end = start + doc_count * dtypes.int64.itemsize
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
+
+    # bin file
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
+
+  def _index(self, idx) -> tuple[int, int]:
+    return self.pointers[idx].item(), self.sizes[idx].item()
+
+  def get(self, idx, offset:int=0, length:int|None=None):
+    ptr, size = self._index(idx)
+    if length is None: length = size - offset
+    ptr += offset * self.dtype.itemsize
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
+
+# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
+class GPTDataset:
+  def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.samples, self.seqlen = samples, seqlen
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)
+
+    self.indexed_dataset = BinIdxDataset(base_path)
+
+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
+    if cache_path.exists():
+      with open(cache_path, "rb") as f:
+        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
+    else:
+      self.doc_idx = self._build_doc_idx()
+      self.sample_idx = self._build_sample_idx()
+      self.shuffle_idx = self._build_shuffle_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
+
+  def __getitem__(self, idx):
+    if idx is None:
+      text = self._get(0)
+    else:
+      text = self._get(idx)
+
+    return text
+
+  def _get(self, idx):
+    idx = self.shuffle_idx[idx]
+
+    doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
+    doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
+
+    doc_ids, sample_parts = [], []
+
+    if doc_idx_beg == doc_idx_end:
+      doc_ids.append(self.doc_idx[doc_idx_beg])
+
+      sample_parts.append(
+          self.indexed_dataset.get(
+            int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
+    else:
+      for i in range(doc_idx_beg, doc_idx_end + 1):
+        doc_ids.append(self.doc_idx[i])
+
+        offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
+        length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
+        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
+
+    # concat all parts
+    text = Tensor.cat(*sample_parts)
+
+    return text
+
+  @functools.cached_property
+  def tokens_per_epoch(self) -> int:
+    return sum(self.indexed_dataset.sizes.tolist())
+
+  @functools.cached_property
+  def num_epochs(self) -> int:
+    # we need enough epochs to cover the requested amount of tokens
+    num_epochs = 1
+    num_tokens = self.tokens_per_epoch
+    while num_tokens < self.samples * self.seqlen:
+      num_epochs += 1
+      num_tokens += self.tokens_per_epoch
+    return num_epochs
+
+  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
+  def _build_doc_idx(self):
+    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
+    doc_idx = doc_idx.reshape(-1)
+    doc_idx = doc_idx.astype(np.int32)
+    if self.shuffle: self.rng.shuffle(doc_idx)
+    return doc_idx
+
+  def _build_sample_idx(self):
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
+
+    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
+    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+    sample_idx_idx += 1
+
+    for _ in tqdm(range(1, self.samples + 1)):
+      remaining_seqlen = self.seqlen + 1
+      while remaining_seqlen > 0:
+        doc_idx = int(self.doc_idx[doc_idx_idx])
+        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
+        remaining_seqlen -= doc_len
+        if remaining_seqlen <= 0:
+          doc_offset += remaining_seqlen + doc_len - 1
+          remaining_seqlen = 0
+        else:
+          if doc_idx_idx == len(self.doc_idx) - 1:
+            assert sample_idx_idx == self.samples
+            doc_idx = int(self.doc_idx[doc_idx_idx])
+            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
+            break
+          doc_idx_idx += 1
+          doc_offset = 0
+
+      sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+      sample_idx_idx += 1
+
+    return sample_idx
+
+  def _build_shuffle_idx(self):
+    shuffle_idx = np.arange(self.samples, dtype=np.int32)
+    if self.shuffle: self.rng.shuffle(shuffle_idx)
+    return shuffle_idx
+
+class BlendedGPTDataset:
+  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.seed = seed
+
+    # normalize weights
+    total_weight = sum(weights)
+    self.weights = [w / total_weight for w in weights]
+
+    self.samples = samples
+    surplus = 0.005
+    samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
+
+    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
+
+  def get(self, idx:int):
+    tokens = self.datasets[0][idx]
+    return tokens
+
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+      base_dir / "c4-train.en_7_text_document",
+    ], [
+      1.0, 1.0
+    ], samples, seqlen, seed, True)
+
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
+    yield Tensor.stack(batch, dim=0)
+
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -538,6 +737,18 @@ if __name__ == "__main__":
      for x in batch_load_retinanet(dataset, val, base_dir):
        pbar.update(x[0].shape[0])

+  def load_llama3(val):
+    bs = 24
+    samples = 5760 if val else 1_200_000
+    seqlen = 512
+
+    max_, min_ = 0, math.inf
+    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
+      max_ = max(max_, tokens.shape[1])
+      min_ = min(min_, tokens.shape[1])
+    print(f"max seq length: {max_}")
+    print(f"min seq length: {min_}")
+
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -1,4 +1,4 @@
-import time
+import time, math
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
@@ -241,6 +241,34 @@ def eval_mrcnn():
  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
  evaluate_predictions_on_coco(mask_output, iou_type='segm')

+def eval_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS
+  from tinygrad.helpers import tqdm
+
+  bs = 4
+  sequence_length = 512
+
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+
+  @TinyJit
+  def eval_step(model, tokens):
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten()
+
+  from examples.mlperf.dataloader import batch_load_llama3
+  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
+
+  losses = []
+  for tokens in tqdm(iter, total=5760//bs):
+    GlobalCounters.reset()
+    losses += eval_step(model, tokens).tolist()
+    tqdm.write(f"loss: {np.mean(losses)}")
+
+  log_perplexity = Tensor(losses).mean()
+  print(f"Log Perplexity: {log_perplexity.item()}")
+
 if __name__ == "__main__":
  # inference only
  Tensor.training = False
@@ -1290,9 +1290,16 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 4)
+  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+  SEED               = config["SEED"]                   = getenv("SEED", 5760)
+  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
+  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
+  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000)
+
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # trains to 7

  opt_adamw_beta_1 = 0.9
  opt_adamw_beta_2 = 0.95
@@ -1300,7 +1307,6 @@ def train_llama3():
  opt_adamw_weight_decay = 0.1

  opt_gradient_clip_norm = 1.0
-  sequence_length = 8192
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
@@ -1308,7 +1314,33 @@ def train_llama3():

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
+  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)
+
+  # TODO: MP
+  # if (GPUS := getenv("GPUS", 1)) > 1:
+  #   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
+  #   for k,v in get_state_dict(model).items():
+  #     if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+  #     # elif '.attention.wq' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wk' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wv' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wo' in k: v.shard_(device, axis=1)
+  #     # elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+  #     # elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+  #     # elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+  #     # elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+  #     elif 'output.weight' in k: v.shard_(device, axis=0)  # 243.32
+  #     else:
+  #       # print(k)
+  #       # attention_norm, ffn_norm, norm
+  #       v.shard_(device, axis=None)

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1316,12 +1348,17 @@ def train_llama3():

  @TinyJit
  @Tensor.train()
-  def train_step(model, x, y):
+  def train_step(model, tokens:Tensor, grad_acc:int):
    optim.zero_grad()
-    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
-    loss = logits.cross_entropy(y)
-    loss.backward()
-
+    # grad acc
+    for batch in tokens.split(tokens.shape[0]//grad_acc):
+      if (DP := getenv("DP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+        batch = batch.shard(device, 0)
+      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
+      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
+      loss.backward()
+      Tensor.realize(*[p.grad for p in optim.params])
    # L2 norm grad clip
    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1340,19 +1377,32 @@ def train_llama3():
    loss.realize(lr)
    return loss, lr

-  # overfitting this example should give cross_entropy log(BS)
-  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
-  fake_label = Tensor(list(range(BS)), dtype="int16")
+  if getenv("FAKEDATA", 0):
+    def fake_data():
+      for _ in range(SAMPLES // GBS):
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+    iter = fake_data()
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))

-  for _ in range(100):
+  i = 0
+  for tokens in tqdm(iter, total=SAMPLES//GBS):
+    t = time.perf_counter()
    GlobalCounters.reset()
-    loss, lr = train_step(model, fake_input, fake_label)
-    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
-    # uses 43% ~= 83GB
-    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
-    # TODO: this OOM
-    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
-    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
+    loss, lr = train_step(model, tokens, grad_acc)
+    # above as tqdm.write f-string
+    tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
+    if (fname:=getenv("LOSS_FILE", "")):
+      with open(fname, "a") as f:
+        f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+
+    if getenv("CKPT") and (i % 200 == 0 or i == 10):
+      tqdm.write("saving checkpoint")
+      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+      fn = f"{ckpt_dir}/{i}.safe"
+      safe_save(get_state_dict(model), fn)
+    i += 1

 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')
@@ -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,42 @@
 # 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])
+  elif getenv("ASM") == -2:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
+    prgfast = replace(prg, name="kernel4_gmem_db", 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,143 @@
+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 __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+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 256
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel4_gmem_db(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};
+
+  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)];
+        }
+      }
+    }
+  }
+}
@@ -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
+  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,335 @@
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.dtype import AddrSpace
+from tinygrad.schedule.kernelize import merge_views, view_left
+from tinygrad.helpers import getenv, colored, prod, unwrap
+from tinygrad.shape.shapetracker import ShapeTracker, View
+from tinygrad.shape.view import strides_for_shape
+from tinygrad.opt.kernel import axis_colors
+
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
+
+N = 4096
+run_count = 5
+
+BN = 128
+BM = 128
+BK = 8
+
+TN = 4
+TM = 4
+
+# NOTE: this is from testgrad
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
+# src->r->view  -->   src->view->r
+def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
+  if r.tag is not None: return None
+  # confirm the input is in order
+  # TODO: replace this with a UOp that allows for nothing else then remove this
+  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
+  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"
+
+  # append the reduce shape to each of the views
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
+
+  # no reshape required with shrinking REDUCE_AXIS
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
+             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))
+
+pm = PatternMatcher([
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])
+
+def top_spec_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  sink = c.schedule()[-1].ast
+  L = 16
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
+  sink = graph_rewrite(sink, view_left+pm)
+  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
+  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
+
+def hl_spec_kernel3():
+  nbIterWaveM = 2
+  nbIterWaveN = 2
+
+  # define buffers
+  # TODO: remove these views once the defines have a shape
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
+
+  # shape buffers. TODO: permutes
+  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
+  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
+  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
+  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
+  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
+  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
+  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
+  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)
+
+  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
+  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
+  assert len(expanded_shape) == 20
+  permute_a = list(range(len(expanded_shape)))
+  permute_b = permute_a[:]
+
+  # this makes all the global loads match
+  # this can also be more simply done by rebinding the RANGEs
+  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
+  permute_a[17:20] = [11,12,13]
+  permute_a[11:14] = [17,18,19]
+  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
+  permute_a[2:7] = [3,4,5,6,2]
+
+  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
+  permute_b[17:20] = [5,6,7]
+
+  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+
+  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
+  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
+  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))
+
+  axis_types = (
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.REDUCE, AxisType.REDUCE)
+
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
+  sink = graph_rewrite(sink, merge_views)
+  return sink
+
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256
+
+  nbWaves = BLOCK_SIZE // 32
+  WN = 128 if kernel5 else 64
+  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=1)
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
+
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
+
+  BM_As_stride = (BM+4) if kernel5 else BM
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
+  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), arg=2)
+
+  i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
+  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
+
+  if kernel4:
+    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
+    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
+
+    # initial load from globals into locals (0)
+    kId = 0
+
+    # load from globals into locals
+    i = UOp.range(dtypes.int, nbReadsB, 0)
+    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, 1)
+    index_x = rAIdx + kId
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
+
+    # iterate over the middle chunk
+    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
+    kId = kId_range*BK
+
+    barrier = UOp.barrier(As_store, Bs_store)
+
+    # load from globals into registers (next round)
+    i = UOp.range(dtypes.int, nbReadsB, 3)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
+
+    i = UOp.range(dtypes.int, nbReadsA, 4)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
+
+    def inner_loop(first_range, inp_dep=()):
+      # inner unroll
+      k = UOp.range(dtypes.int, BK, first_range+0)
+
+      # load from locals into registers
+      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
+      i = UOp.range(dtypes.int, TN, first_range+2)
+      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
+      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
+
+      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
+      i = UOp.range(dtypes.int, TM, first_range+4)
+      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
+      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
+
+      # do the GEMM math
+      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
+      yt = UOp.range(dtypes.int, TM, first_range+6)
+      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
+      xt = UOp.range(dtypes.int, TN, first_range+8)
+      x = iterWaveN * TN + xt
+      y = iterWaveM * TM + yt
+      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+      # sketchy, this should end the kId_range but it doesn't
+      sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                              iterWaveM, iterWaveN, yt, xt, k)
+      return sink
+
+    # TODO: kId_range should endrange after a barrier
+    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
+
+    # load from registers into locals
+    i = UOp.range(dtypes.int, nbReadsB, 14)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
+
+    i = UOp.range(dtypes.int, nbReadsA, 15)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
+
+    # final iteration without the copy
+    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
+  else:
+    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_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
+
+    barrier = UOp.barrier(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_As_stride + index].load(barrier), iterWave, i)
+
+    # do the GEMM math
+    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
+    yt = UOp.range(dtypes.int, TM, 9)
+    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
+    xt = UOp.range(dtypes.int, TN, 12)
+    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(init_store) + 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, 1000)
+  yt = UOp.range(dtypes.int, TM, 1001)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
+  xt = UOp.range(dtypes.int, TN, 1003)
+  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
+  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
+  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__":
+  HL = getenv("HL")
+  if HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
+  prg = get_program(hprg, Device.default.renderer)
+  print(prg.src)
+  if getenv("SRC"): exit(0)
+  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()
+  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
+  ei = ExecItem(hrunner, buffers)
+  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)
@@ -0,0 +1,122 @@
+#!/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, args) -> list[str]:
+  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  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}") and dev_id.startswith(target_dev): 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):
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  to_unload = [m for m in modules if _is_module_loaded(m)]
+  if not to_unload: print("Kernel modules are not loaded")
+  else:
+    print("Removing 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)
+
+  module = "nvidia" if args.backend == "nv" else "amdgpu"
+  if _is_module_loaded(module):
+    print(f"{module} kernel module already loaded")
+    return
+
+  print(f"Inserting kernel module: {module}")
+  if args.backend == "nv":
+    subprocess.run(["nvidia-smi"], check=True)
+  elif args.backend == "amd":
+    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    print(f"Resetting device {dev}")
+    if args.backend != "amd": _do_reset_device(dev)
+    time.sleep(0.2)
+
+def cmd_show_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    for i in range(128):
+      if i > 0: time.sleep(0.2)
+
+      try:
+        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+        except subprocess.CalledProcessError: break
+
+        print(f"Killing process {pid} (which uses {dev})")
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+      except subprocess.CalledProcessError as e:
+        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
+
+def add_common_commands(parent_subparsers):
+  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
+  p_insmod.set_defaults(func=cmd_insert_module)
+
+  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.set_defaults(func=cmd_remove_module)
+
+  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
+  p_reset.set_defaults(func=cmd_reset_devices)
+
+  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_show_pids)
+
+  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_kill_pids)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
+
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
+  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(nv_commands)
+
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(amd_commands)
+
+  args = parser.parse_args()
+  if args.command is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)
@@ -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()
@@ -99,7 +99,9 @@ class FeedForward:
    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit

  def __call__(self, x:Tensor) -> Tensor:
-    return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x.contiguous_backward())  # this fixes a strange fusion that makes tensor cores miss
+    return self.w2(w1 * w3)

 class TransformerBlock:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -111,7 +113,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
    h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()

 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -185,10 +187,10 @@ class Transformer:

    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h)).float()[:, -1, :]
+    logits = self.output(self.norm(h)).float()
    if math.isnan(temperature): return logits

-    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
+    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)

  def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
    # TODO: better way to handle the first call v.s. the rest?
@@ -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;
@@ -1,2 +0,0 @@
-GPU="$1"
-echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null
@@ -1,10 +1,12 @@
+# mypy: disable-error-code="misc, list-item, assignment, 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,17 +713,96 @@ 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]

  def Gather(x:Tensor, indices:Tensor, axis:int=0):
    if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
-      x_sh = list(x.shape)
-      ret_shape = x_sh[:axis] + list(indices.shape) + x_sh[axis+1:]
+      ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
      if indices.ndim > 1: indices = indices.flatten()
-      indices = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
-      indices = [x_sh[axis]+x if x<0 else x for x in indices]
-      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x_sh)] for i in indices] # type: ignore
+      index_consts = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
+      index_consts = [x.shape[axis]+i if i<0 else i for i in index_consts]
+      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x.shape)] for i in index_consts]
      return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
    # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
    return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
@@ -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
@@ -128,6 +128,12 @@ def _linalg_eigh(self, UPLO: str = 'U'):
  w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
  return w.tiny(), v.tiny()

+@torch.library.impl("aten::_linalg_det", "privateuseone")
+# TODO: move to tinygrad
+def _linalg_det(self: torch.Tensor):
+  result = aten._linalg_det(self.cpu())
+  return result[0].tiny(), result[1].tiny(), result[2].tiny()
+
 def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()

 for i in [
@@ -352,6 +358,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 +423,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 +485,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])
@@ -177,6 +198,11 @@ class TestTorchBackend(unittest.TestCase):
    recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
    np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)

+  def test_linalg_det(self):
+    a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
+    b = torch.linalg.det(a)
+    np.testing.assert_equal(b.cpu().numpy(), 120.0)
+
  def test_scalar_assign(self):
    a = torch.tensor([1, 2, 3], device=device)
    a[1] = 4
@@ -9,7 +9,7 @@ with open(directory / 'README.md', encoding='utf-8') as f:

 testing_minimal = [
  "numpy",
-  "torch",
+  "torch==2.7.1",
  "pytest",
  "pytest-xdist",
  "hypothesis",
@@ -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,7 +1,8 @@
 import pathlib
 from tinygrad import Tensor, Device, Context
+from tinygrad.helpers import getenv

 if __name__ == "__main__":
  with Context(DEBUG=2):
-    disk_llama = Tensor(pathlib.Path("/raid/weights/LLaMA-3/8B/consolidated.00.pth"))
+    disk_llama = Tensor(pathlib.Path(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
    device_llama = disk_llama.to(Device.DEFAULT).realize()
@@ -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.from_gguf_kv(kv)
+  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,41 @@
+from transformers import AutoTokenizer
+from datasets import load_dataset
+from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
+from tinygrad.helpers import tqdm, getenv, partition
+
+# 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")
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
+    lambda e: e[1] in base_tokenizer.all_special_ids)
+  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
+  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
+
+  color_codes = [ 91, 92, 94, 93, 95 ]
+  def color_tokens(tids):
+    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([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}

@@ -61,6 +74,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
      warnings.warn(f"detected changes in over {MAX_DIFF_PCT}%. skipping further diff generation.", ProcessReplayWarning)
      early_stop.set()
      break
+    name, loc = "", ""
    try:
      name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
      ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
@@ -77,7 +91,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
        warnings.warn("PROCESS REPLAY DETECTED CHANGE", ProcessReplayWarning)
    except Exception as e:
      changed += 1
-      warnings.warn(e, ProcessReplayWarning)
+      warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
  conn.commit()
  cur.close()

@@ -110,5 +124,5 @@ if __name__ == "__main__":
  logging.info(f"running process replay with {ASSERT_DIFF=}")
  try: _pmap(replayers)
  except Exception as e:
-    logging.info("process replay err", e)
+    logging.info(f"process replay err: {e}")
    exit(int(ASSERT_DIFF))
@@ -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()
@@ -1,10 +1,9 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad import Device
+import unittest, numpy as np
+from tinygrad import Tensor, Device, TinyJit
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory

-N = 4096
+N = 256 if CI else 4096
 class TestCopySpeed(unittest.TestCase):
  @classmethod
  def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -49,6 +48,32 @@ class TestCopySpeed(unittest.TestCase):
      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
        t.to('CPU').realize()

+  def testCopyDefaulttoCPUJit(self):
+    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
+
+    @TinyJit
+    def _do_copy(t): return t.to('CPU').realize()
+
+    t = Tensor.randn(N, N, 4).contiguous().realize()
+    for _ in range(5):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        x = _do_copy(t)
+        Device[Device.DEFAULT].synchronize()
+      np.testing.assert_equal(t.numpy(), x.numpy())
+
+  def testCopytoCPUtoDefaultJit(self):
+    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
+
+    @TinyJit
+    def _do_copy(x): return t.to(Device.DEFAULT).realize()
+
+    for _ in range(5):
+      t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize()
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+        x = _do_copy(t)
+        Device[Device.DEFAULT].synchronize()
+      np.testing.assert_equal(t.numpy(), x.numpy())
+
  @unittest.skipIf(CI, "CI doesn't have 6 GPUs")
  @unittest.skipIf(Device.DEFAULT != "GPU", "only test this on GPU")
  def testCopyCPUto6GPUs(self):
@@ -0,0 +1,32 @@
+import unittest
+from tinygrad import dtypes, Device, Tensor, Context
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.engine.realize import get_program, ExecItem, CompiledRunner
+
+class TestDefineReg(unittest.TestCase):
+  def test_simple(self, at=AxisType.UPCAST):
+    N = 16
+    bout = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+    a_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N), (0,1)))
+
+    out = a_col.load(a_col.store(a.load()))
+    sink = bout.store(out).sink(arg=KernelInfo(name="regcopy", axis_types=(AxisType.LOOP, at)))
+    prg = get_program(sink, Device.default.renderer)
+
+    with Context(DEBUG=0):
+      a = Tensor.randn(N, N).realize()
+      b = Tensor.empty(N, N).realize()
+    hrunner = CompiledRunner(prg)
+    ExecItem(hrunner, [b.uop.buffer, a.uop.buffer]).run(wait=True)
+    with Context(DEBUG=0):
+      self.assertEqual((b-a).mean().item(), 0.0)
+
+  @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
+  def test_simple_loop(self): self.test_simple(AxisType.LOOP)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -4,7 +4,7 @@ import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_dtype, to_dtype, fp8_to_float, float_to_fp8
+from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
 from hypothesis import assume, given, settings, strategies as strat
@@ -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)
@@ -376,30 +384,6 @@ class TestPtrDType(unittest.TestCase):
    self.assertEqual(dt.v, 4)
    self.assertEqual(dt.count, 4)

-class TestImageDType(unittest.TestCase):
-  def test_image_scalar(self):
-    assert dtypes.imagef((10,10)).base.scalar() == dtypes.float32
-    assert dtypes.imageh((10,10)).base.scalar() == dtypes.float32
-  def test_image_vec(self):
-    assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4)
-    assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4)
-
-class TestEqStrDType(unittest.TestCase):
-  def test_image_ne(self):
-    if ImageDType is None: raise unittest.SkipTest("no ImageDType support")
-    assert dtypes.float == dtypes.float32, "float doesn't match?"
-    assert dtypes.imagef((1,2,4)) != dtypes.imageh((1,2,4)), "different image dtype doesn't match"
-    assert dtypes.imageh((1,2,4)) != dtypes.imageh((1,4,2)), "different shape doesn't match"
-    assert dtypes.imageh((1,2,4)) == dtypes.imageh((1,2,4)), "same shape matches"
-    assert isinstance(dtypes.imageh((1,2,4)), ImageDType)
-  def test_ptr_eq(self):
-    assert dtypes.float32.ptr() == dtypes.float32.ptr()
-    assert not (dtypes.float32.ptr() != dtypes.float32.ptr())
-  def test_strs(self):
-    if PtrDType is None: raise unittest.SkipTest("no PtrDType support")
-    self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")
-    self.assertEqual(str(dtypes.float32.ptr(16)), "dtypes.float.ptr(16)")
-
 class TestImplicitFunctionTypeChange(unittest.TestCase):
  def test_functions(self):
    result = []
@@ -430,19 +414,6 @@ class TestDtypeUsage(unittest.TestCase):
        t = Tensor([[1, 2], [3, 4]], dtype=d)
        (t*t).max().item()

-class TestToDtype(unittest.TestCase):
-  def test_dtype_to_dtype(self):
-    dtype = dtypes.int32
-    res = to_dtype(dtype)
-    self.assertIsInstance(res, DType)
-    self.assertEqual(res, dtypes.int32)
-
-  def test_str_to_dtype(self):
-    dtype = "int32"
-    res = to_dtype(dtype)
-    self.assertIsInstance(res, DType)
-    self.assertEqual(res, dtypes.int32)
-
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
  def test_cast(self):
@@ -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
@@ -107,8 +107,9 @@ class TestGraph(unittest.TestCase):
    helper_test_graphs(Device[d0].graph, graphs)

  def skip_if_not_multigraph(self):
-    graph = g.func if isinstance(g:=Device[Device.DEFAULT].graph, functools.partial) else g
+    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
+    if not hasattr(d.allocator, '_transfer'): self.skipTest("device is not supported (no transfers)")

  def test_order_copy_writed(self):
    self.skip_if_not_multigraph()
@@ -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.helpers import getenv, CI, mv_address, DEBUG
 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) \
@@ -510,6 +513,31 @@ class TestHCQ(unittest.TestCase):

      assert buf2.as_buffer()[0] == i

+  def test_map_cpu_buffer_to_device(self):
+    if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")
+
+    sz = 0x2000
+    cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
+    cpu_buffer._buf.cpu_view().view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
+
+    for devid in range(6):
+      if DEBUG >= 2: print(f"Testing map to device {Device.DEFAULT}:{devid}")
+
+      try: d = Device[f"{Device.DEFAULT}:{devid}"]
+      except Exception: break
+
+      local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
+
+      d.allocator.map(cpu_buffer._buf)
+
+      d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
+                         .copy(local_buf._buf.va_addr, cpu_buffer._buf.va_addr, sz) \
+                         .signal(d.timeline_signal, d.timeline_value).submit(d)
+      d.timeline_signal.wait(d.timeline_value)
+      d.timeline_value += 1
+
+      np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
+
  @unittest.skipUnless(MOCKGPU, "Emulate this on MOCKGPU to check the path in CI")
  def test_on_device_hang(self):
    if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
@@ -530,10 +558,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()
@@ -5,9 +5,10 @@ import numpy as np
 from hypothesis import given, settings, strategies as strat
 from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV
 from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit
+from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
+from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
 from tinygrad.device import Device
-from tinygrad.helpers import Context, JIT, GlobalCounters
+from tinygrad.helpers import Context, JIT, GlobalCounters, getenv
 from tinygrad.dtype import dtypes
 from extra.models.unet import ResBlock

@@ -669,5 +670,169 @@ class TestJitFree(unittest.TestCase):
    out = fxn(Tensor([11,1,2,3,4]))
    self.assertEqual(out.item(), 13600)

+class TestJitGraphSplit(unittest.TestCase):
+  def compute(self, device, inp):
+    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
+    return (inp + 1.0).contiguous().realize()
+
+  def copy(self, device, to_device, inp):
+    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
+    return inp.to(to_device).realize()
+
+  def expect(self, f, *args, graph=None, multigraph=None, hcqgraph=None):
+    def _numpies(tpl): return tpl.numpy() if tpl.__class__ is Tensor else tuple([t.numpy() for t in tpl])
+
+    expected = _numpies(f(*args))
+    for i in range(4):
+      res = _numpies(f(*args))
+      np.testing.assert_allclose(res, expected, atol=1e-4, rtol=1e-5)
+
+    dev = Device[Device.DEFAULT]
+    graph_t = graph_class(dev)
+    if graph_t is None: return
+
+    got = f.jit_cache
+    from tinygrad.runtime.graph.hcq import HCQGraph
+    if graph_t is HCQGraph:
+      validate = hcqgraph
+    elif issubclass(graph_t, MultiGraphRunner):
+      validate = multigraph
+    else:
+      validate = graph
+
+    assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
+    for expected, got in zip(validate, got):
+      if expected["type"] == "graph":
+        assert isinstance(got.prg, GraphRunner), f"Expected GraphRunner, got {type(got.prg)}"
+        assert len(got.prg.jit_cache) == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {len(got.prg.jit_cache)}"
+      elif expected["type"] == "comp":
+        assert isinstance(got.prg, CompiledRunner), f"Expected CompiledRunner, got {type(got.prg)}"
+      elif expected["type"] == "copy":
+        assert isinstance(got.prg, BufferCopy), f"Expected BufferCopy, got {type(got.prg)}"
+      elif expected["type"] == "xfer":
+        assert isinstance(got.prg, BufferXfer), f"Expected BufferXfer, got {type(got.prg)}"
+
+  def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
+  def ji_comp(self): return {"type": "comp"}
+  def ji_copy(self): return {"type": "copy"}
+  def ji_xfer(self): return {"type": "xfer"}
+
+  def test_jit_split_simple(self):
+    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")
+
+    @TinyJit
+    def f(inp):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(Device.DEFAULT, op1)
+      return op2
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    self.expect(f, inp,
+      graph=[self.ji_graph(3)],
+      multigraph=[self.ji_graph(3)],
+      hcqgraph=[self.ji_graph(3)])
+
+  def test_jit_cpu_simple(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    @TinyJit
+    def f(inp, inp_cpu):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute("CPU", inp_cpu)
+      op3 = self.compute(Device.DEFAULT, op1)
+      return op2, op3
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
+    self.expect(f, inp, inp_cpu,
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(4)])
+
+  def test_jit_cpu_several(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    @TinyJit
+    def f(inp, inp_cpu):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute("CPU", inp_cpu)
+      op3 = self.compute("CPU", op2)
+      op4 = self.compute(Device.DEFAULT, op1)
+      return op3, op4
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
+    self.expect(f, inp, inp_cpu,
+      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      hcqgraph=[self.ji_graph(5)])
+
+  def test_jit_multidev(self):
+    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
+
+    try: Device[f"{Device.DEFAULT}:1"]
+    except Exception: raise unittest.SkipTest("no multidevice")
+
+    @TinyJit
+    def f(inp, inp_d1):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
+      op3 = self.compute(f"{Device.DEFAULT}:1", op2)
+      op4 = self.compute(Device.DEFAULT, op1)
+      return op3, op4
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
+    self.expect(f, inp, inp_d1,
+      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
+      multigraph=[self.ji_graph(5)],
+      hcqgraph=[self.ji_graph(5)])
+
+  def test_jit_multidev_xfer(self):
+    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
+
+    try: Device[f"{Device.DEFAULT}:1"]
+    except Exception: raise unittest.SkipTest("no multidevice")
+
+    @TinyJit
+    def f(inp, inp_d1):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
+      op3 = self.copy(f"{Device.DEFAULT}:1", Device.DEFAULT, op2)
+      op4 = self.compute(f"{Device.DEFAULT}:1", op2)
+      op5 = self.compute(Device.DEFAULT, op3)
+      return op1, op4, op5
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
+    self.expect(f, inp, inp_d1,
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_xfer(), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(6)],
+      hcqgraph=[self.ji_graph(6)])
+
+  @unittest.skipIf(getenv("MOCKGPU"), "MockGPU does not support parallel copies")
+  def test_jit_multidev_copy(self):
+    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
+    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")
+
+    @TinyJit
+    def f(inp):
+      op0 = self.compute(Device.DEFAULT, inp)
+      op1 = self.compute(Device.DEFAULT, op0)
+      op2 = self.copy(Device.DEFAULT, "CPU", op1)
+      op3 = self.compute("CPU", op2)
+      return op3
+
+    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
+    self.expect(f, inp,
+      graph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(4)])
+
 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),(3,4,4),(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]))
@@ -114,27 +114,6 @@ class TestLinearizer(unittest.TestCase):
        if skip and i in skip: continue
        assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"

-  @unittest.expectedFailure
-  def test_const_alu_indexing(self):
-    st = ShapeTracker.from_shape((4,)).to_uop()
-    load = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=1, src=()), st, dtype=dtypes.float)
-    op = load+UOp.const(dtypes.float, 1.0)*UOp.const(dtypes.float, -1)
-    store = UOp.store(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0, src=()), st, op)
-    Tensor.manual_seed(0)
-    x = Tensor.randn(4,).realize()
-    helper_linearizer_ast(store.sink(), [x], wanna_output=[x.numpy()+1*-1], opts=[])
-
-  # shapeless CONST in AST is not supported
-  @unittest.expectedFailure
-  def test_const_alu_indexing_one_const_fine(self):
-    st = ShapeTracker.from_shape((4,)).to_uop()
-    load = UOp.load(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=1, src=()), st, dtype=dtypes.float)
-    op = load+UOp.const(dtypes.float, 1.0)
-    store = UOp.store(UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0, src=()), st, op)
-    Tensor.manual_seed(0)
-    x = Tensor.randn(4,).realize()
-    helper_linearizer_ast(store.sink(), [x], wanna_output=[x.numpy()+1], opts=[])
-
  @unittest.skipIf(CI and Device.DEFAULT in {"PTX", "AMD", "NV"}, "very slow")
  def test_indexing_multireduce(self):
    dataset = Tensor.rand(16384, 256).realize()
@@ -148,76 +127,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 +215,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 +228,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 +239,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 +255,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 +270,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.src[0].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 +284,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 +298,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 +306,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 +344,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 +435,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 +446,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 +459,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 +470,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 +555,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 +595,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)]
@@ -622,6 +612,7 @@ class TestLinearizer(unittest.TestCase):
    helper(Tensor.arange(255), max_ops=2)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
  def test_grouped_store_phis(self):
    """
    float4 acc0 = float4(0.0,0.0,0.0,0.0);
@@ -635,8 +626,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.src[0].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 +651,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,32 +664,35 @@ 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")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
+  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
  def test_grouped_store_local_only(self):
    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.src[0].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 +710,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 +728,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 +763,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 +778,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 +811,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 +826,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 +847,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 +863,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 +879,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 +893,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 +1051,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 +1061,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 +1081,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 +1338,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 +1443,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,11 +77,12 @@ 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"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
+  @unittest.skip("not applicable")
  def test_expander_new_srcs(self):
    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
      UOp(Ops.STORE, dtypes.void, arg=None, src=(
@@ -93,9 +95,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 +136,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 +165,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 +189,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):
@@ -1,7 +1,6 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad import dtypes, Device
-from tinygrad.helpers import CI
+from tinygrad import dtypes
 from tinygrad.opt.kernel import Kernel
 from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer
@@ -14,7 +13,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)
@@ -163,33 +161,5 @@ class TestLinearizerOverflow(unittest.TestCase):
    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=4), Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=2, arg=4)]
    _test_overflow(ast, opts)

-@unittest.skipIf(Device.DEFAULT not in {"GPU", "HSA", "CUDA", "METAL"}, "only backends with locals")
-@unittest.skipIf(CI, "slow")
-class TestLinearizerOverflowAlt(unittest.TestCase):
-  def test_overflow_1(self):
-    BS = 2
-    g0, g1, g2 = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=i) for i in range(3)]
-    in_st_1 = ShapeTracker(views=(View(shape=(1, BS, 1, 3, 8, 230, 8, 230), strides=(0, 150528, 0, 50176, 0, 224, 0, 1), offset=-675, mask=((0, 1), (0, BS), (0, 1), (0, 3), (0, 8), (3, 227), (0, 8), (3, 227)), contiguous=False),
-                                  View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(10156800, 0, 0, 3680, 2, 3385600, 425040, 231), offset=0, mask=None, contiguous=False))).to_uop()
-    in_st_2 = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)).to_uop()
-    ot_st = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 1, 1, 1), strides=(802816, 0, 12544, 112, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)).to_uop()
-    prod = UOp(Ops.LOAD, dtypes.float, (g1.view(in_st_1.arg),)) * UOp(Ops.LOAD, dtypes.float, (g2.view(in_st_2.arg),))
-    store = UOp(Ops.STORE, src=(g0.view(ot_st.arg), UOp(Ops.REDUCE_AXIS, dtypes.float, (prod,), (Ops.ADD, (7, 6, 5)))))
-    ast = UOp(Ops.SINK, src=(store,))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.LOCAL, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
-    _test_overflow(ast, opts)
-  def test_overflow_2(self):
-    BS = 2
-    g0, g1, g2 = [UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=i) for i in range(3)]
-    in_st_1 = ShapeTracker(views=(View(shape=(1, BS, 1, 3, 8, 230, 8, 230), strides=(0, 150528, 0, 50176, 0, 224, 0, 1), offset=-675, mask=((0, 1), (0, BS), (0, 1), (0, 3), (0, 8), (3, 227), (0, 8), (3, 227)), contiguous=False),
-                                  View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(10156800, 0, 0, 3680, 2, 3385600, 425040, 231), offset=0, mask=None, contiguous=False))).to_uop()
-    in_st_2 = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 3, 7, 7), strides=(0, 0, 147, 0, 0, 49, 7, 1), offset=0, mask=None, contiguous=False),)).to_uop()
-    ot_st = ShapeTracker(views=(View(shape=(BS, 1, 64, 112, 112, 1, 1, 1), strides=(802816, 0, 12544, 112, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)).to_uop()
-    prod = UOp(Ops.LOAD, dtypes.float, (g1.view(in_st_1.arg),)) * UOp(Ops.LOAD, dtypes.float, (g2.view(in_st_2.arg),))
-    store = UOp(Ops.STORE, src=(g0.view(ot_st.arg), UOp(Ops.REDUCE_AXIS, dtypes.float, (prod,), (Ops.ADD, (7, 6, 5)))))
-    ast = UOp(Ops.SINK, src=(store,))
-    opts = [Opt(op=OptOps.LOCAL, axis=3, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=2, arg=16), Opt(op=OptOps.UPCAST, axis=4, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=5, arg=2)]
-    _test_overflow(ast, opts)
-
 if __name__ == '__main__':
  unittest.main()
@@ -2,7 +2,7 @@ import unittest, functools, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import CI, getenv, prod, Context, OSX
+from tinygrad.helpers import CI, getenv, prod, Context
 from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
@@ -374,7 +374,6 @@ class TestMultiTensor(unittest.TestCase):

  # NOTE: this is failing on LLVM CI, no idea why. Works locally.
  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
-  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_data_parallel_resnet(self):
    from extra.models.resnet import ResNet18

@@ -411,7 +410,6 @@ class TestMultiTensor(unittest.TestCase):
    np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)

  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
-  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_data_parallel_resnet_train_step(self):
    from extra.models.resnet import ResNet18
    fake_image = Tensor.rand((2, 3, 224//8, 224//8))
@@ -938,7 +936,6 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
    np.testing.assert_allclose(output.numpy(), expected)

@unittest.skipIf(not_support_multi_device(), "no multi")
-@unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestBatchNorm(unittest.TestCase):
  def test_unsynced_backprop_conv_bn(self):
    with Tensor.train():
@@ -966,7 +963,6 @@ class TestBatchNorm(unittest.TestCase):
      optim.step()
      out.numpy()

-  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_unsynced_backprop_standalone_bn(self):
    from extra.lr_scheduler import OneCycleLR
    GPUS = (d1, d2)
@@ -1126,6 +1122,7 @@ class TestMultiRamUsage(unittest.TestCase):
    # NOTE: the first one on the DEFAULT device should be freed
    self.assertUsed(self.N*self.N*4*2)

+  @unittest.skip("flaky")
  def test_zeros_shard(self, devices=(d1, d2)):
    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices, axis=0).realize()
    self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
@@ -4,7 +4,7 @@ import numpy as np
 import torch
 from tinygrad import Tensor, Device, TinyJit
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import GlobalCounters, CI, Context, OSX
+from tinygrad.helpers import GlobalCounters, CI, Context
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
@@ -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]
@@ -131,105 +108,39 @@ class TestNN(unittest.TestCase):
    _test_linear(Tensor.randn(BS, in_dim), in_dim, out_dim)
    _test_linear(Tensor.randn(BS, T, in_dim), in_dim, out_dim) # test with more dims

-  def test_conv1d(self):
-    BS, C1, W = 4, 16, 224//4
-    C2, K, S, P = 64, 7, 2, 1
-
+  def _test_conv(self, tiny_conv, torch_conv, BS, C1, DIMS, C2, K, S, P, D=1):
    # create in tinygrad
-    layer = Conv1d(C1, C2, kernel_size=K, stride=S, padding=P)
+    layer = tiny_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D)

    # create in torch
    with torch.no_grad():
-      torch_layer = torch.nn.Conv1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+      torch_layer = torch_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D).eval()
      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)

    # test
-    x = Tensor.uniform(BS, C1, W)
+    x = Tensor.uniform(BS, C1, *DIMS)
    z = layer(x)
    torch_x = torch.tensor(x.numpy())
    torch_z = torch_layer(torch_x)
    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)

-  def test_conv2d(self):
-    BS, C1, H, W = 4, 16, 224//4, 224//4
-    C2, K, S, P = 64, 7, 2, 1
-
-    # create in tinygrad
-    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
-
-    # create in torch
-    with torch.no_grad():
-      torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
-      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
-      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
-
-    # test
-    x = Tensor.uniform(BS, C1, H, W)
-    z = layer(x)
-    torch_x = torch.tensor(x.numpy())
-    torch_z = torch_layer(torch_x)
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+  def test_conv1d(self): self._test_conv(Conv1d, torch.nn.Conv1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
+  def test_conv2d(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)

  def test_conv1d_same_padding(self):
-    BS, C1, W = 8, 3, 32
-    C2, K, S, P = 16, 3, 1, 'same'
-
-    # create in tinygrad
-    layer = Conv1d(C1, C2, kernel_size=K, stride=S, padding=P)
-
-    # create in torch
-    with torch.no_grad():
-      torch_layer = torch.nn.Conv1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
-      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
-      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
-
-    # test
-    x = Tensor.uniform(BS, C1, W)
-    z = layer(x)
-    torch_x = torch.tensor(x.numpy())
-    torch_z = torch_layer(torch_x)
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
-
-  def _run_conv2d_same_padding_test(self, BS, C1, C2, H, W, K, S, padding='same', D=1):
-    # create in tinygrad
-    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=padding, dilation=D)
-
-    # create in torch
-    with torch.no_grad():
-      torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=padding, dilation=D).eval()
-      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
-      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
-
-    # test
-    x = Tensor.uniform(BS, C1, H, W)
-    z = layer(x)
-    torch_x = torch.tensor(x.numpy())
-    torch_z = torch_layer(torch_x)
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
-
+    self._test_conv(Conv1d, torch.nn.Conv1d, BS=8, C1=3, DIMS=[32], C2=16, K=3, S=1, P='same')
  def test_conv2d_same_padding_odd_input(self):
-    BS, C1, H, W = 16, 16, 29, 31
-    C2, K, S, P = 32, 5, 1, 'same'
-    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P)
-
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[29, 31], C2=32, K=5, S=1, P='same')
  def test_conv2d_same_padding_large_kernel(self):
-    BS, C1, H, W = 16, 16, 28, 33
-    C2, K, S, P = 32, 9, 1, 'same'
-    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P)
-
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
  def test_conv2d_same_padding_with_dilation(self):
-    BS, C1, H, W = 16, 3, 28, 28
-    C2, K, S, P, D = 32, 3, 1, 'same', 3
-    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P, D)
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)

  def test_conv2d_same_padding_invalid_stride(self):
-    C1, C2, K, S, P = 16, 32, 2, 2, 'same'
-    self.assertRaises(ValueError, Conv2d, C1, C2, kernel_size=K, stride=S, padding=P)
-
+    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
  def test_conv2d_same_padding_invalid_padding_str(self):
-    C1, C2, K, S, P = 16, 32, 2, 1, 'not_same'
-    self.assertRaises(ValueError, Conv2d, C1, C2, kernel_size=K, stride=S, padding=P)
+    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')

  @unittest.skip("Takes too long to compile for Compiled backends")
  def test_conv2d_winograd(self):
@@ -252,12 +163,13 @@ class TestNN(unittest.TestCase):
    with Context(WINO=1):
      z = layer(x)

+    m = z.mean()
+    m.backward()
+
    torch_x = torch.tensor(x.numpy(), requires_grad=True)
    torch_z = torch_layer(torch_x)
    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)

-    m = z.mean()
-    m.backward()
    gw = layer.weight.grad.realize()
    gb = layer.bias.grad.realize()
    gx = x.grad.realize()
@@ -268,46 +180,10 @@ class TestNN(unittest.TestCase):
    np.testing.assert_allclose(gx.numpy(), torch_x.grad.numpy(), atol=5e-4, rtol=1e-5)

  def test_conv_transpose1d(self):
-    BS, C1, W = 4, 16, 224//4
-    C2, K, S, P = 64, 7, 2, 1
-
-    # create in tinygrad
-    layer = ConvTranspose1d(C1, C2, kernel_size=K, stride=S, padding=P)
-
-    # create in torch
-    with torch.no_grad():
-      torch_layer = torch.nn.ConvTranspose1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
-      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
-      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
-
-    # test
-    x = Tensor.uniform(BS, C1, W)
-    z = layer(x)
-    torch_x = torch.tensor(x.numpy())
-    torch_z = torch_layer(torch_x)
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
-
+    self._test_conv(ConvTranspose1d, torch.nn.ConvTranspose1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
  def test_conv_transpose2d(self):
-    BS, C1, H, W = 4, 16, 224//4, 224//4
-    C2, K, S, P = 64, 7, 2, 1
+    self._test_conv(ConvTranspose2d, torch.nn.ConvTranspose2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)

-    # create in tinygrad
-    layer = ConvTranspose2d(C1, C2, kernel_size=K, stride=S, padding=P)
-
-    # create in torch
-    with torch.no_grad():
-      torch_layer = torch.nn.ConvTranspose2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
-      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
-      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
-
-    # test
-    x = Tensor.uniform(BS, C1, H, W)
-    z = layer(x)
-    torch_x = torch.tensor(x.numpy())
-    torch_z = torch_layer(torch_x)
-    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
-
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_groupnorm(self):
    BS, H, W, C, G = 20, 10, 10, 6, 3

@@ -334,7 +210,6 @@ class TestNN(unittest.TestCase):
      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_layernorm(self):
    N, C, H, W = 20, 5, 10, 10

@@ -361,7 +236,6 @@ class TestNN(unittest.TestCase):
      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_layernorm_2d(self):
    N, C, H, W = 20, 5, 10, 10

@@ -388,7 +262,6 @@ class TestNN(unittest.TestCase):
      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_instancenorm_2d(self):
    N, C, H, W = 20, 10, 10, 10

@@ -415,7 +288,6 @@ class TestNN(unittest.TestCase):
      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_instancenorm_3d(self):
    N, C, D, H, W = 20, 10, 10, 10, 10

@@ -442,7 +314,6 @@ class TestNN(unittest.TestCase):
      np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
      np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_rmsnorm(self):
    class TorchRMSNorm(torch.nn.Module):
      # https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L34C1-L77C36
@@ -1,8 +1,8 @@
-import time, math, unittest, functools, warnings
+import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, OSX, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -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)
@@ -2664,7 +2682,6 @@ class TestOps(unittest.TestCase):
    i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
    return a,b,c,d,e,i,j,k,o,p

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU can only run kernels with up to 10 buffers")
  def test_slice_fancy_indexing_no_dim_collapse(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    # no dim collapse from int or dim injection from None
@@ -2716,7 +2733,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,1,-1],[-1,-2,0]]), torch.tensor([2,1,-1])],
                            lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU can only run kernels with up to 10 buffers")
  def test_slice_fancy_indexing_list_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
@@ -2736,7 +2752,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,((2,),(1,),(0,)),c,(2,1,0)], lambda x: x[i,((2,),(1,),(0,)),k,(2,1,0)])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,(2,1,0),None,c,(2,1,0),e], lambda x: x[1,(2,1,0),None,k,(2,1,0),p])

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
  def test_slice_fancy_indexing_list_with_tensors(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a]], lambda x: x[[i]])
@@ -2897,6 +2912,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 +2970,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 +3081,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)
@@ -3,6 +3,7 @@ import numpy as np
 from tinygrad import Tensor, Variable, Device
 from tinygrad.helpers import OSX

+# TODO: still fails with MAX_KERNEL_BUFFERS
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestSample(unittest.TestCase):
  def test_sample(self):
@@ -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)
@@ -34,25 +34,29 @@ class TestBEAM(unittest.TestCase):
    capturing.clear()
    self.assertNotEqual(k_beam_0[-1].prg.p.src, k_beam_1[-1].prg.p.src)

-  def test_get_kernel_actions(self):
+  def test_get_kernel_actions_dedup(self):
    from test.test_linearizer import helper_realized_ast
-    a = Tensor.rand(4, 3)
-    b = Tensor.rand(3)
-    realized_ast, _ = helper_realized_ast(a @ b)
    from tinygrad.opt.search import get_kernel_actions
-    lins = get_kernel_actions(Kernel(realized_ast), False).values()
+    a = Tensor.empty(4, 3)
+    b = Tensor.empty(3)
+    realized_ast, _ = helper_realized_ast(a @ b)
+    candidates = [
+      Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4),
+      Opt(op=OptOps.LOCAL, axis=0, arg=0), Opt(op=OptOps.LOCAL, axis=0, arg=4),
+      Opt(op=OptOps.UNROLL, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=3),
+      Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=3),
+      Opt(op=OptOps.GROUPTOP, axis=0, arg=0), Opt(op=OptOps.GROUPTOP, axis=0, arg=3),
+    ]
+    lins = get_kernel_actions(Kernel(realized_ast), include_0=False, candidates=candidates).values()

    # ensure amt=0 are not duplicated
-    if Opt(OptOps.UPCAST, 0, 0) in actions:
-      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UPCAST, axis=0, arg=4)]) == 0, "did not de-dup UPCAST"
-    if Opt(OptOps.LOCAL, 0, 0) in actions:
-      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.LOCAL, axis=0, arg=4)]) == 0, "did not de-dup LOCAL"
-    if Opt(OptOps.UNROLL, 0, 0) in actions:
-      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UNROLL, axis=0, arg=3)]) == 0, "did not de-dup UNROLL"
-    if Opt(OptOps.GROUP, 0, 0) in actions:
-      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUP, axis=0, arg=3)]) == 0, "did not de-dup GROUP"
-    if Opt(OptOps.GROUPTOP, 0, 0) in actions:
-      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUPTOP, axis=0, arg=3)]) == 0, "did not de-dup GROUPTOP"
+    assert all(len(x.applied_opts) == 1 for x in lins)
+    kernel_actions = [x.applied_opts[0] for x in lins]
+    assert Opt(OptOps.UPCAST, axis=0, arg=4) not in kernel_actions, "did not de-dup UPCAST"
+    assert Opt(OptOps.LOCAL, axis=0, arg=4) not in kernel_actions, "did not de-dup LOCAL"
+    assert Opt(OptOps.UNROLL, axis=0, arg=3) not in kernel_actions, "did not de-dup UNROLL"
+    assert Opt(OptOps.GROUP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUP"
+    assert Opt(OptOps.GROUPTOP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUPTOP"

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_search_over_shape(self):
@@ -158,6 +158,13 @@ class TestSetitem(unittest.TestCase):
    t[:-1] = t[1:]
    self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])

+  def test_setitem_big(self):
+    idx_size, val = 256, 4
+    t = Tensor.arange(0, idx_size+1)
+    idx = Tensor.arange(0, idx_size)
+    t[idx] = val
+    self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
+
 class TestWithGrad(unittest.TestCase):
  def test_no_requires_grad_works(self):
    z = Tensor.rand(8, 8)
@@ -176,5 +183,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()
--- a/Show More
+++ b/Show More