comments

more syntax
2025-10-30 12:27:55 +08:00 · 2025-10-30 12:11:18 +08:00 · 2025-10-30 12:04:26 +08:00 · 2025-10-30 11:22:36 +08:00 · 2025-10-30 11:12:37 +08:00 · 2025-10-30 11:01:53 +08:00
215 changed files with 216318 additions and 141490 deletions
@@ -1,3 +0,0 @@
-[run]
-source = tinygrad
-branch = True
@@ -61,7 +61,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/pyproject.toml') }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}

    # **** Caching downloads ****

@@ -70,13 +70,13 @@ runs:
      uses: actions/cache@v4
      with:
        path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
    - name: Cache downloads (macOS)
      if: inputs.key != '' && runner.os == 'macOS'
      uses: actions/cache@v4
      with:
        path: ~/Library/Caches/tinygrad/downloads/
-        key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: osx-downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}

    # **** Python deps ****

@@ -187,7 +187,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -247,7 +247,7 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -1,7 +1,10 @@
 name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -11,15 +14,15 @@ on:
    branches:
      - master
  pull_request:
-    paths:
+    paths: 
    - 'tinygrad/runtime/autogen/**/*'
  workflow_dispatch:
-    paths:
+    paths: 
    - 'tinygrad/runtime/autogen/**/*'

 jobs:
  autogen:
-    name: In-tree Autogen
+    name: Autogen
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
@@ -31,146 +34,64 @@ jobs:
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
-        llvm: 'true'
        webgpu: 'true'
-        mesa: 'true'
+        llvm: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
    - name: Verify OpenCL autogen
      run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        python3 -c "from tinygrad.runtime.autogen import opencl"
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        ./autogen_stubs.sh opencl
        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
    - name: Verify CUDA autogen
      run: |
-        mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
-        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
+        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
+        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
+        ./autogen_stubs.sh cuda
+        ./autogen_stubs.sh nv
        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
-        diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
-        diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
-        diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
+        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
    - name: Verify AMD autogen
      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
+        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
+        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
+        ./autogen_stubs.sh comgr
+        ./autogen_stubs.sh amd
+        ./autogen_stubs.sh sqtt
        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
+        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-        diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
-        diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
-        diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
-        diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
-        diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
-        diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
-        diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
-        diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
-        diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
    - name: Verify Linux autogen
      run: |
-        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
+        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/kfd.py.bak tinygrad/runtime/autogen/kfd.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-        diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
-        diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
-    - name: Verify LLVM autogen
-      run: |
-        mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        python3 -c "from tinygrad.runtime.autogen import llvm"
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
    - name: Verify WebGPU autogen
      run: |
-        mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        python3 -c "from tinygrad.runtime.autogen import webgpu"
+        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
+        ./autogen_stubs.sh webgpu
        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify Qualcomm autogen
+    - name: Verify LLVM autogen
      run: |
-        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
-        python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
-        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.py
-        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
-    - name: Verify libusb autogen
-      run: |
-        mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libusb"
-        diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
+        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
+        ./autogen_stubs.sh llvm
+        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
    - name: Verify mesa autogen
      run: |
-        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
-        python3 -c "from tinygrad.runtime.autogen import mesa"
+        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
+        ./autogen_stubs.sh mesa
        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
-        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
-  autogen-mac:
-    name: In-tree Autogen (macos)
-    runs-on: macos-14
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        llvm: 'true'
-    - name: Verify macos autogen
-      run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
-  autogen-comgr-3:
-    name: In-tree Autogen (comgr 3)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-    - name: Install autogen support packages
-      run: |
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt -qq update || true
-        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Verify comgr (3) autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr_3"
-        diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
@@ -54,7 +54,7 @@ jobs:
    - name: Print macOS version
      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
@@ -199,7 +199,7 @@ jobs:
    - name: Test speed vs torch
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -320,20 +320,19 @@ jobs:
    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-    # TODO: too slow
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
-    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    # - name: Run full CIFAR training steps w 6 GPUS
-    #   run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
@@ -410,7 +409,7 @@ jobs:
    #    python3 -c "import torch; print(torch.__version__)"
    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
      run: |
        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
@@ -525,18 +524,17 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-    # TODO: too slow
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
-    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    # - name: Run 10 CIFAR training steps w BF16
    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    #- name: Run full CIFAR training steps w 6 GPUS
    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
@@ -625,32 +623,32 @@ jobs:
        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: openpilot compile3 0.9.9 driving_vision
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    # - name: openpilot compile3 0.9.9 driving_policy
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    # - name: openpilot compile3 0.9.9 dmonitoring
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      # TODO: ASSERT_MIN_STEP_TIME=17
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
-    # - name: benchmark MobileNetV2 on DSP
-    #   run: |
-    #     # generate quantized weights
-    #     ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-    #     ln -s /data/home/tiny/tinygrad/testsig-*.so .
-    #     PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
-    #     # benchmark on DSP with NOOPT=1, the devectorizer has issues
-    #     PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+      # TODO: ASSERT_MIN_STEP_TIME=10
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: benchmark MobileNetV2 on DSP
+      run: |
+        # generate quantized weights
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        # benchmark on DSP with NOOPT=1, the devectorizer has issues
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - 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

@@ -706,9 +704,8 @@ jobs:
      run: |
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
    # TODO: enable
    # - 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
@@ -770,9 +767,8 @@ jobs:
        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - 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
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 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)
@@ -22,13 +22,13 @@ jobs:
    - name: Run SDXL with new search
      # TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
      run: |
-        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run SDXL with cached search
      run: |
        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
    - name: Run winograd cifar with cached search
      run: |
        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -20,11 +20,11 @@ jobs:
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
-        pip install setuptools wheel build twine
+        pip install setuptools wheel twine
    - name: Build and publish
      env:
        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
      run: |
-        python -m build
+        python setup.py sdist bdist_wheel
        twine upload dist/*
@@ -1,7 +1,10 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -227,7 +230,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: linting-only
-        python-version: '3.11'
+        python-version: '3.10'
        deps: linting
    - name: Lint bad-indentation and trailing-whitespace with pylint
      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
@@ -240,9 +243,8 @@ jobs:
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-    # broken because of UPatAny
-    #- name: Run TYPED=1
-    #  run: TYPED=1 python -c "import tinygrad"
+    - name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"

  unittest:
    name: Unit Tests
@@ -287,8 +289,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 19000 lines
-      run: MAX_LINE_COUNT=19000 python sz.py
+    - name: Repo line count < 18000 lines
+      run: MAX_LINE_COUNT=18000 python sz.py

  spec:
    strategy:
@@ -306,7 +308,6 @@ jobs:
      with:
        key: spec-unit
        deps: testing_unit
-        python-version: '3.14'
    - name: Test SPEC=2
      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

@@ -342,11 +343,10 @@ jobs:
          key: gpu-image
          deps: testing_minimal
          opencl: 'true'
-      - name: Test CL IMAGE=2 ops
+      - name: Test CL IMAGE=2 ops + training
        run: |
          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-        # TODO: training is broken
-        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -391,7 +391,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp16
        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
@@ -63,5 +63,3 @@ profile_stats
 *.log
 target
 .mypy_cache
-mutants
-.mutmut-cache
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: subset of tests
-        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
@@ -21,38 +21,17 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

-tinygrad is an end-to-end deep learning stack:
+Despite tinygrad's size, it is a fully featured deep learning framework.

- **Tensor library** with autograd
- **IR and compiler** that fuse and lower kernels
- **JIT + graph execution**
- **nn / optim / datasets** for real training
+Due to its extreme simplicity, it is the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.

-It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.
+tinygrad is now beta software, we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.

---
+## Features

-## How tinygrad compares
+### LLaMA and Stable Diffusion

-**PyTorch**
-
- ✅ Similar: eager `Tensor` API, autograd, `optim`, basic datasets and layers.
- ✅ You can write familiar training loops.
- 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable.
-
-**JAX**
-
- ✅ IR-based autodiff over primitives (like JAXPR + XLA).
- ✅ Function-level JIT (`TinyJit`) that captures and replays kernels.
- 🔁 Fewer functional transforms (no full `vmap`/`pmap` yet), but far easier to read.
-
-**TVM**
-
- ✅ Multiple lowering passes, scheduling, and BEAM search over kernels.
- ✅ Device “graphs” for batched execution.
- 🔁 tinygrad also ships the **front-end framework** (tensors, nn, optim), not just the compiler.
-
---
+tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!

 ### Laziness

@@ -0,0 +1,564 @@
+#!/bin/bash -e
+
+# setup instructions for clang2py
+if [[ ! $(clang2py -V) ]]; then
+  pushd .
+  cd /tmp
+  sudo apt-get install -y --no-install-recommends clang
+  pip install --upgrade pip setuptools
+  pip install clang==14.0.6
+  git clone https://github.com/nimlgen/ctypeslib.git
+  cd ctypeslib
+  pip install .
+  clang2py -V
+  popd
+fi
+
+BASE=tinygrad/runtime/autogen/
+
+fixup() {
+  sed -i '1s/^/# mypy: ignore-errors\n/' $1
+  sed -i 's/ *$//' $1
+  grep FIXME_STUB $1 || true
+}
+
+patch_dlopen() {
+  path=$1; shift
+  name=$1; shift
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $path
+PATHS_TO_TRY = [
+$(for p in "$@"; do echo "  $p,"; done)
+]
+def _try_dlopen_$name():
+  library = ctypes.util.find_library("$name")
+  if library: return ctypes.CDLL(library)
+  for candidate in PATHS_TO_TRY:
+    try: return ctypes.CDLL(candidate)
+    except OSError: pass
+  return None
+EOF
+}
+
+generate_opencl() {
+  clang2py /usr/include/CL/cl.h -o $BASE/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
+  fixup $BASE/opencl.py
+  # hot patches
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/opencl.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libOpenCL.so.1')\ctypes.CDLL(ctypes.util.find_library('OpenCL'))\g" $BASE/opencl.py
+  python3 -c "import tinygrad.runtime.autogen.opencl"
+}
+
+generate_hip() {
+  clang2py /opt/rocm/include/hip/hip_ext.h /opt/rocm/include/hip/hiprtc.h \
+  /opt/rocm/include/hip/hip_runtime_api.h /opt/rocm/include/hip/driver_types.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/hip.py -l /opt/rocm/lib/libamdhip64.so
+  echo "hipDeviceProp_t = hipDeviceProp_tR0600" >> $BASE/hip.py
+  echo "hipGetDeviceProperties = hipGetDevicePropertiesR0600" >> $BASE/hip.py
+  fixup $BASE/hip.py
+  # we can trust HIP is always at /opt/rocm/lib
+  #sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhiprtc.so')\ctypes.CDLL(ctypes.util.find_library('hiprtc'))\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamdhip64.so')\ctypes.CDLL(ctypes.util.find_library('amdhip64'))\g" $BASE/hip.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/hip.py
+  sed -i "s\'/opt/rocm/\os.getenv('ROCM_PATH', '/opt/rocm/')+'/\g" $BASE/hip.py
+  python3 -c "import tinygrad.runtime.autogen.hip"
+}
+
+generate_comgr() {
+  clang2py /opt/rocm/include/amd_comgr/amd_comgr.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
+  fixup $BASE/comgr.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
+  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
+  python3 -c "import tinygrad.runtime.autogen.comgr"
+}
+
+generate_kfd() {
+  clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum
+
+  fixup $BASE/kfd.py
+  sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
+  sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
+  sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
+  sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
+  sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
+  sed -i "s/!!/not not /g" $BASE/kfd.py
+  python3 -c "import tinygrad.runtime.autogen.kfd"
+}
+
+generate_cuda() {
+  clang2py /usr/include/cuda.h --clang-args="-D__CUDA_API_VERSION_INTERNAL" -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/cuda.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" $BASE/cuda.py
+  fixup $BASE/cuda.py
+  python3 -c "import tinygrad.runtime.autogen.cuda"
+}
+
+generate_nvrtc() {
+  clang2py /usr/local/cuda/include/nvrtc.h /usr/local/cuda/include/nvJitLink.h -o $BASE/nvrtc.py -l /usr/local/cuda/lib64/libnvrtc.so -l /usr/local/cuda/lib64/libnvJitLink.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvJitLink.so')\ctypes.CDLL(ctypes.util.find_library('nvJitLink'))\g" $BASE/nvrtc.py
+  fixup $BASE/nvrtc.py
+  python3 -c "import tinygrad.runtime.autogen.nvrtc"
+}
+
+generate_nv() {
+  NVKERN_COMMIT_HASH=81fe4fb417c8ac3b9bdcc1d56827d116743892a5
+  NVKERN_SRC=/tmp/open-gpu-kernel-modules-$NVKERN_COMMIT_HASH
+  if [ ! -d "$NVKERN_SRC" ]; then
+    git clone https://github.com/NVIDIA/open-gpu-kernel-modules $NVKERN_SRC
+    pushd .
+    cd $NVKERN_SRC
+    git reset --hard $NVKERN_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    extra/nv_gpu_driver/clc6c0qmd.h \
+    extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc96f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc761.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl83de.h \
+    $NVKERN_SRC/src/nvidia/generated/g_allclasses.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc6c0.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clcdc0.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc6b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc9b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_linux_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/hwref/ampere/ga100/dev_fault.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv_escape.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numbers.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numa.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-unix-nvos-params-wrappers.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/alloc/alloc_channel.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/nvos.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0000/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl2080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl83de/*.h \
+    $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
+  sed -i "s\(0000000001)\1\g" $BASE/nv_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/nv_gpu.py
+  sed -i 's/#\?\s\([A-Za-z0-9_]\+\) = MW ( \([0-9]\+\) : \([0-9]\+\) )/\1 = (\2 , \3)/' $BASE/nv_gpu.py # NVC6C0_QMDV03_00 processing
+  sed -i 's/#\sdef NVC6C0_QMD\([A-Za-z0-9_()]\+\):/def NVC6C0_QMD\1:/' $BASE/nv_gpu.py
+  sed -i 's/#\sdef NVCEC0_QMD\([A-Za-z0-9_()]\+\):/def NVCEC0_QMD\1:/' $BASE/nv_gpu.py
+  sed -E -i -n '/^def (NVCEC0_QMDV05_00_RELEASE)(_ENABLE)\(i\):/{p;s//\1'"0"'\2=\1\2(0)\n\1'"1"'\2=\1\2(1)/;H;b};p;${x;s/^\n//;p}' "$BASE/nv_gpu.py"
+  sed -i 's/#\s*return MW(\([0-9i()*+]\+\):\([0-9i()*+]\+\))/    return (\1 , \2)/' $BASE/nv_gpu.py
+  sed -i 's/#\?\s*\(.*\)\s*=\s*\(NV\)\?BIT\(32\)\?\s*(\s*\([0-9]\+\)\s*)/\1 = (1 << \4)/' $BASE/nv_gpu.py # name = BIT(x) -> name = (1 << x)
+  sed -i "s/UVM_\([A-Za-z0-9_]\+\) = \['i', '(', '\([0-9]\+\)', ')'\]/UVM_\1 = \2/" $BASE/nv_gpu.py # UVM_name = ['i', '(', '<num>', ')'] -> UVM_name = <num>
+
+  # Parse status codes
+  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 \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $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 \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/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/nv.py
+
+  fixup $BASE/nv/nv.py
+  python3 -c "import tinygrad.runtime.autogen.nv.nv"
+}
+
+generate_amd() {
+  # clang2py broken when pass -x c++ to prev headers
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    extra/hip_gpu_driver/nvd.h \
+    extra/hip_gpu_driver/gc_11_0_0_offset.h \
+    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/amd_gpu.py
+
+  fixup $BASE/amd_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/amd_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.amd_gpu"
+}
+
+generate_hsa() {
+  clang2py \
+    /opt/rocm/include/hsa/hsa.h \
+    /opt/rocm/include/hsa/hsa_ext_amd.h \
+    /opt/rocm/include/hsa/amd_hsa_signal.h \
+    /opt/rocm/include/hsa/amd_hsa_queue.h \
+    /opt/rocm/include/hsa/amd_hsa_kernel_code.h \
+    /opt/rocm/include/hsa/hsa_ext_finalize.h /opt/rocm/include/hsa/hsa_ext_image.h \
+    /opt/rocm/include/hsa/hsa_ven_amd_aqlprofile.h \
+    --clang-args="-I/opt/rocm/include" \
+    -o $BASE/hsa.py -l /opt/rocm/lib/libhsa-runtime64.so
+
+  fixup $BASE/hsa.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/hsa.py
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhsa-runtime64.so')\ctypes.CDLL(os.getenv('ROCM_PATH')+'/lib/libhsa-runtime64.so' if os.getenv('ROCM_PATH') else ctypes.util.find_library('hsa-runtime64'))\g" $BASE/hsa.py
+  python3 -c "import tinygrad.runtime.autogen.hsa"
+}
+
+generate_io_uring() {
+  clang2py -k cdefstum \
+    /usr/include/liburing.h \
+    /usr/include/linux/io_uring.h \
+    -o $BASE/io_uring.py
+
+  sed -r '/^#define __NR_io_uring/ s/^#define __(NR_io_uring[^ ]+) (.*)$/\1 = \2/; t; d' /usr/include/asm-generic/unistd.h >> $BASE/io_uring.py # io_uring syscalls numbers
+  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) \
+    $(dpkg -L libc6-dev | grep sys/syscall.h) \
+    /usr/include/string.h \
+    /usr/include/elf.h \
+    /usr/include/unistd.h \
+    /usr/include/asm-generic/mman-common.h \
+    -o $BASE/libc.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
+  sed -i "s\FIXME_STUB\libc\g" $BASE/libc.py
+  sed -i "s\FunctionFactoryStub()\None if (libc_path := ctypes.util.find_library('c')) is None else ctypes.CDLL(libc_path, use_errno=True)\g" $BASE/libc.py
+
+  fixup $BASE/libc.py
+}
+
+generate_llvm() {
+  INC="$(llvm-config-14 --includedir)"
+  clang2py -k cdefstum \
+    $(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
+    "$INC/llvm/Config/Targets.def" \
+    "$INC/llvm/Config/AsmPrinters.def" \
+    "$INC/llvm/Config/AsmParsers.def" \
+    "$INC/llvm/Config/Disassemblers.def" \
+    --clang-args="$(llvm-config-14 --cflags)" \
+    -o "$BASE/llvm.py"
+
+  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
+  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
+
+  fixup "$BASE/llvm.py"
+}
+
+generate_kgsl() {
+  clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
+  fixup $BASE/kgsl.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
+  sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
+  python3 -c "import tinygrad.runtime.autogen.kgsl"
+}
+
+generate_adreno() {
+  clang2py extra/qcom_gpu_driver/a6xx.xml.h -o $BASE/adreno.py -k cestum
+  sed -nE 's/#define ([A-Za-z0-9_]+)__SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1__SHIFT) \& \1__MASK/p' extra/qcom_gpu_driver/a6xx.xml.h >> $BASE/adreno.py
+  fixup $BASE/adreno.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/adreno.py
+  python3 -c "import tinygrad.runtime.autogen.adreno"
+}
+
+generate_qcom() {
+  clang2py -k cdefstum \
+    extra/dsp/include/ion.h \
+    extra/dsp/include/msm_ion.h \
+    extra/dsp/include/adsprpc_shared.h \
+    extra/dsp/include/remote_default.h \
+    extra/dsp/include/apps_std.h \
+    -o $BASE/qcom_dsp.py
+
+  fixup $BASE/qcom_dsp.py
+  python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
+}
+
+generate_pci() {
+  clang2py -k cdefstum \
+    /usr/include/linux/pci_regs.h \
+    -o $BASE/pci.py
+  fixup $BASE/pci.py
+}
+
+generate_vfio() {
+  clang2py -k cdefstum \
+    /usr/include/linux/vfio.h \
+    -o $BASE/vfio.py
+  fixup $BASE/vfio.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
+  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
+  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
+}
+
+generate_am() {
+  AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
+  AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
+  if [ ! -d "$AMKERN_SRC" ]; then
+    git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
+  fi
+  AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
+  AMKERN_INC=$AMKERN_AMD/include/
+
+  clang2py -k cdefstum \
+    extra/amdpci/headers/v11_structs.h \
+    extra/amdpci/headers/v12_structs.h \
+    extra/amdpci/headers/amdgpu_vm.h \
+    extra/amdpci/headers/discovery.h \
+    extra/amdpci/headers/amdgpu_ucode.h \
+    extra/amdpci/headers/psp_gfx_if.h \
+    extra/amdpci/headers/amdgpu_psp.h \
+    extra/amdpci/headers/amdgpu_irq.h \
+    extra/amdpci/headers/amdgpu_doorbell.h \
+    $AMKERN_INC/soc15_ih_clientid.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/am.py
+  fixup $BASE/am/am.py
+  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
+  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/soc15d.h \
+    -o $BASE/am/pm4_soc15.py
+  fixup $BASE/am/pm4_soc15.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/nvd.h \
+    -o $BASE/am/pm4_nv.py
+  fixup $BASE/am/pm4_nv.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_4_0_0.py
+  fixup $BASE/am/sdma_4_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_5_0_0.py
+  fixup $BASE/am/sdma_5_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_6_0_0.py
+  fixup $BASE/am/sdma_6_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    -o $BASE/am/smu_v13_0_0.py
+  fixup $BASE/am/smu_v13_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/smu_v14_0_2.py
+  fixup $BASE/am/smu_v14_0_2.py
+}
+
+generate_sqtt() {
+  clang2py -k cdefstum \
+    extra/sqtt/sqtt.h \
+    -o $BASE/sqtt.py
+  fixup $BASE/sqtt.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
+  python3 -c "import tinygrad.runtime.autogen.sqtt"
+
+  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
+  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
+  if [ ! -d "$ROCPROF_SRC" ]; then
+    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
+    pushd .
+    cd $ROCPROF_SRC
+    git reset --hard $ROCPROF_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
+    $ROCPROF_SRC/include/trace_decoder_instrument.h \
+    $ROCPROF_SRC/include/trace_decoder_types.h \
+    -o extra/sqtt/rocprof/rocprof.py
+  fixup extra/sqtt/rocprof/rocprof.py
+  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
+}
+
+generate_webgpu() {
+  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
+  fixup $BASE/webgpu.py
+  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
+  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
+  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
+  python3 -c "import tinygrad.runtime.autogen.webgpu"
+}
+
+generate_libusb() {
+  clang2py -k cdefstum \
+    /usr/include/libusb-1.0/libusb.h \
+    -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/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"
+  python3 -c "import tinygrad.runtime.autogen.libusb"
+}
+
+generate_mesa() {
+  MESA_TAG="mesa-25.2.4"
+  MESA_SRC=/tmp/mesa-$MESA_TAG
+  TINYMESA_TAG=tinymesa-32dc66c
+  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
+  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
+  if [ ! -d "$MESA_SRC" ]; then
+    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
+    pushd .
+    cd $MESA_SRC
+    git reset --hard $MESA_COMMIT_HASH
+    # clang 14 doesn't support packed enums
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
+    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
+    mkdir -p gen/util/format
+    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
+    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
+    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
+    popd
+  fi
+
+  if [ ! -d "$TINYMESA_DIR" ]; then
+    mkdir $TINYMESA_DIR
+    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
+  fi
+
+  clang2py -k cdefstu \
+    $MESA_SRC/src/compiler/nir/nir.h \
+    $MESA_SRC/src/compiler/nir/nir_builder.h \
+    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
+    $MESA_SRC/src/compiler/nir/nir_serialize.h \
+    $MESA_SRC/gen/nir_intrinsics.h \
+    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
+    $MESA_SRC/src/nouveau/compiler/nak.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
+    $MESA_SRC/src/compiler/glsl_types.h \
+    $MESA_SRC/src/util/blob.h \
+    $MESA_SRC/src/util/ralloc.h \
+    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
+    -l $TINYMESA_SO \
+    -o $BASE/mesa.py
+
+  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
+
+  fixup $BASE/mesa.py
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
+  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
+  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
+  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
+  echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
+  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
+  # bitfield bug in clang2py
+  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
+  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
+  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
+  python3 -c "import tinygrad.runtime.autogen.mesa"
+}
+
+if [ "$1" == "opencl" ]; then generate_opencl
+elif [ "$1" == "hip" ]; then generate_hip
+elif [ "$1" == "comgr" ]; then generate_comgr
+elif [ "$1" == "cuda" ]; then generate_cuda
+elif [ "$1" == "nvrtc" ]; then generate_nvrtc
+elif [ "$1" == "hsa" ]; then generate_hsa
+elif [ "$1" == "kfd" ]; then generate_kfd
+elif [ "$1" == "nv" ]; then generate_nv
+elif [ "$1" == "amd" ]; then generate_amd
+elif [ "$1" == "am" ]; then generate_am
+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
+elif [ "$1" == "adreno" ]; then generate_adreno
+elif [ "$1" == "pci" ]; then generate_pci
+elif [ "$1" == "vfio" ]; then generate_vfio
+elif [ "$1" == "webgpu" ]; then generate_webgpu
+elif [ "$1" == "libusb" ]; then generate_libusb
+elif [ "$1" == "mesa" ]; then generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
+else echo "usage: $0 <type>"
+fi
@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -10,8 +12,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
  pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
  def __init__(self, model_path: str):
-    import tiktoken
-    from tiktoken.load import load_tiktoken_bpe
    mergeable_ranks = load_tiktoken_bpe(model_path)
    self.num_base_tokens = len(mergeable_ranks)
    special_tokens = [
@@ -15,7 +15,7 @@ export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

 # pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
+export LOGMLPERF=1

 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
@@ -1,10 +1,11 @@
 import os, sys, pickle, time, re
 import numpy as np
-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.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
+
+import onnx
 from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -39,7 +40,7 @@ def compile(onnx_file):
  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

-  # check gated read_image usage
+  # checks from compile2
  kernel_count = 0
  read_image_count = 0
  gated_read_image_count = 0
@@ -95,7 +96,6 @@ def test_vs_compile(run, inputs, test_val=None):
  return val

 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
-  import onnx
  import onnxruntime as ort
  
  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
@@ -2,8 +2,7 @@ import os, pathlib, argparse
 from examples.llama3 import Tokenizer
 from tabulate import tabulate
 from tinygrad import fetch
-from tinygrad.helpers import flatten, getenv
-from sz import NONCORE_DIRS
+from tinygrad.helpers import flatten

 # llama 3 tokenizer
 tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
@@ -11,15 +10,19 @@ tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/
 def read_code(base_path):
  ret = []
  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
-    if not getenv("CORE") and any(path.split("./")[1].startswith(x) for x in NONCORE_DIRS): continue
    for name in files:
      if not name.endswith(".py"): continue
      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
      fullpath = os.path.join(path, name)
      code = pathlib.Path(fullpath).read_text()
-      ret.append((fullpath.split("tinygrad/", 1)[1], code))
+      ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
  return ret

+def write_code_to_file(filename, code_list):
+  """Writes the combined code to a specified file."""
+  with open(filename, 'w') as f:
+    f.write('\n'.join(flatten(code_list)))
+
 if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
  parser.add_argument("--output", help="Output file to write the combined code to.")
@@ -29,11 +32,10 @@ if __name__ == "__main__":

  table = []
  for name,code in ret:
-    table.append([name, len(tokenizer.encode(code))])
+    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))

-  banner = "#"*40
-  code_str = ''.join([f"{banner}\n# {name}\n{banner}\n\n{code}\n" for name,code in ret])
+  code_str = '\x00'.join(flatten(ret))
  print(f"code has {len(code_str)} chars")
  newline_count = code_str.count('\n')
  print(f"code has {newline_count} newlines")
@@ -42,5 +44,5 @@ if __name__ == "__main__":
  print(f"code has {len(encoded)} tokens")

  if args.output:
-    with open(args.output, 'w') as f: f.write(code_str)
-    print(f"Combined code written to {args.output}")
+    write_code_to_file(args.output, ret)
+    print(f"Combined code written to {args.output}")
@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

-from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
+from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch

@@ -244,16 +244,15 @@ def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):

  log_spec = prep_audio(waveforms, model.batch_size, truncate)
  nsample = model.decoder.max_tokens_to_sample
-  nctx = model.decoder.max_self_attn_cache_len

  def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
    pos, next_tokens = 0, ctx
-    for i in range(nsample):
-      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range((nsample-len(start_tokens))*2):
+      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
      next_tokens[ctx[:, -1] == eot] = eot
      ctx = np.concatenate((ctx, next_tokens), axis=1)
      pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all() or pos == nctx: break
+      if (next_tokens == eot).all(): break
    return ctx

  def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]
@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
-from tinygrad.helpers import getenv, colored
-
-SZ = 8_000_000_000
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
-
-if __name__ == "__main__":
-  # create tensors
-  tens = [Tensor.ones(SZ, dtype=dtypes.uint8, device=f"{Device.DEFAULT}:{i}").contiguous() for i in range(GPUS)]
-  Tensor.realize(*tens)
-
-  bw = [[0.0]*GPUS for _ in range(GPUS)]
-  for i in range(GPUS):
-    for j in range(GPUS):
-      GlobalCounters.reset()
-      with Context(DEBUG=2):
-        if i == j:
-          # this copy would be optimized out, just add 1
-          (tens[i]+1).realize()
-        else:
-          tens[i].to(f"{Device.DEFAULT}:{j}").realize()
-      t = max(GlobalCounters.time_sum_s, 1e-9)
-      bw[i][j] = SZ / t / 1e9  # GB/s
-
-  def fmt(x):
-    c = "green" if x > 50 else "yellow" if x > 20 else "red"
-    return colored(f"{x:6.1f}", c)
-
-  # header
-  print(" " * 8 + " ".join(f"{'d'+str(j):>6}" for j in range(GPUS)))
-  # rows
-  for i in range(GPUS):
-    print(f"{'s'+str(i):>6} -> " + " ".join(fmt(x) for x in bw[i]))
@@ -1,169 +1,353 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
-from tinygrad.engine.realize import ExecItem, get_runner
+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.helpers import getenv
+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.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.swizzler import merge_views, view_left

-N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])

-# ---------------------------
-# launch/config constants
-# ---------------------------
+N = 4096
+run_count = 5

-WARP_SIZE = 32
+BN = 128
+BM = 128
+BK = 8

-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
+TN = 4
+TM = 4

-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+# 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"

-is_kernel5 = getenv("K5", 0)
-THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+  # 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]

-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
+  # 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)))))

-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+pm = PatternMatcher([
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])

-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
-  assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+def rangeify_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
+  sink = c.schedule()[-1].ast
+  #print(sink)

-def hand_spec_kernel3():
-  # ---------------------------
-  # block indices & placeholders
-  # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
+  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
+  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
+  opts += [Opt(OptOps.UNROLL, 0, 8)]

-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))

-  # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+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(N//BM, 0), 2:UOp.range(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))

-  # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+def hl_spec_kernel3():
+  nbIterWaveM = 2
+  nbIterWaveN = 2

-  # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  # 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,)))

-  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
-  # ---------------------------
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  # 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)

-  # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  #                     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[:]

-  # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  # 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]

-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  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]

-  # open inner k range
-  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
+  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)

-  # ---------------------------
-  # LOCAL -> REG (per-wave tiles)
-  # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  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)

-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  #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))

-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
+  axis_types = (
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.REDUCE, AxisType.REDUCE)

-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  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

-  # ---------------------------
-  # FMA: c_regs += A_col * B_row
-  # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  i = UOp.range(c_regs.size, 16)
-  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256

-  # TODO: why don't these work as upcast?
-  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  nbWaves = BLOCK_SIZE // 32
+  WN = 128 if kernel5 else 64
+  WM = BN * BM // nbWaves // WN

-  # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  nbWaveX = BN // WN
+  nbWaveY = BM // WM

-  # ---------------------------
-  # REG -> GLOBAL (epilogue)
-  # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
-  c = c[waveIdy, :, laneIdy, :,
-        waveIdx, :, laneIdx, :]
-  sink = copy(c, c_regs.after(sink), rng=600)
+  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

-  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+  nbThreadXPerWave = 8
+  nbThreadYPerWave = 4

-def test_matmul(sink:UOp, N=N):
-  with Context(DEBUG=0):
-    a = Tensor.randn(N, N)
-    b = Tensor.randn(N, N)
-    hc = Tensor.empty(N, N)
-    Tensor.realize(a, b, hc)
+  idxInWave = indexInWave % nbThreadXPerWave
+  idyInWave = indexInWave // nbThreadXPerWave

-  ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
+  nbIterWaveN = WN // (nbThreadXPerWave * TN)
+  nbIterWaveM = WM // (nbThreadYPerWave * TM)

-  GlobalCounters.reset()
-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  SUBWN = WN // nbIterWaveN
+  SUBWM = WM // nbIterWaveM

-  if getenv("VERIFY", 1):
-    GlobalCounters.reset()
-    with Context(DEBUG=2):
-      tc = (a @ b).realize()
-    with Context(DEBUG=0):
-      err = (hc - tc).square().mean().item()
-    print(f"mean squared error {err}")
-    if err > 1e-06:
-      raise RuntimeError("matmul is wrong!")
+  # 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(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(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(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(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(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(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(BK, first_range+0)
+
+      # load from locals into registers
+      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      i = UOp.range(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(nbIterWaveM, first_range+3)
+      i = UOp.range(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(nbIterWaveM, first_range+5)
+      yt = UOp.range(TM, first_range+6)
+      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
+      xt = UOp.range(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(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(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(N//BK, 0)
+    kId = kId_range*BK
+
+    # load from globals into locals
+    i = UOp.range(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(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(BK, 3)
+
+    # load from locals into registers
+    iterWave = UOp.range(nbIterWaveN, 4)
+    i = UOp.range(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(nbIterWaveM, 6)
+    i = UOp.range(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(nbIterWaveM, 8)
+    yt = UOp.range(TM, 9)
+    iterWaveN = UOp.range(nbIterWaveN, 10)
+    xt = UOp.range(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(nbIterWaveM, 1000)
+  yt = UOp.range(TM, 1001)
+  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  xt = UOp.range(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__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  HL = getenv("HL")
+  if HL == 3: hprg = rangeify_kernel3()
+  elif HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
+  if HL == 3:
+    prg = get_program(hprg, Device.default.renderer)
+  else:
+    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,11 +1,17 @@
 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 import dtypes
-from tinygrad.dtype import DTYPES_DICT
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
+from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad import Device, dtypes, Tensor
+from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View

 script_dir = os.path.dirname(os.path.abspath(__file__))

@@ -47,6 +53,12 @@ def randoms():
    nc = nc.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
  return na, nb, nc

+def ast_to_cuda_prog(compiler, ast, opts):
+  k = Kernel(ast)
+  k.apply_opts(opts)
+  p = get_program(k.ast, k.opts, k.applied_opts)
+  return CUDAProgram(device, p.function_name, compiler.compile(p.src))
+
 if __name__ == "__main__":
  print(f"gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
  prog, global_size, local_size = None, None, None
@@ -177,11 +189,11 @@ if __name__ == "__main__":

    tms = []
    na, nb, nc = randoms()
-    cudaalloc._copyin(a, memoryview(bytearray(na)))
-    cudaalloc._copyin(b, memoryview(bytearray(nb)))
+    cudaalloc.copyin(a, bytearray(na))
+    cudaalloc.copyin(b, bytearray(nb))
    for i in range(CNT):
      tms.append(prog(*args, **kwargs))
-    cudaalloc._copyout(flat_mv(nc.data), c)
+    cudaalloc.copyout(flat_mv(nc.data), c)
    comp = na.astype(np.float32) @ nb.astype(np.float32)
    result = nc.reshape(M, N).astype(np.float32)

@@ -1,42 +0,0 @@
-from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 2048
-
-# metal has an 8x8 tensor core. this is the indexing
-def mat_idx(buf, g0, g1, warp, u):
-  l = [(warp//2**i)%2 for i in range(5)]
-  return buf[g0, l[4]*4 + l[2]*2 + l[1], g1, l[3]*4 + l[0]*2 + u]
-
-def hand_spec_tc_cores():
-  gx = UOp.special(N // 8, "gidx0")
-  gy = UOp.special(N // 8, "gidx1")
-  warp = UOp.special(32, "lidx0")
-
-  c = UOp.placeholder((N, N), dtypes.float, slot=0).reshape((N//8, 8, N//8, 8))
-  a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
-  b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
-
-  gk = UOp.range(N // 8, 0, AxisType.REDUCE)
-
-  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
-  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
-
-  acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  acc = acc[0].set(0.0)
-  acc = acc[1].set(0.0)
-
-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
-  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
-
-  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
-
-  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
-  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
-
-if __name__ == "__main__":
-  test_matmul(hand_spec_tc_cores(), N=N)
@@ -1,229 +0,0 @@
-import os
-import numpy as np
-np.set_printoptions(linewidth=1000000)
-os.environ["AMD_LLVM"] = "0"
-
-from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
-from tinygrad.helpers import DEBUG, getenv
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-
-WARP_SIZE = 64
-
-# Reg tile sizes (tensor cores)
-TC_M = 16
-TC_N = 16
-TC_K = 32
-
-# 1024 matrix cores
-# 16 cycle mfma
-# 2.2 GHz
-# 16x16x32x2 FLOPS/mma = 16384
-# 2.2*1e9*16384*1024/16*1e-12 TFLOPS = 2306 TFLOPS
-
-#N,M,K = 256,256,64
-N,M,K = 4096,4096,4096
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-#BLOCK_M = 128   # rows of C (M-dim) per block
-#BLOCK_N = 128   # columns of C (N-dim) per block
-#BLOCK_K = 128   # K-slice per block iteration
-
-BLOCK_M = 64
-BLOCK_N = 64
-BLOCK_K = 128
-
-WARPGROUP_SIZE = 1
-BLOCK_M = BLOCK_M * WARPGROUP_SIZE
-
-# TODO: improve the syntax of this. better syntax, faster iteration
-#  -- DONE: add working slice a[gx, :, i] -> shape of the : (aka (16,16,32) becomes (16,))
-#  -- DONE(ish): add argfix to movement (traits shared with Tensor)
-#  -- fix WMMA to not require all the junk
-#  -- improve syntax for vectorized loads/stores (both with DEVECTORIZE and without)
-#  -- DONE: be able to use CONTRACT on a range
-#  -- fix upcasted RANGE on an already vectorized buffer
-#  -- improve "all ranges not ended error" / fix the bug with after on ended ranges (if you are after end of range, range is closed)
-
-CUS_PER_GPU = 256
-assert ((M//BLOCK_M) * (N//BLOCK_N)) >= CUS_PER_GPU, "not enough globals"
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  # A = (M x K)
-  # B = (K x N)
-  # C = (M x N)
-
-  # check it's proper matmul
-  assert C.shape[0] == A.shape[0]
-  assert C.shape[1] == B.shape[1]
-  assert A.shape[1] == B.shape[0]
-
-  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
-  warp = UOp.special(WARP_SIZE, "lidx0")
-  warpgroup = UOp.special(WARPGROUP_SIZE, "lidx1")
-
-  # generic copy logic (not good)
-  def generic_copy(glbl, gargs, lcl, rng):
-    # Fully coalesced 128-bit loads/stores.
-    INNER_SIZE = 8
-    cp_i = UOp.range(lcl.size//(WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE), rng)
-    cp_inner = UOp.range(INNER_SIZE, rng+1, AxisType.UPCAST)
-    idx_i = cp_i*WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE + warpgroup*WARP_SIZE*INNER_SIZE + warp*INNER_SIZE + cp_inner
-    return lcl[idx_i].store(glbl[*gargs, idx_i]).end(cp_i, cp_inner)
-
-  # split out the globals into blocks
-  C = C.reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
-  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))
-  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))
-
-  # this is the big accumulator
-  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
-  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)
-
-  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
-  def make_locals(slot) -> tuple[UOp, UOp]:
-    BM_As_stride = (BLOCK_M + 1)
-    BN_Bs_stride = (BLOCK_N + 0)
-    INNER_SLICE = 8
-    As = UOp.placeholder((BLOCK_K//INNER_SLICE, BM_As_stride, INNER_SLICE), dtypes.half, slot=slot, addrspace=AddrSpace.LOCAL)
-    INNER_SLICE = 1
-    Bs = UOp.placeholder((BLOCK_K//INNER_SLICE, BN_Bs_stride, INNER_SLICE), dtypes.half, slot=slot+1, addrspace=AddrSpace.LOCAL)
-    As = As.permute((0,2,1)).reshape((BLOCK_K, BM_As_stride)).shrink_to((BLOCK_K, BLOCK_M))
-    Bs = Bs.permute((0,2,1)).reshape((BLOCK_K, BN_Bs_stride)).shrink_to((BLOCK_K, BLOCK_N))
-    return As, Bs
-
-  # load from globals into locals (TODO: use the warpgroup)
-
-  def load_to_locals(l_K_outer_loop:UOp, Asl:UOp, Bsl:UOp, rng:int, barrier=True) -> tuple[UOp, UOp]:
-    if getenv("FAKE"):
-      return Asl[0].set(0), Bsl[0].set(0)
-    else:
-      pA = A.permute((0,2,1,3)).reshape((M//BLOCK_M, K//BLOCK_K, BLOCK_M*BLOCK_K))
-      pas = Asl.permute((1,0)).reshape((BLOCK_M*BLOCK_K,))
-      As_store = generic_copy(pA, (gx, l_K_outer_loop), pas, rng)
-
-      pB = B.permute((0,2,1,3)).reshape((K//BLOCK_K, N//BLOCK_N, BLOCK_K*BLOCK_N))
-      pbs = Bsl.reshape((BLOCK_K*BLOCK_N,))
-      Bs_store = generic_copy(pB, (l_K_outer_loop, gy), pbs, rng+2)
-
-      barrier = UOp.barrier(As_store, Bs_store) if barrier else UOp.group(As_store, Bs_store)
-      return Asl.after(barrier), Bsl.after(barrier)
-
-  def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]=()) -> UOp:
-    K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
-
-    # load from locals into registers
-    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
-    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
-
-    M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
-    Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
-    load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
-    A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
-    Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
-
-    N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
-    Bsl = Bsl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N))
-    load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
-    B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
-    Br = Br[N_load_loop].set(B_in, end=N_load_loop)
-
-    M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
-    N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
-
-    # load values
-    acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
-    acc_load = acc_after[N_inner_loop, M_inner_loop]
-
-    # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
-
-    # store back the acc
-    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
-    return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
-
-  # **** START INNER LOOP *****
-  # inner loop -- locals -> regs
-
-  # no pipeline
-  if not getenv("PIPELINE"):
-    As, Bs = make_locals(slot=0)
-
-    K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
-    As, Bs = load_to_locals(K_outer_loop, As, Bs, 1000, barrier=True)
-    acc_store = compute_on_locals(acc, As, Bs, 1500, afters=(K_outer_loop,))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-  else:
-    # this doesn't work
-    As0, Bs0 = make_locals(slot=0)
-    As1, Bs1 = make_locals(slot=2)
-    As0, Bs0 = load_to_locals(0, As0, Bs0, 1000)
-
-    K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.REDUCE)
-    As1, Bs1 = load_to_locals(K_outer_loop+1, As1, Bs1, 2000, barrier=False)
-    acc_store = compute_on_locals(acc, As0, Bs0, 1500, afters=(K_outer_loop,))
-    As0, Bs0 = load_to_locals(K_outer_loop+2, As0, Bs0, 3000, barrier=False)
-    acc_store = compute_on_locals(acc, As1, Bs1, 2500, afters=(acc_store, As0, Bs0))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-
-    #acc_store = compute_on_locals(acc, As0, Bs0, 3500, afters=(acc_store.barrier().end(K_outer_loop)))
-    """
-    As1, Bs1 = load_to_locals(K//BLOCK_K-1, As1, Bs1, 4000)
-    acc_store = compute_on_locals(acc, As1, Bs1, 4500, afters=(acc_store))
-    """
-    #acc = acc.after(acc_store)
-
-  # **** END LOOPS *****
-
-  # store the acc into gmem
-  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
-  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
-  store = store.end(cp_i, cp_j)
-
-  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
-
-# simplest WMMA
-"""
-# init the acc
-acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
-acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
-
-# do the wmma
-acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
-
-# store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
-
-# store the acc into gmem
-store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
-"""
-
-if __name__ == "__main__":
-  a = Tensor.randn(M, K, dtype=dtypes.half)
-  b = Tensor.randn(K, N, dtype=dtypes.half)
-
-  #a = Tensor.zeros(M, K, dtype=dtypes.half).contiguous()
-  #a[0,16] = 1
-  #b = Tensor.ones(K, N, dtype=dtypes.half).contiguous()
-
-  c = Tensor.empty(M, N, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(a,b)
-
-  ref = a.dot(b, dtype=dtypes.float)
-  ref.realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    tst.realize()
-  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
-
-  with Context(DEBUG=0):
-    #print(ref.numpy())
-    #print(tst.numpy())
-    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"
@@ -1,141 +0,0 @@
-import os
-import numpy as np
-np.set_printoptions(linewidth=1000000)
-os.environ["AMD_LLVM"] = "0"
-
-from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
-from tinygrad.helpers import DEBUG, getenv
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
-
-WARP_SIZE = 64
-
-# Reg tile sizes (tensor cores)
-TC_M = 16
-TC_N = 16
-TC_K = 32
-
-N,M,K = 4096,4096,4096
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_M = 64
-BLOCK_N = 64
-BLOCK_K = 64
-
-WARPGROUP_SIZE = 1
-BLOCK_M = BLOCK_M * WARPGROUP_SIZE
-
-TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
-
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
-  assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
-
-def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...], warpgroup, warp) -> UOp:
-  K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
-
-  # load from locals into registers
-  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
-  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
-
-  M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
-  Asl = Asl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M)
-  load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
-  A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
-  Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
-
-  N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
-  Bsl = Bsl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N)
-  load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
-  B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
-  Br = Br[N_load_loop].set(B_in, end=N_load_loop)
-
-  M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
-  N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
-
-  # load values
-  acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
-  acc_load = acc_after[N_inner_loop, M_inner_loop]
-
-  # do WMMA
-  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-  out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
-
-  # store back the acc
-  acc_store = acc[N_inner_loop, M_inner_loop].store(out)
-  return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
-  K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
-
-  # split out the globals into blocks
-  C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
-  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :]
-  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :]
-
-  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
-  # ---------------------------
-  tid = UOp.special(TID_SIZE, "lidx0")
-  warpgroup, warp = tid//WARP_SIZE, tid%WARP_SIZE
-
-  A_view = A.reshape(-1, TID_SIZE, 8)
-  B_view = B.reshape(-1, TID_SIZE, 8)
-
-  # A: read BM x BK tiles (permute on store into locals)
-  As = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_M)
-  As_view = As.reshape(-1, TID_SIZE, 8)
-
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N+4), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL).shrink_to(BLOCK_K, BLOCK_N)
-  Bs_view = Bs.reshape(-1, TID_SIZE, 8)
-
-  outer_copy = UOp.range(A_view.shape[0], 100, AxisType.UPCAST)
-  inner_copy = UOp.range(A_view.shape[2], 101, AxisType.UPCAST)
-  As_store = As_view[outer_copy, tid, inner_copy].store(A_view[outer_copy, tid, inner_copy])
-  Bs_store = Bs_view[outer_copy, tid, inner_copy].store(B_view[outer_copy, tid, inner_copy])
-
-  if getenv("NOLOAD"):
-    As_store = As[0,0].store(0)
-    Bs_store = Bs[0,0].store(0)
-
-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(UOp.group(As_store, Bs_store).end(outer_copy, inner_copy))
-
-  if getenv("COMPUTE"):
-    As, Bs = As.after(barrier), Bs.after(barrier)
-
-    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
-
-    sink = compute_on_locals(acc, As, Bs, 200, afters=(barrier,), warpgroup=warpgroup, warp=warp)
-    sink = sink.end(K_outer_loop)
-
-    C_view = C[gx, :, gy, :].reshape(BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M, BLOCK_N//TC_N, TC_N)[:, warpgroup, warp%16, :, (warp//16)*4]
-    sink = copy(C_view, acc.after(sink), rng=300)
-  else:
-    sink = C.after(barrier.end(K_outer_loop))[0,0,0,0].store(As[0,0]+Bs[0,0])
-
-  return sink.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
-
-if __name__ == "__main__":
-  a = Tensor.randn(M, K, dtype=dtypes.half)
-  b = Tensor.randn(K, N, dtype=dtypes.half)
-  c = Tensor.empty(M, N, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(a,b)
-
-
-  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    tst.realize()
-  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
-
-
-  with Context(DEBUG=0):
-    ref = a.dot(b, dtype=dtypes.float)
-    ref.realize()
-    #print(ref.numpy())
-    #print(tst.numpy())
-    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"
@@ -17,7 +17,7 @@ M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)

-atol, rtol = {dtypes.half:{1e-3, 1e-2}, dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
+atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
 ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)

 INT_LOW = getenv("INT_LOW", 0)
@@ -9,10 +9,9 @@ torch.set_num_threads(1)
 from tinygrad.helpers import getenv
 CUDA = getenv("CUDA", 1)
 MPS = getenv("MPS", 0)
-if getenv("FP16_ACC"): torch.backends.cuda.matmul.allow_fp16_accumulation = True

 for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-  for N in [256, 512, 1024, 2048, 4096] + ([6144, 8192] if getenv("BIG") else []):
+  for N in [256, 512, 1024, 2048, 4096]:
    FLOPS = N*N*N*2

    b = torch.rand((N,N), dtype=dtype)
@@ -1,16 +0,0 @@
-from tinygrad import Tensor, Device, TinyJit, dtypes
-from tinygrad.helpers import getenv
-
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
-N = 6144
-
-@TinyJit
-def many_matmul(A, B):
-  out = A
-  for _ in range(8): out = out@B
-  return out
-
-if __name__ == "__main__":
-  A = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
-  B = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
-  while 1: many_matmul(A, B)
@@ -51,15 +51,11 @@ def create_report(dev, test, result, stdout, stderr):
  dmesg_output = subprocess.check_output(["sudo", "dmesg", "--ctime", "--color=never"], text=True)
  with open(dmesg_path, "w") as f: f.write(dmesg_output)

-  env_vars = " ".join(f"{k}={v}" for k, v in test.env.items())
-  reproduce_cmd = f"{env_vars} {test.cmd}"
-
  summary_path = os.path.join(report_path, "summary.txt")
  with open(summary_path, "w") as f:
    f.write(f"Test: {test.name()}\n")
    f.write(f"Dev params: {vars(dev)}\n")
    f.write(f"Test params: {vars(test)}\n")
-    f.write(f"Reproduce cmd: {reproduce_cmd}\n")
    f.write(f"Exit Code: {result}\n")

  print(f"Crash report saved to {report_path}")
@@ -19,6 +19,5 @@ trap 'rm -f "$TMP"' EXIT
 EOF
  sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
  echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
-) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include || exit 1
+) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0

-printf 'lvp_nir_options = gzip.decompress(base64.b64decode("%s"))' $("$TMP" | gzip | base64 -w0)
@@ -1,10 +1,8 @@
-import os, pathlib
-
-# TODO: there is a timing bug without this
-os.environ["AMD_AQL"] = "1"
-
+import pathlib
 from tinygrad.device import Device
 from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
+import time
+import os

 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -34,7 +32,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra="")
  src = src.replace("DIRECTIVE", DIRECTIVE)
  lib = COMPILER.compile(src)
  fxn = AMDProgram(DEV, "matmul", lib)
-  elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
+  elapsed = fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

@@ -46,9 +44,9 @@ if __name__=="__main__":
    raise RuntimeError("Error while initiating AMD device")

  COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
-    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
-    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 40
+  if DEV.arch in {'gfx1100', 'gfx1103'}:
+    if DEV.arch == 'gfx1103':
+      NUM_WORKGROUPS = 8
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
@@ -86,14 +84,12 @@ if __name__=="__main__":
    NUM_WORKGROUPS = 256
    WAVE_SIZE = 64
    NUM_WAVES = 4
-    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
    FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
    FLOPS_PER_MATMUL = 16*16*128*2
    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
  else:
-    raise RuntimeError(f"arch {DEV.arch} not supported.")
+    raise RuntimeError(f"arch {DEV.arch} not supported.")
@@ -3,14 +3,14 @@
    .p2align 8
    .type matmul,@function
 matmul:
-  s_mov_b32 s1, INTERNAL_LOOP
-  s_mov_b32 s2, 0
-  inner_loop:
-    INSTRUCTION
-    s_sub_u32 s1, s1, 1
-    s_cmp_lg_i32 s1, s2
-    s_cbranch_scc1 inner_loop
-  s_endpgm
+    s_mov_b32 s1, INTERNAL_LOOP
+    s_mov_b32 s2, 0
+    inner_loop:
+        INSTRUCTION
+        s_sub_u32 s1, s1, 1
+        s_cmp_lg_i32 s1, s2
+        s_cbranch_scc1 inner_loop
+    s_endpgm

 .rodata
 .p2align 6
@@ -89,20 +89,6 @@ class Attention:
    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
-    if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
-      from tinygrad.uop.ops import UOp, KernelInfo
-      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-        return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
-      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
-      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
-        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
-        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
-      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
    attn = attn.reshape(bsz, seqlen, -1)
    return self.wo(attn)

@@ -1,53 +0,0 @@
-/*
- * NVIDIA_COPYRIGHT_BEGIN
- *
- * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
- *
- * NVIDIA CORPORATION and its licensors retain all intellectual property
- * and proprietary rights in and to this software, related documentation
- * and any modifications thereto.  Any use, reproduction, disclosure or
- * distribution of this software and related documentation without an express
- * license agreement from NVIDIA CORPORATION is strictly prohibited.
- *
- * NVIDIA_COPYRIGHT_END
- */
-
-#include <stdint.h>
-#include <stdlib.h>
-
-typedef enum {
-  NVJITLINK_SUCCESS = 0,
-  NVJITLINK_ERROR_UNRECOGNIZED_OPTION,
-  NVJITLINK_ERROR_MISSING_ARCH,
-  NVJITLINK_ERROR_INVALID_INPUT,
-  NVJITLINK_ERROR_PTX_COMPILE,
-  NVJITLINK_ERROR_NVVM_COMPILE,
-  NVJITLINK_ERROR_INTERNAL
-} nvJitLinkResult;
-
-typedef enum {
-  NVJITLINK_INPUT_NONE = 0,
-  NVJITLINK_INPUT_CUBIN = 1,
-  NVJITLINK_INPUT_PTX,
-  NVJITLINK_INPUT_LTOIR,
-  NVJITLINK_INPUT_FATBIN,
-  NVJITLINK_INPUT_OBJECT,
-  NVJITLINK_INPUT_LIBRARY
-} nvJitLinkInputType;
-
-typedef struct nvJitLink* nvJitLinkHandle;
-
-nvJitLinkResult nvJitLinkCreate(nvJitLinkHandle *handle, uint32_t numOptions, const char **options);
-nvJitLinkResult nvJitLinkDestroy(nvJitLinkHandle *handle);
-nvJitLinkResult nvJitLinkAddData(nvJitLinkHandle handle, nvJitLinkInputType inputType, const void *data, size_t size, const char *name);
-nvJitLinkResult nvJitLinkAddFile(nvJitLinkHandle handle, nvJitLinkInputType inputType, const char *fileName);
-nvJitLinkResult nvJitLinkComplete(nvJitLinkHandle handle);
-nvJitLinkResult nvJitLinkGetLinkedCubinSize(nvJitLinkHandle handle, size_t *size);
-nvJitLinkResult nvJitLinkGetLinkedCubin(nvJitLinkHandle handle, void *cubin);
-nvJitLinkResult nvJitLinkGetLinkedPtxSize(nvJitLinkHandle handle, size_t *size);
-nvJitLinkResult nvJitLinkGetLinkedPtx(nvJitLinkHandle handle, char *ptx);
-nvJitLinkResult nvJitLinkGetErrorLogSize(nvJitLinkHandle handle, size_t *size);
-nvJitLinkResult nvJitLinkGetErrorLog(nvJitLinkHandle handle, char *log);
-nvJitLinkResult nvJitLinkGetInfoLogSize(nvJitLinkHandle handle, size_t *size);
-nvJitLinkResult nvJitLinkGetInfoLog(nvJitLinkHandle handle, char *log);
-nvJitLinkResult nvJitLinkVersion(unsigned int *major, unsigned int *minor);
@@ -65,8 +65,6 @@
 #define NVCEC0_QMDV05_00_GRID_HEIGHT_RESUME                        MW(271:256)
 #define NVCEC0_QMDV05_00_GRID_DEPTH_RESUME                         MW(287:272)
 #define NVCEC0_QMDV05_00_RELEASE_ENABLE(i)                         MW((288+(i)*16):(288+(i)*16))
-#define NVCEC0_QMDV05_00_RELEASE0_ENABLE                           NVCEC0_QMDV05_00_RELEASE_ENABLE(0)
-#define NVCEC0_QMDV05_00_RELEASE1_ENABLE                           NVCEC0_QMDV05_00_RELEASE_ENABLE(1)
 #define NVCEC0_QMDV05_00_RELEASE_ENABLE_FALSE                      0x00000000
 #define NVCEC0_QMDV05_00_RELEASE_ENABLE_TRUE                       0x00000001
 #define NVCEC0_QMDV05_00_RELEASE_STRUCTURE_SIZE(i)                 MW((290+(i)*16):(289+(i)*16))
@@ -58,23 +58,20 @@ def install_hook(c_function, python_function):
  return orig_func

 # *** ioctl lib end ***
-from tinygrad.runtime.autogen import nv_570 as nv_gpu
+import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
 nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
          x.startswith("NV") and x[6:].startswith("_CTRL_") and isinstance(getattr(nv_gpu, x), int)}

 def get_classes():
-  res = {}
-  known_classes = {"NV01_DEVICE_0", "NV01_ROOT", "NV1_MEMORY_SYSTEM", "NV01_MEMORY_VIRTUAL", "NV1_MEMORY_USER", "NV50_MEMORY_VIRTUAL", "NV_FERMI_VASPACE_A",
-                   "NV20_SUBDEVICE_0"}
-  for nm,val in nv_gpu.__dict__.items():
-    if not isinstance(val, int): continue
-    if 0x3000 < val < 0xffff: res[val] = nm
-    if nm in known_classes: res[val] = nm
-  return res
+  hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_gpu.py").read_text()
+  clss = re.search(r'NV01_ROOT.*?NV_SEMAPHORE_SURFACE = \(0x000000da\) # macro', hdrpy, re.DOTALL).group()
+  pattern = r'([0-9a-zA-Z_]*) = +\((0x[0-9a-fA-F]+)\)'
+  matches = re.findall(pattern, clss, re.MULTILINE)
+  return {int(num, base=16):name for name, num in matches}
 nvclasses = get_classes()
 nvuvms = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("UVM_") and nv_gpu.__dict__.get(x+"_PARAMS")}
-nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC9B0_", "NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
+nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}

 global_ioctl_id = 0
 gpus_user_modes = []
@@ -275,4 +272,4 @@ def compare_launch_state(states, good_states):

  return True, "PASS"

-# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
+# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
@@ -1,134 +0,0 @@
-/*
- * SPDX-FileCopyrightText: Copyright (c) 1993-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: MIT
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
- * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
- * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
- * DEALINGS IN THE SOFTWARE.
- */
-
-#ifndef PCIEXPTBL_H
-#define PCIEXPTBL_H
-
-#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE   0x00
-#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT    0xE0
-
-//
-// The VBIOS object comes from walking the PCI expansion code block
-// The following structure holds the expansion code format.
-//
-#define PCI_EXP_ROM_SIGNATURE     0xaa55
-#define PCI_EXP_ROM_SIGNATURE_NV  0x4e56 // "VN" in word format
-#define PCI_EXP_ROM_SIGNATURE_NV2 0xbb77
-#define IS_VALID_PCI_ROM_SIG(sig) ((sig == PCI_EXP_ROM_SIGNATURE) ||    \
-                                   (sig == PCI_EXP_ROM_SIGNATURE_NV) || \
-                                   (sig == PCI_EXP_ROM_SIGNATURE_NV2))
-
-#define OFFSETOF_PCI_EXP_ROM_SIG                    0x0
-#define OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET       0x16
-#define OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR    0x18
-
-#pragma pack(1)
-typedef struct _PCI_EXP_ROM_STANDARD
-{
-    NvU16       sig;                //  00h: ROM Signature 0xaa55
-    NvU8        reserved [0x16];    //  02h: Reserved (processor architecture unique data)
-    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
-    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
-} PCI_EXP_ROM_STANDARD, *PPCI_EXP_ROM_STANDARD;
-#pragma pack()
-
-#pragma pack(1)
-typedef struct _PCI_EXP_ROM_NBSI
-{
-    NvU16       sig;                //  00h: ROM Signature 0xaa55
-    NvU8        reserved [0x14];    //  02h: Reserved (processor architecture unique data)
-    NvU16       nbsiDataOffset;     //  16h: Offset from header to NBSI image
-    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
-    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
-} PCI_EXP_ROM_NBSI, *PPCI_EXP_ROM_NBSI;
-#pragma pack()
-
-typedef union _PCI_EXP_ROM {
-    PCI_EXP_ROM_STANDARD standard;
-    PCI_EXP_ROM_NBSI nbsi;
-} PCI_EXP_ROM, *PPCI_EXP_ROM;
-
-#define PCI_DATA_STRUCT_SIGNATURE     0x52494350 // "PCIR" in dword format
-#define PCI_DATA_STRUCT_SIGNATURE_NV  0x5344504E // "NPDS" in dword format
-#define PCI_DATA_STRUCT_SIGNATURE_NV2 0x53494752 // "RGIS" in dword format
-#define IS_VALID_PCI_DATA_SIG(sig) ((sig == PCI_DATA_STRUCT_SIGNATURE) ||    \
-                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV) || \
-                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV2))
-
-#define PCI_LAST_IMAGE NVBIT(7)
-#define PCI_ROM_IMAGE_BLOCK_SIZE 512U
-
-#define OFFSETOF_PCI_DATA_STRUCT_SIG        0x0
-#define OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID  0x4
-#define OFFSETOF_PCI_DATA_STRUCT_LEN        0xa
-#define OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE 0xd
-#define OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE  0x14
-#define OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN  0x10
-#define OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE 0x15
-
-#pragma pack(1)
-typedef struct _PCI_DATA_STRUCT
-{
-    NvU32       sig;                //  00h: Signature, the string "PCIR" or NVIDIA's alternate "NPDS"
-    NvU16       vendorID;           //  04h: Vendor Identification
-    NvU16       deviceID;           //  06h: Device Identification
-    NvU16       deviceListPtr;      //  08h: Device List Pointer
-    NvU16       pciDataStructLen;   //  0Ah: PCI Data Structure Length
-    NvU8        pciDataStructRev;   //  0Ch: PCI Data Structure Revision
-    NvU8        classCode[3];       //  0Dh: Class Code
-    NvU16       imageLen;           //  10h: Image Length (units of 512 bytes)
-    NvU16       vendorRomRev;       //  12h: Revision Level of the Vendor's ROM
-    NvU8        codeType;           //  14h: holds NBSI_OBJ_CODE_TYPE (0x70) and others
-    NvU8        lastImage;          //  15h: Last Image Indicator: bit7=1 is lastImage
-    NvU16       maxRunTimeImageLen; //  16h: Maximum Run-time Image Length (units of 512 bytes)
-} PCI_DATA_STRUCT, *PPCI_DATA_STRUCT;
-#pragma pack()
-
-#define NV_PCI_DATA_EXT_SIG 0x4544504E // "NPDE" in dword format
-#define NV_PCI_DATA_EXT_REV_10 0x100      // 1.0
-#define NV_PCI_DATA_EXT_REV_11 0x101      // 1.1
-
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_SIG            0x0
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_LEN            0x6
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_REV            0x4
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN   0x8
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE     0xa
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS          0xb
-
-#define PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED 0x04
-
-#pragma pack(1)
-typedef struct _NV_PCI_DATA_EXT_STRUCT
-{
-    NvU32   signature;          //  00h: Signature, the string "NPDE"
-    NvU16   nvPciDataExtRev;    //  04h: NVIDIA PCI Data Extension Revision
-    NvU16   nvPciDataExtLen;    //  06h: NVIDIA PCI Data Extension Length
-    NvU16   subimageLen;        //  08h: Sub-image Length
-    NvU8    privLastImage;      //  0Ah: Private Last Image Indicator
-    NvU8    flags;              //  0Bh: Private images enabled if bit0=1
-} NV_PCI_DATA_EXT_STRUCT, *PNV_PCI_DATA_EXT_STRUCT;
-#pragma pack()
-
-#endif // PCIEXPTBL_H
-
-
@@ -10,8 +10,6 @@ import ctypes


 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
    @classmethod
    def as_dict(cls, self):
        result = {}
@@ -2,6 +2,8 @@

 ## Getting SQ Thread Trace

+Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)
+
 SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
@@ -1,131 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["PROFILE"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-from dataclasses import replace
-import atexit, contextlib
-from tinygrad import Tensor
-from tinygrad.helpers import system, OSX
-from tinygrad.runtime.ops_amd import AMDProgram
-from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
-from tinygrad.device import Device, ProfileDeviceEvent
-
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-# TODO: should really check for AM driver / USB
-if not OSX:
-  def set_power(x): system(f"sudo /opt/rocm/bin/amd-smi set -l {x}")
-  @atexit.register
-  def reset_power(): set_power("auto")
-  set_power("stable_std")
-
-dev = Device["AMD"]
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[str, list[WaveExec]] = {}
-  yield sqtt
-  events = dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())]
-
-  rctx = decode(events)
-  assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  sqtt.update(rctx.inst_execs)
-
-  for e in events:
-    if isinstance(e, ProfileSQTTEvent):
-      print(replace(e, blob=b''))
-      if e.se == 0:
-        parse_sqtt_print_packets(e.blob)
-
-template = """.text
-.globl matmul
-.p2align 8
-.type matmul,@function
-matmul:
-  INSTRUCTION
-  s_endpgm
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 32
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
-
-def run_asm(src):
-  NUM_WORKGROUPS = 1
-  WAVE_SIZE = 32
-  NUM_WAVES = 1
-  t = Tensor.empty(0x1000).realize()
-  buf = t.uop.buffer.ensure_allocated()
-  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
-  dev.compiler.disassemble(lib)
-  fxn = AMDProgram(dev, "matmul", lib)
-  fxn(buf._buf, global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
-
-if __name__ == "__main__":
-  with save_sqtt() as sqtt:
-    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
-    Tensor.empty(1).elu().realize()
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    # what's in v0?
-    run_asm([
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v1, 0",
-      "s_clause 0x1",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-    ]+[
-      "global_load_b32 v1, v0, s[0:1]",
-    ]*10+[
-      "global_load_b32 v10, v1, s[0:1]",
-      "s_waitcnt vmcnt(0)",
-
-      #"v_rcp_f32 v1, v0"
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v5 v4 v4",
-      #"v_add_f32_e32 v7 v6 v6",
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v2 v1 v1",
-      #"s_nop 1"
-    ]*5+[
-      "v_add_f32_e32 v3 v2 v2",
-    ]*5+[
-      "v_mul_f32_e32 v3 v2 v2",
-    ]*7)
@@ -1,543 +0,0 @@
-import pickle
-from tinygrad.helpers import getenv
-from extra.sqtt.roc import decode, ProfileSQTTEvent
-
-# Instruction packets (one per ISA op)
-# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
-# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
-
-OPCODE_NAMES = {
-  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
-  0x02: "VMEMEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
-  0x03: "ALUEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
-  0x01: "VALUINST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
-  0x06: "WAVERDY",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
-  0x08: "WAVEEND",
-  0x09: "WAVESTART",
-  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
-  0x04: "IMMEDIATE_4",
-  0x05: "IMMEDIATE_5",
-  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there
-  0x14: "REG",
-  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
-  0x12: "EVENT",
-  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
-  0x18: "INST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
-  0x19: "UTILCTR",
-
-  # ------------------------------------------------------------------------
-  # 0x07–0x0F: pure timestamp-ish deltas
-  # ------------------------------------------------------------------------
-  0x07: "TS_DELTA_S8_W3",           # shift=8,  width=3  (small delta)
-  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
-  0x0B: "TS_DELTA_S5_W3_A",         # shift=5,  width=3
-  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
-  0x0D: "TS_DELTA_S5_W3_C",         # shift=5,  width=3
-  0x0E: "TS_DELTA_S7_W2",           # shift=7,  width=2
-  0x0F: "TS_DELTA_SHORT_PLUS4",     # short delta; ROCm adds +4 before accumulate
-
-  # ------------------------------------------------------------------------
-  # 0x10–0x19: timestamps, layout headers, events, perf
-  # ------------------------------------------------------------------------
-  0x10: "PSEUDO_NEED_MORE_BITS",    # not a real packet; decoder refill hint
-
-  0x11: "TS_WAVE_STATE_SAMPLE",     # wave stall/termination sample (byte at +10)
-  0x13: "EVT_SMALL_GENERIC",        # same structural family as 0x08/0x12/0x19
-
-  0x15: "PERFCOUNTER_SNAPSHOT",     # small delta + 50-ish bits of snapshot
-  0x16: "TS_DELTA36_OR_MARK",       # 36-bit long delta or 36-bit marker
-  0x17: "LAYOUT_MODE_HEADER",       # layout/mode/group + selectors A/B
-}
-
-# these tables are from rocprof trace decoder
-# rocprof_trace_decoder_parse_data-0x11c6a0
-# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
-
-# ---------- 1. local_138: 256-byte state->token table ----------
-
-STATE_TO_TOKEN: bytes = bytes([
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-])
-
-
-# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
-
-NIBBLE_BUDGET = [
-  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40,
-  0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
-  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40,
-  0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-]
-assert len(NIBBLE_BUDGET) == 32
-
-
-# ---------- 3. delta_map from your hash nodes ----------
-
-# opcode -> (shift, width)
-DELTA_MAP_DEFAULT = {
-  0x01: (3,  3),   # shift=3,  end=6
-  0x02: (4,  2),   # shift=4,  end=6
-  0x03: (4,  2),   # shift=4,  end=6
-  0x04: (4,  3),   # shift=4,  end=7
-  0x05: (5,  3),   # shift=5,  end=8
-  0x06: (5,  3),   # shift=5,  end=8
-  0x07: (8,  3),   # shift=8,  end=11
-  0x08: (5,  3),   # shift=5,  end=8
-  0x09: (5,  2),   # shift=5,  end=7
-  0x0A: (5,  2),   # shift=5,  end=7
-  0x0B: (5,  3),   # shift=5,  end=8
-  0x0C: (5,  3),   # shift=5,  end=8
-  0x0D: (5,  3),   # shift=5,  end=8
-  0x0E: (7,  2),   # shift=7,  end=9
-  0x0F: (4,  4),   # shift=4,  end=8
-  0x10: (0,  0),   # shift=0,  end=0  (no delta)
-  0x11: (7,  9),   # shift=7,  end=16
-  0x12: (8,  3),   # shift=8,  end=11
-  0x13: (8,  3),   # shift=8,  end=11
-  0x14: (4,  3),   # shift=4,  end=7
-  0x15: (7,  3),   # shift=7,  end=10
-  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
-  0x17: (0,  0),   # shift=0,  end=0  (no delta)
-  0x18: (4,  3),   # shift=4,  end=7
-  0x19: (7,  2),   # shift=7,  end=9
-}
-
-# ---------- 4. One-line-per-packet parser ----------
-
-def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
-  """
-  Decode packet payloads conservatively, using:
-    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
-    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
-    - Per-opcode layouts derived from rocprof's decompiled consumers.
-  """
-  # --- 0. Restrict to real packet bits ---------------------------------
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
-  if nb_bits <= 0 or nb_bits >= 64:
-    pkt = reg & ((1 << 64) - 1)
-  else:
-    pkt = reg & ((1 << nb_bits) - 1)
-
-  fields: list[str] = []
-
-  shift, width = DELTA_MAP_DEFAULT.get(opcode, (0, 0))
-  if width:
-    field_mask = (1 << width) - 1
-    shaped_field = (pkt >> shift) & field_mask
-  else:
-    field_mask = 0
-    shaped_field = 0
-
-  # =====================================================================
-  # 1. Timestamp-centric opcodes (actually drive 'time')
-  # =====================================================================
-
-  if opcode == 0x0F:  # TS_DELTA_SHORT_PLUS4
-    # In the caller, delta already has +4 applied.
-    raw_delta = shaped_field
-    fields.append(f"raw_delta={raw_delta}")
-    fields.append(f"ts_short_plus4={delta}")
-    return ", ".join(fields)
-
-  if opcode == 0x11:  # TS_WAVE_STATE_SAMPLE
-    # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
-    raw_delta = shaped_field
-    coarse    = (pkt >> (shift + width)) & 0xFF  # matches byte at +10 in C
-    fields.append(f"raw_delta={raw_delta}")
-    if coarse:
-      fields.append(f"coarse_state=0x{coarse:02x}")
-    # From decomp:
-    #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
-    #  - when coarse&8, it marks all live waves as "terminated"
-    if coarse & 0x01:
-      fields.append("flag_wave_interest=1")
-    if coarse & 0x08:
-      fields.append("flag_terminate_all=1")
-    return ", ".join(fields)
-
-  if opcode == 0x16:  # TS_DELTA36_OR_MARK
-    # Bits:
-    #   bit8  -> 0x100
-    #   bit9  -> 0x200
-    #   bits 12..47 -> 36-bit field used as delta or marker
-    bit8 = bool(pkt & 0x100)
-    bit9 = bool(pkt & 0x200)
-    if not bit9:
-      mode = "delta"
-    elif not bit8:
-      mode = "marker"
-    else:
-      mode = "other"
-    val36 = (pkt >> 12) & ((1 << 36) - 1)
-    fields.append(f"mode={mode}")
-    if mode != "delta":
-      fields.append(f"val36=0x{val36:x}")
-    return ", ".join(fields)
-
-  # For 0x07, 0x0A–0x0E, we know they drive time (via DELTA_MAP_DEFAULT),
-  # but we don't see any other fields used in the decomp.
-  if opcode in (0x07, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E):
-    if width:
-      raw_delta = shaped_field
-      leftover  = pkt & ~(field_mask << shift)
-      fields.append(f"raw_delta={raw_delta}")
-      if leftover:
-        fields.append(f"payload=0x{leftover:x}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 2. Small "meta + tiny delta" packets (0x01–0x06)
-  # =====================================================================
-
-  if opcode == 0x01:  # META_ID12_TS_SMALL
-    id12 = pkt & 0xFFF
-    fields.append(f"id12=0x{id12:03x}")
-    if width:
-      fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    return ", ".join(fields)
-
-  if opcode == 0x02:  # META_FLAG8_TS_SMALL
-    flag8 = pkt & 0xFF
-    fields.append(f"flag8=0x{flag8:02x}")
-    if width:
-      fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    return ", ".join(fields)
-
-  if opcode == 0x03:  # META_SUBEVENT8_TS_SMALL
-    sub8 = pkt & 0xFF
-    fields.append(f"subevent8=0x{sub8:02x}")
-    if width:
-      fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    return ", ".join(fields)
-
-  if opcode == 0x04:  # META_BASE_INDEX12_TS
-    idx12 = pkt & 0xFFF
-    fields.append(f"base_index12=0x{idx12:03x}")
-    if width:
-      fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    return ", ".join(fields)
-
-  if opcode in (0x05, 0x06):  # META_DESC24_TS_A/B
-    desc24 = pkt & 0xFFFFFF
-    fields.append(f"desc24=0x{desc24:06x}")
-    if width:
-      fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 3. Opcode 0x14: exec/config record (+ COR marker)
-  # =====================================================================
-
-  if opcode == 0x14:  # INST_EXEC_OR_CFG
-    subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
-    val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
-    slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
-    hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
-
-    fields.append(f"subop=0x{subop:04x}")
-    fields.append(f"slot={slot}")
-    fields.append(f"val32=0x{val32:08x}")
-
-    if hi_byte & 0x80:
-      # Config flavour: writes config words into per-slot state arrays.
-      fields.append("kind=config")
-      if subop == 0x000C:
-        fields.append("cfg_target=local_168[slot].lo")
-      elif subop == 0x000D:
-        fields.append("cfg_target=local_168[slot].hi")
-    else:
-      # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
-      if subop == 0xC342:
-        fields.append("kind=cor_stream")
-        if val32 == 0x434F5200:
-          fields.append("cor_magic='COR\\0'")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 4. Opcode 0x17: layout / mode header
-  # =====================================================================
-
-  if opcode == 0x17:  # LAYOUT_MODE_HEADER
-    # From decomp (two sites with identical logic):
-    #   layout = (w >> 7) & 0x3f
-    #   mode   = (w >> 0xd) & 3
-    #   group  = (w >> 0xf) & 7
-    #   sel_a  = (w >> 0x1c) & 0xf
-    #   sel_b  = (w >> 0x21) & 7
-    #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
-    layout = (pkt >> 7)  & 0x3F
-    mode   = (pkt >> 13) & 0x3
-    group  = (pkt >> 15) & 0x7
-    sel_a  = (pkt >> 0x1C) & 0xF
-    sel_b  = (pkt >> 0x21) & 0x7
-    flag4  = (pkt >> 0x3B) & 0x1
-
-    fields.append(f"layout={layout}")
-    fields.append(f"group={group}")
-    fields.append(f"mode={mode}")
-    fields.append(f"sel_a={sel_a}")
-    fields.append(f"sel_b={sel_b}")
-    if layout == 4:
-      fields.append(f"layout4_flag={flag4}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 5. Opcode 0x09: state / route config record
-  # =====================================================================
-
-  if opcode == 0x09:  # PERF_ROUTE_CONFIG
-    # From case 9 in multiple consumers:
-    #   flag7  = (w >> 7) & 1        (low bit of uVar41)
-    #   cls2   = (w >> 8) & 3        (class / group)
-    #   slot4  = (w >> 10) & 0xf     (slot / group index)
-    #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
-    #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
-    #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
-    flag7   = (pkt >> 7) & 0x1
-    cls2    = (pkt >> 8) & 0x3
-    slot4   = (pkt >> 10) & 0xF
-    idx_lo  = (pkt >> 13) & 0x1F
-    idx_hi  = (pkt >> 15) & 0x1F
-    id7     = (pkt >> 0x19) & 0x7F
-
-    fields.append(f"flag7={flag7}")
-    fields.append(f"cls2={cls2}")
-    fields.append(f"slot4=0x{slot4:x}")
-    fields.append(f"idx_lo5=0x{idx_lo:x}")
-    fields.append(f"idx_hi5=0x{idx_hi:x}")
-    fields.append(f"id7=0x{id7:x}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 6. Opcode 0x18: perf/event selector (FUN_0010aba0)
-  # =====================================================================
-
-  if opcode == 0x18:  # PERF_EVENT_SELECT
-    # From case 0x18:
-    #   low3   = w & 7
-    #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
-    #   flags  = bits 6 (B6) and 7 (B7)
-    #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
-    #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
-    #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
-    low3     = pkt & 0x7
-    grp3_a   = (pkt >> 3) & 0x7
-    grp3_b   = (pkt >> 4) & 0x7
-    flag_b6  = (pkt >> 6) & 0x1
-    flag_b7  = (pkt >> 7) & 0x1
-    idx5_a   = (pkt >> 7) & 0x1F
-    idx5_b   = (pkt >> 8) & 0x1F
-    hi8      = (pkt >> 12) & 0xFF
-    hi7      = (pkt >> 13) & 0x7F
-
-    fields.append(f"low3=0x{low3:x}")
-    fields.append(f"grp3_a=0x{grp3_a:x}")
-    fields.append(f"grp3_b=0x{grp3_b:x}")
-    fields.append(f"flag_b6={flag_b6}")
-    fields.append(f"flag_b7={flag_b7}")
-    fields.append(f"idx5_a=0x{idx5_a:x}")
-    fields.append(f"idx5_b=0x{idx5_b:x}")
-    fields.append(f"hi8=0x{hi8:02x}")
-    fields.append(f"hi7=0x{hi7:02x}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 7. Opcode 0x15: perfcounter snapshot
-  # =====================================================================
-
-  if opcode == 0x15:  # PERFCOUNTER_SNAPSHOT
-    # NIBBLE_BUDGET gives full 64 bits here.
-    # DELTA_MAP_DEFAULT: shift=7, width=3 → tiny delta field.
-    raw_delta = shaped_field if width else 0
-    # low bits below the delta field
-    snap_low  = pkt & ((1 << shift) - 1) if shift else 0
-    # everything above delta field
-    snap_hi   = pkt >> (shift + width) if width else (pkt >> shift)
-
-    fields.append(f"raw_delta={raw_delta}")
-    fields.append(f"snap_low_s{shift}=0x{snap_low:x}")
-    fields.append(f"snap_hi=0x{snap_hi:x}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 8. Small event-ish packets (0x08 / 0x12 / 0x13 / 0x19)
-  # =====================================================================
-
-  if opcode in (0x08, 0x12, 0x13, 0x19):
-    # These are all "small event / metric" style tokens. The exact semantics
-    # depend on layout (0x17) and accumulated state (local_500 etc), so we
-    # expose:
-    #   - low 8 bits as kind byte
-    #   - rest as opaque payload.
-    kind    = pkt & 0xFF
-    payload = pkt >> 8
-    fields.append(f"kind_byte=0x{kind:02x}")
-    if payload:
-      fields.append(f"payload=0x{payload:x}")
-    return ", ".join(fields)
-
-  # =====================================================================
-  # 9. Pseudo opcode 0x10: never a "real" packet
-  # =====================================================================
-
-  if opcode == 0x10:  # PSEUDO_NEED_MORE_BITS
-    # The main loop never prints these; they're just a control token.
-    return ""
-
-  # =====================================================================
-  # 10. Generic fallback: expose the DELTA_MAP_DEFAULT field + leftover
-  # =====================================================================
-
-  if width:
-    fields.append(f"field_s{shift}_w{width}={shaped_field}")
-    leftover = pkt & ~(field_mask << shift)
-    if leftover:
-      fields.append(f"payload=0x{leftover:x}")
-
-  return ", ".join(fields)
-
-# 0xb is time something
-# 0xd is time something
-# 0xf is small time advance
-# 0x11 is time advance
-# 0x16 is big time advance + markers
-# 0x14 is REG
-DEFAULT_FILTER = (0xb, 0xd, 0xf, 0x11, 0x16, 0x14) if getenv("FILTER", 1) else None
-
-def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000, filter=DEFAULT_FILTER) -> None:
-  """
-  Minimal debug: print ONE LINE per decoded token (packet).
-
-  Now prints only the actual nibbles that belong to each packet, instead of
-  the full 64-bit shift register.
-  """
-  n = len(data)
-  time = 0
-  reg = 0          # shift register
-  offset = 0       # bit offset, in steps of 4 (one nibble)
-  nib_budget = 0x40
-  flags = 0
-  token_index = 0
-
-  while (offset >> 3) < n and token_index < max_tokens:
-    # Remember where we started refilling for this step (bit offset),
-    # but the *logical* start of the current packet is last_real_offset.
-    refill_start = offset
-
-    # 1) Fill register with nibbles according to nib_budget
-    if nib_budget != 0:
-      target = refill_start + 4 + ((nib_budget - 1) & ~3)
-      cur = refill_start
-      while cur != target and (cur >> 3) < n:
-        byte_index = cur >> 3
-        byte = data[byte_index]
-        shift = 4 if (cur & 4) else 0  # low then high nibble
-        nib = (byte >> shift) & 0xF
-        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
-        cur += 4
-      offset = cur
-
-    # 2) Decode token from low 8 bits
-    state = reg & 0xFF
-    opcode = STATE_TO_TOKEN[state]
-
-    # 3) Handle pseudo-token 0x10: need more bits, don't print. Looks like a NOP.
-    if opcode == 0x10:
-      # "need more bits" pseudo-token: adjust nibble budget and continue
-      nib_budget = 4
-      if (offset >> 3) >= n:
-        break
-      # Do NOT count this as a real packet; do not update last_real_offset.
-      continue
-
-    # 4) Set next nibble budget
-    nb_index = opcode & 0x1F
-    nib_budget = NIBBLE_BUDGET[nb_index]
-    time_before = time
-    note = ""
-    # 5) Special opcode 0x16 (timestamp / marker)
-    if opcode == 0x16:
-      two_bits = (reg >> 8) & 0x3
-      if two_bits == 1:
-        flags |= 0x01
-
-      # Common 36-bit field at bits [12..47]
-
-      if (reg & 0x200) == 0:
-        # delta mode: add 36-bit delta to time
-        delta = (reg >> 12) & ((1 << 36) - 1)
-        time += delta
-      else:
-        # marker / other modes: no time advance
-        if (reg & 0x100) == 0:
-          # real marker: bit9=1, bit8=0, non-zero payload
-          # "other" 0x16 variants, ignored for timing
-          delta = 0
-    else:
-      # 6) Generic opcode (including 0x0F)
-      shift, width = DELTA_MAP_DEFAULT[opcode]
-      mask = (1 << width) - 1
-      delta = (reg >> shift) & mask
-
-      # TODO: add more opcode parsers here that add notes to other opcodes
-      if opcode == 0x0F:
-        delta_with_fix = delta + 4
-        time += delta_with_fix
-        delta = delta_with_fix
-      else:
-        time += delta
-
-    # Append extra decoded fields into the note string
-    note = decode_packet_fields(opcode, reg, delta)
-
-    if filter is None or opcode not in filter:
-      my_reg = reg
-      my_reg &= (1 << nib_budget) - 1
-      print(
-        f"{token_index:4d}  "
-        f"off={offset//4:5d}  "
-        f"op=0x{opcode:02x} "
-        f"{OPCODE_NAMES[opcode]:24s} "
-        f" time={time_before:8d}+{delta:8d}  "
-        f"{my_reg:16X} "
-        f"{note}"
-      )
-
-    token_index += 1
-
-  # Optional summary at the end
-  print(f"# done: tokens={token_index}, final_time={time}, flags=0x{flags:02x}")
-
-def parse(fn:str):
-  dat = pickle.load(open(fn, "rb"))
-  ctx = decode(dat)
-  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
-  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
-  return dat_sqtt
-
-if __name__ == "__main__":
-  #dat_sqtt = parse("extra/sqtt/examples/profile_empty_run_0.pkl")
-  #dat_sqtt = parse("extra/sqtt/examples/profile_plus_run_0.pkl")
-  dat_sqtt = parse("extra/sqtt/examples/profile_gemm_run_0.pkl")
-  blob_0 = dat_sqtt[0].blob
-  parse_sqtt_print_packets(blob_0[8:])
@@ -0,0 +1,68 @@
+import ctypes
+from dataclasses import dataclass
+import tinygrad.runtime.autogen.comgr as comgr
+from tinygrad.runtime.support.compiler_amd import check
+
+@dataclass
+class InstrCtx:
+  pc:int=0
+  inst:str=""
+
+@comgr.amd_comgr_create_disassembly_info.argtypes[2]
+def instr_cb(text, user_data):
+  c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
+  c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
+  return comgr.AMD_COMGR_STATUS_SUCCESS
+
+# nop callback
+@comgr.amd_comgr_create_disassembly_info.argtypes[3]
+def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
+
+def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
+  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
+  lib_buf = ctypes.create_string_buffer(lib, len(lib))
+  check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
+  check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
+
+  @comgr.amd_comgr_create_disassembly_info.argtypes[1]
+  def memory_cb(from_addr, to, size, _):
+    base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
+    start = int(from_addr) - base
+    if start < 0 or start >= buf_len: return 0
+    ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
+    return n
+
+  info_src = comgr.amd_comgr_disassembly_info_t()
+  check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
+
+  @comgr.amd_comgr_iterate_symbols.argtypes[1]
+  def sym_callback(sym, udata):
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
+    if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
+    check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
+                                                                        ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
+    check(nobits.value)
+    base = ctypes.addressof(lib_buf)
+    pc = base + offset.value
+    end = pc + size.value
+    addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
+    instr_ref = ctypes.py_object(ctx:=InstrCtx())
+    instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
+    while pc < end:
+      size_read = ctypes.c_uint64(0)
+      ctx.pc = pc
+      st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
+      if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
+        rel = (pc - base) - offset.value
+        addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
+        pc += size_read.value
+      else: # don't inf loop if comgr fails
+        b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
+        addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
+        pc += 1
+    return comgr.AMD_COMGR_STATUS_SUCCESS
+  addr_table:dict[int, tuple[str, int]] = {}
+  check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
+  return addr_table
@@ -4,7 +4,7 @@ import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
 import tinygrad.runtime.autogen.sqtt as sqtt
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from tinygrad.helpers import round_up, flatten, all_same, temp
+from tinygrad.helpers import round_up, flatten, all_same
 from dataclasses import dataclass

 CHUNK_CLASSES = {
@@ -154,22 +154,8 @@ class RGP:
      if device not in device_events: raise RuntimeError(f"Device {device} not found in profile, devices in profile: {', '.join(device_events.keys())} ")
      device_event = device_events[device]
    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
-    device_props = device_event.props
-    # merge events per SE
-    merged_sqtt_events:dict[int, ProfileSQTTEvent] = {}
-    for ev in sqtt_events:
-      if ev.se not in merged_sqtt_events: merged_sqtt_events[ev.se] = ev
-      else:
-        merged_sqtt_events[ev.se] = ProfileSQTTEvent(
-          device=ev.device,
-          kern=ev.kern,
-          se=ev.se,
-          itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
-          blob=merged_sqtt_events[ev.se].blob + ev.blob,
-        )
-    sqtt_events = list(merged_sqtt_events.values())
-
    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
+    device_props = sqtt_events[0].props
    gfx_ver = device_props['gfx_target_version'] // 10000
    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
@@ -185,7 +171,9 @@ class RGP:
      magic_number=sqtt.SQTT_FILE_MAGIC_NUMBER,
      version_major=sqtt.SQTT_FILE_VERSION_MAJOR,
      version_minor=sqtt.SQTT_FILE_VERSION_MINOR,
-      flags=sqtt.struct_sqtt_file_header_flags(value=1,),
+      flags=sqtt.struct_sqtt_file_header_flags(
+        _0=sqtt.union_sqtt_file_header_flags_0(value=1),
+      ),
      chunk_offset=ctypes.sizeof(sqtt.struct_sqtt_file_header),
    )
    chunks = [
@@ -208,7 +196,7 @@ class RGP:
        flags=0,
        trace_shader_core_clock=0x93f05080,
        trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
        device_revision_id=0xc8,
        vgprs_per_simd=1536,
        sgprs_per_simd=128*16,
@@ -263,7 +251,7 @@ class RGP:
        profiling_mode=sqtt.SQTT_PROFILING_MODE_PRESENT,
        instruction_trace_mode=sqtt.SQTT_INSTRUCTION_TRACE_FULL_FRAME if sqtt_itrace_enabled else sqtt.SQTT_INSTRUCTION_TRACE_DISABLED,
        instruction_trace_data=sqtt.union_sqtt_instruction_trace_data(
-          shader_engine_filter=sqtt.union_sqtt_instruction_trace_data_shader_engine_filter(mask=sqtt_itrace_se_mask),
+          shader_engine_filter=sqtt.struct_sqtt_instruction_trace_data_shader_engine_filter(mask=sqtt_itrace_se_mask),
        ),
      )),
      *flatten([(
@@ -274,11 +262,13 @@ class RGP:
          ),
          shader_engine_index=sqtt_event.se,
          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
-          v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
-            instrumentation_spec_version=1,
-            instrumentation_api_version=0,
-            compute_unit_index=0,
-          )
+          _0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
+            v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
+              instrumentation_spec_version=1,
+              instrumentation_api_version=0,
+              compute_unit_index=0,
+            )
+          ),
        )),
        RGPChunk(sqtt.struct_sqtt_file_chunk_sqtt_data(
          header=sqtt.struct_sqtt_file_chunk_header(
@@ -320,7 +310,7 @@ class RGP:
 if __name__ == '__main__':
  parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
  parser.add_argument('command')
-  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
+  parser.add_argument('input')
  parser.add_argument('-d', '--device')
  parser.add_argument('-o', '--output')
  args = parser.parse_args()
@@ -342,4 +332,3 @@ if __name__ == '__main__':

  if args.output is not None:
    with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
-    print(f"Saved to {args.output}")
@@ -1,119 +1,91 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
-from tinygrad.helpers import temp, unwrap, DEBUG
-from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+from extra.sqtt.rocprof import rocprof
+from extra.sqtt.disasm import comgr_get_address_table
+from tinygrad.helpers import temp, DEBUG
+from tinygrad.device import ProfileEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
-from tinygrad.runtime.autogen import llvm, rocprof
-from tinygrad.runtime.support.elf import elf_loader

-# to pass NULL to callbacks
-llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
-def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
-  llvm.LLVMInitializeAMDGPUTargetInfo()
-  llvm.LLVMInitializeAMDGPUTargetMC()
-  llvm.LLVMInitializeAMDGPUAsmParser()
-  llvm.LLVMInitializeAMDGPUDisassembler()
-  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
+@dataclasses.dataclass
+class InstInfo:
+  typ:str=""
+  inst:str=""
+  hit:int=0
+  lat:int=0
+  stall:int=0
+  def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"

-  image, sections, relocs = elf_loader(lib)
-  text = next((sh.header for sh in sections if sh.name == ".text"), None)
-  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
-
-  addr_table:dict[int, tuple[str, int]] = {}
-  out = ctypes.create_string_buffer(128)
-  cur_off = off
-  while cur_off < sz + off:
-    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
-    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
-    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
-    cur_off += instr_sz
-  return addr_table
-
-@dataclasses.dataclass(frozen=True)
-class InstExec:
-  typ:str
-  inst:str
-  stall:int
-  dur:int
-  time:int
-
-@dataclasses.dataclass(frozen=True)
-class WaveExec:
-  wave_id:int
-  cu:int
-  simd:int
-  se:int
-  begin_time:int
-  end_time:int
-  insts:list[InstExec]
+  def on_ev(self, ev):
+    self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall

 class _ROCParseCtx:
-  def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
-    self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
-    self.disasms:dict[tuple[str, int], tuple[str, int]] = {}
-    self.inst_execs:dict[str, list[WaveExec]] = {}
+  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
+    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
+    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}

    for prog in prog_evs:
-      arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
-      for addr, info in llvm_disasm(arch, unwrap(prog.lib)).items():
-        self.disasms[(prog.name, unwrap(prog.base) + addr)] = info
+      for addr, info in comgr_get_address_table(prog.lib).items():
+        self.disasms[prog.base + addr] = info
+        self.addr2prg[prog.base + addr] = prog

  def next_sqtt(self):
    x = next(self.sqtt_evs, None)
-    self.active_kern = x.kern if x is not None else None
    self.active_se = x.se if x is not None else None
-    self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
-    return self.active_blob
+    return x

-  def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
-    if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
+  def find_program(self, addr): return self.addr2prg[addr]

-  def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
-    if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+  def on_occupancy_ev(self, ev):
+    if DEBUG >= 4: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)

-    inst_execs:list[InstExec] = []
+  def on_wave_ev(self, ev):
+    if DEBUG >= 4: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+
+    asm = {}
    for j in range(ev.instructions_size):
      inst_ev = ev.instructions_array[j]
-      inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst_ev.category)
-      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
-      inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
-      if DEBUG >= 8: print(inst_execs[-1])
+      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
+      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
+      asm[inst_ev.pc.address].on_ev(inst_ev)

-    if ev.instructions_size > 0:
-      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
-                                                                               ev.end_time, inst_execs))
+    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm

-def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
-  dev_events:dict[str, ProfileDeviceEvent] = {}
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  args = parser.parse_args()
+
+  with args.profile.open("rb") as f: profile = pickle.load(f)
  sqtt_events:list[ProfileSQTTEvent] = []
+  pmc_events:list[ProfilePMCEvent] = []
  prog_events:list[ProfileProgramEvent] = []
  for e in profile:
-    if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
    if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
+    if isinstance(e, ProfilePMCEvent): pmc_events.append(e)
    if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)

-  ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
+  ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)

  @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, _):
-    if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
-    buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
-    buf_size[0] = len(prof_info)
-    return len(prof_info)
+  def copy_cb(buf, buf_size, data_ptr):
+    if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
+    buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
+    buf_size[0] = len(prof.blob)
+    return len(prof.blob)

  @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, _):
+  def trace_cb(record_type, events_ptr, n, data_ptr):
    match record_type:
      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
        for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
      case _:
-        if DEBUG >= 5: print(rocprof.enum_rocprofiler_thread_trace_decoder_record_type_t.get(record_type), events_ptr, n)
+        if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS

  @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]

    # this is the number of bytes to next instruction, set to 0 for end_pgm
    if instr == "s_endpgm": mem_size_ptr[0] = 0
@@ -126,29 +98,15 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:

    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS

-  def worker():
-    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
-  (t:=threading.Thread(target=worker, daemon=True)).start()
-  t.join()
-  return ROCParseCtx
+  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+  print('SQTT:', ROCParseCtx.wave_events.keys())

-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  args = parser.parse_args()
-
-  with args.profile.open("rb") as f: profile = pickle.load(f)
-  rctx = decode(profile)
-  print('SQTT:', rctx.inst_execs.keys())
-
-  for ev in profile:
-    if not isinstance(ev, ProfilePMCEvent): continue
+  for ev in pmc_events:
    print(f"PMC Event: dev={ev.device} kern={ev.kern}")
    ptr = 0
    for s in ev.sched:
      view = memoryview(ev.blob).cast('Q')
      print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
+      for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        print(f"\t\tInst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]}")
        ptr += 1
@@ -12,9 +12,7 @@ if __name__ == "__main__":
    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
    os.chmod(fp, 0o755)
    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
-    shutil.copy2(lib, DEST)
  else:
-    lib = DEST/"librocprof-trace-decoder.so"
-    os.system("sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o"+str(lib))
-    os.system("sudo ldconfig")
+    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
+  shutil.copy2(lib, DEST)
  print(f"Installed {lib.name} to", DEST)
@@ -0,0 +1,656 @@
+# pylint: skip-file
+# mypy: ignore-errors
+# -*- coding: utf-8 -*-
+#
+# TARGET arch is: []
+# WORD_SIZE is: 8
+# POINTER_SIZE is: 8
+# LONGDOUBLE_SIZE is: 16
+#
+import ctypes, ctypes.util
+
+
+class AsDictMixin:
+    @classmethod
+    def as_dict(cls, self):
+        result = {}
+        if not isinstance(self, AsDictMixin):
+            # not a structure, assume it's already a python object
+            return self
+        if not hasattr(cls, "_fields_"):
+            return result
+        # sys.version_info >= (3, 5)
+        # for (field, *_) in cls._fields_:  # noqa
+        for field_tuple in cls._fields_:  # noqa
+            field = field_tuple[0]
+            if field.startswith('PADDING_'):
+                continue
+            value = getattr(self, field)
+            type_ = type(value)
+            if hasattr(value, "_length_") and hasattr(value, "_type_"):
+                # array
+                if not hasattr(type_, "as_dict"):
+                    value = [v for v in value]
+                else:
+                    type_ = type_._type_
+                    value = [type_.as_dict(v) for v in value]
+            elif hasattr(value, "contents") and hasattr(value, "_type_"):
+                # pointer
+                try:
+                    if not hasattr(type_, "as_dict"):
+                        value = value.contents
+                    else:
+                        type_ = type_._type_
+                        value = type_.as_dict(value.contents)
+                except ValueError:
+                    # nullptr
+                    value = None
+            elif isinstance(value, AsDictMixin):
+                # other structure
+                value = type_.as_dict(value)
+            result[field] = value
+        return result
+
+
+class Structure(ctypes.Structure, AsDictMixin):
+
+    def __init__(self, *args, **kwds):
+        # We don't want to use positional arguments fill PADDING_* fields
+
+        args = dict(zip(self.__class__._field_names_(), args))
+        args.update(kwds)
+        super(Structure, self).__init__(**args)
+
+    @classmethod
+    def _field_names_(cls):
+        if hasattr(cls, '_fields_'):
+            return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
+        else:
+            return ()
+
+    @classmethod
+    def get_type(cls, field):
+        for f in cls._fields_:
+            if f[0] == field:
+                return f[1]
+        return None
+
+    @classmethod
+    def bind(cls, bound_fields):
+        fields = {}
+        for name, type_ in cls._fields_:
+            if hasattr(type_, "restype"):
+                if name in bound_fields:
+                    if bound_fields[name] is None:
+                        fields[name] = type_()
+                    else:
+                        # use a closure to capture the callback from the loop scope
+                        fields[name] = (
+                            type_((lambda callback: lambda *args: callback(*args))(
+                                bound_fields[name]))
+                        )
+                    del bound_fields[name]
+                else:
+                    # default callback implementation (does nothing)
+                    try:
+                        default_ = type_(0).restype().value
+                    except TypeError:
+                        default_ = None
+                    fields[name] = type_((
+                        lambda default_: lambda *args: default_)(default_))
+            else:
+                # not a callback function, use default initialization
+                if name in bound_fields:
+                    fields[name] = bound_fields[name]
+                    del bound_fields[name]
+                else:
+                    fields[name] = type_()
+        if len(bound_fields) != 0:
+            raise ValueError(
+                "Cannot bind the following unknown callback(s) {}.{}".format(
+                    cls.__name__, bound_fields.keys()
+            ))
+        return cls(**fields)
+
+
+class Union(ctypes.Union, AsDictMixin):
+    pass
+
+
+
+c_int128 = ctypes.c_ubyte*16
+c_uint128 = c_int128
+void = None
+if ctypes.sizeof(ctypes.c_longdouble) == 16:
+    c_long_double_t = ctypes.c_longdouble
+else:
+    c_long_double_t = ctypes.c_ubyte*16
+
+def string_cast(char_pointer, encoding='utf-8', errors='strict'):
+    value = ctypes.cast(char_pointer, ctypes.c_char_p).value
+    if value is not None and encoding is not None:
+        value = value.decode(encoding, errors=errors)
+    return value
+
+
+def char_pointer_cast(string, encoding='utf-8'):
+    if encoding is not None:
+        try:
+            string = string.encode(encoding)
+        except AttributeError:
+            # In Python3, bytes has no encode attribute
+            pass
+    string = ctypes.c_char_p(string)
+    return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
+
+
+
+class FunctionFactoryStub:
+    def __getattr__(self, _):
+      return ctypes.CFUNCTYPE(lambda y:y)
+
+# libraries['FIXME_STUB'] explanation
+# As you did not list (-l libraryname.so) a library that exports this function
+# This is a non-working stub instead.
+# You can either re-run clan2py with -l /path/to/library.so
+# Or manually fix this by comment the ctypes.CDLL loading
+_libraries = {}
+_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')
+
+
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_info_t'
+rocprofiler_thread_trace_decoder_info_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = 0
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = 1
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = 2
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = 3
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = 4
+rocprofiler_thread_trace_decoder_info_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_pc_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_pc_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
+    ('address', ctypes.c_uint64),
+    ('code_object_id', ctypes.c_uint64),
+]
+
+rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
+class struct_rocprofiler_thread_trace_decoder_perfevent_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_perfevent_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
+    ('time', ctypes.c_int64),
+    ('events0', ctypes.c_uint16),
+    ('events1', ctypes.c_uint16),
+    ('events2', ctypes.c_uint16),
+    ('events3', ctypes.c_uint16),
+    ('CU', ctypes.c_ubyte),
+    ('bank', ctypes.c_ubyte),
+    ('PADDING_0', ctypes.c_ubyte * 6),
+]
+
+rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
+class struct_rocprofiler_thread_trace_decoder_occupancy_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_occupancy_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
+    ('pc', rocprofiler_thread_trace_decoder_pc_t),
+    ('time', ctypes.c_uint64),
+    ('reserved', ctypes.c_ubyte),
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('start', ctypes.c_uint32, 1),
+    ('_rsvd', ctypes.c_uint32, 31),
+]
+
+rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_wstate_type_t'
+rocprofiler_thread_trace_decoder_wstate_type_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = 0
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = 1
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = 2
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = 3
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = 4
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = 5
+rocprofiler_thread_trace_decoder_wstate_type_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_wave_state_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_wave_state_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
+    ('type', ctypes.c_int32),
+    ('duration', ctypes.c_int32),
+]
+
+rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_inst_category_t'
+rocprofiler_thread_trace_decoder_inst_category_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
+    6: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
+    7: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
+    8: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
+    9: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
+    10: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
+    11: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
+    12: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
+    13: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = 0
+ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = 1
+ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = 2
+ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = 3
+ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = 4
+ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = 5
+ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = 6
+ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = 7
+ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = 8
+ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = 9
+ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = 10
+ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = 11
+ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = 12
+ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = 13
+rocprofiler_thread_trace_decoder_inst_category_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_inst_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_inst_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
+    ('category', ctypes.c_uint32, 8),
+    ('stall', ctypes.c_uint32, 24),
+    ('duration', ctypes.c_int32),
+    ('time', ctypes.c_int64),
+    ('pc', rocprofiler_thread_trace_decoder_pc_t),
+]
+
+rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
+class struct_rocprofiler_thread_trace_decoder_wave_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_wave_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('contexts', ctypes.c_ubyte),
+    ('_rsvd1', ctypes.c_uint32),
+    ('_rsvd2', ctypes.c_uint32),
+    ('_rsvd3', ctypes.c_uint32),
+    ('begin_time', ctypes.c_int64),
+    ('end_time', ctypes.c_int64),
+    ('timeline_size', ctypes.c_uint64),
+    ('instructions_size', ctypes.c_uint64),
+    ('timeline_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_wave_state_t)),
+    ('instructions_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_inst_t)),
+]
+
+rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
+class struct_rocprofiler_thread_trace_decoder_realtime_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_realtime_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
+    ('shader_clock', ctypes.c_int64),
+    ('realtime_clock', ctypes.c_uint64),
+    ('reserved', ctypes.c_uint64),
+]
+
+rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_shaderdata_flags_t'
+rocprofiler_thread_trace_decoder_shaderdata_flags_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = 0
+ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = 1
+rocprofiler_thread_trace_decoder_shaderdata_flags_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_shaderdata_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
+    ('time', ctypes.c_int64),
+    ('value', ctypes.c_uint64),
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('flags', ctypes.c_ubyte),
+    ('reserved', ctypes.c_uint32),
+]
+
+rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_record_type_t'
+rocprofiler_thread_trace_decoder_record_type_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
+    6: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
+    7: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
+    8: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
+    9: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = 0
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = 1
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = 2
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = 3
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = 4
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = 5
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = 6
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = 7
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = 8
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = 9
+rocprofiler_thread_trace_decoder_record_type_t = ctypes.c_uint32 # enum
+
+# values for enumeration 'c__EA_rocprofiler_thread_trace_decoder_status_t'
+c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = 0
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = 1
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = 2
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = 3
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = 4
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = 5
+c__EA_rocprofiler_thread_trace_decoder_status_t = ctypes.c_uint32 # enum
+rocprofiler_thread_trace_decoder_status_t = c__EA_rocprofiler_thread_trace_decoder_status_t
+rocprofiler_thread_trace_decoder_status_t__enumvalues = c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues
+rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, rocprofiler_thread_trace_decoder_record_type_t, ctypes.POINTER(None), ctypes.c_uint64, ctypes.POINTER(None))
+rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.POINTER(None))
+rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(None))
+try:
+    rocprof_trace_decoder_parse_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_parse_data
+    rocprof_trace_decoder_parse_data.restype = rocprofiler_thread_trace_decoder_status_t
+    rocprof_trace_decoder_parse_data.argtypes = [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.POINTER(None)]
+except AttributeError:
+    pass
+try:
+    rocprof_trace_decoder_get_info_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_info_string
+    rocprof_trace_decoder_get_info_string.restype = ctypes.POINTER(ctypes.c_char)
+    rocprof_trace_decoder_get_info_string.argtypes = [rocprofiler_thread_trace_decoder_info_t]
+except AttributeError:
+    pass
+try:
+    rocprof_trace_decoder_get_status_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_status_string
+    rocprof_trace_decoder_get_status_string.restype = ctypes.POINTER(ctypes.c_char)
+    rocprof_trace_decoder_get_status_string.argtypes = [rocprofiler_thread_trace_decoder_status_t]
+except AttributeError:
+    pass
+rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(None))
+uint64_t = ctypes.c_uint64
+try:
+    rocprof_trace_decoder_dump_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_dump_data
+    rocprof_trace_decoder_dump_data.restype = rocprofiler_thread_trace_decoder_status_t
+    rocprof_trace_decoder_dump_data.argtypes = [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.POINTER(None)]
+except AttributeError:
+    pass
+class union_rocprof_trace_decoder_gfx9_header_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_gfx9_header_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_gfx9_header_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
+    ('legacy_version', ctypes.c_uint64, 13),
+    ('gfx9_version2', ctypes.c_uint64, 3),
+    ('DSIMDM', ctypes.c_uint64, 4),
+    ('DCU', ctypes.c_uint64, 5),
+    ('reserved1', ctypes.c_uint64, 1),
+    ('SEID', ctypes.c_uint64, 6),
+    ('reserved2', ctypes.c_uint64, 32),
+]
+
+union_rocprof_trace_decoder_gfx9_header_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_gfx9_header_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_gfx9_header_t_0),
+    ('raw', ctypes.c_uint64),
+]
+
+rocprof_trace_decoder_gfx9_header_t = union_rocprof_trace_decoder_gfx9_header_t
+class union_rocprof_trace_decoder_instrument_enable_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_instrument_enable_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_instrument_enable_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
+    ('char1', ctypes.c_uint32, 8),
+    ('char2', ctypes.c_uint32, 8),
+    ('char3', ctypes.c_uint32, 8),
+    ('char4', ctypes.c_uint32, 8),
+]
+
+union_rocprof_trace_decoder_instrument_enable_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_instrument_enable_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_instrument_enable_t_0),
+    ('u32All', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_instrument_enable_t = union_rocprof_trace_decoder_instrument_enable_t
+class union_rocprof_trace_decoder_packet_header_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_packet_header_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_packet_header_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_packet_header_t_0._fields_ = [
+    ('opcode', ctypes.c_uint32, 8),
+    ('type', ctypes.c_uint32, 4),
+    ('data20', ctypes.c_uint32, 20),
+]
+
+union_rocprof_trace_decoder_packet_header_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_packet_header_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_packet_header_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_packet_header_t_0),
+    ('u32All', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_packet_header_t = union_rocprof_trace_decoder_packet_header_t
+
+# values for enumeration 'rocprof_trace_decoder_packet_opcode_t'
+rocprof_trace_decoder_packet_opcode_t__enumvalues = {
+    4: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
+    5: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
+    6: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
+}
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = 4
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = 5
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = 6
+rocprof_trace_decoder_packet_opcode_t = ctypes.c_uint32 # enum
+
+# values for enumeration 'rocprof_trace_decoder_agent_info_type_t'
+rocprof_trace_decoder_agent_info_type_t__enumvalues = {
+    0: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
+    1: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
+    2: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
+}
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = 0
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = 1
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = 2
+rocprof_trace_decoder_agent_info_type_t = ctypes.c_uint32 # enum
+class union_rocprof_trace_decoder_codeobj_marker_tail_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_codeobj_marker_tail_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
+    ('isUnload', ctypes.c_uint32, 1),
+    ('bFromStart', ctypes.c_uint32, 1),
+    ('legacy_id', ctypes.c_uint32, 30),
+]
+
+union_rocprof_trace_decoder_codeobj_marker_tail_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_codeobj_marker_tail_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_codeobj_marker_tail_t_0),
+    ('raw', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_codeobj_marker_tail_t = union_rocprof_trace_decoder_codeobj_marker_tail_t
+
+# values for enumeration 'rocprof_trace_decoder_codeobj_marker_type_t'
+rocprof_trace_decoder_codeobj_marker_type_t__enumvalues = {
+    0: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
+    1: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
+    2: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
+    3: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
+    4: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
+    5: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
+    6: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
+    7: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
+}
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = 0
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = 1
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = 2
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = 3
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = 4
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = 5
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = 6
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = 7
+rocprof_trace_decoder_codeobj_marker_type_t = ctypes.c_uint32 # enum
+__all__ = \
+    ['ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
+    'c__EA_rocprofiler_thread_trace_decoder_status_t',
+    'rocprof_trace_decoder_agent_info_type_t',
+    'rocprof_trace_decoder_codeobj_marker_tail_t',
+    'rocprof_trace_decoder_codeobj_marker_type_t',
+    'rocprof_trace_decoder_dump_data',
+    'rocprof_trace_decoder_get_info_string',
+    'rocprof_trace_decoder_get_status_string',
+    'rocprof_trace_decoder_gfx9_header_t',
+    'rocprof_trace_decoder_instrument_enable_t',
+    'rocprof_trace_decoder_isa_callback_t',
+    'rocprof_trace_decoder_packet_header_t',
+    'rocprof_trace_decoder_packet_opcode_t',
+    'rocprof_trace_decoder_parse_data',
+    'rocprof_trace_decoder_se_data_callback_t',
+    'rocprof_trace_decoder_trace_callback_t',
+    'rocprofiler_thread_trace_decoder_debug_callback_t',
+    'rocprofiler_thread_trace_decoder_info_t',
+    'rocprofiler_thread_trace_decoder_inst_category_t',
+    'rocprofiler_thread_trace_decoder_inst_t',
+    'rocprofiler_thread_trace_decoder_occupancy_t',
+    'rocprofiler_thread_trace_decoder_pc_t',
+    'rocprofiler_thread_trace_decoder_perfevent_t',
+    'rocprofiler_thread_trace_decoder_realtime_t',
+    'rocprofiler_thread_trace_decoder_record_type_t',
+    'rocprofiler_thread_trace_decoder_shaderdata_flags_t',
+    'rocprofiler_thread_trace_decoder_shaderdata_t',
+    'rocprofiler_thread_trace_decoder_status_t',
+    'rocprofiler_thread_trace_decoder_status_t__enumvalues',
+    'rocprofiler_thread_trace_decoder_wave_state_t',
+    'rocprofiler_thread_trace_decoder_wave_t',
+    'rocprofiler_thread_trace_decoder_wstate_type_t',
+    'struct_rocprof_trace_decoder_codeobj_marker_tail_t_0',
+    'struct_rocprof_trace_decoder_gfx9_header_t_0',
+    'struct_rocprof_trace_decoder_instrument_enable_t_0',
+    'struct_rocprof_trace_decoder_packet_header_t_0',
+    'struct_rocprofiler_thread_trace_decoder_inst_t',
+    'struct_rocprofiler_thread_trace_decoder_occupancy_t',
+    'struct_rocprofiler_thread_trace_decoder_pc_t',
+    'struct_rocprofiler_thread_trace_decoder_perfevent_t',
+    'struct_rocprofiler_thread_trace_decoder_realtime_t',
+    'struct_rocprofiler_thread_trace_decoder_shaderdata_t',
+    'struct_rocprofiler_thread_trace_decoder_wave_state_t',
+    'struct_rocprofiler_thread_trace_decoder_wave_t', 'uint64_t',
+    'union_rocprof_trace_decoder_codeobj_marker_tail_t',
+    'union_rocprof_trace_decoder_gfx9_header_t',
+    'union_rocprof_trace_decoder_instrument_enable_t',
+    'union_rocprof_trace_decoder_packet_header_t']
@@ -1,133 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["VIZ"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-import unittest
-import sys, contextlib
-from tinygrad import Tensor, dtypes
-from tinygrad.helpers import getenv
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.device import Device, ProfileDeviceEvent
-
-from extra.sqtt.roc import decode, WaveExec
-
-dev = Device[os.environ["DEV"]]
-
-def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
-
-def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
-  name = sys._getframe(1).f_code.co_name
-  def fxn(_):
-    L = UOp.special(l, "lidx0")
-    G = UOp.special(g, "gidx0")
-    op = custom("asm volatile (")
-    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
-    op = custom(");", op)
-    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
-  k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-  return k
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[str, list[WaveExec]] = {}
-  yield sqtt
-  # decode sqtt
-  if os.environ["DEV"] != "AMD": return
-  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
-  assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  sqtt.update(rctx.inst_execs)
-
-class TestTiming(unittest.TestCase):
-  def test_v_add(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([f"v_add_f32 v{10+i} v{10+i+1} {10+i}" for i in range(3)]).realize()
-    wave = list(sqtt.values())[0][:-1]
-    assert all(s.dur == 1 for s in wave)
-    assert all(s.stall == 0 for s in wave)
-
-  def test_chain_v_add_1l(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_add_f32_e32 v1 v0 v0",
-        "v_add_f32_e32 v2 v1 v1",
-      ]).realize()
-    wave = list(sqtt.values())[0][:-1]
-    assert all(s.dur == 1 for s in wave)
-    assert all(s.stall == 0 for s in wave)
-
-  def test_multi_cycle_inst(self):
-    def custom_vrcp(A, B):
-      op = custom("float a = 0.0;")
-      op = custom("float b = (*(data1_1+0));", op)
-      #op = custom('asm volatile("v_mul_f32_e32 %2 %2 %1" : "+v"(a) : "v"(b));', op)
-      op = custom('asm volatile("v_rcp_f32_e32 %2 %1" : "+v"(a) : "v"(b));', op)
-      op = custom('asm volatile("v_add_f32_e64 %1 %1 1.0" : "+v"(a));', op)
-      op = custom("*(data0_1+0) = a;", op)
-      return UOp.sink(op, A, B, arg=KernelInfo(name="custom_vrcp"))
-    out = Tensor([0.]).realize()
-    inp = Tensor([-2.0]).realize()
-    with save_sqtt() as sqtt:
-      Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
-    wave = list(sqtt.values())[0][0]
-    for i in range(len(wave.insts)):
-      if wave.insts[i].inst.startswith("global_store"):
-        print(f"store diff {wave.insts[i].time-(wave.insts[i-1].time)}")
-    self.assertEqual(out.item(), 0.5)
-
-  def test_wmma(self):
-    with save_sqtt() as sqtt:
-      for tc in dev.renderer.get_tensor_cores(dev.arch):
-        M, K, N = tc.dims
-        s = 32
-        a = Tensor.empty(M*s, K*s, dtype=tc.dtype_in)@Tensor.empty(K*s, N*s, dtype=tc.dtype_in)
-        a.realize()
-        print(a)
-    for p,waves in sqtt.items():
-      for e in waves[0].insts:
-        if (e.inst.startswith("v_wmma")):
-          instruction = e.inst.split(" ")[0]
-          print(f"{instruction:<29} : {e.dur} cycles")
-
-  def test_sleep(self):
-    n = 1
-    def sleep_kernel(data0):
-      assert data0.dtype.base == dtypes.ulong
-      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
-      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom("unsigned long long t1 = __builtin_readcyclecounter();", op)
-      op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
-      return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
-    diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
-    diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
-    with save_sqtt() as sqtt:
-      diff_hw_reg.realize()
-    sleep = next((e for e in sqtt[f"sleep_{n}"][0].insts if e.inst.startswith("s_sleep")))
-    # cycles = sleep dur + overhead of storing hi/lo REG_SHADER_CYCLES
-    self.assertGreaterEqual(diff_hw_reg.item(), sleep.dur)
-
-  def test_nop(self):
-    with save_sqtt() as sqtt:
-      asm_kernel(["s_nop 1"]*10).realize()
-    wave = list(sqtt.values())[0][0]
-    for e in wave.insts:
-      print(f"{e.inst} {e.dur=} {e.stall=}")
-
-  def test_wave_sched(self):
-    num_waves = getenv("NUM_WAVES", 16)
-    num_wgps = getenv("NUM_WGPS", 2)
-    num_vgpr = getenv("NUM_VGPR", 256)
-    with save_sqtt() as sqtt:
-      # 1 cycle decode, no stall
-      asm_kernel([f"v_mov_b32_e32 v{i} {i}" for i in range(num_vgpr)], l=32*num_waves, g=num_wgps).realize()
-    waves = list(sqtt.values())[0]
-    print(len(waves), "waves decoded")
-    for w in waves:
-      print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
-
-if __name__ == "__main__":
-  unittest.main()
@@ -2,7 +2,7 @@

 using namespace kittens;

-constexpr int NUM_WORKERS = 4;
+constexpr int NUM_WORKERS = 2;
 constexpr int PIPE_STAGES = 3;

 constexpr int ATTN_B = 16;
@@ -10,7 +10,7 @@ constexpr int ATTN_N = 1024;
 constexpr int ATTN_H = 16;
 constexpr int ATTN_D = 64;

-template<int D> constexpr size_t ROWS = 16*(64/D); // height of each worker tile (rows)
+template<int D> constexpr size_t ROWS = 16*(128/D); // height of each worker tile (rows)
 template<int D, typename T=bf16, typename L=row_l> using qkvo_tile = rt<T, ROWS<D>, D, L>;
 template<int D, typename T=float> using attn_tile = rt<T, ROWS<D>, ROWS<D>>;
 template<int D> using shared_tile = st_bf<ROWS<D>, D>;
@@ -13,7 +13,7 @@ if __name__ == "__main__":
  print(pretty_ptx(lib.decode()))

  prg = device.runtime(kernel_name, lib)
-  prg.smem = 16384 * 3
+  prg.smem = 16384 * 2

  B, N, H, D = 16, 1024, 16, 64
  q = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
@@ -22,8 +22,8 @@ if __name__ == "__main__":
  out = Tensor.empty(B, N, H, D, device='CUDA', dtype="bfloat16")
  Tensor.realize(q, k, v, out)

-  NUM_WORKERS = 4
-  ROWS = 16 * (64 // D)
+  NUM_WORKERS = 2
+  ROWS = 16 * (128 // D)

  gsz = (N // (ROWS*NUM_WORKERS), H, B)
  for _ in range(5):
@@ -5,11 +5,11 @@ using namespace kittens;
 constexpr int g_N = 8192;
 constexpr int BLOCK_SIZE = 32;
 #define NUM_WORKERS (1)
+#define NUM_THREADS (NUM_WORKERS*kittens::WARP_THREADS)

 using sub_tile = st_bf<BLOCK_SIZE,BLOCK_SIZE>;
-using tile_gl =  gl<bf16, 1, 1, g_N, g_N>;
+using tile_gl =  gl<bf16,  1, 1, g_N, g_N, sub_tile>;

-__launch_bounds__(NUM_WORKERS*WARP_THREADS, 1)
 __global__ void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
  tile_gl g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
  tile_gl g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -1,14 +1,10 @@
 import pathlib
 from tinygrad import Device, Tensor
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context
 from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler

 if __name__ == "__main__":
-  if getenv("MATMUL2"):
-    code = (pathlib.Path(__file__).parent / "matmul2.cu").read_text()
-  else:
-    code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
-
+  code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
  device = Device["CUDA"]
  kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr"]
  lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
@@ -17,10 +13,7 @@ if __name__ == "__main__":
  print(pretty_ptx(lib.decode()))

  prg = device.runtime(kernel_name, lib)
-  if getenv("MATMUL2"):
-    prg.smem = 16384 * 2
-  else:
-    prg.smem = 10000
+  prg.smem = 10000

  N = 8192
  a = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
@@ -28,25 +21,14 @@ if __name__ == "__main__":
  c = Tensor.empty(N, N, device='CUDA', dtype="bfloat16")
  Tensor.realize(a, b, c)

-  WARP_THREADS = 32
-  if getenv("MATMUL2"):
-    SUPER_N = 2
-    SUPER_M = 2
-    NUM_WORKERS = SUPER_N * SUPER_M
-    BLOCK_SIZE = 32
-    gsz = (N // (BLOCK_SIZE * SUPER_N), N // (BLOCK_SIZE * SUPER_M), 1)
-  else:
-    NUM_WORKERS = 1
-    BLOCK_SIZE = 32
-    gsz = (N // (BLOCK_SIZE), N // (BLOCK_SIZE), 1)
+  BLOCK_SIZE = 32

+  gsz = (N // BLOCK_SIZE, N // BLOCK_SIZE, 1)
  for _ in range(5):
    et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
-             global_size=gsz, local_size=(NUM_WORKERS*WARP_THREADS,1,1), wait=True)
+             global_size=gsz, local_size=(32,1,1), wait=True)
    print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")

-  # print(c.tolist())
-
  for _ in range(5):
    with Context(DEBUG=2):
      ref = (a@b).realize()
@@ -1,105 +0,0 @@
-#include "kittens.cuh"
-using namespace kittens;
-
-constexpr int g_N = 8192;
-
-constexpr int SUPER_N = 2;
-constexpr int SUPER_M = 2;
-constexpr int NUM_WORKERS = SUPER_N * SUPER_M;
-constexpr int LOAD_TASKS = SUPER_N + SUPER_M;
-
-constexpr int WORKER_M = 32;
-constexpr int WORKER_N = 32;
-
-constexpr int BLOCK_K = 32;
-constexpr int BLOCK_M = WORKER_M * SUPER_M;
-constexpr int BLOCK_N = WORKER_N * SUPER_N;
-
-constexpr int PIPE_STAGES = 2;
-
-using reg_tile_A = rt_bf<WORKER_M, BLOCK_K>;
-using reg_tile_B_col = rt_bf<BLOCK_K, WORKER_N, ducks::rt_layout::col>;
-using reg_tile_C = rt_fl<WORKER_M, WORKER_N>;
-
-using shared_tile_A = st_bf<WORKER_M, BLOCK_K>;
-using shared_tile_B = st_bf<BLOCK_K, WORKER_N>;
-using shared_tile_C = st_bf<WORKER_M, WORKER_N>;
-
-using gl_tile_A = gl<bf16, 1, 1, g_N, g_N, shared_tile_A>;
-using gl_tile_B = gl<bf16, 1, 1, g_N, g_N, shared_tile_B>;
-using gl_tile_C = gl<bf16, 1, 1, g_N, g_N, shared_tile_C>;
-
-__launch_bounds__(NUM_WORKERS *WARP_THREADS, 1) __global__
-    void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
-  gl_tile_C g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
-  gl_tile_A g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
-  gl_tile_B g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
-
-  extern __shared__ alignment_dummy __shm[];
-  shared_allocator al((int *)&__shm[0]);
-
-  shared_tile_A(&As)[SUPER_M][PIPE_STAGES] =
-      al.allocate<shared_tile_A, SUPER_M, PIPE_STAGES>();
-  shared_tile_B(&Bs)[SUPER_N][PIPE_STAGES] =
-      al.allocate<shared_tile_B, SUPER_N, PIPE_STAGES>();
-
-  reg_tile_A A_reg;
-  reg_tile_B_col B_reg_col;
-  reg_tile_C C_accum;
-
-  int warpid = kittens::warpid();
-  int warp_m = warpid % SUPER_M;
-  int warp_n = warpid / SUPER_M;
-
-  int load_group_id = warpgroup::groupid();
-
-  int block_row = blockIdx.y * SUPER_M;
-  int block_col = blockIdx.x * SUPER_N;
-
-  warp::zero(C_accum);
-  int num_tiles = (g_N + BLOCK_K - 1) / BLOCK_K;
-
-  for (int load_tile = 0; load_tile < (PIPE_STAGES - 1); load_tile++) {
-    if (load_tile < num_tiles) {
-      int load_smem_idx = load_tile % PIPE_STAGES;
-      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
-        if (task_id < SUPER_M) {
-          warp::load_async(As[task_id][load_smem_idx], g_A, {0, 0, block_row + task_id, load_tile});
-        } else {
-          int n_index = task_id - SUPER_M;
-          warp::load_async(Bs[n_index][load_smem_idx], g_B, {0, 0, load_tile, block_col + n_index});
-        }
-      }
-    }
-  }
-
-  for (int tile = 0; tile < num_tiles; tile++) {
-    int compute_smem_idx = tile % PIPE_STAGES;
-
-    int load_tile = tile + PIPE_STAGES - 1;
-    int load_smem_idx = load_tile % PIPE_STAGES;
-
-    if (load_tile < num_tiles) {
-      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
-        if (task_id < SUPER_M) {
-          warp::load_async(As[task_id][load_smem_idx], g_A,
-                           {0, 0, block_row + task_id, load_tile});
-        } else {
-          int n_index = task_id - SUPER_M;
-          warp::load_async(Bs[n_index][load_smem_idx], g_B,
-                           {0, 0, load_tile, block_col + n_index});
-        }
-      }
-      load_async_wait<1>();
-    } else
-      load_async_wait();
-    __syncthreads();
-
-    warp::load(A_reg, As[warp_m][compute_smem_idx]);
-    warp::load(B_reg_col, Bs[warp_n][compute_smem_idx]);
-
-    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
-    __syncthreads();
-  }
-  warp::store(g_C, C_accum, {0, 0, block_row + warp_m, block_col + warp_n});
-}
@@ -1,6 +0,0 @@
-from tinygrad.device import Device
-
-if Device.DEFAULT == "AMD":
-  WARP_THREADS = 64
-else:
-  WARP_THREADS = 32
@@ -1,290 +0,0 @@
-import math, functools
-from typing import cast, Callable
-from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
-from tinygrad.engine.realize import ExecItem, get_runner
-from tinygrad.dtype import AddrSpace, PtrDType
-from tinygrad.helpers import getenv, prod
-
-from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, ST, RT, RV
-
-class Group:
-  def __init__(self, warps:int, ker):
-    self.warps = warps
-    self.group_threads = warps * WARP_THREADS
-    self.threadIdx_x = ker.threadIdx_x
-    self.ker = ker
-
-  # helpers
-  @property
-  def laneid(self): return self.threadIdx_x % self.group_threads
-  @property
-  def warpid(self): return self.laneid // WARP_THREADS
-  @property
-  def groupid(self): return self.threadIdx_x // self.group_threads
-
-  # ops that only work on a single warp
-
-  clear_rid = 1000
-  def clear(self, reg:ALL_TILES, value:float=0):
-    reg = cast(UOp, reg)
-    assert self.warps == 1
-
-    rngs_for_shape = tuple(UOp.range(dim, Group.clear_rid + i) for i, dim in enumerate(reg.shape))
-    Group.clear_rid += len(reg.shape)
-
-    reg_store = reg[*rngs_for_shape].store(value).end(*rngs_for_shape)
-
-    self.ker.push_store(reg_store, reg)
-    return reg.after(reg_store).reshape(reg.shape)
-
-  def zero(self, reg:ALL_TILES): return self.clear(reg, 0)
-  def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
-
-  copy_rid = 300
-  def copy(self, dst:ALL_TILES, src:ALL_TILES):
-    dst, src = cast(UOp, dst), cast(UOp, src)
-    assert self.warps == 1
-    assert dst.shape == src.shape
-
-    rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
-    Group.copy_rid += len(dst.shape)
-
-    dst_store = dst[*rngs_for_shape].store(src[*rngs_for_shape].cast(dst.dtype.base)).end(*rngs_for_shape)
-
-    self.ker.push_store(dst_store, dst)
-    return dst.after(dst_store).reshape(dst.shape)
-
-  def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
-    c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
-    assert self.warps == 1
-
-    for height in self.ker.range(c.shape[-3], track=False):
-      for width in self.ker.range(c.shape[-2], track=False):
-        for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
-          wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
-
-          a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
-          b_in1 = UOp.vectorize(*([b[inner, width, i] for i in range(2)] + [b[inner, width, 4+i] for i in range(2)]))
-          c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
-          b_in2 = UOp.vectorize(*([b[inner, width, 2+i] for i in range(2)] + [b[inner, width, 6+i] for i in range(2)]))
-          c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
-
-          out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
-          out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
-          c_store = UOp.group(*c_i).end(height, width, inner)
-
-    self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape)
-
-  def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
-    c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
-    assert self.warps == 1
-
-    for height in self.ker.range(c.shape[-3], track=False):
-      for width in self.ker.range(c.shape[-2], track=False):
-        for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
-          wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
-
-          a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
-          b_in1 = UOp.vectorize(*([b[width, inner, i] for i in range(2)] + [b[width, inner, 4+i] for i in range(2)]))
-          c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
-          b_in2 = UOp.vectorize(*([b[width, inner, 2+i] for i in range(2)] + [b[width, inner, 6+i] for i in range(2)]))
-          c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
-
-          out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
-          out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
-          c_store = UOp.group(*c_i).end(height, width, inner)
-
-    self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape)
-
-  map_rid = 400
-  def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
-    a = cast(UOp, a)
-    assert self.warps == 1
-
-    rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
-    Group.map_rid += len(a.shape)
-
-    if op.__code__.co_argcount == 1:
-      to_store = op(a[*rngs_for_shape])
-    else:
-      to_store = op(a[*rngs_for_shape], rngs_for_shape)
-
-    a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
-
-    self.ker.push_store(a_store, a)
-    return a.after(a_store).reshape(a.shape)
-
-  def row_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp]):
-    vec, src = cast(UOp, vec), cast(UOp, src)
-    assert self.warps == 1
-
-    red_local = self.ker.alloc((self.group_threads, 2), src.dtype.base, AddrSpace.LOCAL)
-    red_reg = self.ker.alloc((2,), src.dtype.base, AddrSpace.REG)
-
-    for height in self.ker.range(src.shape[-3], track=False):
-      i = UOp.range(red_reg.size, Group.clear_rid)
-      Group.clear_rid += 1
-      red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
-      reg_store = red_reg.flatten()[i].store(0.).end(i)
-      red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
-
-      for outer in self.ker.range(2, track=False):
-        for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
-          for inner in self.ker.range(4, AxisType.REDUCE, track=False):
-            elem_index = inner + 2 * (inner // 2) + outer * 2
-            reg_store = red_reg[outer].store(op(red_reg[outer], src[height, width, elem_index])).end(inner, width, outer)
-            red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
-
-      # store to shared memory
-      for outer in self.ker.range(2, track=False):
-        red_local_store = red_local[self.laneid, outer].store(red_reg[outer]).end(outer)
-        red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
-
-      # reduce from shared memory
-      for outer in self.ker.range(2, track=False):
-        for inner in self.ker.range(3, AxisType.REDUCE, track=False):
-          offset = (self.laneid // 4) * 4 + ((self.laneid + inner + 1) % 4)
-          reg_store = red_reg[outer].store(op(red_reg[outer], red_local[offset, outer])).end(inner, outer)
-          red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
-
-      # reduce with vec
-      for outer in self.ker.range(2, track=False):
-        vec_store = vec[height, 0, outer].store(op(vec[height, 0, outer], red_reg[outer])).end(outer, height)
-
-    self.ker.push_store(vec_store, vec)
-    return vec.after(vec_store).reshape(vec.shape)
-
-  # ops that can work across multiple warps
-
-  LOAD_INNER = 4
-  def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
-    dst, src = cast(UOp, dst), cast(UOp, src)
-    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
-    if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
-      srcf = src.flatten(-2)
-
-      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
-      else: local_warpid = self.warpid
-      warp_laneid = self.threadIdx_x % WARP_THREADS
-
-      for height in self.ker.range(dst.shape[-3], track=False):
-        for width in self.ker.range(dst.shape[-2], track=False):
-          for inner in self.ker.range(RT.BASE_TILE_NEPT, track=False):
-            base_row = (local_warpid * dst.shape[-3] + height) * RT.BASE_TILE_ROWS
-            base_col = width * RT.BASE_TILE_COLS
-
-            if not transpose:
-              row = base_row + (warp_laneid // 4)
-              col = base_col + 2 * (warp_laneid % 4)
-
-              row_offset = ((inner % 4) // 2) * 8
-              col_offset = (inner % 2) + (inner // 4) * 8
-            else:
-              row = base_row + 2 * (warp_laneid % 4)
-              col = base_col + (warp_laneid // 4)
-
-              row_offset = (inner % 2) + (inner // 4) * 8
-              col_offset = ((inner % 4) // 2) * 8
-
-            src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
-
-            dst_store = dst[*dst_idxs, height, width, inner].store(srcf[*idxs[:-2], src_i_last])
-            dst_store = dst_store.end(height, width, inner)
-    elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
-      dstf = dst.flatten(-2)
-
-      srcf = src.flatten()
-      row_stride = prod(src.shape[axis+1:])
-
-      idxs = tuple(idx * dst.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
-      idxs = tuple(idx * dst.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
-      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
-
-      memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
-      total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
-
-      for outer in self.ker.range(total_calls, track=False):
-        for inner in self.ker.range(Group.LOAD_INNER, track=False):
-          load_idx = outer * self.group_threads + self.laneid
-          row = load_idx // memcpy_per_row
-          col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
-
-          dst_i = row * dst.shape[-1] + col + inner
-          src_i += row * row_stride + col + inner
-
-          dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(outer, inner)
-    else:
-      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
-
-    return dst.after(dst_store.barrier()).reshape(dst.shape)
-
-  STORE_INNER = 4
-  def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
-    dst, src = cast(UOp, dst), cast(UOp, src)
-    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
-      dstf = dst.flatten(-2)
-
-      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
-      else: local_warpid = self.warpid
-      warp_laneid = self.threadIdx_x % WARP_THREADS
-
-      for height in self.ker.range(src.shape[-3], track=False):
-        for width in self.ker.range(src.shape[-2], track=False):
-          for inner in self.ker.range(RT.BASE_TILE_NEPT, track=False):
-            base_row = (local_warpid * src.shape[-3] + height) * RT.BASE_TILE_ROWS
-            base_col = width * RT.BASE_TILE_COLS
-
-            if not transpose:
-              row = base_row + (warp_laneid // 4)
-              col = base_col + 2 * (warp_laneid % 4)
-
-              row_offset = ((inner % 4) // 2) * 8
-              col_offset = (inner % 2) + (inner // 4) * 8
-            else:
-              row = base_row + 2 * (warp_laneid % 4)
-              col = base_col + (warp_laneid // 4)
-
-              row_offset = (inner % 2) + (inner // 4) * 8
-              col_offset = ((inner % 4) // 2) * 8
-
-            dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
-
-            dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, height, width, inner])
-            dst_store = dst_store.end(height, width, inner)
-    elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
-      dstf = dst.flatten()
-      row_stride = prod(dst.shape[axis+1:])
-
-      idxs = tuple(idx * src.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
-      idxs = tuple(idx * src.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
-      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
-
-      srcf = src.flatten(-2)
-
-      memcpy_per_row = src.shape[-1] // Group.STORE_INNER
-      total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
-
-      for outer in self.ker.range(total_calls, track=False):
-        for inner in self.ker.range(Group.STORE_INNER, track=False):
-          load_idx = outer * self.group_threads + self.laneid
-          row = load_idx // memcpy_per_row
-          col = (load_idx * Group.STORE_INNER) % src.shape[-1]
-
-          src_i = row * src.shape[-1] + col + inner
-          dst_i += row * row_stride + col + inner
-
-          dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(outer, inner)
-    else:
-      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
-
-    self.ker.push_store(dst_store, dst)
-    return dst.after(dst_store.barrier()).reshape(dst.shape)
@@ -1,92 +0,0 @@
-from contextlib import AbstractContextManager
-from tinygrad.uop.ops import UOp, KernelInfo, AxisType, AddrSpace
-from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.group import Group
-from extra.thunder.tiny.tk.tiles import GL, ST, RT, RV
-
-class _tk_range:
-  user_rid = 0
-  def __init__(self, end:int, axis_type:AxisType): self.end, self.axis_type, self.done = end, axis_type, False
-  def __iter__(self): return self
-  def __next__(self):
-    if not self.done:
-      self.done = True
-      _tk_range.user_rid += 1
-      self._rng = UOp.range(self.end, _tk_range.user_rid-1, axis_type=self.axis_type)
-      return self._rng
-    raise StopIteration
-
-class Kernel(AbstractContextManager):
-  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
-    self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
-    self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
-    self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
-    self.threadIdx_x = UOp.special(block_size, "lidx0")
-
-    self.range_stack = []
-    self.store_stack = []
-
-    self.global_slot = 0
-    self.shared_slot = 0
-    self.register_slot = 0
-    self.allocs = {}
-
-  @property
-  def warpid(self): return self.threadIdx_x // WARP_THREADS
-
-  def __enter__(self): return self
-  def __exit__(self, exc_type, exc_value, traceback): pass
-
-  def group(self, size:int): return Group(size, self)
-  @property
-  def warp(self): return self.group(1)
-  @property
-  def warpgroup(self): return self.group(4)
-
-  def range(self, end:int, axis_type:AxisType=AxisType.LOOP, track:bool=True):
-    rng = _tk_range(end, axis_type)
-    if track: self.range_stack.append(rng)
-    return rng
-
-  def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
-    match addrspace:
-      case AddrSpace.GLOBAL:
-        slot = self.global_slot
-        self.global_slot += 1
-      case AddrSpace.LOCAL:
-        slot = self.shared_slot
-        self.shared_slot += 1
-      case AddrSpace.REG:
-        slot = self.register_slot
-        self.register_slot += 1
-
-    uop = UOp.placeholder(shape, dtype, slot=slot, addrspace=addrspace)
-
-    if name:
-      if (name, shape) in self.allocs: return self.allocs[(name, shape)]
-      self.allocs[(name, shape)] = uop
-
-    return uop
-
-  def gl(self, shape, dtype): return GL.create(shape, dtype, self)
-  def st(self, shape, dtype): return ST.create(shape, dtype, self)
-  def rt(self, shape, dtype): return RT.create(shape, dtype, self)
-  def rv(self, length, dtype, layout="naive"): return RV.create(length, dtype, layout, self)
-
-  def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
-
-  def finish(self):
-    # end all ranges
-    rngs = []
-    while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
-
-    last_store = self.store_stack.pop()[0]
-    if hasattr(last_store, '_uop'): uop = last_store._uop
-    else: uop = last_store
-
-    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
-
-  def endrange(self):
-    last_store = self.store_stack.pop()
-    last_range = self.range_stack.pop()
-    return last_store[1].after(last_store[0].barrier().end(last_range._rng)).reshape(last_store[1].shape)
@@ -1,143 +0,0 @@
-import functools
-from tinygrad.dtype import AddrSpace
-from tinygrad.mixin import MathMixin
-from tinygrad.uop.ops import UOp, Ops
-
-from extra.thunder.tiny.tk import WARP_THREADS
-
-def unwrap(x):
-  if hasattr(x, "_uop"): return x._uop
-  if isinstance(x, (list, tuple)): return type(x)(unwrap(y) for y in x)
-  if isinstance(x, dict): return {k: unwrap(v) for k,v in x.items()}
-  return x
-
-def wrap(x, ker, cls):
-  if isinstance(x, UOp): return cls(x, ker)
-  if isinstance(x, (list, tuple)): return type(x)(wrap(y, ker, cls) for y in x)
-  return x
-
-def autowrap(source_cls, blacklist=None):
-  if blacklist is None:
-    blacklist = {
-      "__init__", "__new__", "__str__", "__del__", "__repr__", "__dict__", "__getattribute__",
-      "__setattr__", "__delattr__", "__weakref__", "__slots__", "__class__",
-      "__reduce__", "__reduce_ex__", "__getstate__", "__setstate__", "__hash__"
-    }
-
-  def decorator(cls):
-    def __getattr__(self, name):
-      uop = object.__getattribute__(self, "_uop")
-      val = getattr(uop, name)
-      if callable(val):
-        @functools.wraps(val)
-        def proxy(*args, **kwargs):
-          return wrap(val(*unwrap(args), **unwrap(kwargs)), self.ker, cls)
-        return proxy
-      if name in UOp.__slots__: return val
-      return wrap(val, self.ker, cls)
-    cls.__getattr__ = __getattr__
-
-    for name in dir(source_cls):
-      if name in blacklist or not name.startswith("__"): continue
-
-      for base in cls.mro():
-        if base is source_cls: break
-        if name in base.__dict__: break
-      else:
-        original = getattr(source_cls, name)
-        if callable(original):
-          def make_proxy(op_name, func):
-            def proxy(self, *args, **kwargs):
-              return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self.ker, cls)
-            return proxy
-          setattr(cls, name, make_proxy(name, original))
-
-    return cls
-  return decorator
-
-class TileMathMixin(MathMixin):
-  def alu(self, op, *src, inner_op=lambda x:x):
-    assert isinstance(self, (RT, RV))
-    if len(src) == 0:
-      if self._uop._shape is None: uop = UOp.alu(self._uop, op)
-      else: uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op))
-    elif len(src) == 1:
-      if self._uop._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
-      elif isinstance(src[0], (int,float,bool)): uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op, inner_op(x.ufix(src[0]))))
-      elif src[0]._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
-      else:
-        if isinstance(self, RT) and isinstance(src[0], RV): uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0, (idx[2]%4)//2])))
-        else: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[*idx])))
-    else: raise NotImplementedError
-    return type(self)(uop, self.ker)
-  def const_like(self, b): return b
-
-  # override ops that do compute on the src uop
-  def sub(self, x, reverse=False):
-    return self.ufix(x).alu(Ops.ADD, self, inner_op=lambda y: -y) if reverse else self.alu(Ops.ADD, self.ufix(x), inner_op=lambda y: -y)
-  def div(self, x, reverse=False):
-    return self.ufix(x).alu(Ops.MUL, self, inner_op=lambda y: 1/y) if reverse else self.alu(Ops.MUL, self.ufix(x), inner_op=lambda y: 1/y)
-
-@autowrap(UOp)
-class GL:
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
-
-  @classmethod
-  def create(cls, shape, dtype, ker):
-    uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
-    return cls(uop, ker)
-
-@autowrap(UOp)
-class ST:
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
-
-  @classmethod
-  def create(cls, shape, dtype, ker):
-    uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
-    return cls(uop, ker)
-
-@autowrap(UOp)
-class RT(TileMathMixin):
-  BASE_TILE_ROWS, BASE_TILE_COLS = 16, 16
-  BASE_TILE_NE = BASE_TILE_ROWS * BASE_TILE_COLS
-  BASE_TILE_NEPT = BASE_TILE_NE // WARP_THREADS
-
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
-
-  @classmethod
-  def create(cls, shape, dtype, ker):
-    assert len(shape) == 2
-    assert shape[0] % RT.BASE_TILE_ROWS == 0
-    assert shape[1] % RT.BASE_TILE_COLS == 0
-
-    height = shape[0] // RT.BASE_TILE_ROWS
-    width = shape[1] // RT.BASE_TILE_COLS
-
-    uop = ker.alloc((height, width, RT.BASE_TILE_NEPT), dtype, AddrSpace.REG)
-    return cls(uop, ker)
-
-@autowrap(UOp)
-class RV(TileMathMixin):
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
-
-  @classmethod
-  def create(cls, length, dtype, layout, ker):
-    tiles = length // RT.BASE_TILE_ROWS
-
-    match layout:
-      case "naive":
-        inner_dim = 1
-        outer_dim = (tiles + 1) // 2
-      case "ortho":
-        inner_dim = 1
-        outer_dim = tiles
-      case _: raise NotImplementedError(f"rv layout {layout} not implemented")
-
-    uop = ker.alloc((outer_dim, inner_dim, 2), dtype, AddrSpace.REG)
-    return RV(uop, ker)
-
-ALL_TILES = UOp | GL | ST | RT | RV
@@ -3,9 +3,9 @@ import argparse

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
-  parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
-  parser.add_argument("--len", type=int, required=True, help="file length to fetch")
-  parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
+  parser.add_argument("hash", type=str, required=True, help="file hash to fetch")
+  parser.add_argument("len", type=int, required=True, help="file length to fetch")
+  parser.add_argument("dest", type=str, required=True, help="destination path to save the file")
  args = parser.parse_args()

  Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import sys, os, zlib, struct, hashlib
+from hexdump import hexdump
 from tinygrad.helpers import DEBUG, getenv, fetch
 from tinygrad.runtime.support.usb import USB3

@@ -119,7 +119,14 @@ extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {

 		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)

+		// Add appropriate logic here to determine whether to replace the extension
+		// with the new extension. Common things to check for include
+		// testing whether the new extension's version number is newer than
+		// the current version number, or whether the bundleIdentifier is different.
+		// For simplicity, this sample always replaces the current extension
+		// with the new one.
 		replacementAction = .replace
+
 		self.state = .activating
 		return replacementAction
 	}
@@ -7,48 +7,30 @@
 struct TinyGPUDriverUserClient_IVars
 {
 	OSSharedPtr<TinyGPUDriver> provider = nullptr;
-
-	TinyGPUCreateDMAResp *dmas = nullptr;
-	size_t dmaCount = 0;
-	size_t dmaCap = 0;
-
-	int ensureDMACap(size_t need)
-	{
-		// not thread-safe
-		if (need <= dmaCap) return 0;
-
-		size_t newCap = dmaCap ? dmaCap * 2 : 16;
-		while (newCap < need) newCap *= 2;
-
-		auto *newArr = IONewZero(TinyGPUCreateDMAResp, newCap);
-		if (!newArr) return -kIOReturnNoMemory;
-
-		if (dmas && dmaCount) {
-			memcpy(newArr, dmas, dmaCount * sizeof(TinyGPUCreateDMAResp));
-		}
-
-		IOSafeDeleteNULL(dmas, TinyGPUCreateDMAResp, dmaCap);
-		dmas = newArr;
-		dmaCap = newCap;
-		return 0;
-	}
 };

 bool TinyGPUDriverUserClient::init()
 {
-	auto ok = super::init();
-	if (!ok) return false;
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}

 	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
-	if (!ivars) return false;
+	if (ivars == nullptr) {
+		return false;
+	}
+
 	return true;
 }

 void TinyGPUDriverUserClient::free()
 {
-	if (ivars) {
-		IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	if (ivars != nullptr) {
+		ivars->provider.reset();
 	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
 	super::free();
 }

@@ -77,22 +59,6 @@ error:

 kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
 {
-	// release all DMA allocations for this client
-	if (ivars) {
-		for (size_t i = 0; i < ivars->dmaCount; i++) {
-			auto &d = ivars->dmas[i];
-			if (d.dmaCmd) {
-				d.dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
-				d.dmaCmd->release();
-				d.dmaCmd = nullptr;
-			}
-		}
-		ivars->dmaCount = 0;
-		IOSafeDeleteNULL(ivars->dmas, TinyGPUCreateDMAResp, ivars->dmaCap);
-		ivars->dmas = nullptr;
-		ivars->provider.reset();
-	}
-
 	return Stop(in_provider, SUPERDISPATCH);
 }

@@ -136,26 +102,26 @@ kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserC

 kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
 {
-	if (!memory) return kIOReturnBadArgument;
-	if (!ivars->provider.get()) return kIOReturnNotAttached;
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}

-	// bar handling, type is bar num
 	if (type < 6) {
 		uint32_t bar = (uint32_t)type;
 		return ivars->provider->MapBar(bar, memory);
 	}

-	// dma handling, type is size
-	if (ivars->ensureDMACap(ivars->dmaCount + 1)) {
-		os_log(OS_LOG_DEFAULT, "tinygpu: cannot grow dma array");
-		return kIOReturnNoMemory;
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
 	}

-	TinyGPUCreateDMAResp buf{};
-	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
-	if (err) return err;
-
-	ivars->dmas[ivars->dmaCount++] = buf;
 	*memory = buf.sharedBuf;
 	return 0;
 }
@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt

-NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","wozeparrot"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
@@ -40,4 +40,4 @@ for d in days:
 print("** Commits by day (last 7) **")
 print("```")
 print("\n".join([header, rule] + rows))
-print("```")
+print("```")
@@ -0,0 +1,10 @@
+[mypy]
+warn_unused_configs = True
+files = tinygrad
+ignore_missing_imports = True
+check_untyped_defs = True
+explicit_package_bases = True
+warn_unreachable = True
+warn_redundant_casts = True
+# NOTE: had to comment this out to make mypy pass on both CI and OSX
+#warn_unused_ignores = True
@@ -1,232 +0,0 @@
-[project]
-name = "tinygrad"
-version = "0.11.0"
-description = "You like pytorch? You like micrograd? You love tinygrad! <3"
-authors = [{ name = "George Hotz" }]
-
-classifiers = ["Programming Language :: Python :: 3"]
-
-license = 'MIT'
-readme = "README.md"
-requires-python = ">=3.11"
-dependencies = []
-
-[build-system]
-requires = ["setuptools"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools]
-include-package-data = true
-packages = [
-  'tinygrad',
-  'tinygrad.apps',
-  'tinygrad.codegen',
-  'tinygrad.codegen.opt',
-  'tinygrad.codegen.late',
-  'tinygrad.engine',
-  'tinygrad.mixin',
-  'tinygrad.nn',
-  'tinygrad.renderer',
-  'tinygrad.runtime',
-  'tinygrad.runtime.autogen',
-  'tinygrad.runtime.autogen.am',
-  'tinygrad.runtime.graph',
-  'tinygrad.runtime.support',
-  'tinygrad.runtime.support.am',
-  'tinygrad.runtime.support.nv',
-  'tinygrad.schedule',
-  'tinygrad.uop',
-  'tinygrad.viz',
-]
-
-[tool.setuptools.package-data]
-tinygrad = ["py.typed"]
-"tinygrad.viz" = ["index.html", "assets/**/*", "js/*"]
-
-
-[project.optional-dependencies]
-arm = ["unicorn"]
-triton = ["triton-nightly>=2.1.0.dev20231014192330"]
-linting = [
-  "pylint",
-  "mypy==1.18.1",
-  "typing-extensions",
-  "pre-commit",
-  "ruff",
-  "numpy",
-  "typeguard",
-]
-# mlperf = [
-#   "mlperf-logging @ git+https://github.com/mlperf/[email protected]",
-# ]
-testing_minimal = [
-  "numpy",
-  "torch==2.9.0",
-  "pytest",
-  "pytest-xdist",
-  "pytest-timeout",
-  "pytest-split",
-  "hypothesis",
-  "z3-solver",
-]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate"]
-testing = [
-  "tinygrad[testing_minimal]",
-  "pillow",
-  "onnx==1.18.0",
-  "onnx2torch",
-  "onnxruntime",
-  "opencv-python",
-  "tabulate",
-  "tqdm",
-  "safetensors",
-  "transformers",
-  "sentencepiece",
-  "tiktoken",
-  "blobfile",
-  "librosa",
-  # librosa needs numba but uv ignores python upper bounds and some numba versions require <python3.10
-  "numba>=0.55",
-  "networkx",
-  "nibabel",
-  "bottle",
-  "ggml-python",
-  "capstone",
-  "pycocotools",
-  "boto3",
-  "pandas",
-  "influxdb3-python",
-]
-docs = [
-  "mkdocs",
-  "mkdocs-material",
-  "mkdocstrings[python]",
-  "markdown-callouts",
-  "markdown-exec[ansi]",
-  "black",
-  "numpy",
-]
-
-
-[tool.mutmut]
-paths_to_mutate = ["tinygrad/"]
-do_not_mutate = [
-  "tinygrad/apps/*",
-  "tinygrad/codegen/*",
-  "tinygrad/engine/*",
-  "tinygrad/nn/*",
-  "tinygrad/renderer/*",
-  "tinygrad/runtime/*",
-  "tinygrad/schedule/*",
-  "tinygrad/uop/*",
-  "tinygrad/viz/*",
-  "tinygrad/device.py",
-  "tinygrad/dtype.py",
-  "tinygrad/gradient.py",
-  "tinygrad/helpers.py",
-  "tinygrad/tensor.py",
-]
-tests_dir = ["test/test_tiny.py", "test/test_ops.py"]
-debug = true
-
-
-[tool.mypy]
-warn_unused_configs = true
-files = ["tinygrad"]
-ignore_missing_imports = true
-check_untyped_defs = true
-explicit_package_bases = true
-warn_unreachable = true
-warn_redundant_casts = true
-# NOTE: had to comment this out to make mypy pass on both CI and OSX
-#warn_unused_ignores = true
-
-[tool.pytest.ini_options]
-norecursedirs = [
-  "extra",
-  ".hypothesis",
-  ".git",
-]
-timeout = 300
-timeout_method = "thread"
-timeout_func_only = true
-testpaths = ["test"]
-
-[tool.ruff]
-preview = true
-target-version = "py311"
-line-length = 150
-indent-width = 2
-exclude = [
-  ".git/",
-  "docs/",
-  "extra/",
-  "test/external/mlperf_resnet",
-  "test/external/mlperf_unet3d",
-]
-
-[tool.ruff.lint]
-select = [
-  "F",  # Pyflakes
-  "W6",
-  "E71",
-  "E72",
-  "E112",   # no-indented-block
-  "E113",   # unexpected-indentation
-  # "E124",
-  "E203",   # whitespace-before-punctuation
-  "E272",   # multiple-spaces-before-keyword
-  "E275",   # missing-whitespace-after-keyword
-  "E303",   # too-many-blank-lines
-  "E304",   # blank-line-after-decorator
-  "E501",   # line-too-long
-  # "E502",
-  "E702",   # multiple-statements-on-one-line-semicolon
-  "E703",   # useless-semicolon
-  "E731",   # lambda-assignment
-  "W191",   # tab-indentation
-  "W291",   # trailing-whitespace
-  "W293",   # blank-line-with-whitespace
-  "UP039",  # unnecessary-class-parentheses
-  "C416",   # unnecessary-comprehension
-  "RET506", # superfluous-else-raise
-  "RET507", # superfluous-else-continue
-  "A",      # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing
-  "FURB110",# if-exp-instead-of-or-operator
-  "RUF018", # assignment-in-assert
-]
-
-# detect unused imports in examples
-[tool.ruff.lint.per-file-ignores]
-"examples/**/*.py" = [
-  "W6",
-  "E71",
-  "E72",
-  "E112",
-  "E113",
-  "E203",
-  "E272",
-  "E275",
-  "E303",
-  "E304",
-  "E501",
-  "E702",
-  "E703",
-  "E731",
-  "W191",
-  "W291",
-  "W293",
-  "UP039",
-  "C416",
-  "RET506",
-  "RET507",
-  "A",
-  "FURB110",
-  "RUF018",
-  "F541",
-  "F841",
-]
-"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E722", "E731", "F821", "A006"]
-
-[tool.ruff.format]
-exclude = ["*"]
@@ -0,0 +1,6 @@
+[pytest]
+norecursedirs = extra
+timeout = 300
+timeout_method = thread
+timeout_func_only = true
+testpaths = test
@@ -0,0 +1,56 @@
+indent-width = 2
+preview = true
+target-version = "py310"
+
+lint.select = [
+  "F",  # Pyflakes
+  "W6",
+  "E71",
+  "E72",
+  "E112",   # no-indented-block
+  "E113",   # unexpected-indentation
+  # "E124",
+  "E203",   # whitespace-before-punctuation
+  "E272",   # multiple-spaces-before-keyword
+  "E275",   # missing-whitespace-after-keyword
+  "E303",   # too-many-blank-lines
+  "E304",   # blank-line-after-decorator
+  "E501",   # line-too-long
+  # "E502",
+  "E702",   # multiple-statements-on-one-line-semicolon
+  "E703",   # useless-semicolon
+  "E731",   # lambda-assignment
+  "W191",   # tab-indentation
+  "W291",   # trailing-whitespace
+  "W293",   # blank-line-with-whitespace
+  "UP039",  # unnecessary-class-parentheses
+  "C416",   # unnecessary-comprehension
+  "RET506", # superfluous-else-raise
+  "RET507", # superfluous-else-continue
+  "A",      # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing
+  "FURB110",# if-exp-instead-of-or-operator
+  "RUF018", # assignment-in-assert
+]
+
+line-length = 150
+
+exclude = [
+  ".git/",
+  "docs/",
+  "extra/",
+  "tinygrad/runtime/autogen",
+  "test/external/mlperf_resnet",
+  "test/external/mlperf_unet3d",
+]
+
+# detect unused imports in examples
+[lint.per-file-ignores]
+"examples/**/*.py" = [
+  "W6", "E71", "E72", "E112", "E113", "E203", "E272", "E275",
+  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
+  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
+  "FURB110", "RUF018", "F541", "F841"
+]
+
+[format]
+exclude = ["*"]
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+
+from pathlib import Path
+from setuptools import setup
+
+directory = Path(__file__).resolve().parent
+with open(directory / 'README.md', encoding='utf-8') as f:
+  long_description = f.read()
+
+testing_minimal = [
+  "numpy",
+  "torch==2.9.0",
+  "pytest",
+  "pytest-xdist",
+  "pytest-timeout",
+  "pytest-split",
+  "hypothesis",
+  "z3-solver",
+]
+
+setup(name='tinygrad',
+      version='0.11.0',
+      description='You like pytorch? You like micrograd? You love tinygrad! <3',
+      author='George Hotz',
+      license='MIT',
+      long_description=long_description,
+      long_description_content_type='text/markdown',
+      packages = [
+        'tinygrad',
+        'tinygrad.apps',
+        'tinygrad.codegen',
+        'tinygrad.codegen.opt',
+        'tinygrad.codegen.late',
+        'tinygrad.engine',
+        'tinygrad.nn',
+        'tinygrad.renderer',
+        'tinygrad.runtime',
+        'tinygrad.runtime.autogen',
+        'tinygrad.runtime.autogen.am',
+        'tinygrad.runtime.autogen.nv',
+        'tinygrad.runtime.graph',
+        'tinygrad.runtime.support',
+        'tinygrad.runtime.support.am',
+        'tinygrad.runtime.support.nv',
+        'tinygrad.schedule',
+        'tinygrad.uop',
+        'tinygrad.viz',
+      ],
+      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
+      classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License"
+      ],
+      install_requires=[],
+      python_requires='>=3.10',
+      extras_require={
+        'arm': ["unicorn"],
+        'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
+        'linting': [
+            "pylint",
+            "mypy==1.18.1",
+            "typing-extensions",
+            "pre-commit",
+            "ruff",
+            "numpy",
+            "typeguard",
+        ],
+        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
+        'testing_minimal': testing_minimal,
+        'testing_unit': testing_minimal + [
+            "tqdm",
+            "safetensors",
+            "tabulate",  # for sz.py
+        ],
+        'testing': testing_minimal + [
+            "pillow",
+            "onnx==1.18.0",
+            "onnx2torch",
+            "onnxruntime",
+            "opencv-python",
+            "tabulate",
+            "tqdm",
+            "safetensors",
+            "transformers",
+            "sentencepiece",
+            "tiktoken",
+            "blobfile",
+            "librosa",
+            "numba>=0.55",  # librosa needs numba but uv ignores python upper bounds and some numba versions require <python3.10
+            "networkx",
+            "nibabel",
+            "bottle",
+            "ggml-python",
+            "capstone",
+            "pycocotools",
+            "boto3",
+            "pandas",
+            "influxdb3-python"
+        ],
+        'docs': [
+            "mkdocs",
+            "mkdocs-material",
+            "mkdocstrings[python]",
+            "markdown-callouts",
+            "markdown-exec[ansi]",
+            "black",
+            "numpy",
+        ],
+      },
+      include_package_data=True)
@@ -1,106 +0,0 @@
-# ruff: noqa: E501 E712 F401
-from tinygrad import dtypes, Device
-from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
-from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup, getenv
-from tinygrad.device import Buffer
-from tinygrad.dtype import ImageDType, Invalid
-
-# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-
-def vision_conv_143():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
-  c2 = UOp.range(32, 3, AxisType.LOOP)
-  c5 = UOp.range(128, 4, AxisType.LOOP)
-  c8 = UOp.range(16, 2, AxisType.LOOP)
-  c16 = UOp.range(7, 0, AxisType.REDUCE)
-  c17 = c8*2+c16
-  c24 = ((c17<3)!=True)&(c17<35)
-  c26 = UOp.range(7, 1, AxisType.REDUCE)
-  c27 = c2*2+c26
-  c32 = ((c27<3)!=True)&(c27<67)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
-  c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
-  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
-  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
-  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
-
-  opts = None
-  # JITBEAM=2
-  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
-  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
-
-def vision_conv_153():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
-  c2 = UOp.range(16, 3, AxisType.LOOP)
-  c5 = UOp.range(256, 4, AxisType.LOOP)
-  c8 = UOp.range(8, 2, AxisType.LOOP)
-  c16 = UOp.range(7, 0, AxisType.REDUCE)
-  c17 = c8*2+c16
-  c24 = ((c17<3)!=True)&(c17<19)
-  c26 = UOp.range(7, 1, AxisType.REDUCE)
-  c27 = c2*2+c26
-  c32 = ((c27<3)!=True)&(c27<35)
-  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
-  c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
-  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
-  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
-  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
-  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
-
-  opts = None
-  # JITBEAM=2
-  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
-  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
-
-def dm_conv_172():
-  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 240, 4)), (), 0)
-  c2 = UOp.range(960, 4, AxisType.LOOP)
-  c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 384, 4)), (), 1)
-  c7 = UOp.range(32, 0, AxisType.REDUCE)
-  c10 = UOp.range(4, 1, AxisType.REDUCE)
-  c13 = UOp.range(12, 3, AxisType.REDUCE)
-  c18 = UOp.range(8, 2, AxisType.REDUCE)
-  c23 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((240, 128, 4)), (), 2)
-  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(960), (), 3)
-  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
-  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
-  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
-  c55 = c0.index(c2, ptr=True).store(c53).end(c2)
-
-  opts = None
-  # JITBEAM=2
-  # (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
-  return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
-
-ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
-
-compiler = Device.default.compiler
-renderer = Device.default.renderer
-allocator = Device.default.allocator
-
-uops = full_rewrite(ast, renderer)
-src = renderer.render(uops)
-
-lib = compiler.compile(src)
-ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
-cr = CompiledRunner(ps, precompiled=lib)
-
-gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
-# print(len(gs))
-# print([g.dtype for g in gs])
-bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
-
-t = cr(bufs, wait=True)
-print(f"{t*1e6:.2f} us")
@@ -0,0 +1,39 @@
+import subprocess
+import random
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+def run_test(i, full_run=False):
+  print(f"\rRunning iteration {i}...", end=" ", flush=True)
+
+  p = subprocess.Popen(['python3', 'test/test_tiny.py', 'TestTiny.test_plus'],  stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+  if not full_run:
+    time.sleep(random.uniform(0, 1200) / 1000)
+    p.kill()
+    _, stderr = p.communicate()
+  else:
+    _, stderr = p.communicate()
+
+  if full_run:
+    stderr_text = stderr.decode()
+    print(stderr_text)
+    assert "Ran 1 test in" in stderr_text and "OK" in stderr_text
+
+max_workers = 4
+with ThreadPoolExecutor(max_workers=max_workers) as executor:
+  futures = []
+  for i in range(1000000):
+    if i % 100 == 0:
+      for future in as_completed(futures):
+        try: future.result()
+        except Exception as e:
+          print(f"\nError in iteration: {e}")
+      futures = []
+
+      run_test(i, True)
+    else:
+      future = executor.submit(run_test, i, False)
+      futures.append(future)
+
+    if len(futures) > max_workers * 2: futures = [f for f in futures if not f.done()]
@@ -1,44 +0,0 @@
-import subprocess
-import random
-import time
-from concurrent.futures import ProcessPoolExecutor, as_completed
-from tinygrad.helpers import getenv
-
-# checks that HCQ drivers can be killed during operation without causing issues
-
-def run_test(i, full_run=False, force_ok=False):
-  print(f"\rRunning iteration {i}...", end=" ", flush=True)
-
-  p = subprocess.Popen(["python3", "test/test_tiny.py", "TestTiny.test_plus"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-  if not full_run:
-    time.sleep(random.uniform(0, 1200) / 1000.0)
-    p.kill()
-    _, stderr = p.communicate()
-  else:
-    _, stderr = p.communicate()
-    stderr_text = stderr.decode()
-    assert ("Ran 1 test in" in stderr_text and "OK" in stderr_text) or (not force_ok and "Failed to take lock file" in stderr_text), stderr_text
-
-if __name__ == "__main__":
-  max_workers = getenv("MAX_WORKERS", 4)
-  with ProcessPoolExecutor(max_workers=max_workers) as executor:
-    futures = []
-    for i in range(1000000):
-      if i % 100 == 0:
-        # wait for everything we launched so far
-        for f in as_completed(futures):
-          try:
-            f.result()
-          except Exception as e:
-            print(f"\nError in iteration: {e}")
-        futures = []
-
-        # do a full run in the main proc
-        run_test(i, True, force_ok=True)
-      else:
-        futures.append(executor.submit(run_test, i, bool(getenv("FULL_RUN", 0))))
-
-      # keep list small
-      if len(futures) > max_workers * 2:
-        futures = [f for f in futures if not f.done()]
@@ -1,55 +0,0 @@
-import time
-from tinygrad.tensor import Tensor, Device
-
-MODEL_WIDTH = 512
-MODEL_HEIGHT = 256
-MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
-IMG_INPUT_SHAPE = (1, 12, 128, 256)
-
-def tensor_arange(end): return Tensor([float(i) for i in range(end)])
-def tensor_round(tensor:Tensor): return (tensor + 0.5).floor()
-
-h_src, w_src = 1208, 1928
-h_dst, w_dst = MODEL_HEIGHT, MODEL_WIDTH
-x = tensor_arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
-y = tensor_arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
-ones = Tensor.ones_like(x)
-dst_coords = x.reshape((1,-1)).cat(y.reshape((1,-1))).cat(ones.reshape((1,-1)))
-
-def warp_perspective_tinygrad(src:Tensor, M_inv:Tensor) -> Tensor:
-  src_coords = M_inv @ dst_coords
-  src_coords = src_coords / src_coords[2:3, :]
-
-  x_src = src_coords[0].reshape(h_dst, w_dst)
-  y_src = src_coords[1].reshape(h_dst, w_dst)
-
-  x_nearest = tensor_round(x_src).clip(0, w_src - 1).cast('int')
-  y_nearest = tensor_round(y_src).clip(0, h_src - 1).cast('int')
-
-  # TODO: make 2d indexing fast
-  idx = y_nearest*src.shape[1] + x_nearest
-  dst = src.flatten()[idx]
-  return dst.reshape(h_dst, w_dst)
-
-if __name__ == "__main__":
-  from tinygrad.engine.jit import TinyJit
-  update_img_jit = TinyJit(warp_perspective_tinygrad, prune=True)
-
-  step_times = []
-  for _ in range(10):
-    # regenerate inputs
-    inputs = [Tensor.randn(1928,1208), Tensor.randn(3,3)]
-    Tensor.realize(*inputs)
-    Device.default.synchronize()
-
-    # do the warp
-    st = time.perf_counter()
-    out = update_img_jit(*inputs)
-    mt = time.perf_counter()
-    val = out.contiguous().realize()
-    Device.default.synchronize()
-    et = time.perf_counter()
-
-    # measure the time
-    step_times.append((et-st)*1e3)
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -1,14 +1,13 @@
 import unittest
 from tinygrad import Device
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, CI, OSX
+from tinygrad.helpers import getenv, CI

 def multidevice_test(fxn):
  exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
  def ret(self):
    for device in Device._devices:
-      # broken on OSX USB AMD, why?
-      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
+      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"]: continue
      if not CI: print(device)
      if device in exclude_devices:
        if not CI: print(f"WARNING: {device} test is excluded")
@@ -1,20 +0,0 @@
-import os
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-
-import unittest, time
-from tinygrad import Device
-
-class TestOpen(unittest.TestCase):
-  def generate_test_open(n):
-    def test(self):
-      dev = Device[Device.DEFAULT]
-      for i in range(10):
-        dev.allocator.alloc(10 << 20)
-        time.sleep(0.5)
-    test.__name__ = f'test_open_{n}'
-    return test
-
-  for i in range(64): locals()[f'test_open_{i}'] = generate_test_open(i)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -112,6 +112,7 @@ backend_test.exclude('test_dequantizelinear_e5m2_cpu')
 backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')

 # we don't support indexes
+backend_test.exclude('test_nonzero_*')

 # no support for int pow
 backend_test.exclude('test_pow_types_int32_int32_cpu')
@@ -36,9 +36,7 @@ def trunc_log(x):
  logging.info("\n".join(lines))

 # user config
-# NOTE: process replay is slow so it's now disabled by default. add [pr] to enable it
-#SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
-SKIP_PROCESS_REPLAY = not ASSERT_DIFF
+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

@@ -85,11 +85,12 @@ class TestKernelSpeed(unittest.TestCase):
    gbs = mems / tm / 1e9
    self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)

+  # NOTE: tiny7 was slower than tiny12
  # TODO: why are convs so slow?!?
  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=14)

  # theoretical is nv_tflops=165, amd_tflops=123
-  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=110, amd_tflops=65)
+  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, amd_tflops=65)
  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=115, amd_tflops=60)

  # theoretical is nv_gbs=1008, amd_gbs=960
@@ -33,7 +33,7 @@ remu = _try_dlopen_remu()
 def create_sdma_packets():
  # TODO: clean up this, if we want to keep it
  structs = {}
-  for name,pkt in [(name,s) for name,s in amd_gpu.__dict__.items() if name.startswith("rocr_AMD_SDMA_PKT_") and name.endswith("_TAG")]:
+  for name,pkt in [(name,s) for name,s in amd_gpu.__dict__.items() if name.startswith("struct_SDMA_PKT_") and name.endswith("_TAG")]:
    names = set()
    fields = []
    for pkt_fields in pkt._fields_:
@@ -47,7 +47,7 @@ def create_sdma_packets():
          # merge together 64-bit fields, otherwise just append them
          if fname.endswith("_63_32") and fields[-1][0].endswith("_31_0"): fields[-1] = tuple([fname[:-6], ctypes.c_ulong, 64])
          else: fields.append(tuple([fname, *union_fields[1:]]))
-    new_name = name[18:-4].lower()
+    new_name = name[16:-4].lower()
    structs[new_name] = init_c_struct_t(tuple(fields))
    assert ctypes.sizeof(structs[new_name]) == ctypes.sizeof(pkt), f"{ctypes.sizeof(structs[new_name])} != {ctypes.sizeof(pkt)}"
  return type("SDMA_PKTS", (object, ), structs)
@@ -124,7 +124,6 @@ class PM4Executor(AMDQueue):
    elif mem_data_sel == 3:
      if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
      else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
-    elif mem_data_sel == 0: pass # no write
    else: raise RuntimeError(f"Unknown {mem_data_sel=}")

  def _exec_copy_data(self, n):
@@ -164,7 +164,7 @@ def cuStreamWaitEvent(stream: Any, event, flags: int) -> int: return orig_cuda.C
 def cuCtxSynchronize() -> int: return orig_cuda.CUDA_SUCCESS

 def cuGetErrorString(error: int, pStr) -> int:
-  error_str = orig_cuda.enum_cudaError_enum.get(error, "Unknown CUDA error").encode()
+  error_str = orig_cuda.cudaError_enum__enumvalues.get(error, "Unknown CUDA error").encode()
  buf = ctypes.create_string_buffer(error_str)
  # Set the pointer to point to our error string buffer
  pStr._obj.value = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char))
@@ -1,5 +1,5 @@
 import ctypes, mmap, collections, functools, os
-from tinygrad.runtime.autogen import nv_570 as nv_gpu
+import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from typing import Any
 from tinygrad.helpers import to_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, VirtFile
@@ -153,10 +153,8 @@ class NVDriver(VirtDriver):
                 51059, 51069, 51071, 51632, 51639, 51639, 51706, 52019, 222, 50287, 50273, 50031, 50017] # from ada102
      params.numClasses = len(classes)
      if struct.cmd == nv_gpu.NV0080_CTRL_CMD_GPU_GET_CLASSLIST:
-        if params.classList and params.numClasses > 0:
-          clslist = to_mv(params.classList, params.numClasses * 4).cast('I')
-          for i,c in enumerate(classes): clslist[i] = c
-        else: params.numClasses = len(classes)
+        clslist = to_mv(params.classList, params.numClasses * 4).cast('I')
+        for i,c in enumerate(classes): clslist[i] = c
      else:
        for i,c in enumerate(classes): params.classList[i] = c
    elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GR_GET_INFO:
@@ -194,9 +192,6 @@ class NVDriver(VirtDriver):
      params.mmuFaultInfoList[0].faultAddress = int(os.environ['MOCKGPU_EMU_FAULTADDR'], base=16)
      params.mmuFaultInfoList[0].faultType = 1
      params.mmuFaultInfoList[0].accessType = 1
-    elif struct.cmd == nv_gpu.NV0000_CTRL_CMD_SYSTEM_GET_BUILD_VERSION_V2:
-      params = nv_gpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_V2_PARAMS.from_address(params_ptr)
-      params.driverVersionBuffer = b"570.00.00\0"
    else: raise RuntimeError(f"Unknown {struct.cmd} to rm_control")
    return 0

@@ -259,4 +254,4 @@ class NVDriver(VirtDriver):
      for gpu in self.gpus.values():
        for q in gpu.queues:
          if q.ctrl.GPGet != q.ctrl.GPPut:
-            any_progress |= q.execute()
+            any_progress |= q.execute()
@@ -1,5 +1,5 @@
 import ctypes, time
-from tinygrad.runtime.autogen import nv_570 as nv_gpu
+import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from enum import Enum, auto
 from test.mockgpu.gpu import VirtGPU
 from test.mockgpu.helpers import _try_dlopen_gpuocelot
@@ -5,40 +5,26 @@ from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.uop.ops import Ops
-from tinygrad.renderer import Estimates
-from tinygrad.renderer.ptx import PTXRenderer

 class TestArange(unittest.TestCase):
-  def _get_flops(self, tensor, desired):
+  def _get_flops(self, N):
    GlobalCounters.reset()
-    sched = tensor.schedule()
+    tt = Tensor.arange(N)
+    sched = tt.schedule()
    self.assertEqual(len(sched), 1)
    p = get_program(sched[-1].ast)
-    ExecItem(CompiledRunner(p), [tensor.uop.buffer]).run()
-    np.testing.assert_equal(tensor.numpy(), desired)
+    ExecItem(CompiledRunner(p), [tt.uop.buffer]).run()
+    np.testing.assert_equal(tt.numpy(), np.arange(N))
    return p.estimates.ops

-  def test_arange_complexity(self):
-    self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
-    self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
+  def test_complexity(self):
+    self.assertEqual(self._get_flops(256), 0)
+    self.assertEqual(self._get_flops(2560), 0)

  def test_arange_cat(self):
    t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
    self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])

-  def test_eye_complexity(self):
-    with Context(NOOPT=1):
-      # NOTE: not every backend supports CMPEQ
-      self.assertLessEqual(self._get_flops(Tensor.eye(2560).contiguous(), np.eye(2560)), 2*2560*2560)
-
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
-  def test_tri_complexity(self):
-    with Context(NOOPT=1):
-      t = Tensor.ones(256, 256).contiguous().realize()
-      sched = t.triu().schedule()
-      p = get_program(sched[-1].ast)
-      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)
-
 DSET, DDIM = 2048, 32

 class TestIndexing(unittest.TestCase):
@@ -1,226 +0,0 @@
-import unittest
-from tinygrad import Tensor, UOp, Context
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import KernelInfo, AxisType
-
-# **** kernels ****
-
-def custom_arange_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.size, 0)
-  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
-
-def custom_eye_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.shape[0], 0)
-  j = UOp.range(C.shape[1], 1)
-  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
-
-def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
-  A,B = A.flatten(), B.flatten()
-  assert B.size == A.size
-  i = UOp.range(A.size, 0)
-  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.size}"))
-
-def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
-  C,A,B = C.flatten(), A.flatten(), B.flatten()
-  i = UOp.range(C.size, 0)
-  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.size}")).simplify()
-
-def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
-  C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
-  assert C.size == D.size
-  i = UOp.range(C.size, 0)
-  store_c = C[i].store(A[i]+B[i])
-  store_d = D[i].store(A[i]*B[i])
-  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  assert A.shape[1] == B.shape[0]
-  i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
-  C = C[i, j].set(0.0)
-  C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
-  prog = C.end(i, j)
-  return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
-
-def custom_sum(B:UOp, A:UOp) -> UOp:
-  i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
-  B = B[0].set(0.0)
-  B = B[0].set(B.after(i)[0] + A[i], end=i)
-  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
-
-def flip_contract_kernel(dest:UOp, src:UOp):
-  i = UOp.range(dest.shape[0], 0)
-  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
-  vec = src[i, j].contract(j)
-  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
-
-def slice_sum_kernel(dest:UOp, src:UOp):
-  G = UOp.range(src.shape[0], 0)
-  slice_src = src[G, :]
-  reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
-  reg = reg.after(G)[0].set(0)
-  R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
-  reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
-  ast = dest[G].set(reg[0], end=G)
-  return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
-
-def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
-  # attention without softmax
-  N, d = Q.shape[0], Q.shape[1]
-
-  i = UOp.range(N, 0)  # output row
-  d_out = UOp.range(d, 1)  # output column
-  j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
-
-  k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
-  qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
-  qk_acc = qk_acc.after(i, j)[0].set(0.0)
-  qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
-  qk_score = qk_acc[0] / (d ** 0.5)
-
-  out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
-  out_acc = out_acc.after(i, d_out)[0].set(0.0)
-  out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
-
-  store = O[i, d_out].store(out_acc[0])
-  return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
-
-# **** backward callbacks ****
-
-def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
-  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
-  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
-  return (None, grad_a, grad_b)
-
-def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
-  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
-  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
-  return (None, grad_a, grad_b)
-
-# **** tests ****
-
-class TestCustomKernel(unittest.TestCase):
-  def test_empty(self):
-    a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink())[0]
-    a.realize()
-
-  def test_simple(self):
-    a = Tensor.ones(16, 16).contiguous()
-    b = Tensor.ones(16, 16).contiguous()
-    c = Tensor.empty(16, 16)
-
-    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
-
-    out = c.flatten().tolist()
-    assert all(x == 2 for x in out), "all 2"
-
-  def test_multioutput(self):
-    a = Tensor.full((16, 16), 3.).contiguous()
-    b = Tensor.full((16, 16), 3.).contiguous()
-    c = Tensor.empty(16, 16)
-    d = Tensor.empty(16, 16)
-
-    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
-    Tensor.realize(c,d)
-
-    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
-    assert all(x == 9 for x in d.flatten().tolist()), "all 9"
-
-  def test_arange(self):
-    ref = Tensor.arange(100)
-    tst = Tensor.empty_like(ref)
-    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
-    self.assertTrue((ref == tst).all().item())
-
-  def test_eye(self):
-    ref = Tensor.eye(1024).contiguous().realize()
-    tst = Tensor.empty_like(ref)
-    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
-    self.assertTrue((ref == tst).all().item())
-
-  def test_flip_contract(self):
-    a = Tensor.randn(10,4)
-    b = Tensor.empty_like(a)
-    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
-    self.assertTrue((a.flip(1) == b).all().item())
-
-  def test_noncontig(self):
-    a = Tensor.ones(16, 16).contiguous()
-    tst = Tensor.empty_like(a)
-    b = a+1
-    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
-    self.assertTrue((b_p1 == 3).all().item())
-
-  def test_sum(self):
-    # TODO: this only works for float, and silently fails with int
-    a = Tensor([1.0, 2, 3, 4, 5])
-    tst = Tensor.empty(1)
-    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
-    self.assertEqual(b.item(), 15)
-
-  def test_slice_sum(self):
-    A = Tensor.randn(16, 16).contiguous()
-    B = Tensor.empty(16)
-    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
-    self.assertTrue(B.allclose(A.sum(1)))
-
-  def test_gemm(self):
-    N = 16
-    a = Tensor.randn(N, N)
-    b = Tensor.randn(N, N)
-    c = Tensor.empty(N, N)
-
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    err = (tst - (a@b)).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
-  # NOTE: grad_fxn doesn't work with pyrender
-  @Context(SPEC=1)
-  def test_gemm_backward(self, custom_backward_gemm=False):
-    N = 4
-    a_rand = Tensor.randn(N, 8)
-    b_rand = Tensor.randn(8, N)
-    Tensor.realize(a_rand, b_rand)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    c = Tensor.empty(N, N)
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
-    tst.sum().backward()
-    grad_a, grad_b = a.grad, b.grad
-    Tensor.realize(tst, grad_a, grad_b)
-
-    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
-    ref = (a@b)
-    ref.sum().backward()
-    real_grad_a, real_grad_b = a.grad, b.grad
-    Tensor.realize(ref, real_grad_a, real_grad_b)
-
-    err = (tst - ref).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_a - real_grad_a).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_b - real_grad_b).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-  def test_simple_qkv(self):
-    N, d = 8, 4
-    Q = Tensor.randn(N, d)
-    K = Tensor.randn(N, d)
-    V = Tensor.randn(N, d)
-    O = Tensor.empty(N, d)
-
-    O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
-    O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
-
-    Tensor.realize(O_custom, O_ref)
-    err = (O_custom - O_ref).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -4,7 +4,7 @@ from dataclasses import replace

 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, PatternMatcher, graph_rewrite, UPat
+from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
@@ -14,8 +14,6 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 MOCKGPU = getenv("MOCKGPU")

-from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
-
 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
    # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
@@ -40,38 +38,6 @@ class TestLinearizer(unittest.TestCase):
    np.testing.assert_equal(a.numpy(), ta)
    np.testing.assert_equal(b.numpy(), tb)

-  @unittest.skip("TODO: some backends insert more casts")
-  def test_cast_there_and_back(self):
-    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
-    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
-    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
-
-  @unittest.expectedFailure
-  def test_cast_back_and_there(self):
-    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
-    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
-    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
-
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
-  def test_late_bias_load(self):
-    img = Tensor.empty(1, 3, 16, 16)
-    w = Tensor.empty(16, 3, 3, 3)
-    b = Tensor.empty(16)
-    out = img.conv2d(w, b)
-    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
-    # slice at the last loop end
-    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
-    # only valid test if outermost range is the reduce
-    if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
-      load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
-      # assert that there is a global load after the reduce ends
-      assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
-
  def _test_no_nested_ranges(self, lins, skip=None):
    for l in lins:
      range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
@@ -96,7 +62,6 @@ class TestLinearizer(unittest.TestCase):
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

-  @unittest.expectedFailure # TODO: investigate
  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()
@@ -297,8 +262,6 @@ class TestLinearizer(unittest.TestCase):
    _assert_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
    #                              2             ->             3
    _assert_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
-    #                              2             ->             2
-    _assert_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
    # test when the only divisor is the square root of dim
    _assert_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)

@@ -323,27 +286,6 @@ class TestLinearizer(unittest.TestCase):
    with self.assertRaises(RuntimeError):
      get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))

-    # TODO: In the above cases we only test if the shape after reshape is correct, never the indices.
-    # We should check if the returned indices are correct, for all cases.
-    # (65536, 2) -> (32768, 4)
-    dims, expected_limited_dims = (65536,2), (32768, 4)
-    idxs = get_grouped_dims("gidx", dims, (65535,65535,65535))
-    def match_div(): raise RuntimeError("match_div")
-    def match_mod(): raise RuntimeError("match_mod")
-    flat_idx_pattern = UPat(Ops.SPECIAL, arg='gidx0')*expected_limited_dims[1]+UPat(Ops.SPECIAL, arg='gidx1')
-    pm = PatternMatcher([
-      (flat_idx_pattern//dims[1], match_div),
-      (flat_idx_pattern%dims[1], match_mod)
-    ])
-
-    with self.assertRaises(RuntimeError) as error:
-      graph_rewrite(idxs[0], pm)
-    self.assertIn("match_div", str(error.exception))
-
-    with self.assertRaises(RuntimeError) as error:
-      graph_rewrite(idxs[1], pm)
-    self.assertIn("match_mod", str(error.exception))
-
    # # variable too large
    # with self.assertRaises(AssertionError):
    #   get_grouped_dims("gidx", (Variable("start_pos",0,16),3,4), (16,16,16), False,)
@@ -490,8 +432,6 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  # now all input buffers in s[-1] should be realized
  # create fresh buffers for the outputs
  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
-  # ensure buffers are allocated
-  for b in bufs: b.ensure_allocated()
  return s[-1].ast, bufs

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -509,7 +449,7 @@ def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
  return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]

 def reset_bufs(bufs:list[Buffer]):
-  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
+  for buf in bufs: buf.copyin(np.zeros((buf.size, ), dtype=_to_np_dtype(buf.dtype)).data) # Zero to check that all values are filled

 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
@@ -596,12 +596,6 @@ class TestMultiTensor(unittest.TestCase):
    # ast are the same on devices
    self.assertEqual(len(set(asts)), 1)

-  def test_flip(self):
-    rng = Tensor.rand((10, 10, 10))
-    t0 = rng.shard(devices_2, axis=1)
-    out = t0.flip(0) + 1
-    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))
-
  def test_reshape_on_axis(self):
    t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)

@@ -1551,10 +1551,8 @@ class TestOps(unittest.TestCase):
                                lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))

  def test_std_mean_loaded_nan(self):
-    with warnings.catch_warnings():
-      warnings.filterwarnings("ignore", message="std_mean\\(\\): degrees of freedom is <= 0")
-      helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
-                                    lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
+    helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
+                                  lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
  def test_softmax(self):
    helper_test_op([(45,65)], torch.nn.Softmax(dim=1), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
    helper_test_op([(45)], torch.nn.Softmax(dim=0), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
@@ -2822,13 +2820,13 @@ class TestOps(unittest.TestCase):
  @slow_test
  def test_slice_fancy_indexing_list_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(0,),b,c,d,:], lambda x: x[(0,),j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[0],b,c,d,:], lambda x: x[[0],j,k,o,:])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[[0]]],b,c,d,[[1]]], lambda x: x[[[[0]]],j,k,o,[[1]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(1,0,-1),b,c,d,:], lambda x: x[(1,0,-1),j,k,o,:])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,(1,2,3),...], lambda x: x[i,j,k,(1,2,3),...])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[1,0,-1],b,c,d,:], lambda x: x[[1,0,-1],j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[1,2,3],...], lambda x: x[i,j,k,[1,2,3],...])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[[1],[2],[3]],...], lambda x: x[i,j,k,[[1],[2],[3]],...])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[-2,1,0],e], lambda x: x[i,[2,1,0],k,[-2,1,0],p])

  @slow_test
  def test_slice_fancy_indexing_tuple_indices(self):
@@ -2843,10 +2841,11 @@ class TestOps(unittest.TestCase):
  @slow_test
  def test_slice_fancy_indexing_list_with_tensors(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,)], lambda x: x[(i,)])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,1)], lambda x: x[(i,1)])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,(1,1))], lambda x: x[(i,(1,1))])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,b,c,d,e)], lambda x: x[(i,j,k,o,p)])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a]], lambda x: x[[i]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,1]], lambda x: x[[i,1]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,[1,1]]], lambda x: x[[i,[1,1]]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,(1,1)]], lambda x: x[[i,(1,1)]])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a,b,c,d,e]], lambda x: x[[i,j,k,o,p]])

  def test_slice_fancy_indexing_errors(self):
    a = Tensor.ones(10,11,12)
@@ -1,6 +1,5 @@
 import unittest
-import numpy as np
-from tinygrad import Tensor, UOp, nn
+from tinygrad import Tensor, UOp
 from tinygrad.uop.ops import AxisType, Ops

 class TestOuterworldReduce(unittest.TestCase):
@@ -12,81 +11,6 @@ class TestOuterworldReduce(unittest.TestCase):
    t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, a), arg=Ops.ADD))
    self.assertListEqual(t.tolist(), [5.,5.,5.,5.,5.])

-# TODO: delete test_outerworld_range?
-class TestOuterRange(unittest.TestCase):
-  def test_simple_range(self):
-    a = Tensor.ones(10).contiguous()
-    acc = Tensor.zeros().contiguous()
-    Tensor.realize(a, acc)
-
-    # this is fold
-    i = UOp.range(10, -100, AxisType.OUTER)
-    acc_i = acc.uop.after(i)
-    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
-    out = Tensor(acc.uop.after(acc_i.store(acc_i + a[vi].uop).end(i)))
-    out.realize()
-    assert out.item() == 10.0
-
-  def test_inner_range(self):
-    a = Tensor.ones(10, 10).contiguous()
-    acc = Tensor.zeros(10).contiguous()
-    Tensor.realize(a, acc)
-
-    # this is fold
-    i = UOp.range(10, -100, AxisType.OUTER)
-    acc_i = acc.uop.after(i)
-    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
-    out = Tensor(acc.uop.after(acc_i.store(acc_i + a[:, vi].uop).end(i)))
-    out.realize()
-    assert all(x == 10.0 for x in out.tolist())
-
-  def test_range_matmul(self):
-    vec = Tensor.randn(1, 10).realize()
-    mats = Tensor.randn(3, 10, 10).realize()
-
-    # 3 matmuls in "scan"
-    ref = ((vec @ mats[0]) @ mats[1]) @ mats[2]
-    ref.realize()
-
-    # 3 matmuls with outer world range
-    i = UOp.range(3, -100, AxisType.OUTER)
-    vec_i = Tensor(vec.uop.after(i))
-    comp = vec_i.contiguous() @ mats[i]
-    store = vec_i.uop.store(comp.uop).end(i)
-    out = Tensor(vec.uop.after(store))
-    out.realize()
-
-    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
-
-class TestOuterScan(unittest.TestCase):
-  def _test_scan(self):
-    vec = Tensor.randn(1, 10).realize()
-    mats = Tensor.randn(3, 10, 10).realize()
-
-    # 3 matmuls in "scan"
-    vec1 = vec @ mats[0]
-    vec2 = vec1 @ mats[1]
-    vec3 = vec2 @ mats[2]
-    ref = Tensor.stack(vec1, vec2, vec3)
-    ref.realize()
-    return vec, mats, ref
-
-  def test_uop_scan_matmul(self):
-    vec, mats, ref = self._test_scan()
-
-    # 3 matmuls with SCAN
-    i = UOp.range(3, -100, AxisType.OUTER)
-    out = Tensor.empty(3, 1, 10)
-    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
-    comp = phi @ mats[i]
-    store = out[i].uop.store(comp.uop).end(i)
-    out = Tensor(out.uop.after(store))
-    out.realize()
-
-    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
-
 class TestOuterworld(unittest.TestCase):
  def test_range_plus_1(self):
    t = Tensor.arange(100).reshape(10,10).realize()
@@ -146,85 +70,5 @@ class TestOuterworld(unittest.TestCase):
    out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
    self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())

-class TestVmap(unittest.TestCase):
-  def test_vmap_inner(self, axis_type=AxisType.LOOP, fuse=False, grad=False):
-    x = Tensor.ones(1, 10).contiguous().requires_grad_()
-    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
-
-    ref = x @ mats
-    if fuse: ref = ref * 2
-
-    # vmap across axis 0
-    a = UOp.range(3, -1, axis_type)
-    out = x @ mats[a]
-    out = out.reshape(1, 10).pad(((a,(3-a)-1), None))
-    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
-    if fuse: out = out * 2
-    if grad:
-      out.mean().backward()
-      np.testing.assert_allclose(mats.grad.numpy(), (2./30) if fuse else (1./30))
-    out.realize()
-
-    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
-  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
-  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
-  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
-
-  def test_vmap_inner_grad(self): self.test_vmap_inner(grad=True)
-  def test_vmap_inner_fuse_grad(self): self.test_vmap_inner(fuse=True, grad=True)
-  def test_vmap_outer_grad(self): self.test_vmap_inner(AxisType.OUTER, grad=True)
-
-  def test_vmap_convs(self):
-    layers = [
-      nn.Conv2d(1, 8, 3), Tensor.relu,
-      nn.Conv2d(8, 8, 3), Tensor.relu]
-    img = Tensor.randn(4, 1, 16, 16).realize(*nn.state.get_parameters(layers))
-    a = UOp.range(4, -1, AxisType.OUTER)
-    out = img[a:a+1].sequential(layers)
-    out = out.pad(((a,(4-a)-1), None, None, None))
-    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
-    out.realize()
-    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
-
-  def test_vmap_gemm(self):
-    layers = [
-      nn.Linear(16, 16, bias=False), Tensor.relu,
-      nn.Linear(16, 16, bias=False), Tensor.relu]
-    img = Tensor.randn(4, 16).realize(*nn.state.get_parameters(layers))
-    a = UOp.range(4, -1, AxisType.OUTER)
-    out = img[a:a+1].sequential(layers)
-    out = out.pad(((a,(4-a)-1), None))
-    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
-    out.realize()
-    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
-
-  @unittest.skip("this is broken, we need to lower the outer reduce in the outer graph")
-  def test_vmap_gemm_grad(self):
-    layers = [
-      nn.Linear(16, 16, bias=False), Tensor.relu,
-      nn.Linear(16, 16, bias=False), Tensor.relu]
-    layer_tensors = nn.state.get_parameters(layers)
-    img = Tensor.randn(4, 16).realize(*layer_tensors)
-    for l in layer_tensors: l.requires_grad_()
-    a = UOp.range(4, -1, AxisType.OUTER)
-    out = img[a:a+1].sequential(layers)
-    out = out.pad(((a,(4-a)-1), None))
-    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
-    out.mean().backward()
-    grads = [l.grad for l in layer_tensors]
-    out.realize(*grads)
-    out_grads = [x.numpy() for x in grads]
-
-    # compute reference grads
-    for l in layer_tensors: l.grad = None
-    img.sequential(layers).mean().backward()
-    grads = [l.grad for l in layer_tensors]
-    out.realize(*grads)
-    ref_grads = [x.numpy() for x in grads]
-
-    # compare
-    for o,r in zip(out_grads, ref_grads): np.testing.assert_allclose(o, r, atol=1e-6)
-
 if __name__ == '__main__':
  unittest.main()
@@ -17,7 +17,7 @@ def helper_collect_profile(*devs):
  cpu_events.clear()

  profile_list = []
-  with Context(VIZ=1, PROFILE=1):
+  with Context(VIZ=1):
    yield profile_list
    for dev in devs: dev.synchronize()
    for dev in devs: dev._at_profile_finalize()
@@ -300,12 +300,6 @@ class TestRangeify(unittest.TestCase):
    w2 = Tensor.empty(12, 8, 3, 3)
    x.conv2d(w1).conv2d(w2).realize()

-  def test_resnet_conv2d(self):
-    x = Tensor.empty(1, 8, 32, 32)
-    w1 = Tensor.empty(8, 8, 3, 3)
-    w2 = Tensor.empty(8, 8, 1, 1)
-    x.conv2d(w1).conv2d(w2).realize()
-
  def test_xception_conv2d(self):
    # NOTE: this fusion is bad, it's recomputing the inner many times
    x = Tensor.empty(1, 4, 32, 32)
@@ -711,7 +711,7 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(b.buffer.numpy(), [12])

  # unlike schedule, kernelize can be called multiple times on a Tensor
-  def test_double_kernelize(self):
+  def test_double_kerenlize(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10)
    c = (a+b)
@@ -1042,12 +1042,13 @@ class TestSchedule(unittest.TestCase):
      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
      np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)

-    out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 4)) # TODO: should be 1?
-    if getenv("CHECK", 1):
-      import torch
-      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
-      np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
+    with Context(FUSE_ATTENTION=1):
+      out = Tensor.scaled_dot_product_attention(q,k,v)
+      run_schedule(check_schedule(out, 4)) # TODO: should be 1?
+      if getenv("CHECK", 1):
+        import torch
+        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
+        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)

  def test_ugly_reduceop_pairing(self):
    Tensor.manual_seed(0)
@@ -1502,18 +1503,6 @@ class TestSchedule(unittest.TestCase):
    run_schedule(sched)
    np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)

-  def test_fuse_arange_avg_pool2d_ceil_mode(self):
-    x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True)
-    sched = check_schedule(x, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1)
-
-  def test_fuse_arange_pad_circular_mode_bw(self):
-    x = Tensor.empty(1,1,5,5,5)
-    out = x.pad((1,2,3,5,1,2), mode="circular")
-    g = out.sum().gradient(x)[0]
-    sched = check_schedule(g, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0)
-
  # TODO like openpilot with imagef
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_base_change_expand_expand(self):
@@ -1573,13 +1562,6 @@ class TestSchedule(unittest.TestCase):
  def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
  def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)

-  def test_resnet_conv2d(self):
-    x = Tensor.empty(1, 8, 32, 32)
-    w1 = Tensor.empty(8, 8, 3, 3)
-    w2 = Tensor.empty(8, 8, 1, 1)
-    out = x.conv2d(w1).conv2d(w2)
-    check_schedule(out, 2)
-
  @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
  def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -2285,7 +2267,7 @@ class TestContiguous(unittest.TestCase):
  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 2) # TODO: should be 1?

  def test_view_does_not_realize(self):
    a = Tensor.empty(4)
@@ -32,7 +32,7 @@ def run_one_schedule_item(out): lower_schedule_item(get_single_element(out.sched
 class TestFuse(unittest.TestCase):
  def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
    GlobalCounters.reset()
-    out_single = fxn(*args, **kwargs)
+    out_single = fxn(*args, **kwargs).fuse()
    if not allow_multiple: run_one_schedule_item(out_single)
    np_single = out_single.numpy()
    GlobalCounters.reset()
@@ -100,7 +100,7 @@ class TestFuse(unittest.TestCase):
    q = (x @ wq).contiguous()
    k = (x @ wk).contiguous()
    v = (x @ wv).contiguous()
-    attn = q.scaled_dot_product_attention(k, v)
+    attn = q.scaled_dot_product_attention(k, v).fuse()
    s = attn.schedule()
    self.assertEqual(len(s), 4) # 3 matmul and 1 attention

@@ -121,7 +121,7 @@ class TestFuse(unittest.TestCase):
  def test_mismatch_reduce(self):
    a = Tensor.ones(16, 10).contiguous().realize()
    b = Tensor.ones(16, 20).contiguous().realize()
-    c = (a.sum(axis=1) + b.sum(axis=1))
+    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
    self.assertListEqual(c.tolist(), [30]*16)

  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
@@ -129,7 +129,7 @@ class TestFuse(unittest.TestCase):
    A = Tensor.randn(8, 8).realize()
    B = Tensor.randn(8, 8).realize()
    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
-    out = (C + (A @ B))
+    out = (C + (A @ B)).fuse()
    out.realize()

 class TestSoftmaxFusion(unittest.TestCase):
@@ -180,7 +180,7 @@ class TestSoftmaxFusion(unittest.TestCase):

    print("*** auto single kernel softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
-      out = self.test.contiguous().softmax(-1)
+      out = self.test.contiguous().softmax(-1).fuse()
      run_one_schedule_item(out)

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
@@ -3,7 +3,7 @@ import torch
 import unittest, copy, mmap, random, math, array
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _METADATA
-from tinygrad.helpers import Context, getenv, temp, mv_address
+from tinygrad.helpers import getenv, temp, mv_address
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
@@ -846,16 +846,6 @@ class TestTensorMetadata(unittest.TestCase):
    #self.assertEqual(len(bw), 1)
    #self.assertEqual(bw[0].name, "sigmoid")

-  def test_tracemeta_0(self):
-    with Context(TRACEMETA=0):
-      x = Tensor.rand(3, requires_grad=True)
-      y = Tensor.rand(3, requires_grad=True)
-      out = (x.relu() * y.sigmoid()).sum()
-      self.assertIsNone(out.uop.metadata)
-      self.assertIsNone(out.uop.src[0].metadata)
-      si = out.schedule()[-1]
-      self.assertEqual(si.metadata, ())
-
 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
    if ast.op is op: return ast
@@ -32,8 +32,8 @@ class TestTiny(unittest.TestCase):
    self.assertListEqual(out.tolist(), [2]*16)

  def test_cat(self):
-    out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.zeros(8).contiguous())
-    self.assertListEqual(out.tolist(), [1]*8+[0]*8)
+    out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.ones(8).contiguous())
+    self.assertListEqual(out.tolist(), [1]*16)

  def test_sum(self):
    out = Tensor.ones(256).contiguous().sum()
@@ -62,7 +62,7 @@ class TestTiny(unittest.TestCase):
    out = Tensor.rand(10)
    for x in out.tolist():
      self.assertGreaterEqual(x, 0.0)
-      self.assertLess(x, 1.0)
+      self.assertLessEqual(x, 1.0)

  # *** JIT (for Python speed) ***

@@ -138,7 +138,9 @@ class TestTiny(unittest.TestCase):
      nn.Conv2d(8, 8, 5), Tensor.relu]

    # replace random weights with ones
-    Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
+    # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
+    #Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
+    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))

    # realize gradients
    for x in nn.state.get_parameters(layers): x.requires_grad_()
@@ -569,21 +569,11 @@ class TestUOpPrograms(unittest.TestCase):
  def _run(self, prog:UOp, *tensors:Tensor):
    ExecItem(get_runner(Device.DEFAULT, prog), [t.uop.buffer for t in tensors]).run(wait=True)

-  def test_simple(self):
-    out = Tensor.empty(10,10,dtype=dtypes.int)
-
-    ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
-    i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
-    prog = ptr[i,j].set(42).end(i,j)
-    self._run(prog.sink(), out)
-
-    with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
-
  def test_matmul(self):
-    a = Tensor.randn(10,10)
-    b = Tensor.randn(10,10)
+    a = Tensor.rand(10,10)
+    b = Tensor.rand(10,10)
    c = Tensor.empty(10,10)
-    ref = (a@b)
+    ref = a@b
    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)

    # C[i,j] = sum_k A[i,k] * B[k,j]
@@ -591,16 +581,16 @@ class TestUOpPrograms(unittest.TestCase):
    M = N = K = 10
    DT = dtypes.float32

-    # Placeholders (bind slots explicitly)
-    A = UOp.placeholder((M, K), DT, slot=0)
-    B = UOp.placeholder((K, N), DT, slot=1)
-    C = UOp.placeholder((M, N), DT, slot=2)
-
    # Axes: i,j are spatial; k is a reduction axis over the shared dim K
    i = UOp.range(M, axis_id=0)                             # rows of A/C
    j = UOp.range(N, axis_id=1)                             # cols of B/C
    k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE)  # reduction over K

+    # Placeholders (bind slots explicitly)
+    A = UOp.placeholder(DT, (M, K), slot=0)
+    B = UOp.placeholder(DT, (K, N), slot=1)
+    C = UOp.placeholder(DT, (M, N), slot=2)
+
    # Zero-init: write a scalar 0 to each (i,j).
    C = C[i, j].set(0.0)

@@ -611,27 +601,9 @@ class TestUOpPrograms(unittest.TestCase):
    prog = C.end(i, j, k)

    # run program
-    # TODO: make this work with opts_to_apply
    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)

    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)

-  def test_matmul_relu(self):
-    a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
-    ref = (a@b).relu()
-    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
-
-    A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
-    i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
-
-    C = C[i, j].set(0.0)
-    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
-    C = C[i, j].set(C[i, j].maximum(0.0))
-
-    prog = C.end(i, j)
-
-    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
-    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
-
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -1,443 +0,0 @@
-import unittest, math
-
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.engine.realize import ExecItem, get_runner
-from tinygrad.helpers import CI
-from tinygrad.renderer.ptx import PTXRenderer
-import numpy as np
-
-from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.kernel import Kernel
-
-@unittest.skipIf(CI and Device.DEFAULT not in ["CUDA", "NV"], "only cuda")
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
-class TestTK(unittest.TestCase):
-  @unittest.skipIf(CI, "no wmma in ci")
-  def test_simple_matmul(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      c = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.bfloat16)
-      b = ker.gl((1, 1, N, N), dtypes.bfloat16)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      col, row = ker.blockIdx_x, ker.blockIdx_y
-
-      c_reg = warp.zero(c_reg)
-      for tile in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
-        b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
-
-        a_reg = warp.load(a_reg, a_smem)
-        b_reg = warp.load(b_reg, b_smem, transpose=True)
-
-        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
-      c_reg = ker.endrange()
-
-      c_smem = warp.store(c_smem, c_reg)
-      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
-      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
-      c = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b, c)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
-    for _ in range(5): ei.run(wait=True)
-    c = c.float()
-
-    ref = a.matmul(b, dtype=dtypes.float32).float()
-
-    np.testing.assert_allclose(c.numpy(), ref.numpy())
-
-  @unittest.skipIf(CI, "no wmma in ci")
-  def test_simple_matmul_transposed(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      c = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.bfloat16)
-      b = ker.gl((1, 1, N, N), dtypes.bfloat16)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      col, row = ker.blockIdx_x, ker.blockIdx_y
-
-      c_reg = warp.zero(c_reg)
-      for tile in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
-        b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
-
-        a_reg = warp.load(a_reg, a_smem)
-        b_reg = warp.load(b_reg, b_smem)
-
-        c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
-      c_reg = ker.endrange()
-
-      c_smem = warp.store(c_smem, c_reg)
-      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
-      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
-      c = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b, c)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
-    for _ in range(5): ei.run(wait=True)
-    c = c.float()
-
-    ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float()
-
-    np.testing.assert_allclose(c.numpy(), ref.numpy())
-
-  def test_load_store(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      col, row = ker.blockIdx_x, ker.blockIdx_y
-
-      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
-      a_reg = warp.load(a_reg, a_smem)
-      b_reg = warp.copy(b_reg, a_reg)
-      b_smem = warp.store(b_smem, b_reg)
-      b = warp.store(b, b_smem, (0, 0, row, col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float()
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy())
-
-  def test_add(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      for tile_row in ker.range(N // BLOCK_SIZE):
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
-          a_reg = warp.load(a_reg, a_smem)
-
-          a_reg += 1
-
-          a_smem = warp.store(a_smem, a_reg)
-          b = warp.store(b, a_smem, (0, 0, tile_row, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-      with Context(DEBUG=0):
-        a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
-        b = Tensor.empty(1, 1, N, N, dtype="float32")
-        Tensor.realize(a, b)
-
-      ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-      for _ in range(5): ei.run(wait=True)
-      b = b.float()
-
-      ref = a.float() + 1
-
-      np.testing.assert_allclose(b.numpy(), ref.numpy())
-
-  def test_max(self):
-    N = 16
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      max_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-
-      for tile_row in ker.range(N // BLOCK_SIZE):
-        max_reg = warp.neg_inf(max_reg.after(tile_row))
-
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
-          a_reg = warp.load(a_reg, a_smem)
-          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
-        max_reg = ker.endrange()
-
-        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
-
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy())
-
-  def test_max_nonsquare(self):
-    N, M = 16, 64
-    BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, M), dtypes.float32)
-      a = ker.gl((1, 1, N, M), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-
-      max_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
-
-      for tile_row in ker.range(N // BLOCK_N):
-        max_reg = warp.neg_inf(max_reg.after(tile_row))
-
-        for tile_col in ker.range(M // BLOCK_M):
-          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
-          a_reg = warp.load(a_reg, a_smem)
-          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
-        max_reg = ker.endrange()
-
-        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
-
-        for tile_col in ker.range(M // BLOCK_M):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, M, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy())
-
-  def test_sum(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-
-      for tile_row in ker.range(N // BLOCK_SIZE):
-        sum_reg = warp.zero(sum_reg.after(tile_row))
-
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
-          a_reg = warp.load(a_reg, a_smem)
-          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
-        sum_reg = ker.endrange()
-
-        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
-
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
-
-  def test_sum_nonsquare(self):
-    N, M = 16, 64
-    BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, M), dtypes.float32)
-      a = ker.gl((1, 1, N, M), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-
-      sum_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
-
-      for tile_row in ker.range(N // BLOCK_N):
-        sum_reg = warp.zero(sum_reg.after(tile_row))
-
-        for tile_col in ker.range(M // BLOCK_M):
-          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
-          a_reg = warp.load(a_reg, a_smem)
-          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
-        sum_reg = ker.endrange()
-
-        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
-
-        for tile_col in ker.range(M // BLOCK_M):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, M, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
-
-  @unittest.skip("fake range not ended")
-  def test_softmax(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
-      a = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-      max_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-      norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-
-      max_vec = warp.neg_inf(max_vec)
-      norm_vec = warp.zero(norm_vec)
-
-      for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
-
-        a_reg *= 1.0 / math.log(2)
-
-        max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
-        max_vec = warp.row_reduce(max_vec, a_reg, lambda a, b: a.maximum(b))
-        a_reg = (a_reg - max_vec).exp2()
-        max_vec_last = (max_vec_last - max_vec).exp2()
-        norm_vec *= max_vec_last
-        norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
-      norm_vec = ker.endrange()
-
-      for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
-
-        a_reg *= 1.0 / math.log(2)
-        a_reg = (a_reg - max_vec).exp2()
-        a_reg /= norm_vec
-
-        a_smem = warp.store(a_smem, a_reg)
-        b = warp.store(b, a_smem, (0, 0, 0, tile_col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, BLOCK_SIZE, N, dtype="float32")
-      b = Tensor.empty(1, 1, BLOCK_SIZE, N, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float().softmax(axis=3)
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,85 +0,0 @@
-import ctypes, subprocess, tempfile, unittest
-from tinygrad.helpers import WIN
-from tinygrad.runtime.support.c import Struct
-
-class TestAutogen(unittest.TestCase):
-  def test_packed_struct_sizeof(self):
-    layout = [('a', ctypes.c_char), ('b', ctypes.c_int, 5), ('c', ctypes.c_char)]
-    class Y(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Z(Struct): pass
-    Z._packed_, Z._fields_ = True, layout
-    self.assertEqual(ctypes.sizeof(Y), 6)
-    self.assertEqual(ctypes.sizeof(Z), 3)
-    layout = [('a', ctypes.c_int, 31), ('b', ctypes.c_int, 31), ('c', ctypes.c_int, 1), ('d', ctypes.c_int, 1)]
-    class Foo(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
-    class Bar(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Baz(Struct): pass
-    Baz._packed_, Baz._fields_ = True, layout
-    self.assertEqual(ctypes.sizeof(Foo), 12)
-    self.assertEqual(ctypes.sizeof(Bar), 12)
-    self.assertEqual(ctypes.sizeof(Baz), 8)
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_packed_struct_interop(self):
-    class Baz(Struct): pass
-    Baz._packed_ = True
-    Baz._fields_ = [('a', ctypes.c_int, 30), ('b', ctypes.c_int, 30), ('c', ctypes.c_int, 2), ('d', ctypes.c_int, 2)]
-    src = '''
-      struct __attribute__((packed)) baz {
-        int a:30;
-        int b:30;
-        int c:2;
-        int d:2;
-      };
-
-      int test(struct baz x) {
-        return x.a + x.b + x.c + x.d;
-      }
-    '''
-    args = ('-x', 'c', '-fPIC', '-shared')
-    with tempfile.NamedTemporaryFile(suffix=".so") as f:
-      subprocess.check_output(('clang',) + args + ('-', '-o', f.name), input=src.encode('utf-8'))
-      b = Baz(0xAA000, 0x00BB0, 0, 1)
-      test = ctypes.CDLL(f.name).test
-      test.argtypes = [Baz]
-      self.assertEqual(test(b), b.a + b.b + b.c + b.d)
-
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_packed_structs(self):
-    NvU32 = ctypes.c_uint32
-    NvU64 = ctypes.c_uint64
-    class FWSECLIC_READ_VBIOS_DESC(Struct): pass
-    FWSECLIC_READ_VBIOS_DESC._packed_ = True
-    FWSECLIC_READ_VBIOS_DESC._fields_ = [
-      ('version', NvU32),
-      ('size', NvU32),
-      ('gfwImageOffset', NvU64),
-      ('gfwImageSize', NvU32),
-      ('flags', NvU32),
-    ]
-    class FWSECLIC_FRTS_REGION_DESC(Struct): pass
-    FWSECLIC_FRTS_REGION_DESC._packed_ = True
-    FWSECLIC_FRTS_REGION_DESC._fields_ = [
-      ('version', NvU32),
-      ('size', NvU32),
-      ('frtsRegionOffset4K', NvU32),
-      ('frtsRegionSize', NvU32),
-      ('frtsRegionMediaType', NvU32),
-    ]
-    class FWSECLIC_FRTS_CMD(Struct): pass
-    FWSECLIC_FRTS_CMD._packed_ = True
-    FWSECLIC_FRTS_CMD._fields_ = [
-      ('readVbiosDesc', FWSECLIC_READ_VBIOS_DESC),
-      ('frtsRegionDesc', FWSECLIC_FRTS_REGION_DESC),
-    ]
-    read_vbios_desc = FWSECLIC_READ_VBIOS_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_READ_VBIOS_DESC), flags=2)
-    frst_reg_desc = FWSECLIC_FRTS_REGION_DESC(version=0x1, size=ctypes.sizeof(FWSECLIC_FRTS_REGION_DESC),
-      frtsRegionOffset4K=0xdead, frtsRegionSize=0x100, frtsRegionMediaType=2)
-    frts_cmd = FWSECLIC_FRTS_CMD(readVbiosDesc=read_vbios_desc, frtsRegionDesc=frst_reg_desc)
-    assert int.from_bytes(frts_cmd, 'little') == 0x2000001000000dead0000001400000001000000020000000000000000000000000000001800000001
-    assert int.from_bytes(frts_cmd.readVbiosDesc, 'little') == int.from_bytes(read_vbios_desc, 'little')
-    assert int.from_bytes(frts_cmd.frtsRegionDesc, 'little') == int.from_bytes(frst_reg_desc, 'little')
-    assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
-    assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
-
-if __name__ == "__main__": unittest.main()
@@ -62,7 +62,6 @@ class TestConv(unittest.TestCase):
    np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
    np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)

-  @unittest.skip("this test is flaky")
  def test_two_overlapping_binops_no_rerun_wino(self):
    with Context(WINO=1):
      x = Tensor.randn(1,4,16,16)
@@ -30,7 +30,7 @@ class TestDevice(unittest.TestCase):

  @unittest.skipIf(WIN and CI, "skipping windows test") # TODO: subproccess causes memory violation?
  def test_env_overwrite_default_compiler(self):
-    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept Exception: pass"
+    expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept RuntimeError: pass"

    if Device.DEFAULT == "CPU":
      from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
@@ -81,20 +81,20 @@ class TestCompiler(unittest.TestCase):
  def test_compile_cached(self):
    diskcache_put("key", "123", None) # clear cache
    getenv.cache_clear()
-    with Context(CCACHE=1):
+    with Context(DISABLE_COMPILER_CACHE=0):
      self.assertEqual(MockCompiler("key").compile_cached("123"), str.encode("123"))
      self.assertEqual(diskcache_get("key", "123"), str.encode("123"))

  def test_compile_cached_disabled(self):
    diskcache_put("disabled_key", "123", None) # clear cache
    getenv.cache_clear()
-    with Context(CCACHE=0):
+    with Context(DISABLE_COMPILER_CACHE=1):
      self.assertEqual(MockCompiler("disabled_key").compile_cached("123"), str.encode("123"))
      self.assertIsNone(diskcache_get("disabled_key", "123"))

  def test_device_compile(self):
    getenv.cache_clear()
-    with Context(CCACHE=0):
+    with Context(DISABLE_COMPILER_CACHE=1):
      a = Tensor([0.,1.], device=Device.DEFAULT).realize()
      (a + 1).realize()

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	9d4e74a52a	comments	2025-10-30 12:27:55 +08:00
geohot	e887ad6317	more syntax	2025-10-30 12:11:18 +08:00
geohot	98b76aa3aa	more syntax	2025-10-30 12:04:26 +08:00
George HotzandGitHub	2f80994ce8	Merge branch 'master' into uop_prg	2025-10-30 11:22:36 +08:00
geohot	2ef52db298	work	2025-10-30 11:12:37 +08:00
geohot	80f9347e53	work	2025-10-30 11:01:53 +08:00
geohot	79f98a6624	uops programs	2025-10-30 10:52:44 +08:00