notes

uc warp
Merge branch 'master' into warp_fun
2025-07-01 10:36:05 -07:00 · 2025-06-30 20:54:29 -07:00 · 2025-06-30 18:18:52 -07:00 · 2025-06-30 15:54:57 -07:00 · 2025-06-29 09:26:32 -07:00 · 2025-06-29 09:13:16 -07:00
187 changed files with 2966 additions and 13221 deletions
@@ -112,16 +112,7 @@ 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
@@ -144,11 +135,14 @@ 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: Compute Package List + Hash
+    - name: apt-get update + install
      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
@@ -159,7 +153,7 @@ runs:
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
        fi
        # **** CUDA ****
        if [[ "${{ inputs.cuda }}" == "true" ]]; then
@@ -174,31 +168,14 @@ 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 "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
-          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
+        if [[ -n "$pkgs" ]]; then
+          sudo apt-get -y --allow-unauthenticated --no-install-recommends install $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
@@ -251,7 +228,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 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
+        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
    - 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)
@@ -637,125 +637,3 @@ 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 amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
-    - 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/nvpci/nv_smi.py rmmod
-    - 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,7 +10,6 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 100

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

    steps:
    - name: Checkout Code
@@ -132,13 +132,10 @@ 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
@@ -242,8 +239,8 @@ jobs:
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 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 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 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
@@ -255,8 +252,8 @@ jobs:
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 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 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 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
@@ -329,21 +326,16 @@ 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
-        python -m mypy --strict-equality extra/onnx_parser.py
-        python -m mypy --strict-equality extra/onnx.py
+      run: python -m mypy --strict-equality --lineprecision-report . && cat lineprecision.txt

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

    steps:
    - name: Checkout Code
@@ -376,8 +368,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 < 15500 lines
-      run: MAX_LINE_COUNT=15500 python sz.py
+    - name: Repo line count < 14600 lines
+      run: MAX_LINE_COUNT=14600 python sz.py

  fuzzing:
    name: Fuzzing
@@ -430,29 +422,6 @@ 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
@@ -471,13 +440,11 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          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
+          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
      - 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
@@ -508,10 +475,15 @@ 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

@@ -535,6 +507,10 @@ 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
@@ -548,20 +524,6 @@ 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
@@ -601,7 +563,7 @@ jobs:
      with:
        key: dsp-minimal
        deps: testing_minimal
-        pydeps: "onnx==1.18.0 onnxruntime pillow"
+        pydeps: "onnx==1.17.0 onnxruntime pillow"
        llvm: "true"
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
@@ -689,9 +651,7 @@ 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 --durations=20
-      - name: Run pytest (amd)
-        run: python -m pytest 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 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
@@ -844,7 +804,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: osx-webgpu
-        deps: testing
+        deps: testing_minimal
        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
@@ -867,8 +827,6 @@ 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)
@@ -895,11 +853,12 @@ jobs:
          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

  amdremote:
-    name: Linux (remote)
+    name: Linux (remote amd)
    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
@@ -907,58 +866,38 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: linux-remote
+          key: linux-remote-amd
          deps: testing_minimal
          amd: 'true'
          llvm: 'true'
-          opencl: 'true'
      - name: Start remote server
        run: |
          start_server() {
            systemd-run --user \
              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
+              --setenv=REMOTEDEV=AMD \
              --setenv=MOCKGPU=1 \
              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
+              --setenv=PORT="$2" \
              --working-directory="$(pwd)" \
              python tinygrad/runtime/ops_remote.py
          }

-          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
+          start_server "remote-server-1" 6667
+          start_server "remote-server-2" 6668
      - 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 (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+      - 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_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-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
+          journalctl --user -u remote-server-1 --no-pager
+          journalctl --user -u remote-server-2 --no-pager

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

    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
@@ -1012,12 +951,11 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: windows-${{ matrix.backend }}-minimal
+          key: windows-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' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=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,8 +39,9 @@ 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.empty(N, N), Tensor.empty(N, N);
-(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
+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())"
 ```

 And we can change `DEBUG` to `4` to see the generated code.
@@ -149,7 +149,6 @@ 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
@@ -167,7 +166,6 @@ 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 \
@@ -182,7 +180,6 @@ 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 \
@@ -190,7 +187,7 @@ nv_status_codes = {}
    -o $BASE/nv/nv.py

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

 generate_amd() {
@@ -240,21 +237,6 @@ 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) \
@@ -459,7 +441,7 @@ generate_libusb() {
    -o $BASE/libusb.py

  fixup $BASE/libusb.py
-  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, 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"
@@ -480,7 +462,6 @@ 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,30 +7,28 @@

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

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

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

 # load in some values (little endian)
-cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
-cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+MallocAllocator._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 = cpu.runtime("add", lib)
+fxn = CPUProgram("add", lib)

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

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


@@ -48,7 +46,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 cpu.allocator.alloc
+# NOTE: a._buf is the same as the return from MallocAllocator.alloc

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

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.schedule.kernelize import get_kernelize_map
+from tinygrad.kernelize.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.new_signal(value=0)
+signal = your_device.signal_t()

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

 ---

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

 Group UOps into kernels.

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

 ## Interoperability

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

 def load_onnx_model(onnx_file):
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx_load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)
  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,7 +10,6 @@ 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):
@@ -26,7 +25,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(Tensor(onnx_model.SerializeToString(), device="PYTHON"))
+    self.run_onnx = OnnxRunner(onnx_model)

  def forward(self, x):
    return self.run_onnx({"x": x}, debug=False)['predictions']
@@ -48,8 +47,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}, {cl._buf.size});")
-    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl._buf)*4});")
+    weights += bytes(cl._buf)
  cprog.append("}")

  # write the weights to disk
@@ -9,7 +9,6 @@ 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()
@@ -109,7 +108,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:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
+      ops = (prg:=lin.to_program()).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))

@@ -122,7 +121,7 @@ if __name__ == "__main__":
    if getenv("SRC"):
      print(si.ast)
      print(lin.applied_opts)
-      print(get_program(lin.get_optimized_ast(), lin.opts).src)
+      print(lin.to_program().src)
    total_tm += tm
    running_gflops += gflops * tm
    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
@@ -5,26 +5,29 @@ 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, dtypes
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
 from tinygrad.helpers import DEBUG, getenv
+from tinygrad.tensor import _from_np_dtype
 from tinygrad.engine.realize import CompiledRunner

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

 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):
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx_load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)
  print("loaded model")

-  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()}
+  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}
  # Float inputs and outputs to tinyjits for openpilot are always float32
-  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
+  input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  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_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

@@ -54,11 +57,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"different kernels! {kernel_count=}, {allowed_kernel_count=}"
+    assert kernel_count <= allowed_kernel_count, f"too many 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"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+    assert gated_read_image_count <= allowed_gated_read_image, f"too many 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
+import sys, onnx
 from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.schedule.kernelize import get_kernelize_map
+from tinygrad.kernelize.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule

@@ -12,8 +12,10 @@ 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)
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx.load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)

  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).realize()
+      out = model.decoder(Tensor([lst]), 0, encoded_audio, streaming=True).realize()
      idx = int(out[0,-1].argmax().numpy().item())
      lst.append(idx)
      dec = enc.decode(lst)
@@ -2,12 +2,13 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 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])
-run_onnx = OnnxRunner("yolov8n-seg.onnx")
+onnx_model = onnx_load(open("yolov8n-seg.onnx", "rb"))
+run_onnx = OnnxRunner(onnx_model)
 run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)
@@ -8,6 +8,8 @@ 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:
@@ -71,12 +73,21 @@ class AMSMI(AMDev):
    self._run_discovery()
    self._build_regs()

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

-    self.is_booting = True
-    self.init_sw(smi_dev=True)
+    self.is_booting, self.smi_dev = True, 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()
@@ -125,7 +136,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() != AMDev.Version:
+      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AM_VERSION:
        self.devs.remove(d)
        self.opened_pcidevs.remove(d.pcibus)
        os.system('clear')
@@ -284,8 +295,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.1)
-            if stopped_pids[pid] == 64:
+            if stopped_pids[pid] > 0: time.sleep(0.5)
+            if stopped_pids[pid] == 10:
              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_at=uops.uops.index(u))
-      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_at=uops.uops.index(u))
+      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))
      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, to_mv
+from tinygrad.helpers import Context
 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(to_mv(cl._buf.va_addr, cl._buf.size))])
+      weight = ''.join(["\\x%02X"%x for x in bytes(cl._buf)])
      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,42 +1,96 @@
 # 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, GlobalCounters
+from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+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

 N = 4096
+LN = 16
 run_count = 5

-if __name__ == "__main__":
-  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])
+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:
-    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)
+    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))

-  a = Tensor.randn(N, N).realize()
-  b = Tensor.randn(N, N).realize()
-  c = Tensor.zeros(N, N).contiguous().realize()
+local_loads_pm = PatternMatcher([
+  (UPat(Ops.LOAD, name="x"), transform_load),
+])

-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    for _ in range(run_count): tc = (a@b).realize()
+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

-  GlobalCounters.reset()
-  ei = ExecItem(runner, [a.uop.buffer, b.uop.buffer, c.uop.buffer])
+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)
  with Context(DEBUG=2):
    for _ in range(run_count): ei.run(wait=True)
-  print(f"custom  {(c-tc).square().mean().item()}")
+  nc = c.numpy()
+  np.testing.assert_allclose(na@nb, nc, rtol=1e-5)
@@ -1,143 +0,0 @@
-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)];
-        }
-      }
-    }
-  }
-}
@@ -1,172 +0,0 @@
-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)];
-        }
-      }
-    }
-  }
-}
@@ -1,172 +0,0 @@
-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,335 +0,0 @@
-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!")
@@ -1,5 +1,6 @@
 	.text
 	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
+	;.amdhsa_code_object_version 5
 	.protected	kernel                  ; -- Begin function kernel
 	.globl	kernel
 	.p2align	8
@@ -8,7 +9,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], 0x0 ; Matrix A and B
+	s_load_b128 s[20:23], s[0:1], 0x8 ; Matrix A and B
 	s_waitcnt lgkmcnt(0)

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


 	s_clause 0x1
-	; 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_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_lshl_b32 s2, s14, 7
 	v_lshrrev_b32_e32 v4, 3, v0
 	v_or_b32_e32 v1, s2, v0
@@ -90,7 +93,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], 0x10
+	s_load_b64 s[0:1], s[0:1], 0  ; Matrix C
 	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]
@@ -460,7 +463,7 @@ kernel:                                 ; @kernel

 	v_mov_b32_e32 v5, 0
 	v_mov_b32_e32 v3, 0
-	s_add_i32 s7, s4, -8
+	s_add_i32 s7, s4, -1
 	s_add_u32 s8, s8, 32
 	s_addc_u32 s9, s9, 0
 	s_mov_b32 s12, 0
@@ -2395,9 +2398,18 @@ 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: 24
+    .kernarg_segment_size: 36
    .language:       OpenCL C
    .language_version:
      - 2
@@ -2,7 +2,6 @@ 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
@@ -56,7 +55,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
  k = Kernel(ast)
  k.apply_opts(opts)
-  p = get_program(k.get_optimized_ast(), k.opts)
+  p = k.to_program()
  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, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem
 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 = get_program(k.get_optimized_ast(), k.opts)
+  prg = k.to_program()
  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.to_program()
+lin.linearize()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)
 print(src)
@@ -5,7 +5,7 @@ start = time.perf_counter()
 # *** ioctl lib ***
 libc = ctypes.CDLL(ctypes.util.find_library("c"))
 # platform.processor calls `uname -p` which can return `unknown` on some systems
-processor = os.getenv("IOCTL_PROCESSOR") or platform.processor() or platform.machine()
+processor = os.getenv("IOCTL_PROCESSOR") or platform.processor()
 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
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs

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

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

@@ -45,7 +46,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(model_path)
+    onnx_runner = OnnxRunner(onnx.load(model_path))
    for node in onnx_runner.graph_nodes:
      op_counter[node.op] += 1
      if node.op not in supported_ops:
@@ -7,7 +7,6 @@ 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):
@@ -111,7 +110,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
      seen_asts[opt_ast.key] = node

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

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

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

 def has_field(onnx_type: TypeProto|SimpleNamespace, field):
  if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
@@ -49,19 +48,30 @@ 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.")
-  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
+  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)

 def type_parse(onnx_type: TypeProto):
  elem_type = onnx_type
@@ -84,24 +94,10 @@ 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]
@@ -136,27 +132,17 @@ 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:
-  """
-  `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)
+  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)
    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)
-    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.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
    self.variable_dims: dict[str, int] = {}

    self.onnx_ops = onnx_ops
@@ -169,7 +155,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 is_numpy_ndarray(value) else spec.dtype
+    dtype = _from_np_dtype(value.dtype) if str(type(value)) == "<class 'numpy.ndarray'>" 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)):
@@ -178,25 +164,20 @@ 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 _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 _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 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}")
@@ -212,7 +193,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._select_op(node.op, node.opset_id)(*inps, **opts)
+      ret = self._dispatch_op(node.op, 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)))

@@ -227,10 +208,8 @@ class OnnxRunner:
 ####################
 ##### ONNX OPS #####
 ####################
-def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
+def get_onnx_ops():
  # ***** 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, ...)
@@ -301,8 +280,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    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|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 Range(start:float|int, limit:float|int, delta:float|int): return Tensor.arange(start=start, stop=limit, step=delta)

  def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
    try: import PIL.Image
@@ -329,13 +307,13 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  # ***** 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)  # noqa: A002
+  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 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 = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 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 = {1:Softmax_1, 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)
@@ -472,7 +450,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
                        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)
@@ -490,16 +468,14 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

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

-  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
-    axis = X._resolve_dim(_resolve_const(axis))
+  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)
    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|list[int]=0, upper:int=1):
-    k_ = _resolve_const(k)
-    return x.triu(k_) if upper else x.tril(k_)
+  def Trilu(x:Tensor, k:int=0, upper:int=1): 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,
@@ -560,8 +536,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    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):  # noqa: A002
-    val, idx = X.topk(_resolve_const(K), axis, largest, sorted)
+  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)
    return val, idx.cast(dtypes.int64)

  # ***** Neural Network Ops *****
@@ -623,9 +599,9 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  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[int|float], values:Tensor, axis:int=-1):
+  def OneHot(indices:Tensor, depth:float|int|list, values:Tensor, axis:int=-1):
    # Scalar or Rank 1 tensor containing exactly one element
-    depth = int(_resolve_const(depth))
+    depth = int(depth[0] if isinstance(depth, list) else 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])
@@ -636,8 +612,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    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):
-    import numpy as np
+  def Dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
    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)
@@ -646,8 +621,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    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 = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 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 = {6:Dropout_6, 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)
@@ -662,17 +637,16 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
    N, _, *spatial_dims = size
    def generate_grid(steps):
-      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)
+      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)
    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_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):
+  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):
    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)
@@ -713,96 +687,17 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    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
-      ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
+      x_sh = list(x.shape)
+      ret_shape = x_sh[:axis] + list(indices.shape) + x_sh[axis+1:]
      if indices.ndim > 1: indices = indices.flatten()
-      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]
+      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
      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
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 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_file)
+  run_onnx = OnnxRunner(onnx_load(onnx_file))

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

 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
+from tinygrad.tensor import Tensor, dtypes

 # 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
@@ -21,13 +22,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: dict[str, dict] = {
+PB_INFOS = {
  "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),
-    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
+    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},
  "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),
@@ -36,16 +37,16 @@ PB_INFOS: dict[str, dict] = {
    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": None, "int32_data": None, "string_data": None,
-                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
+    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [],
+                                                                  "int64_data": [], "double_data": [], "uint64_data": []})),
    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": None, "int32_data": None, "string_data": None, "int64_data": None,
-                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
+    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [], "int64_data": [],
+                                                        "double_data": [], "uint64_data": []}))},
  "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)},
@@ -54,7 +55,7 @@ PB_INFOS: dict[str, dict] = {
  "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
 }

-def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
+def onnx_load(fn: Union[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)
@@ -70,8 +71,8 @@ def dict_to_namespace(d):
  return d

 class OnnxParser:
-  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
-    self.file_path: pathlib.Path|None = None
+  def __init__(self, inp: Union[Tensor, str, pathlib.Path], load_external_data: bool=True):
+    self.file_path: Union[pathlib.Path, None] = None
    self.load_external_data = load_external_data
    if not isinstance(inp, Tensor):
      self.file_path = pathlib.Path(inp)
@@ -89,6 +90,7 @@ 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

@@ -129,19 +131,16 @@ 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) -> Tensor|bytes:
+  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tuple[bytes, Tensor]:
    str_len = self.decode_varint(reader)
    if not use_tensor: return reader.read(str_len)
-    raw = reader.raw
-    assert isinstance(raw, TensorIO)
-    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
+    res = reader.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):
@@ -167,17 +166,16 @@ 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] = values
+    obj[key_name] = Tensor(values, dtype=dtypes.int64)

  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
+    obj[key_name] = value.bitcast(dtypes.float32)

  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)))
@@ -196,7 +194,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 = pathlib.Path(self.tensor.device[5:])
+        self.file_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,7 +6,6 @@ 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

@@ -116,7 +115,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 = get_program(lin.get_optimized_ast(), lin.opts)
+  p = lin.to_program()
  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 ProfileDeviceEvent, ProfileGraphEvent
-from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent
+from tinygrad.helpers import tqdm, temp

 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, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 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 = get_program(k.get_optimized_ast(), k.opts)
+        p2 = k.to_program()
        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 multiple gpus are used you can select which one to export with `-d` like this:
+If multiplle gpus are used you can select which one to export with `-d` like this:

 ```bash
 extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
@@ -352,11 +352,6 @@ 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 = [
@@ -417,7 +412,6 @@ 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]
@@ -479,7 +473,6 @@ 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,27 +103,6 @@ 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])
@@ -0,0 +1,73 @@
+# play with upcasted warps
+from tinygrad import Tensor, Device
+from tinygrad.uop.ops import KernelInfo
+from tinygrad.opt import get_optimized_ast
+from tinygrad.opt.kernel import OptOps, Opt
+from tinygrad.engine.realize import get_program
+
+if __name__ == "__main__":
+  renderer = Device.default.renderer
+  N = 64
+
+  """
+  a = Tensor.empty(N,N)
+
+  out = (a + 1) #.sum(axis=2)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 32),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # how you split the store determines everything if you don't allow cross warp comms.
+  # actually not everything, there's also the split before the horizontal (unrolled) reduces
+
+  # new flow
+  #  - pull out any dimensions from the store that you want to upcast.
+  #  - decide how you want to assign them to registers. GPUs have a 512-byte memory LOAD/STORE which loads into 4 regs. see BUFFER_LOAD_B128
+  #    - the loads and stores can be shuffled, but only in restrictive ways. in kernels without reduces, the store determines everything
+  #    - it loads 16 bytes from up 32 different places = 512 bytes
+  #  - in kernels with reduces, you now have more flexibility. the final target of the reduce must be what is stored
+  #    - warp dimensions can be in the reduce (this is GROUP)
+
+  # every dimension can be assigned to <global, local, loop, upcast, warp>
+
+  """
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # gemm
+  b = Tensor.empty(N,N)
+  # metal TC
+  #opts = (Opt(OptOps.UPCAST, 0, 2), # not the warp
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UPCAST, 1, 2),
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2))
+  # new TC should just be able to extract from this and swizzle as needed
+  opts = (Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 1, 8), Opt(OptOps.UNROLL, 0, 8))
+  c = (a@b)
+  ast = c.schedule()[-1].ast
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
@@ -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.schedule',
+                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize',
                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
+                  'tinygrad.runtime.support.nv'],
      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.18.0",
+            "onnx==1.17.0",
            "onnx2torch",
            "onnxruntime",
            "opencv-python",
@@ -1,8 +1,7 @@
 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(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
+    disk_llama = Tensor(pathlib.Path("/raid/weights/LLaMA-3/8B/consolidated.00.pth"))
    device_llama = disk_llama.to(Device.DEFAULT).realize()
@@ -1,21 +1,24 @@
 import time, sys, hashlib
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from onnx.helper import tensor_dtype_to_np_dtype
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 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__":
-  run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))
+  onnx_model = onnx_load(onnx_path := fetch(OPENPILOT_MODEL))
+  run_onnx = OnnxRunner(onnx_model)

  Tensor.manual_seed(100)
-  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()}
+  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()}
  new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}

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

 if __name__ == "__main__":
  t = Tensor.empty(81920, 4096, dtype=dtypes.half)
@@ -24,5 +23,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:=get_program(k.get_optimized_ast(), k.opts))
+  run = CompiledRunner(prg:=k.to_program())
  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, get_program
+from tinygrad.engine.realize import CompiledRunner
 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(get_program(k.get_optimized_ast(), k.opts))
+prg = CompiledRunner(k.to_program())

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

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

  fn = fetch(MODELS[m])
-  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()}
+  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()}
  assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"

  # print input names
-  if DEBUG >= 2: print(list(runner.graph_inputs))
+  if DEBUG >= 2: print([inp.name for inp in onnx_model.graph.input if inp.name not in excluded])
  for device in devices:
    Device.DEFAULT = device
    inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
-    tinygrad_model = runner.to(device)
+    tinygrad_model = OnnxRunner(onnx_model)
    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
@@ -106,12 +107,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 rely on automatic dtype promotion for converting half weights inside the graph
+      # this will reply 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 = runner.to(device)
+      tinygrad_model = OnnxRunner(onnx_model)
      tinygrad_out = tinygrad_model(inputs)

      ort_sess = ort.InferenceSession(str(fn), ort_options, ["CPUExecutionProvider"])
@@ -21,13 +21,12 @@ class FakeAM:
  def __init__(self):
    self.is_booting, self.smi_dev = True, False
    self.pcidev = FakePCIDev()
-    self.vram_size = (512 << 20)
-    self.vram_mv = memoryview(bytearray(self.vram_size))
+    self.vram_mv = memoryview(bytearray(4 << 30))
    self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
    self.gmc = FakeGMC(self)
-    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.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.is_booting = False
    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
  def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
@@ -56,8 +55,6 @@ 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):
@@ -68,9 +65,10 @@ class TestAMPageTable(unittest.TestCase):

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

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

      total_covered = 0
@@ -87,7 +85,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(helper_va(va), sz)
+      mm.unmap_range(va, sz)

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

    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      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)
+      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)

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

    for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
-      exteranl_va = helper_va(va)
+      exteranl_va = va + AMMemoryManager.va_allocator.base
      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(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)

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

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

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

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

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

    # offset from start
    for off in [0, 0x3000, 0x10000]:
-      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)
+      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)

    # offset from end
    for off in [0x1000, 0x20000]:
-      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)
+      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)

  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.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
+  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

  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.to_program()
+    lin.linearize()

    reference = """
 #include <hip/hip_common.h>
@@ -3,7 +3,8 @@ 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 tinygrad.engine.realize import get_runner, CompiledRunner, get_program
+from test.test_linearizer_failures import helper_test_lin
+from tinygrad.engine.realize import get_runner, CompiledRunner
 from test.external.fuzz_linearizer import get_fuzz_rawbufs

 from tinygrad.opt.kernel import Kernel
@@ -23,6 +24,11 @@ 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
@@ -30,7 +36,7 @@ class TestNV(unittest.TestCase):
      lin = Kernel(ast)
      lin.apply_opts(opts)
      rawbufs = get_fuzz_rawbufs(lin)
-      prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+      prg = CompiledRunner(lin.to_program())
      prg(rawbufs, {}, wait=True)
    self.assertEqual(str(cm.exception), "This is a runtime error message")

@@ -1,4 +1,4 @@
-import unittest
+import tempfile, unittest
 from typing import Any, Tuple
 from onnx.backend.base import Backend, BackendRep
 import onnx.backend.test
@@ -6,11 +6,12 @@ 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__()
@@ -29,8 +30,11 @@ 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)
-    model = Tensor(model.SerializeToString(), device="PYTHON")
-    run_onnx = OnnxRunner(model)
+    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)
    return TinygradModel(run_onnx, net_feed_input)

  @classmethod
@@ -40,6 +44,9 @@ 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')

@@ -87,7 +94,6 @@ 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')
@@ -99,12 +105,10 @@ 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_*')
@@ -184,11 +188,6 @@ 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,17 +4,12 @@

 from typing import Any
 import unittest, onnx, tempfile
-from tinygrad import dtypes, Tensor
+from tinygrad import dtypes
 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
@@ -23,14 +18,7 @@ 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()}")
-    # 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)])
+    model = onnx.helper.make_model(graph, producer_name=f"test_{op.lower()}")
    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):
@@ -100,8 +88,7 @@ class TestMainOnnxOps(TestOnnxOps):
    attributes = {"detect_negative":1, "detect_positive":1}
    outputs = ["y"]
    model = self.helper_build_model("IsInf", inputs, attributes, outputs)
-    runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
-    outputs = runner(inputs)
+    outputs = OnnxRunner(model)(inputs)
    assert outputs["y"].dtype is dtypes.bool

  def test_quantize_linear(self):
@@ -216,7 +203,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(Tensor(model.SerializeToString(), device="PYTHON"))
+    runner = OnnxRunner(model)
    tiny_out = runner(inps)
    onnx_out = onnx_fxn(**inps, **opts)
    for (nm, t_out), o_out in  zip(tiny_out.items(), onnx_out):
@@ -251,7 +238,6 @@ 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,139 +1,77 @@
-import unittest, onnx, tempfile, pathlib
-import numpy as np
-from tinygrad import dtypes, Tensor
-from tinygrad.uop.ops import Ops
+import unittest, onnx, tempfile
+from tinygrad import dtypes
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from tinygrad.device import is_dtype_supported
 from extra.onnx import data_types
-from tinygrad.frontend.onnx import OnnxRunner
-from hypothesis import given, strategies as st
+from hypothesis import given, settings, strategies as st
+import numpy as np

-# 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)}
+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)]

 class TestOnnxRunnerDtypes(unittest.TestCase):
-  """
-  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_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)

-  @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_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_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)
+  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_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)
+  @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)

 if __name__ == '__main__':
  unittest.main()
@@ -1,39 +0,0 @@
-from transformers import AutoTokenizer
-from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer
-from tinygrad.helpers import tqdm, getenv
-
-# use ALLOW_FAILED=-1 to go over the entire dataset without printing.
-if __name__ == "__main__":
-  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  vocab_words = [ word for word, _ in sorted(base_tokenizer.get_vocab().items(), key=lambda t: t[1]) ]
-  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
-  simple_tokenizer = SimpleTokenizer(vocab_words)
-
-  color_codes = [ 91, 92, 94, 93, 95 ]
-  def color_tokens(tids): return "".join(f"\033[{color_codes[i%len(color_codes)]}m{inv_vocab[t]}" for i, t in enumerate(tids)) + "\033[0m"
-
-  ds = load_dataset("OpenAssistant/oasst1")
-  allow_failed = getenv("ALLOW_FAILED", 10)
-
-  fail_count, total = 0, 0
-
-  for idx, el in enumerate(tqdm(ds["train"])):
-    total += 1
-
-    try: simple_tokens = tuple(simple_tokenizer.encode(el["text"]))
-    except RuntimeError: simple_tokens = ()
-    base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False))
-
-    if simple_tokens != base_tokens:
-      fail_count += 1
-      allow_failed -= 1
-
-      if allow_failed >= 0:
-        print(f"tokens mismatch at index: {idx}.\n")
-
-        print("simple:  ", color_tokens(simple_tokens))
-        print("official:", color_tokens(base_tokens) + "\n")
-
-      if allow_failed == 0: break
-  print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
@@ -0,0 +1,121 @@
+# 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.optional import fast_idiv
+from tinygrad.codegen.devectorizer import fast_idiv
 random.seed(42)

 z3_renderer = PatternMatcher([
@@ -3,7 +3,6 @@ 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:
@@ -94,7 +93,7 @@ def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]:

  # TODO: images needs required_optimization
  try:
-    prg = CompiledRunner(get_program(lin.get_optimized_ast(), lin.opts))
+    prg = CompiledRunner(lin.to_program())
  except KeyboardInterrupt: raise
  except Exception:
    traceback.print_exc()
@@ -115,7 +114,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.bufs)
+  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.membufs)

  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
  try:
@@ -207,7 +206,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(get_program(test_lin.get_optimized_ast(), test_lin.opts).uops)
+        try: tuops = tuplize_uops(test_lin.linearize().uops)
        except KeyboardInterrupt: raise
        except BaseException as e:
          print(test_lin.ast)
@@ -1,20 +1,12 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools
 from typing import Callable, Any
-
-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))
+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

 # *** process replay settings

@@ -28,11 +20,12 @@ early_stop = multiprocessing.Event()
 logging.basicConfig(level=logging.INFO, format="%(message)s")
 MAX_LINES = 500
 def trunc_log(x):
-  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"]
+  if len(lines:=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
@@ -48,15 +41,9 @@ 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, ...]]:
-  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)
+  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)

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

@@ -4,7 +4,7 @@ from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.opt.search import bufs_from_lin
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.engine.realize import CompiledRunner
 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.bufs)
+      has_bf16 = any(b.dtype == dtypes.bfloat16 for b in amdlin.membufs)

-      amd_prg = CompiledRunner(get_program(amdlin.get_optimized_ast(), amdlin.opts))
+      amd_prg = CompiledRunner(amdlin.to_program())
      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(get_program(amlin.get_optimized_ast(), amlin.opts))
+      am_prg = CompiledRunner(amlin.to_program())
      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(get_program(cpulin.get_optimized_ast(), cpulin.opts))
+      cpu_prg = CompiledRunner(cpulin.to_program())
      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, get_program
+from tinygrad.engine.realize import CompiledRunner
 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.bufs)
+    has_bf16 = any(b.dtype == dtypes.bfloat16 for b in culin.membufs)

-    cuda_prg = CompiledRunner(get_program(culin.get_optimized_ast(), culin.opts))
+    cuda_prg = CompiledRunner(culin.to_program())
    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(get_program(nvlin.get_optimized_ast(), nvlin.opts))
+    nv_prg = CompiledRunner(nvlin.to_program())
    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, get_program
+from tinygrad.engine.realize import CompiledRunner
 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(get_program(lin.get_optimized_ast(), lin.opts))
+    cuda_prg = CompiledRunner(lin.to_program())

    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(get_program(lin.get_optimized_ast(), lin.opts))
+    ptx_prg = CompiledRunner(lin.to_program())

    # 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
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 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 = fetch(OPENPILOT_MODEL)
+    onnx_model = onnx_load(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 = fetch(OPENPILOT_MODEL)
+    onnx_model = onnx_load(fetch(OPENPILOT_MODEL))
    run_onnx = OnnxRunner(onnx_model)
    print("got run_onnx")
    inputs = {
@@ -121,9 +121,10 @@ class TestOnnxModel(unittest.TestCase):
      input_name, input_new)

  def _test_model(self, fn, input_name, input_new, debug=False):
-    run_onnx = OnnxRunner(fn)
+    onnx_model = onnx_load(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, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem
 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 = get_program(k.get_optimized_ast(), k.opts)
+    p = k.to_program()
    print(p.name)
    #print(p.src)
    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
@@ -1,9 +1,10 @@
-import unittest, numpy as np
-from tinygrad import Tensor, Device, TinyJit
+import unittest
+from tinygrad import Tensor
+from tinygrad import Device
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory

-N = 256 if CI else 4096
+N = 4096
 class TestCopySpeed(unittest.TestCase):
  @classmethod
  def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -48,32 +49,6 @@ 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):
@@ -1,32 +0,0 @@
-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()
@@ -96,14 +96,6 @@ class TestDType(unittest.TestCase):
     get_available_cast_dtypes(self.DTYPE)
    ))

-  @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
-  @unittest.skipIf(getenv("PTX"), "skip for now")
-  def test_uint_overflow(self):
-    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = dtypes.max(self.DTYPE)
-    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
-    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
-
  def test_dtypes_fields(self):
    fields = dtypes.fields()
    self.assertIn("float", fields)
@@ -94,20 +94,42 @@ 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 self.assertRaises(ValueError):
-      with Tensor.train():
-        Tensor.ones(10).dropout(-0.1)
+    with Tensor.train():
+      out = Tensor.ones(10).dropout(-0.1)
+      np.testing.assert_allclose(out.numpy(), np.ones(10))

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

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

      with self.subTest(name=str(queue_type)):
-        fake_signal = TestHCQ.d0.new_signal()
+        fake_signal = TestHCQ.d0.signal_t()
        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 or Device.DEFAULT in {"CPU", "LLVM"}, "Can't handle async update on MOCKGPU for now")
+  @unittest.skipIf(MOCKGPU, "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.new_signal()
+        fake_signal = TestHCQ.d0.signal_t()
        queue_type().wait(fake_signal, 1) \
                    .signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)

@@ -101,7 +100,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.new_signal()
+        fake_signal = TestHCQ.d0.signal_t()
        q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)

        fake_signal.value = 0x30
@@ -137,7 +136,6 @@ 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:])
@@ -155,7 +153,6 @@ 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)]
@@ -166,7 +163,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(get_program(k.get_optimized_ast(), k.opts))
+    runner = CompiledRunner(k.to_program())

    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()
@@ -295,7 +292,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.new_signal()
+        fake_signal = TestHCQ.d0.signal_t()
        q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
        q.bind(TestHCQ.d0)

@@ -312,7 +309,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.new_signal(value=0), value=0xfff) \
+    TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0.signal_t(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) \
@@ -326,7 +323,7 @@ class TestHCQ(unittest.TestCase):

  # Test profile api
  def test_speed_exec_time(self):
-    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
+    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
    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) \
@@ -338,7 +335,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 or Device.DEFAULT in {"CPU", "LLVM"} else 100)
+    assert 0.1 <= et <= (15000 if MOCKGPU 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")
@@ -348,7 +345,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.new_signal(), TestHCQ.d0.new_signal()
+    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf.va_addr, b._buf.va_addr, SZ) \
                                .timestamp(sig_en) \
@@ -375,7 +372,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.new_signal(), TestHCQ.d0.new_signal()
+    sig_st, sig_en = TestHCQ.d0.signal_t(), TestHCQ.d0.signal_t()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf.va_addr, b._buf.va_addr, SZ) \
                                .timestamp(sig_en) \
@@ -513,31 +510,6 @@ 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")
@@ -558,28 +530,10 @@ class TestHCQ(unittest.TestCase):
    try: nv_dev = Device["NV"]
    except Exception: self.skipTest("no NV device, test skipped")

-    x = amd_dev.new_signal()
-    y = nv_dev.new_signal()
+    x = amd_dev.signal_t()
+    y = nv_dev.signal_t()
    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()
@@ -1,66 +0,0 @@
-import numpy as np
-import unittest
-from tinygrad import Tensor
-from typing import List
-import functools
-
-def orthogonality_helper(A:Tensor,tolerance=1.0e-5):
-  b_shape,m = A.shape[0:-2],A.shape[-2]  #outer dimension should be the dim along orthogonality
-  A_identity = (Tensor.eye(m).reshape((1,) * len(b_shape)+(m,m)).expand(b_shape+(m,m)))
-  np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)
-
-def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
-  reconstructed_tensor = functools.reduce(Tensor.matmul, A)
-  np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
-
-class TestLinAlg(unittest.TestCase):
-
-  def test_svd_general(self):
-    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
-    for size in sizes:
-      a = Tensor.randn(size).realize()
-      U,S,V = Tensor.svd(a)
-      b_shape,m,n = size[0:-2],size[-2],size[-1]
-      k = min(m,n)
-      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-      s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
-      orthogonality_helper(U)
-      orthogonality_helper(V)
-      reconstruction_helper([U,s_diag,V],a)
-
-  def test_svd_nonfull(self):
-    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
-    for size in sizes:
-      a = Tensor.randn(size).realize()
-      U,S,V = Tensor.svd(a,full_matrices=False)
-      b_shape,m,n = size[0:-2],size[-2],size[-1]
-      k = min(m,n)
-      s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
-      #reduced U,V is only orthogonal along smaller dim
-      if (m < n): orthogonality_helper(U),orthogonality_helper(V)
-      else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
-      reconstruction_helper([U,s_diag,V],a)
-
-  @unittest.skip("very big. recommend wrapping with TinyJit around inner function")
-  def test_svd_large(self):
-    size = (1024,1024)
-    a = Tensor.randn(size).realize()
-    U,S,V = Tensor.svd(a)
-    b_shape,m,n = size[0:-2],size[-2],size[-1]
-    k = min(m,n)
-    s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
-    s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
-    orthogonality_helper(U,tolerance=1.0e-3)
-    orthogonality_helper(V,tolerance=1.0e-3)
-    reconstruction_helper([U,s_diag,V],a,tolerance=1.0e-3)
-
-  def test_qr_general(self):
-    sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
-    for size in sizes:
-      a = Tensor.randn(size).realize()
-      Q,R = Tensor.qr(a)
-      orthogonality_helper(Q)
-      reconstruction_helper([Q,R],a)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,8 +1,9 @@
+from typing import Union
 import numpy as np
 import unittest
 from dataclasses import replace

-from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
+from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel
 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
@@ -12,9 +13,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, AddrSpace
+from tinygrad.dtype import DType, dtypes

-def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
+def helper_realized_ast(r:Union[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
@@ -32,8 +33,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(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"
+  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"
  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
@@ -99,9 +100,8 @@ 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 uops if u.op is Ops.STORE]
+    stores = [u for u in lin.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,6 +114,27 @@ 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()
@@ -127,85 +148,76 @@ 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]
-    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    ranges = [i for i,u in enumerate(lin.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 uops[ranges[0]:ranges[1]])
+    #assert any(x.op is Ops.LOAD for x in lin.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]
-    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    ranges = [i for i,u in enumerate(lin.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 uops[ranges[1]:ranges[2]])
+    #assert any(x.op is Ops.LOAD for x in lin.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]
-    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> ASSIGN
-    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]:])
+    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]:])

  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]
-    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    ranges = [i for i,u in enumerate(lin.uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> ASSIGN
-    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
+    assert len([x for x in lin.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]
-    uops = get_program(lin.get_optimized_ast(), lin.opts).uops
-    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
+    ranges = [i for i,u in enumerate(lin.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 uops[ranges[0]-2].op is Ops.LOAD
+    #  assert lin.uops[ranges[0]-2].op is Ops.LOAD
    #  assert ranges[1] == ranges[0]+6
-    #  assert [x.op for x in uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
+    #  assert [x.op for x in lin.uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
    # LOAD -> RANGE -> LOAD -> ALU -> RANGE -> ASSIGN
    #else:
-    #  assert uops[ranges[0]-2].op is Ops.LOAD
+    #  assert lin.uops[ranges[0]-2].op is Ops.LOAD
    #  assert ranges[1] == ranges[0]+3
-    #  assert [x.op for x in uops[ranges[1]-2:ranges[1]]] == [Ops.LOAD, Ops.ALU]
+    #  assert [x.op for x in lin.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(uops) if u.op is Ops.ENDRANGE)
+    end = max(i for i,u in enumerate(lin.uops) if u.op is Ops.ENDRANGE)
    # the INDEX can be first
-    assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU
+    assert lin.uops[end+1].op in GroupOp.ALU or lin.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(uops) if u.op is Ops.ENDRANGE)
+    end = max(i for i,u in enumerate(lin.uops) if u.op is Ops.ENDRANGE)
    # the INDEX can be first
-    assert uops[end+1].op in GroupOp.ALU or uops[end+2].op in GroupOp.ALU
+    assert lin.uops[end+1].op in GroupOp.ALU or lin.uops[end+2].op in GroupOp.ALU

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

    k = Kernel(r.schedule()[-1].ast)
-    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])
+    k.upcast()
+    k.linearize()
+    num_loads = len([uop for uop in k.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?"

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

    k = Kernel(r.schedule()[-1].ast)
-    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])
+    k.upcast()
+    k.linearize()
+    num_ops = len([uop for uop in k.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")
@@ -239,14 +251,14 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.conv2d(x,w,padding=1).relu()

    k = Kernel(r.schedule()[-1].ast)
-    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]
+    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]
    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
@@ -255,7 +267,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    self.assertEqual(len(prg.src.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -270,13 +282,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 and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in program.uops if u.op is Ops.STORE]

    # 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):
@@ -284,9 +296,9 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.stack(a, b)

    k = Kernel(r.schedule()[-1].ast)
-    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])
+    k.upcast()
+    k.linearize()
+    num_ops = len([uop for uop in k.uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

  def test_sum_acc_dtype(self):
@@ -298,7 +310,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.base == acc_dtype
+        assert local[0].dtype == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
@@ -306,7 +318,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]
-      self.assertEqual(local[0].dtype.base, expected_dtype)
+      assert local[0].dtype == expected_dtype

    tests = (
      (dtypes.float16, None, dtypes.float),
@@ -344,7 +356,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 = get_program(kernel.get_optimized_ast(), kernel.opts)
+      prg = kernel.to_program()
      if Device.DEFAULT == "LLVM":
        assert "0x201000" in prg.src
      elif Device.DEFAULT == "AMD" and AMD_LLVM:
@@ -435,7 +447,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 get_program(k.get_optimized_ast(), k.opts).uops:
+    for u in k.uops:
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.ASSIGN

@@ -446,7 +458,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 get_program(k.get_optimized_ast(), k.opts).uops:
+    for u in k.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
@@ -459,7 +471,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 get_program(k.get_optimized_ast(), k.opts).uops:
+    for u in k.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
@@ -470,14 +482,13 @@ 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
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
-    for u in uops:
+    for u in k.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(uops) if x.op is Ops.ENDRANGE][0]
-        assert end_range < uops.index(u)
+        end_range = [i for i, x in enumerate(k.uops) if x.op is Ops.ENDRANGE][0]
+        assert end_range < k.uops.index(u)

  def test_grouped_dims(self):
    def _assert_grouped_dims(prefix, dims, max_sizes, reverse_dims, expected_sizes, assert_same_length = True):
@@ -555,8 +566,7 @@ 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]
-    uops = get_program(k.get_optimized_ast(), k.opts).uops
-    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
+    idxs = dedup([uop for uop in k.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
@@ -595,7 +605,7 @@ class TestLinearizer(unittest.TestCase):
  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      k = helper_linearizer_opt(t)[-1]
-      uops = get_program(k.get_optimized_ast(), k.opts).uops
+      uops = list(k.linearize().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)]
@@ -612,7 +622,6 @@ 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);
@@ -626,9 +635,8 @@ 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 uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    store_vals = [u.src[-1] for u in k.uops if u.op is Ops.STORE]
    for val in store_vals:
      assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE

@@ -651,7 +659,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 get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[-1] for u in k.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")
@@ -664,35 +672,32 @@ 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])
-    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]
+    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]
    # 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 uops if u.op is Ops.IF and u.src[-1] == barrier]) == 1
+    assert len([u for u in k.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]
-    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]
+    stores = [u for u in k.uops if u.op is Ops.STORE]

    # 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")
@@ -710,8 +715,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 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)
+    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)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -728,19 +733,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 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
+    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

@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()
@@ -763,11 +768,13 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    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
+    k.shift_to(0, 4)  # float4 dimension
+    k.shift_to(0, 2, insert_before=k.shape_len-1)
+    k.upcast()
+    k.upcast()
+    k.linearize()

-    assert TestFloat4.count_float4(uops) == (4, 2)
+    assert TestFloat4.count_float4(k.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):
@@ -778,17 +785,20 @@ class TestFloat4(unittest.TestCase):

      s = c.schedule()[0]
      k = Kernel(s.ast)
-      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
+      k.shift_to(0, 4)
+      k.shift_to(0, shift, insert_before=k.shape_len-1)
+      k.upcast()
+      k.upcast()
+      k.linearize()
+      return k

    sizes = [12, 8, 16]
    shifts = [3, 2, 4]
-    expected_upcast_size = [4, 8, 16]
+    excepted_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]), expected_upcast_size[i]) == expected_output[i]
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]), excepted_upcast_size[i]) == expected_output[i]

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

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    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
+    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()

-    assert TestFloat4.count_float4(uops) == (0, 2)
+    assert TestFloat4.count_float4(k.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):
@@ -826,17 +838,20 @@ class TestFloat4(unittest.TestCase):

      s = c.schedule()[0]
      k = Kernel(s.ast)
-      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
+      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

    sizes = [13, 9, 17]
    shifts = [3, 2, 4]
-    expected_upcast_size = [4, 8, 16]
+    excepted_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]), expected_upcast_size[i]) == expected_output[i]
+      assert TestFloat4.count_float4(kernel_for_shape(sizes[i], shifts[i]).uops, excepted_upcast_size[i]) == expected_output[i]

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

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

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

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

    s = c.schedule()[0]
    k = Kernel(s.ast)
-    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
+    k.upcast()
+    k.upcast()
+    k.linearize()

-    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
+    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)

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

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

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

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

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

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

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

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

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -1061,7 +1094,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:Tensor|list[Tensor], *args, **kwargs):
+def helper_linearizer_opt(r:Union[Tensor, list[Tensor]], *args, **kwargs):
  realized_ast, real_bufs = helper_realized_ast(r)
  return _helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)

@@ -1081,7 +1114,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(get_program(k.get_optimized_ast(), k.opts), device=device))
+  def get_prg(k:Kernel): return CompiledRunner(replace(k.to_program(), device=device))

  def check_opt(opts, create_k, expected_color_size):
    k = create_k()
@@ -1338,9 +1371,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, 1, 8)]])
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 0, 0), Opt(OptOps.PADTO, 2, 8)]])
    with self.assertRaises(KernelOptError):
-      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 1, 8)]])
+      helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 2, 8)]])
    with self.assertRaises(KernelOptError):
      helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])

@@ -1443,7 +1476,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 the first reduce
+      # check to ensure local_dims are stable for full UNROLL of 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,7 +10,6 @@ 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")
@@ -38,12 +37,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    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 prg.uops if u.op is Ops.IF]) == 1, "must have a single IF"
+    assert len([u for u in k.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):
@@ -77,12 +76,11 @@ 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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    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=(
@@ -95,9 +93,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    print(prg.src)
-    if_uops = [u for u in prg.uops if u.op is Ops.IF]
+    if_uops = [u for u in k.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)
@@ -136,7 +134,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    print(prg.src)

  @unittest.expectedFailure
@@ -165,7 +163,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    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}!"
@@ -189,7 +187,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 = get_program(k.get_optimized_ast(), k.opts)
+    prg = k.to_program()
    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,6 +1,7 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad import dtypes
+from tinygrad import dtypes, Device
+from tinygrad.helpers import CI
 from tinygrad.opt.kernel import Kernel
 from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer
@@ -13,6 +14,7 @@ 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)
@@ -161,5 +163,33 @@ 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()
@@ -1126,7 +1126,6 @@ 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
@@ -13,6 +13,29 @@ from test.helpers import not_support_multi_device

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

  def test_meshgrid(self):
    x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.], requires_grad=True)
@@ -548,7 +546,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(Device.DEFAULT == "AMD" and AMD_LLVM, "AMD with LLVM backend generate rcp in FP division causes trunc/floor errors")
+  @unittest.skipIf(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
@@ -576,7 +574,8 @@ 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"), forward_only=True, vals=[[5], [0]], expected=RuntimeError)
+    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)

  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]])
@@ -588,6 +587,12 @@ 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)
@@ -736,7 +741,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)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], torch.bitwise_xor, Tensor.bitwise_xor, expected=RuntimeError)

  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
@@ -753,7 +758,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)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], torch.bitwise_and, Tensor.bitwise_and, expected=RuntimeError)

  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
@@ -768,7 +773,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)], lambda x,y: x.bitwise_or(y), expected=RuntimeError)
+    self.helper_test_exception([(4), (4)], torch.bitwise_or, Tensor.bitwise_or, expected=RuntimeError)

  def test_bitwise_not(self):
    data = [[1,-8,1],[32,1,6]]
@@ -783,7 +788,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)], lambda x: x.bitwise_not(), expected=RuntimeError)
+    self.helper_test_exception([(4)], torch.bitwise_not, Tensor.bitwise_not, expected=RuntimeError)

  def test_lshift(self):
    data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -821,7 +826,6 @@ 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())
@@ -829,7 +833,6 @@ 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)
@@ -1105,12 +1108,6 @@ 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)
@@ -1130,7 +1127,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), expected=(RuntimeError, ValueError))
+    self.helper_test_exception([(4)], lambda x: x.topk(5), lambda x: x.topk(5), expected=(RuntimeError, ValueError))

  def test_einsum(self):
    # matrix transpose
@@ -1284,8 +1281,7 @@ 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
-  or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
+  @unittest.skipIf(CI and Device.DEFAULT in ["NV", "LLVM", "GPU", "CUDA"] or IMAGE, "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):
@@ -1333,9 +1329,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), expected=IndexError)
-    self.helper_test_exception([()], lambda x: x.sum(1), expected=IndexError)
-    self.helper_test_exception([()], lambda x: x.sum((1,)), expected=IndexError)
+    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)

  def test_sum_dtype_arg(self):
    helper_test_op([(45,3)], lambda x: x.sum(), lambda x: x.sum(dtype=dtypes.float32))
@@ -1850,9 +1846,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)), 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)
+    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)

  def test_reshape(self):
    helper_test_op([(4,3,6,6)], lambda x: x.reshape((12,6,6)))
@@ -1863,8 +1859,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)), expected=RuntimeError)
-    self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,-1,2)), expected=RuntimeError)
+    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)

    with self.assertRaises(ValueError):
      x = Tensor.ones((4,3,6,6))
@@ -1889,16 +1885,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)), 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)
+    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)

  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: x.squeeze(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: 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)
    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())
@@ -1906,9 +1902,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: 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)
+    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)

  def test_unsqueeze(self):
    helper_test_op([(4,3,6,6)], lambda x: x.unsqueeze(0))
@@ -1931,31 +1927,20 @@ 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: 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)
+    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)))

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

  def test_stack(self):
    for dim in range(-1, 3):
-      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))
+      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))

    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())
@@ -2654,7 +2636,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), expected=RuntimeError)
+    self.helper_test_exception([(45,65)], lambda x: x.clip(None, None), lambda x: x.clip(None, None), RuntimeError)

  def test_matvecmat(self):
    helper_test_op([(1,128), (128,128), (128,128)], lambda x,y,z: (x@y).relu()@z)
@@ -2915,17 +2897,6 @@ 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)))
@@ -2973,39 +2944,6 @@ 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)],
@@ -3084,20 +3022,6 @@ 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, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem

 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(get_program(k.get_optimized_ast(), k.opts))
+    run = CompiledRunner(k.to_program())
    ExecItem(run, si.bufs).run()
    test = si.bufs[0].numpy().reshape(self.res.shape)
    np.testing.assert_allclose(self.res, test, atol=1e-4)
@@ -1,148 +0,0 @@
-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, time, gc
+import unittest, struct, contextlib, statistics
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events
-from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
+from tinygrad.helpers import CI, getenv, Context
+from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileRangeEvent, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner

@@ -10,11 +10,7 @@ MOCKGPU = getenv("MOCKGPU")
@contextlib.contextmanager
 def helper_collect_profile(*devs):
  for dev in devs: dev.synchronize()
-  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()
+  Compiled.profile_events = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]

  profile_list = []
  with Context(PROFILE=1):
@@ -22,7 +18,6 @@ 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"
@@ -30,10 +25,7 @@ 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]

-# 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")
+@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled) or Device.DEFAULT in {"METAL"}, "HCQ device required to run")
 class TestProfiler(unittest.TestCase):
  @classmethod
  def setUpClass(self):
@@ -81,15 +73,13 @@ class TestProfiler(unittest.TestCase):
    evs = [x for x in profile if isinstance(x, ProfileRangeEvent)]

    assert len(evs) == 3, "3 kernel runs are expected"
-    # 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)
+    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"

-    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"
+    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"]
@@ -169,45 +159,5 @@ 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, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item
 from tinygrad.opt.search import bufs_from_lin
 from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import

@@ -32,9 +32,7 @@ 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")
-  # 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)])
+  model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
  onnx.save_model(model_def, model_path)
  return model_path

@@ -43,7 +41,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 = get_program(k.get_optimized_ast(), k.opts)
+  prg = k.to_program()
  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"))
@@ -65,7 +63,6 @@ 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):
@@ -73,9 +70,10 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
      import onnx # noqa: F401 # pylint: disable=unused-import
    except ImportError:
      raise unittest.SkipTest()
-    from tinygrad.frontend.onnx import OnnxRunner
+    from tinygrad.frontend.onnx import OnnxRunner, onnx_load
    out_file = get_quantized_model(sz)
-    run_onnx = OnnxRunner(out_file)
+    onnx_model = onnx_load(out_file)
+    run_onnx = OnnxRunner(onnx_model)
    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()
@@ -299,7 +297,7 @@ class TestDSPCache(unittest.TestCase):
    with Context(DEVECTORIZE=0, QUANTIZE=1):
      k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
      k.apply_opts(opts)
-      prg = get_program(k.get_optimized_ast(), k.opts)
+      prg = k.to_program()
      #print(prg.src)

    new_src = """
@@ -308,7 +306,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_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(
+__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(
 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,9 +251,6 @@ 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,9 +1,8 @@
 import numpy as np, unittest, string
 from hypothesis import given, strategies as st
-from tinygrad import Device, Tensor, TinyJit, dtypes
+from tinygrad import Device, Tensor, TinyJit
 from tinygrad.runtime.ops_remote import RemoteDevice, parse_hosts
-from tinygrad.runtime.graph.remote import RemoteGraph
-from tinygrad.helpers import LazySeq, all_same, Context
+from tinygrad.helpers import LazySeq, all_same

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

-  @Context(JIT_BATCH_SIZE=2**32)
-  def test_multihost_matmul_jit_graph(self):
+  # 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):
    @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)
@@ -29,38 +27,6 @@ class TestRemoteMultiHost(unittest.TestCase):
      c = do(a, b)
      np.testing.assert_allclose(nc, c.numpy(), rtol=3e-2, atol=1e-4) # tolerances from extra/gemm/simple_matmul.py

-    # Verify that everything is in one big cross-host graph
-    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
-
-  @Context(JIT_BATCH_SIZE=2**32)
-  def test_multihost_aware_schedule(self):
-    @TinyJit
-    def do(*ts:Tensor):
-      acc = Tensor.zeros(1, dtype=dtypes.float32)
-      for t in ts: acc += t.sum()
-      return acc.realize()
-
-    def do_np(*ts:np.ndarray):
-      acc = np.zeros(1, np.float32)
-      for t in ts: acc += t.sum()
-      return acc
-
-    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
-    TS = 64
-    for _ in range(3):
-      inp_np = [np.random.rand(256).astype(np.float32) for _ in range(TS)]
-      inp = [Tensor(inp).shard(ds, 0).contiguous().realize() for inp in inp_np]
-      out_np = do_np(*inp_np)
-      out = do(*inp)
-      np.testing.assert_allclose(out_np, out.numpy(), rtol=3e-2, atol=1e-4)
-
-    # Verify that everything is in one big cross-host graph and that the scheduling is reasonable
-    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
-    # At the time of writing this: 2050 graph breaks without multihost aware scheduling, 14 with it. I've set fail threshold to 28 to not fail on
-    # unrelated scheduling changes. Maybe 2x is a bit too pessimistic, but remote should perform just fine as long as this is not like a half hundred
-    # or more here.
-    self.assertLess(len(do.captured._jit_cache[0].prg.template), 28, "Very bad scheduling! Many unnecesary graph breaks!")
-
 class TestParseHosts(unittest.TestCase):
  def assert_seq(self, result:LazySeq, host:str):
    self.assertIsInstance(result, LazySeq)
@@ -5,7 +5,7 @@
 import unittest
 import numpy as np
 import functools
-from typing import cast
+from typing import List, Optional, Union, 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.schedule.kernelize import merge_views, get_kernelize_map, Kernel
+from tinygrad.kernelize.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:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
+def check_schedule(t:Union[Tensor, List[Tensor], UOp], allowed:int, to_prerealize:Optional[List[Tensor]]=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:Tensor|list[Tensor], cnt:int):
+  def check_schedule(self, xt:Union[Tensor,List[Tensor]], cnt:int):
    with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):
      lst = [xt] if isinstance(xt, Tensor) else xt
      s = Tensor.schedule(*lst)
@@ -176,12 +176,5 @@ 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()
@@ -86,6 +86,7 @@ class TestFuse(unittest.TestCase):
      return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
    self._test_fuse(embedding, a, atol=1e-5)

+  @unittest.skip("still broken")
  def test_flash_attention(self):
    BS = 4
    HEADS = 2
@@ -97,7 +98,7 @@ class TestFuse(unittest.TestCase):
      v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
    # TODO: OPT is breaking things. NOOPT isn't linearizing
    with Context(NOOPT=1):
-      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
+      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v)

 class TestSoftmaxFusion(unittest.TestCase):
  @classmethod
@@ -512,24 +512,6 @@ class TestTinygrad(unittest.TestCase):
    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
                    shell=True, check=True)

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

-  def test_const_bitcast(self):
-    bf = UOp(Ops.CONST, dtypes.float, arg=1.0)
-    out = UOp(Ops.BITCAST, dtypes.uint32, (bf,))
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 1)
-    out = uops[-1]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.arg, 0x3F800000)
-
-  @unittest.expectedFailure
-  def test_const_shape_change_bitcast(self):
-    bf = UOp(Ops.CONST, dtypes.uint8, arg=0x3F)
-    out = UOp(Ops.BITCAST, dtypes.half, (bf,))
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 1)
-
  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
@@ -317,7 +300,6 @@ class TestUOpGraph(unittest.TestCase):
      for uop, const in zip(uops, consts):
        self.assertEqual(uop, const)

-  @unittest.skip("no longer testable standalone")
  def test_wmma_vectorize_fold(self):
    for i in [2, 4, 8]:
      vec = UOp(Ops.VECTORIZE, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
@@ -441,7 +423,6 @@ class TestUOpGraph(unittest.TestCase):
      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-  @given(matchers)
-  def test_no_inf_loop_bottom_up(self, PatternMatcher):
+  def test_no_inf_loop_bottom_up(self):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
    graph_rewrite(a, pm, bottom_up=True)
@@ -1,72 +0,0 @@
-from typing_extensions import Callable
-import hashlib, random, unittest
-from tinygrad import Tensor, Device, getenv, dtypes
-from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import CI
-
-@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
-class TestHashing(unittest.TestCase):
-  def _python_hash_1mb(self, data:bytes):
-    chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
-    chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
-    return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
-
-  @unittest.skipIf(CI, "very slow")
-  def test_abc(self):
-    expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
-    out = Tensor(b"abc").hash()
-    self.assertEqual(bytes(out.data()), expected)
-
-@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
-class TestKeccak(unittest.TestCase):
-  def setUp(self) -> None: random.seed(1337)
-
-  def test_shape_keeping(self):
-    s = (1, 2, 3, 4)
-    for i in range(len(s)):
-      out_shape = Tensor.randint(*s[i:], high=255, dtype=dtypes.uint8).keccak().shape
-      self.assertTupleEqual(s[i:-1], out_shape[:-1])
-
-  def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
-  def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
-  def test_shake_128(self): self._test_preset("shake_128", [167, 168], lambda d: hashlib.shake_128(d).digest(16))
-
-  def _test_preset(self, name: str, special_sizes: list[int], hasher: Callable[[bytes], bytes] | None = None):
-    def default_hasher(d: bytes) -> bytes: return getattr(hashlib, name)(d).digest()
-    if hasher is None: hasher = default_hasher
-
-    for n in (special_sizes + [special_sizes[0] - 1]):
-      a, b = random.randbytes(n), random.randbytes(n)
-
-      ha_ref, hb_ref = hasher(a), hasher(b)
-      tres = Tensor.stack(*(Tensor(d) for d in (a, b))).keccak(name)
-      ha, hb = tres[0].data(), tres[1].data()
-
-      self.assertEqual(ha_ref, ha)
-      self.assertEqual(ha_ref, Tensor(a).keccak(name).data())
-      self.assertEqual(hb_ref, hb)
-
-  def test_referenced(self):
-    # https://www.di-mgt.com.au/sha_testvectors.html
-    self.assertEqual(bytes(Tensor(b"abc").keccak().tolist()),
-                     bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-    self.assertEqual(bytes(Tensor(b"").keccak().tolist()),
-                     bytearray.fromhex("a7ffc6f8bf1ed766 51c14756a061d662 f580ff4de43b49fa 82d80a4b80f8434a"))
-    t = Tensor(b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu").keccak()
-    self.assertEqual(bytes(t.tolist()),
-                     bytearray.fromhex("916f6061fe879741 ca6469b43971dfdb 28b1a32dc36cb325 4e812be27aad1d18"))
-    # TODO: this does not run or very slow
-    # self.assertEqual(bytes(Tensor(b"a" * 1000000).keccak().tolist()),
-    #                  bytearray.fromhex("5c8875ae474a3634 ba4fd55ec85bffd6 61f32aca75c6d699 d0cdcb6c115891c1"))
-
-  def test_long(self):
-    data = b"\x00" * 4
-    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
-
-    data = b"\x00" * (1000 if CI else 4096)
-    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
-
-if __name__ == "__main__":
-  unittest.main()
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	383c0bf05e	notes	2025-07-01 10:36:05 -07:00
geohot	fd2c0e2626	uc warp	2025-06-30 20:54:29 -07:00
George HotzandGitHub	5907f2d443	Merge branch 'master' into warp_fun	2025-06-30 18:18:52 -07:00
George HotzandGitHub	fc59db0aaa	Merge branch 'master' into warp_fun	2025-06-30 15:54:57 -07:00
geohot	256d4403c5	what does order change	2025-06-29 09:26:32 -07:00
George HotzandGitHub	ad823a5199	Merge branch 'master' into warp_fun	2025-06-29 09:13:16 -07:00
geohot	0df4355cd8	upcasted warp experiments	2025-06-29 09:04:37 -07:00