split pm_substitute_recurse

2025-10-06 09:22:01 +08:00
593 changed files with 204191 additions and 188639 deletions
@@ -1,3 +0,0 @@
 [run]
 source = tinygrad
 branch = True
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
  mesa:
    description: "Install mesa"
    required: false
    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -61,7 +57,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 +66,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 +183,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 +243,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
@@ -293,13 +289,3 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
    # **** mesa ****
    - name: Install mesa (linux)
      if: inputs.mesa == 'true' && runner.os == 'Linux'
      shell: bash
      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
    - name: Install mesa (macOS)
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -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,58 @@ 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
    - name: Verify OpenCL autogen
      run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        python3 -c "from tinygrad.runtime.autogen import opencl"
+        ./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
+        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
+        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
+        ./autogen_stubs.sh cuda
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
+        ./autogen_stubs.sh nv
        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
+        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.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
    - name: Verify AMD autogen
      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
+        ./autogen_stubs.sh hsa
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
+        ./autogen_stubs.sh kfd
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
+        ./autogen_stubs.sh comgr
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
+        ./autogen_stubs.sh amd
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
+        ./autogen_stubs.sh sqtt
        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
        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
+        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
+        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
+        ./autogen_stubs.sh libc
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
+        ./autogen_stubs.sh io_uring
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
+        ./autogen_stubs.sh ib
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
        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
+        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        python3 -c "from tinygrad.runtime.autogen import webgpu"
+        ./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
+        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
+        ./autogen_stubs.sh llvm
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
+        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
        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
    - name: Verify mesa autogen
      run: |
        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
        python3 -c "from tinygrad.runtime.autogen import 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
@@ -51,18 +51,15 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
    - 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=500 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
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=700 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      # TODO: very slow step time
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=1600 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
@@ -102,7 +99,7 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=8 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
@@ -111,19 +108,14 @@ jobs:
      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-
+    - name: Run 10 CIFAR training steps
-    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=330 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
-    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
+    - name: Run 10 CIFAR training steps w HALF
-    #- name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=385 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
    #- name: Run 10 CIFAR training steps w HALF
    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #- name: Run 10 CIFAR training steps w BF16
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
+    - name: Run 10 CIFAR training steps w winograd
-    # - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
@@ -131,7 +123,7 @@ jobs:
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -199,7 +191,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
@@ -211,7 +203,6 @@ jobs:
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
    - name: Run Tensor Core GEMM (PTX)
      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
@@ -222,9 +213,8 @@ jobs:
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
+    - name: Run SDXL
-    # - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -239,8 +229,6 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
    - name: Run quantized LLaMA3
      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -250,9 +238,9 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
@@ -274,7 +262,6 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -312,28 +299,24 @@ 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
-    # TODO: too slow
+    - name: Fuzz Padded Tensor Core GEMM (NV)
-    # - name: Fuzz Padded Tensor Core GEMM (NV)
+      run: NV=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
-    #   run: NV=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: Fuzz Padded Tensor Core GEMM (PTX)
-    # TODO: too slow
+      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: Fuzz Padded Tensor Core GEMM (PTX)
    #   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
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
+    - 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
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=68 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
+    - 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_bf16 ASSERT_MIN_STEP_TIME=75 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
-    # - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=35 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    #   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
-    # - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    #   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
-    # - 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.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    #   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 +393,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
@@ -432,10 +415,9 @@ jobs:
    - name: Test AM warm start time
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=450 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
+    - name: Run SDXL
-    # - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=1400 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -525,22 +507,20 @@ 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
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
+    - 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_half ASSERT_MIN_STEP_TIME=188 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
-    # - 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
-    #   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
-    # - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    #   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
+    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD 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_remote.txt
+    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -625,34 +605,42 @@ 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
+    - name: benchmark openpilot 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
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision ASSERT_MIN_STEP_TIME=30 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    # - name: openpilot compile3 0.9.9 driving_policy
+    - name: benchmark openpilot 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
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy ASSERT_MIN_STEP_TIME=45 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    # - name: openpilot compile3 0.9.9 dmonitoring
+    - name: benchmark openpilot 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
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring ASSERT_MIN_STEP_TIME=70 PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 0.10.0 driving_policy
+    - name: openpilot compile3 0.9.9 driving_vision
-      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: PYTHONPATH="." 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.10.0 dmonitoring
+    - name: openpilot compile3 0.9.9 driving_policy
-      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
+      run: PYTHONPATH="." 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: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+    - name: openpilot compile3 0.9.9 dmonitoring
-      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: PYTHONPATH="." 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.1 driving_vision
+    - name: openpilot compile3 Space Lab policy + 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
+      run: |
-    - name: openpilot compile3 0.10.1 driving_policy
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-      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
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
-    - name: openpilot compile3 0.10.1 dmonitoring
+    - name: benchmark MobileNetV2 on DSP
-      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
+      run: |
-    # - name: benchmark MobileNetV2 on DSP
+        # generate quantized weights
-    #   run: |
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-    #     # generate quantized weights
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-    #     ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        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
-    #     ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-    #     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
+        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /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
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (comma)
        path: |
          openpilot_compile_0_9_4.txt
          openpilot_compile_0_9_7.txt
          openpilot_0_9_4.txt
          openpilot_0_9_7.txt
          openpilot_image_0_9_4.txt
          openpilot_image_0_9_7.txt
  testreddriverbenchmark:
    name: AM Benchmark
@@ -706,9 +694,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
-    # - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
    #   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 +757,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
-    # - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
    #   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
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 720
+    timeout-minutes: 360
    steps:
    - name: Checkout Code
@@ -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 }}
@@ -86,65 +89,64 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
-  # TODO: fix the torch backend and reenable
+  torchbackend:
-  # torchbackend:
+    name: Torch Backend Tests
-  #   name: Torch Backend Tests
+    runs-on: ubuntu-latest
-  #   runs-on: ubuntu-latest
+    timeout-minutes: 15
-  #   timeout-minutes: 15
+    steps:
-  #   steps:
+    - name: Checkout Code
-  #   - name: Checkout Code
+      uses: actions/checkout@v4
-  #     uses: actions/checkout@v4
+    - name: Setup Environment
-  #   - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
-  #     uses: ./.github/actions/setup-tinygrad
+      with:
-  #     with:
+        key: torch-backend-pillow-torchvision-et-pt
-  #       key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
-  #       deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
-  #       pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
-  #       llvm: 'true'
+    - name: Install ninja
-  #   - name: Install ninja
+      run: |
-  #     run: |
+        sudo apt update || true
-  #       sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
-  #       sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
-  #   - name: Lint with ruff
+      run: |
-  #     run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
-  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
-  #       python3 -m ruff check extra/torch_backend/backend.py
+    - name: Test one op
-  #   - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
-  #   - name: Test ResNet-18
+      run: DEBUG=2 python3 extra/torch_backend/example.py
-  #     run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: My (custom) tests
-  #   - name: My (custom) tests
+      run: python3 extra/torch_backend/test.py
-  #     run: python3 extra/torch_backend/test.py
+    - name: Test one op in torch tests
-  #   - name: Test one op in torch tests
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
-  #   - name: Test Ops with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
-  #   - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
-  #   - name: Test multi-gpu
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
-  # torchbackendmore:
+  torchbackendmore:
-  #   name: Torch Backend Tests More
+    name: Torch Backend Tests More
-  #   runs-on: ubuntu-latest
+    runs-on: ubuntu-latest
-  #   timeout-minutes: 15
+    timeout-minutes: 15
-  #   steps:
+    steps:
-  #   - name: Checkout Code
+    - name: Checkout Code
-  #     uses: actions/checkout@v4
+      uses: actions/checkout@v4
-  #   - name: Setup Environment
+    - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
+      uses: ./.github/actions/setup-tinygrad
-  #     with:
+      with:
-  #       key: torch-backend-pillow-torchvision-et-pt
+        key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
+        deps: testing_minimal
-  #       llvm: 'true'
+        llvm: 'true'
-  #   - name: Install ninja
+    - name: Install ninja
-  #     run: |
+      run: |
-  #       sudo apt update || true
+        sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
+        sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
-  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-  #   - name: Test some torch tests (expect failure)
+    - name: Test some torch tests (expect failure)
-  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
  bepython:
    name: Python Backend
@@ -158,8 +160,10 @@ jobs:
      with:
        key: be-minimal
        deps: testing_minimal
-    - name: Test dtype with Python emulator
+    - name: Test dtype with Python emulator (with RANGEIFY)
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: |
        RANGEIFY=0 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
        RANGEIFY=1 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
    - name: Test ops with Python emulator
      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
    - name: Test uops with Python emulator
@@ -201,7 +205,7 @@ jobs:
        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
@@ -227,7 +231,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 .
@@ -236,13 +240,14 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
    - name: Run mypy
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-    # broken because of UPatAny
+    - name: Run TYPED=1
-    #- name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"
    #  run: TYPED=1 python -c "import tinygrad"
  unittest:
    name: Unit Tests
@@ -264,13 +269,10 @@ jobs:
      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
-    # TODO: too slow
+    - name: Run SDXL on NULL backend
-    # - name: Run SDXL on NULL backend
+      run: MAX_BUFFER_SIZE=0 NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+      run: MAX_BUFFER_SIZE=0 NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    - name: Run AMD emulated BERT training on NULL backend
      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -287,28 +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
+    - name: Repo line count < 18000 lines
-      run: MAX_LINE_COUNT=19000 python sz.py
+      run: MAX_LINE_COUNT=18000 python sz.py
  spec:
    strategy:
      fail-fast: false
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      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 }}
  fuzzing:
    name: Fuzzing
@@ -326,6 +308,10 @@ jobs:
      run: python test/external/fuzz_symbolic.py
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
      run: |
        python test/external/fuzz_shapetracker.py
        python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py
@@ -342,11 +328,14 @@ 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: Test CL IMAGE=2 ops + training (rangeify)
        run: |
          RANGEIFY=1 CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
          RANGEIFY=1 CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -368,7 +357,7 @@ jobs:
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -391,13 +380,20 @@ 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=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 RANGEIFY=0 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot CL compile fp16
+      - name: Test openpilot model with rangeify
-        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
+        run: |
-      - name: Test openpilot CL compile fp32 (test correctness)
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 RANGEIFY=1 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+      - name: Test openpilot alt model correctness (float32)
-      - name: Test openpilot LLVM compile fp16
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
-        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      - name: Test openpilot fastvits model correctness (float32)
        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      # - name: Test openpilot simple_plan vision model correctness (float32)
      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot compile4
        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -457,11 +453,11 @@ jobs:
      - name: Test MLPerf stuff
        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
      - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+        run: MAX_BUFFER_SIZE=0 NULL=1 python examples/beautiful_mnist_multigpu.py
      - name: Test Bert training
-        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: MAX_BUFFER_SIZE=0 NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: MAX_BUFFER_SIZE=0 NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -524,6 +520,88 @@ jobs:
 # ****** Feature Tests ******
  testrangeifycpu:
    name: Linux (rangeify) CPU
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
        key: rangeify-minimal-llvm
        deps: testing_minimal
        opencl: 'true'
        llvm: "true"
    - name: Test CPU=1 RANGEIFY=1
      # TODO: add more passing tests here
      run: |
        CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \
          test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \
          test/test_setitem.py test/test_assign.py test/test_multitensor.py test/test_const_folding.py
    - name: Test CPU=1 DEVECTORIZE=0 (RANGEIFY=1)
      run: CPU=1 CPU_LLVM=0 RANGEIFY=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
      run: |
        CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
    - name: Test Docs RANGEIFY=1
      run: |
        RANGEIFY=1 python docs/abstractions2.py
    # RANGEIFY=2 isn't supported
    #- name: Test CPU=1 RANGEIFY=2
    #  run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
    # slow (and still wrong on beautiful_mnist)
    #- name: Test LLVM RANGEIFY=1 (slow tests)
    #  run: CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
  testrangeifycl:
    name: Linux (rangeify) CL
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
        key: rangeify-cl
        deps: testing
        opencl: 'true'
        llvm: "true"
    - name: Test CL=1 RANGEIFY=1
      run: CL=1 RANGEIFY=1 pytest -n auto test/test_ops.py test/test_schedule.py test/test_symbolic_ops.py test/test_jit.py test/unit/test_disk_tensor.py test/models/test_mnist.py test/unit/test_mnist_dataset.py test/test_optim.py --durations 20
    - name: Test Fuse
      run: CL=1 RANGEIFY=2 python3 -m pytest --durations 20 test/test_softmax_fusion.py -k "not test_auto_softmax"
    - name: Test ONNX
      run: CL=1 RANGEIFY=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
  testrangeifymacos:
    name: MacOS (rangeify)
    runs-on: macos-14
    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
        key: metal
        deps: testing
    - name: some unit tests
      run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20
    - name: Test METAL=1 RANGEIFY=1
      run: |
        METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py test/test_multitensor.py --durations=20
        METAL=1 MAX_KERNEL_BUFFERS=6 RANGEIFY=1 PYTHONPATH=. python test/test_multitensor.py TestBatchNorm.test_batchnorm
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
  testdevectorize:
    name: Linux (devectorize)
    runs-on: ubuntu-24.04
@@ -539,11 +617,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
  testdsp:
    name: Linux (DSP)
@@ -644,9 +722,10 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run pytest (amd) with RANGEIFY
        run: RANGEIFY=1 python -m pytest test/test_linearizer.py::TestLinearizer::test_where_fold
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -688,7 +767,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl, lvp]
+        backend: [llvm, cpu, opencl]
    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -702,10 +781,9 @@ jobs:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -907,7 +985,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [metal, llvm, cpu, lvp]
+        backend: [metal, llvm, cpu]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
@@ -920,13 +998,12 @@ jobs:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
          pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -968,3 +1045,9 @@ jobs:
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
      - name: Run pytest (${{ matrix.backend }}) with RANGEIFY
        if: matrix.backend=='webgpu'
        env:
          RANGEIFY: 1
        shell: bash
        run: python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
@@ -38,7 +38,6 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
 venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -63,5 +62,3 @@ profile_stats
 *.log
 target
 .mypy_cache
 mutants
 .mutmut-cache
@@ -20,15 +20,21 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
      - id: tests
        name: subset of tests
        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
      - id: example
-        name: test all devices
+        name: multi device tests
        entry: python3 test/external/external_test_example.py
        language: system
        always_run: true
        pass_filenames: false
-      - id: tests
+      - id: pylint
-        name: subset of tests
+        name: pylint
-        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: python3 -m pylint tinygrad/
        language: system
        always_run: true
-        pass_filenames: false
+        pass_filenames: false
@@ -30,6 +30,10 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=
 # When enabled, pylint would attempt to guess common misconfiguration and emit
 # user-friendly hints instead of false-positive error messages
 suggestion-mode=yes
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
@@ -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
+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.
 - **IR and compiler** that fuse and lower kernels
 - **JIT + graph execution**
 - **nn / optim / datasets** for real training
-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**
+tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!
 - ✅ 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.
 ---
 ### Laziness
@@ -0,0 +1,470 @@
 #!/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"
 }
 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"
 }
 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" == "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
 else echo "usage: $0 <type>"
 fi
@@ -42,6 +42,7 @@ import struct
 from tinygrad.dtype import dtypes
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
 # allocate some buffers + load in values
 out = Buffer(DEVICE, 1, dtypes.int32).allocate()
@@ -50,12 +51,13 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 # NOTE: a._buf is the same as the return from cpu.allocator.alloc
 # describe the computation
 idx = UOp.const(dtypes.index, 0)
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-alu = buf_1.index(idx) + buf_2.index(idx)
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
 ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
 alu = ld_1 + ld_2
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
-st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))
 # convert the computation to a "linearized" format (print the format)
@@ -78,6 +80,8 @@ print("******** third, the UOp ***********")
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.helpers import RANGEIFY
 from tinygrad.schedule.kernelize import get_kernelize_map
 from tinygrad.schedule.rangeify import get_rangeify_map
 # allocate some values + load in values
@@ -91,7 +95,7 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))
 # group the computation into kernels
-becomes_map = get_rangeify_map(s)
+becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s)
 # the compute maps to an assign
 assign = becomes_map[a+b].base
@@ -0,0 +1,109 @@
 # Kernel Creation
 Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
 - Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
 - Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
 - Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
 `Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
 ## Kernelize flow
 Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
 ```py
 # initialize 3 input buffers on the device
 a = Tensor([1]).realize()
 b = Tensor([2]).realize()
 c = Tensor([3]).realize()
 # create the Tensor graph
 mul = a*b
 out = mul+c
 print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
 print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
 out.kernelize()
 print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
 print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
 ```
 The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
 ```py
 print(out.uop)
 ```
 The first source is the output BUFFER:
 ```
 UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
 ```
 And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
 ```
 UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
     x1,
    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
     x1,
    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
     x1,
    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
 ```
 KERNEL describes the compute AST, metadata and memory dependencies.
 BUFFER holds a reference to the device memory where the output will be stored.
 Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
 ```py
 child = out+2
 child.kernelize()
 print(child.uop.src[1].arg.ast)
 ```
 ```
 UOp(Ops.SINK, dtypes.void, arg=None, src=(
  UOp(Ops.STORE, dtypes.void, arg=None, src=(
    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
    UOp(Ops.ADD, dtypes.int, arg=None, src=(
      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
         x2,)),
      UOp(Ops.CONST, dtypes.int, arg=2, src=(
         x2,)),)),)),))
 ```
 `Tensor.realize` will execute the kernels and write outputs to memory:
 ```py
 Tensor.realize(out)
 print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
 print(out.item()) # 5
 ```
 <hr />
 **Summary**
 - The large Tensor graph is built from a mix of data, compute and movement Ops.
 - `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
 - `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
 - Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.
 Group UOps into kernels.
-::: tinygrad.schedule.rangeify.get_rangeify_map
+::: tinygrad.schedule.kernelize.get_kernelize_map
    options:
        members: false
        show_labels: false
@@ -41,7 +41,7 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -10,7 +10,7 @@ GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
 # override tinygrad defaults
 dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(FUSE_ARANGE=1, FUSE_OPTIM=1).__enter__()
 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -1,6 +1,6 @@
 import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate
 def load_onnx_model(onnx_file):
@@ -8,7 +8,7 @@ import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)
  if args.benchmark != -1:
-    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -145,6 +145,7 @@ hyp = {
  },
 }
@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():
  def set_seed(seed):
@@ -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 = [
@@ -145,41 +145,6 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear
 def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
  scale = fp8_max / x.abs().max()
  x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
  return x_scl_sat.cast(dtype), scale.float().reciprocal()
 class FP8Linear:
  def __init__(self, in_features, out_features, bias=True):
    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
    self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
    self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
  def __call__(self, x:Tensor):
    y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
    if self.bias is not None: y = y + self.bias.cast(y.dtype)
    return y.cast(x.dtype)
  @staticmethod
  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
    assert not quantize_embeds
    new_tensors = {}
    for name,v in tensors.items():
      if "feed_forward" in name or "attention.w" in name:
        assert "weight" in name, name
        fp8_weight, scale = quantize_to_fp8(v)
        new_tensors[name] = fp8_weight
        new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
        if isinstance(device, tuple):
          new_tensors[name].shard_(device, axis=-1)
          new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
      else:
        new_tensors[name] = v
    return new_tensors
 MODEL_PARAMS = {
  "1B": {
    "args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
@@ -202,7 +167,6 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  # build model
  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
  elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)
@@ -278,7 +242,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
@@ -707,7 +707,7 @@ def train_unet3d():
  ```BASEDIR=<folder_path> ./examples/mlperf/scripts/setup_kits19_dataset.sh```
  2) To start training the model, run the following:
-  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
+  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
  """
  from examples.mlperf.losses import dice_ce_loss
  from examples.mlperf.metrics import dice_score
@@ -749,6 +749,7 @@ def train_unet3d():
    "train_beam": TRAIN_BEAM,
    "eval_beam": EVAL_BEAM,
    "wino": WINO.value,
    "fuse_conv_bw": FUSE_CONV_BW.value,
    "gpus": GPUS,
    "default_float": dtypes.default_float.name
  }
@@ -1188,9 +1189,7 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1224,7 +1223,7 @@ def train_bert():
          return
      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1310,7 +1309,7 @@ def train_llama3():
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)
-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
  # trains to 7
  opt_adamw_beta_1 = 0.9
@@ -1,17 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
 export IGNORE_OOB=1
 export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 # export BEAM_LOG_SURPASS_MAX=1
 # export BASEDIR="/raid/datasets/wiki"
 export RESET_STEP=1
 export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
 python3 examples/mlperf/model_train.py
@@ -1,69 +0,0 @@
 # 1. Problem
 This problem uses BERT for NLP.
 ## Requirements
 Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
 pip install gdown numpy tqdm tensorflow
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 # 2. Directions
 ## Steps to download and verify data
 ### 1. Download raw data
 ```
 BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
 ```
 ### 2. Preprocess train and validation data
 Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
 #### Training:
 ```
 BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
 ```
 Generating a specific topic (Between 0 and 499)
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
 ```
 #### Validation:
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
 ```
 ## Running
 ### tinybox_green
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```
 ### tinybox_red
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
 ### tinybox_8xMI300X
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,17 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
 export BENCHMARK=10 BERT_LAYERS=2
 python3 examples/mlperf/model_train.py
@@ -1,20 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 # similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
 export WANDB=1 PARALLEL=0
 RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_8xMI300X"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 # similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
 # pip install -e ".[mlperf]"
 export LOGMLPERF=1
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
 # 1. Problem
 This problem uses BERT for NLP.
 ## Requirements
 Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
 pip install gdown numpy tqdm tensorflow
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 # 2. Directions
 ## Steps to download and verify data
 ### 1. Download raw data
 ```
 BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
 ```
 ### 2. Preprocess train and validation data
 Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
 #### Training:
 ```
 BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
 ```
 Generating a specific topic (Between 0 and 499)
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
 ```
 #### Validation:
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
 ```
 ## Running
 ### tinybox_green
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```
 ### tinybox_red
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
 ### tinybox_8xMI300X
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,17 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"
 export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
 python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 export WANDB=1 PARALLEL=0
 RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,28 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 # pip install -e ".[mlperf]"
 export LOGMLPERF=1
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_green_${DATETIME}_${SEED}.log"
 # init
 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
 # 1. Problem
 This problem uses BERT for NLP.
 ## Requirements
 Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
 pip install gdown numpy tqdm tensorflow
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 # 2. Directions
 ## Steps to download and verify data
 ### 1. Download raw data
 ```
 BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
 ```
 ### 2. Preprocess train and validation data
 Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
 #### Training:
 ```
 BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
 ```
 Generating a specific topic (Between 0 and 499)
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
 ```
 #### Validation:
 ```
 BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
 ```
 ## Running
 ### tinybox_green
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```
 ### tinybox_red
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
 ### tinybox_8xMI300X
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,18 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"
 export RESET_STEP=1
 export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
 python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 export WANDB=1 PARALLEL=0
 RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 # pip install -e ".[mlperf]"
 export LOGMLPERF=${LOGMLPERF:-1}
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_red_${DATETIME}_${SEED}.log"
 export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
 # init
 sleep 5 && sudo rmmod amdgpu || true
 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
 # 1. Problem
 This problem uses the ResNet-50 CNN to do image classification.
 ## Requirements
 Install tinygrad and mlperf-logging from master.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 ### tinybox_red
 Disable cwsr
 This is the default on production tinybox red.
 ```
 sudo vi /etc/modprobe.d/amdgpu.conf
 cat <<EOF > /etc/modprobe.d/amdgpu.conf
 options amdgpu cwsr_enable=0
 EOF
 sudo update-initramfs -u
 sudo reboot
 # validate
 sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
 ```
 # 2. Directions
 ## Steps to download and verify data
 ```
 IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
 ```
 ## Steps for one time setup
 ### tinybox_red
 ```
 examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
 ```
 ## Steps to run benchmark
 ```
 examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
 ```
@@ -1,13 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 export BENCHMARK=10 DEBUG=2
 python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 export EVAL_START_EPOCH=3 EVAL_FREQ=4
 export WANDB=1 PARALLEL=0
 python3 examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 # pip install -e ".[mlperf]"
 export LOGMLPERF=${LOGMLPERF:-1}
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
 # init
 BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
 # 1. Problem
 This problem uses the ResNet-50 CNN to do image classification.
 ## Requirements
 Install tinygrad and mlperf-logging from master.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 ### tinybox_red
 Disable cwsr
 This is the default on production tinybox red.
 ```
 sudo vi /etc/modprobe.d/amdgpu.conf
 cat <<EOF > /etc/modprobe.d/amdgpu.conf
 options amdgpu cwsr_enable=0
 EOF
 sudo update-initramfs -u
 sudo reboot
 # validate
 sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
 ```
 # 2. Directions
 ## Steps to download and verify data
 ```
 IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
 ```
 ## Steps for one time setup
 ### tinybox_red
 ```
 examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
 ```
 ## Steps to run benchmark
 ```
 examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
 ```
@@ -1,13 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export BENCHMARK=10 DEBUG=${DEBUG:-2}
 python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export EVAL_START_EPOCH=3 EVAL_FREQ=4
 export WANDB=1 PARALLEL=0
 python3 examples/mlperf/model_train.py
@@ -1,26 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 export RESET_STEP=0
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 # pip install -e ".[mlperf]"
 export LOGMLPERF=${LOGMLPERF:-1}
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
 # init
 sleep 5 && sudo rmmod amdgpu || true
 BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,8 +0,0 @@
 #!/bin/bash
 rocm-smi --setprofile compute
 rocm-smi --setmclk 3
 rocm-smi --setperflevel high
 # power cap to 350W
 echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -1,38 +0,0 @@
 # 1. Problem
 This problem uses RetinaNet for SSD.
 ## Requirements
 Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
 Also install the following dependencies:
 ```
 pip install tqdm numpy pycocotools boto3 pandas torch torchvision
 ```
 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.
 # 2. Directions
 ## Steps to download data
 Run the following:
 ```
 BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
 ```
 ## Running
 ### tinybox_green
 #### Steps to run benchmark
 ```
 examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
 ```
@@ -1,14 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
 # export RESET_STEP=0
 export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export BENCHMARK=5 DEBUG=2
 python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
 # export RESET_STEP=0
 export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export WANDB=1 PARALLEL=0
 export RUNMLPERF=1
 python examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
 #!/bin/bash
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/openimages"
 # pip install -e ".[mlperf]"
 export LOGMLPERF=1
 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
 # init
 BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,14 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
 # export RESET_STEP=0
 export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export BENCHMARK=5 DEBUG=2
 python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
 #!/bin/bash
 export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
 # export RESET_STEP=0
 export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 export WANDB=1 PARALLEL=0
 export RUNMLPERF=1
 python examples/mlperf/model_train.py
@@ -1,38 +0,0 @@
 {
    "submitter": "tinycorp",
    "division": "closed",
    "status": "Available on-premise",
    "system_name": "tinybox 8xMI300X",
    "number_of_nodes": "1",
    "host_processors_per_node": "2",
    "host_processor_model_name": "AMD EPYC 9354",
    "host_processor_core_count": "32",
    "host_processor_vcpu_count": "64",
    "host_processor_frequency": "",
    "host_processor_caches": "",
    "host_processor_interconnect": "",
    "host_memory_capacity": "2304GB",
    "host_storage_type": "NVMe SSD",
    "host_storage_capacity": "3x 4TB raid array",
    "host_networking": "",
    "host_networking_topology": "",
    "host_memory_configuration": "24x 96GB DDR5",
    "accelerators_per_node": "8",
    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
    "accelerator_host_interconnect": "PCIe 5.0 x16",
    "accelerator_frequency": "",
    "accelerator_on-chip_memories": "",
    "accelerator_memory_configuration": "HBM3",
    "accelerator_memory_capacity": "192GB",
    "accelerator_interconnect": "",
    "accelerator_interconnect_topology": "",
    "cooling": "air",
    "hw_notes": "",
    "framework": "tinygrad, branch mlperf_training_v5.0",
    "other_software_stack": {
        "python": "3.10.16",
        "ROCm": "3.0.0+94441cb"
    },
    "operating_system": "Ubuntu 24.04.1 LTS",
    "sw_notes": ""
  }
@@ -1,38 +0,0 @@
 {
  "submitter": "tinycorp",
  "division": "closed",
  "status": "Available on-premise",
  "system_name": "tinybox green",
  "number_of_nodes": "1",
  "host_processors_per_node": "1",
  "host_processor_model_name": "AMD EPYC 7532",
  "host_processor_core_count": "32",
  "host_processor_vcpu_count": "64",
  "host_processor_frequency": "",
  "host_processor_caches": "",
  "host_processor_interconnect": "",
  "host_memory_capacity": "128GB",
  "host_storage_type": "NVMe SSD",
  "host_storage_capacity": "4 TB raid array + 1 TB boot",
  "host_networking": "",
  "host_networking_topology": "",
  "host_memory_configuration": "8x 16GB DDR4",
  "accelerators_per_node": "6",
  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
  "accelerator_host_interconnect": "PCIe 4.0 x16",
  "accelerator_frequency": "",
  "accelerator_on-chip_memories": "",
  "accelerator_memory_configuration": "GDDR6X",
  "accelerator_memory_capacity": "24GB",
  "accelerator_interconnect": "",
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
    "python": "3.10.12",
    "CUDA": "12.4"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
 }
@@ -1,37 +0,0 @@
 {
  "submitter": "tinycorp",
  "division": "closed",
  "status": "Available on-premise",
  "system_name": "tinybox red",
  "number_of_nodes": "1",
  "host_processors_per_node": "1",
  "host_processor_model_name": "AMD EPYC 7532",
  "host_processor_core_count": "32",
  "host_processor_vcpu_count": "64",
  "host_processor_frequency": "",
  "host_processor_caches": "",
  "host_processor_interconnect": "",
  "host_memory_capacity": "128GB",
  "host_storage_type": "NVMe SSD",
  "host_storage_capacity": "4 TB raid array + 1 TB boot",
  "host_networking": "",
  "host_networking_topology": "",
  "host_memory_configuration": "8x 16GB DDR4",
  "accelerators_per_node": "6",
  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
  "accelerator_host_interconnect": "PCIe 4.0 x16",
  "accelerator_frequency": "",
  "accelerator_on-chip_memories": "",
  "accelerator_memory_configuration": "GDDR6",
  "accelerator_memory_capacity": "24GB",
  "accelerator_interconnect": "",
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
    "python": "3.10.12"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
 }
@@ -1,11 +1,16 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time
 import numpy as np
 if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
 if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.nn.onnx import OnnxRunner
+
 import onnx
 from tinygrad.frontend.onnx import OnnxRunner
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -16,14 +21,11 @@ def compile(onnx_file):
  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
  # Float inputs and outputs to tinyjits for openpilot are always float32
  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
-  if not getenv("NPY_IMG"):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")
  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -31,6 +33,8 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
@@ -39,7 +43,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
@@ -48,8 +52,6 @@ def compile(onnx_file):
      kernel_count += 1
      read_image_count += ei.prg.p.src.count("read_image")
      gated_read_image_count += ei.prg.p.src.count("?read_image")
      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -65,77 +67,85 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return inputs, test_val
+  return test_val
-def test_vs_compile(run, inputs, test_val=None):
+def test_vs_compile(run, new_inputs, test_val=None):
  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
  # run 20 times
  step_times = []
  for _ in range(20):
    st = time.perf_counter()
    out = run(**inputs)
    mt = time.perf_counter()
    val = out.numpy()
    et = time.perf_counter()
-    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+  print(out, val.shape, val.dtype)
  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")
  # test that changing the numpy changes the model outputs
-  inputs_2x = {k: Tensor(v.numpy()*2, device=v.device) for k,v in inputs.items()}
+  if any([x.device == 'NPY' for x in inputs.values()]):
-  out = run(**inputs_2x)
+    for v in new_inputs_numpy.values(): v *= 2
-  changed_val = out.numpy()
+    out = run(**inputs)
-  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+    changed_val = out.numpy()
    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val
-def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
-  import onnx
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  import onnxruntime as ort
  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)
  ORT_TO_NP_DTYPES: dict[str, np.dtype] = {
    'tensor(float)': np.dtype('float32'),
    'tensor(float16)': np.dtype('float16'),
    'tensor(uint8)': np.dtype('uint8'),
  }
  timings = []
-  onnx_session = ort.InferenceSession(onnx_file)
+  if ort:
-  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
+    # test with onnxruntime
-  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+    import onnxruntime as ort
    onnx_session = ort.InferenceSession(onnx_file)
    for _ in range(1 if test_val is not None else 5):
      st = time.perf_counter()
      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
      timings.append(time.perf_counter() - st)
    new_torch_out = onnx_output[0]
  else:
    # test with torch
    import torch
    from onnx2torch import convert
    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
    torch_model = convert(onnx_model).float()
    with torch.no_grad():
      for _ in range(1 if test_val is not None else 5):
        st = time.perf_counter()
        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
        timings.append(time.perf_counter() - st)
      new_torch_out = torch_out.numpy()
-  for _ in range(1 if test_val is not None else 5):
+  if test_val is not None:
-    st = time.perf_counter()
+    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
-    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
+    print("test vs onnx passed")
    timings.append(time.perf_counter() - st)
  np.testing.assert_allclose(onnx_output[0].reshape(test_val.shape), test_val, atol=tol, rtol=tol)
  print("test vs onnx passed")
  return timings
 def bench(run, inputs):
  from extra.bench_log import WallTimeEvent, BenchEvent
  for _ in range(10):
    with WallTimeEvent(BenchEvent.STEP):
      run(**inputs).numpy()
 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
-  inputs, outputs = compile(onnx_file)
+  test_val = compile(onnx_file) if not getenv("RUN") else None
  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
-  test_vs_compile(pickle_loaded, inputs, outputs)
+  # same randomness as compile
-  if getenv("SELFTEST"):
+  Tensor.manual_seed(100)
-    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
  if getenv("BENCHMARK"):
    for be in ["torch", "ort"]:
      try:
        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
        print(f"timing {be}: {min(timings)*1000:.2f} ms")
      except Exception as e:
        print(f"{be} fail with {e}")
  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))
  if getenv("BENCHMARK_LOG", ""):
    bench(pickle_loaded, inputs)
@@ -0,0 +1,47 @@
 import sys
 from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.schedule.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule
 # NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
  run_onnx = OnnxRunner(onnx_file)
  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
  root = out.uop
  targets = [x.uop for x in inputs.values()]
  print(targets)
  # TODO: abstract this from gradient?
  # compute the target path (top down)
  in_target_path: dict[UOp, bool] = {}
  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
  independent_set = {}
  for u in root.toposort():
    if in_target_path[u]:
      for s in u.src:
        if not in_target_path[s]:
          independent_set[s] = None
  independent = UOp.sink(*independent_set.keys())
  kernelized = get_kernelize_map(independent)
  independent = independent.substitute(kernelized)
  schedule, var_vals = create_schedule_with_vars(independent)
  run_schedule(schedule)
  print("**** real ****")
  GlobalCounters.reset()
  out.uop = root.substitute(kernelized)
  out.kernelize()
  # realize
  out.realize()
@@ -27,7 +27,7 @@ class Model(nn.Module):
 if __name__ == "__main__":
  if getenv("TINY_BACKEND"):
-    import tinygrad.nn.torch  # noqa: F401
+    import tinygrad.frontend.torch  # noqa: F401
    device = torch.device("tiny")
  else:
    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
@@ -99,7 +99,6 @@ if __name__ == "__main__":
  parser.add_argument('--timing',      action='store_true', help="Print timing per step")
  parser.add_argument('--noshow',      action='store_true', help="Don't show the image")
  parser.add_argument('--fp16',        action='store_true', help="Cast the weights to float16")
  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()
  N = 1
@@ -113,22 +112,19 @@ if __name__ == "__main__":
  model = StableDiffusionV2(**params)
-  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+  default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
-    if not args.fakeweights:
+  weights_fn = args.weights_fn
-      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+  if not weights_fn:
-      weights_fn = args.weights_fn
+    weights_url = args.weights_url if args.weights_url else default_weights_url
-      if not weights_fn:
+    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
        weights_url = args.weights_url if args.weights_url else default_weights_url
        weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
-        load_state_dict(model, safe_load(weights_fn), strict=False)
+  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
    load_state_dict(model, safe_load(weights_fn), strict=False)
    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16))
+          v.replace(v.cast(dtypes.float16).realize())
    Tensor.realize(*get_state_dict(model).values())
  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -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 tinygrad.helpers import flatten
 from sz import NONCORE_DIRS
 # 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 = '\x00'.join(flatten(ret))
  code_str = ''.join([f"{banner}\n# {name}\n{banner}\n\n{code}\n" for name,code in 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)
+    write_code_to_file(args.output, ret)
-    print(f"Combined code written to {args.output}")
+    print(f"Combined code written to {args.output}")
@@ -263,23 +263,18 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()
  model = StableDiffusion()
  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
+    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
      model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16))
+          v.replace(v.cast(dtypes.float16).realize())
    Tensor.realize(*get_state_dict(model).values())
  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv
 # QUANT=1 python3 examples/test_onnx_imagenet.py
 #   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -32,7 +32,7 @@ if __name__ == "__main__":
  lr = 5e-3
  transform = ComposeTransforms([
-    lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
+    lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
    lambda x: np.stack([np.asarray(xx) for xx in x], 0),
    lambda x: x / 255.0,
    lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
@@ -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):
+    for i in range((nsample-len(start_tokens))*2):
-      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+      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):]
@@ -2,7 +2,7 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs
 os.chdir("/tmp")
@@ -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,4 +0,0 @@
 # source extra/cl_android.sh
 export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
 export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
@@ -1,169 +1,355 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
-from tinygrad.engine.realize import ExecItem, get_runner
+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)
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
 M = K = N
 run_count = getenv("CNT", 5)
-# ---------------------------
+N = 4096
-# launch/config constants
+run_count = 5
 # ---------------------------
-WARP_SIZE = 32
+BN = 128
 BM = 128
 BK = 8
-# Threadblock tile sizes (block-level tile of C that a block computes)
+TN = 4
-BLOCK_N = 128   # columns of C (N-dim) per block
+TM = 4
 BLOCK_M = 128   # rows of C (M-dim) per block
 BLOCK_K = 8     # K-slice per block iteration
-# Register tile sizes (per-thread accumulator tile of C)
+# NOTE: this is from testgrad
-TN = 4     # columns per thread
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
-TM = 4     # rows per thread
+# 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)
+  # append the reduce shape to each of the views
-THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
+  rstrides = strides_for_shape(rshape)
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
 assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
+  # no reshape required with shrinking REDUCE_AXIS
-WAVE_TILE_N = 128 if is_kernel5 else 64
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
+             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))
 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"
-LANES_PER_WAVE_X = 8
+pm = PatternMatcher([
-LANES_PER_WAVE_Y = 4
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
-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"
-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 rangeify_kernel3():
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+  a = Tensor.empty(N,N)
-  assert dest.shape == src.shape
+  b = Tensor.empty(N,N)
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  c = a@b
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  return dest.after(copy) if set else copy
+  with Context(RANGEIFY=1):
    sink = c.schedule()[-1].ast
  #print(sink)
-def hand_spec_kernel3():
+  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)]
-  # block indices & placeholders
+  opts += [Opt(OptOps.UNROLL, 0, 8)]
  # ---------------------------
  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
+  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
  b = UOp.placeholder((N, N), dtypes.float, slot=2)
  c = UOp.placeholder((N, N), dtypes.float, slot=0)
-  # index the output with the globals
+def top_spec_kernel3():
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  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
+def hl_spec_kernel3():
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
+  nbIterWaveM = 2
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
+  nbIterWaveN = 2
  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
-  # globals are no longer used, they are already in the indexes
+  # define buffers
-  del blockIdx_y, blockIdx_x
+  # TODO: remove these views once the defines have a shape
  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
-  # ---------------------------
+  # shape buffers. TODO: permutes
-  # GLOBAL -> LOCAL (As, Bs)
+  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)
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  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)
+  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  assert len(expanded_shape) == 20
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  permute_a = list(range(len(expanded_shape)))
  permute_b = permute_a[:]
-  # B: read BK x BN tiles
+  # this makes all the global loads match
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  # this can also be more simply done by rebinding the RANGEs
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  # 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?
+  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
-  barrier = UOp.barrier(As_store, Bs_store)
+  permute_b[17:20] = [5,6,7]
  As, Bs = As.after(barrier), Bs.after(barrier)
-  # open inner k range
+  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
-  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
+  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
-  # ---------------------------
+  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
-  # LOCAL -> REG (per-wave tiles)
+  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
  # ---------------------------
  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  #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)
+  axis_types = (
-  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)
+    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)
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
-  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 = graph_rewrite(sink, merge_views)
  return sink
-  # ---------------------------
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
-  # FMA: c_regs += A_col * B_row
+  BLOCK_SIZE = 128 if kernel5 else 256
  # ---------------------------
  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))
-  # TODO: why don't these work as upcast?
+  nbWaves = BLOCK_SIZE // 32
-  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
+  WN = 128 if kernel5 else 64
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
+  WM = BN * BM // nbWaves // WN
  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
-  # Close k, sync, and close K tiles
+  nbWaveX = BN // WN
-  sink = sink.end(k).barrier().end(k_tile_range)
+  nbWaveY = BM // WM
-  # ---------------------------
+  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
-  # REG -> GLOBAL (epilogue)
+  waveIndex = threadIdx_x // 32
-  # ---------------------------
+  waveIdx = waveIndex % nbWaveX
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+  waveIdy = waveIndex // nbWaveX
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  indexInWave = threadIdx_x % 32
  c = c[waveIdy, :, laneIdy, :,
        waveIdx, :, laneIdx, :]
  sink = copy(c, c_regs.after(sink), rng=600)
-  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+  nbThreadXPerWave = 8
  nbThreadYPerWave = 4
-def test_matmul(sink:UOp, N=N):
+  idxInWave = indexInWave % nbThreadXPerWave
-  with Context(DEBUG=0):
+  idyInWave = indexInWave // nbThreadXPerWave
    a = Tensor.randn(N, N)
    b = Tensor.randn(N, N)
    hc = Tensor.empty(N, N)
    Tensor.realize(a, b, hc)
-  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()
+  SUBWN = WN // nbIterWaveN
-  ets = []
+  SUBWM = WM // nbIterWaveM
  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}")
-  if getenv("VERIFY", 1):
+  # Thread mapping to read BKxBN block from A
-    GlobalCounters.reset()
+  rAIdx = threadIdx_x % BK
-    with Context(DEBUG=2):
+  rAIdy = threadIdx_x // BK
-      tc = (a @ b).realize()
+  # Thread mapping to read BNxBK block from B
-    with Context(DEBUG=0):
+  rBIdx = threadIdx_x % BN
-      err = (hc - tc).square().mean().item()
+  rBIdy = threadIdx_x // BN
-    print(f"mean squared error {err}")
+
-    if err > 1e-06:
+  strideReadB = BLOCK_SIZE // BN
-      raise RuntimeError("matmul is wrong!")
+  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:
    with Context(RANGEIFY=1, BLOCK_REORDER=0):
      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.codegen.opt.kernel import Kernel, KernelOptError
-from tinygrad.dtype import DTYPES_DICT
+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(a, bytearray(na))
-    cudaalloc._copyin(b, memoryview(bytearray(nb)))
+    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"
@@ -5,10 +5,8 @@ from tinygrad.dtype import _to_np_dtype
 from tinygrad.codegen.opt import OptOps
 from tinygrad.engine.realize import lower_schedule
-dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
-            dtypes.fp8e4m3 if getenv("FP8E4M3") else dtypes.fp8e5m2 if getenv("FP8E5M2") else dtypes.float)
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
 acc_dtype = (dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else
            dtypes.fp8e4m3 if getenv("ACC_FP8E4M3") else dtypes.fp8e5m2 if getenv("ACC_FP8E5M2") else None)
 if getenv("INT"):  dtype_in, acc_dtype = dtypes.int8, dtypes.int32
 if getenv("UINT"): dtype_in, acc_dtype = dtypes.uint8, dtypes.int32
@@ -16,10 +14,8 @@ N = getenv("N", 4096)
 M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
-
+ATOL = getenv("ATOL", 1e-4)
-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))
+RTOL = getenv("RTOL", 3e-2)
 ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
 INT_LOW = getenv("INT_LOW", 0)
 INT_HIGH = getenv("INT_HIGH", 10)
@@ -8,23 +8,19 @@ import torch
 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 dtype in [torch.float32, torch.float16]:
-  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)
    c = torch.rand((N,N), dtype=dtype)
    if CUDA: b,c = b.cuda(),c.cuda()
    if MPS: b,c = b.to('mps'),c.to('mps')
    def torch_prog(b, c):
      st = time.perf_counter()
      a = b@c
      if CUDA: torch.cuda.synchronize()
      if MPS: torch.mps.synchronize()
      return time.perf_counter() - st
    tm = min([torch_prog(b, c) for _ in range(20)])
    print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS {N:4d}x{N:4d}x{N:4d} matmul in {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)
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
-import argparse, glob, os, time, subprocess, sys
+from tinygrad.runtime.support.system import System
 import argparse, glob, os, re, time, subprocess, sys
 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
@@ -11,7 +12,7 @@ def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
  return devs
-def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
 def cmd_remove_module(args):
@@ -7,6 +7,7 @@ bert_train_params = {
  "GPUS": 6,
  "BS": 96,
  "EVAL_BS": 96,
  "FUSE_ARANGE": 1,
  "BASEDIR": "/raid/datasets/wiki",
 }
@@ -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}")
@@ -50,7 +50,7 @@ def ioctls_from_header():
  hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
  matches = re.findall(pattern, hdr, re.MULTILINE)
-  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname, None)) for name, nr, sname in matches}
+  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname)) for name, nr, sname in matches}
 nrs = ioctls_from_header()
@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
@@ -1,7 +1,7 @@
 import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
 from typing import Any
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import validate, get_example_inputs
 from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
@@ -1,24 +0,0 @@
 #!/bin/sh
 if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
  echo "usage: $0 MESA_PREFIX"
  exit 1
 fi
 TMP=$(mktemp)
 trap 'rm -f "$TMP"' EXIT
 (
  cat <<EOF
 #define HAVE_ENDIAN_H
 #define HAVE_STRUCT_TIMESPEC
 #define HAVE_PTHREAD
 #include <unistd.h>
 #include "nir_shader_compiler_options.h"
 #include "compiler/shader_enums.h"
 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
 printf 'lvp_nir_options = gzip.decompress(base64.b64decode("%s"))' $("$TMP" | gzip | base64 -w0)
@@ -1,42 +1,37 @@
-import os, pathlib
+import pathlib
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"
 from tinygrad.device import Device
 from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
 import time
 import os
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
 NUM_WAVES = 2
-FLOPS_PER_MATMUL = 16*16*16*2
+FLOPS_PER_MATMUL =  16*16*16*2
-INTERNAL_LOOP = 1_000_00
+INTERNAL_LOOP  = 1_000_000
-INSTRUCTIONS_PER_LOOP = 200
+INSTRUCTIONS_PER_LOOP = 1_000
 DIRECTIVE = ".amdhsa_wavefront_size32 1"
 assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
-def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
+def launchBenchmark(instruction, vgprIndices, dense = True):
-  if accum:
+  if dense:
    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
                                                                vgprIndices[1], vgprIndices[2],
                                                                vgprIndices[1], vgprIndices[2], extra)
  elif dense:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2])
+                                                                vgprIndices[1], vgprIndices[2]) * INSTRUCTIONS_PER_LOOP
  else:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                  vgprIndices[1], vgprIndices[2],
+                                                                vgprIndices[1], vgprIndices[2],
-                                                                  vgprIndices[3], vgprIndices[4],
+                                                                vgprIndices[3], vgprIndices[4],
-                                                                  vgprIndices[5])
+                                                                vgprIndices[5]) * INSTRUCTIONS_PER_LOOP
-  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
+  src = assemblyTemplate.replace("INSTRUCTION", instructions)
  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)])
+  start = time.perf_counter()
  fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) #For some reason the returned time is very small after the first kernel execution
  end = time.perf_counter()
  elapsed = end-start
  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")
+  print("{:<29} : {} T(FL)OPS".format(instruction, round(FLOPs/elapsed/10**12, 2)))
 if __name__=="__main__":
  DEVICENUM = os.getenv("DEVICENUM", "0")
@@ -45,17 +40,18 @@ if __name__=="__main__":
  except:
    raise RuntimeError("Error while initiating AMD device")
-  COMPILER = HIPCompiler(DEV.arch)
+  if (ARCH := DEV.arch) not in ['gfx1100', 'gfx1201']:
-  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    raise RuntimeError("only gfx1100 and gfx1201 supported")
-    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
+  COMPILER = HIPCompiler(ARCH)
-    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 40
+
  if ARCH == 'gfx1100':
    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))
    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  elif DEV.arch == 'gfx1201':
+  if ARCH == 'gfx1201':
    NUM_WORKGROUPS = 64
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
@@ -80,20 +76,4 @@ if __name__=="__main__":
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
    FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
+    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
  elif DEV.arch == 'gfx950':
    DIRECTIVE = ".amdhsa_accum_offset 4"
    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.")
@@ -1,23 +1,23 @@
    .text
    .globl matmul
-    .p2align 8
+    .p2align 8 
    .type matmul,@function
 matmul:
-  s_mov_b32 s1, INTERNAL_LOOP
+    s_mov_b32 s1, 1000000
-  s_mov_b32 s2, 0
+    s_mov_b32 s2, 0
-  inner_loop:
+    inner_loop:
-    INSTRUCTION
+        INSTRUCTION
-    s_sub_u32 s1, s1, 1
+        s_sub_u32 s1, s1, 1
-    s_cmp_lg_i32 s1, s2
+        s_cmp_lg_i32 s1, s2
-    s_cbranch_scc1 inner_loop
+        s_cbranch_scc1 inner_loop
-  s_endpgm
+    s_endpgm
 .rodata
 .p2align 6
 .amdhsa_kernel matmul
  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  DIRECTIVE
+  .amdhsa_wavefront_size32 1
 .end_amdhsa_kernel
 .amdgpu_metadata
@@ -28,7 +28,7 @@ amdhsa.version:
 amdhsa.kernels:
  - .name: matmul
    .symbol: matmul.kd
-    .kernarg_segment_size: 0
+    .kernarg_segment_size:  0
    .group_segment_fixed_size: 0
    .private_segment_fixed_size: 0
    .kernarg_segment_align: 4
@@ -36,5 +36,6 @@ amdhsa.kernels:
    .sgpr_count: 8
    .vgpr_count: 32
    .max_flat_workgroup_size: 1024
    .args:
 ...
 .end_amdgpu_metadata
@@ -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 = {}
+  hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_gpu.py").read_text()
-  known_classes = {"NV01_DEVICE_0", "NV01_ROOT", "NV1_MEMORY_SYSTEM", "NV01_MEMORY_VIRTUAL", "NV1_MEMORY_USER", "NV50_MEMORY_VIRTUAL", "NV_FERMI_VASPACE_A",
+  clss = re.search(r'NV01_ROOT.*?NV_SEMAPHORE_SURFACE = \(0x000000da\) # macro', hdrpy, re.DOTALL).group()
-                   "NV20_SUBDEVICE_0"}
+  pattern = r'([0-9a-zA-Z_]*) = +\((0x[0-9a-fA-F]+)\)'
-  for nm,val in nv_gpu.__dict__.items():
+  matches = re.findall(pattern, clss, re.MULTILINE)
-    if not isinstance(val, int): continue
+  return {int(num, base=16):name for name, num in matches}
    if 0x3000 < val < 0xffff: res[val] = nm
    if nm in known_classes: res[val] = nm
  return res
 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
@@ -1,23 +1,10 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.nn.onnx import OnnxRunner, OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  """
  Generate example input tensors based on the provided ONNX graph input specifications.
  NOTE: This is not guaranteed to be reliable. It's a best-effort helper
  that uses heuristics to guess input shapes and values.
  Example:
    from tinygrad.nn.onnx import OnnxRunner
    from extra.onnx_helpers import get_example_inputs
    inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
  """
  def _get_shape(onnx_shape: tuple[str|int]):
    shape = []
    for onnx_dim in onnx_shape:
@@ -57,9 +44,11 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
    ret.update({name:value})
  return ret
-def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
+def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  run_onnx = OnnxRunner(onnx_file)
  ort_options = ort.SessionOptions()
  ort_options.log_severity_level = 3
  ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
  np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
  out_names = list(run_onnx.graph_outputs)
@@ -67,101 +56,9 @@ def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
  ort_out = dict(zip(out_names, out_values))
  tinygrad_out = run_onnx(inputs)
  Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
  tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
  return tinygrad_out, ort_out
 def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  """
  Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
  """
  tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)
  assert tinygrad_out.keys() == ort_out.keys()
  for k in tinygrad_out.keys():
    tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
    if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
-    else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+    else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
 def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  """
  Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
  """
  report = generate_node_output_report(onnx_file, inputs)
  for i, node in enumerate(report):
    node_name = node["node"]
    op = node["op"]
    outputs = node["outputs"]
    for output in outputs:
      output_name = output["name"]
      tinygrad_out = output["tinygrad"]
      ort_out = output["onnxruntime"]
      try:
        if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
        else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
        print(f"Validated {i}: {op=} {node_name=} {output_name=}")
      except AssertionError as e:
        print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
        print(str(e).strip() + "\n")
 def generate_node_output_report(onnx_file, inputs):
  """
  Build a report of all ONNX node outputs from tinygrad and onnxruntime
  Returns:
    A list of dictionaries, where each entry corresponds to one
    node in the ONNX graph. The structure is as follows:
    [
      {
        "node": str,  # The name of the ONNX node.
        "op": str,    # The operation type of the ONNX node.
        "outputs": [
          {
            "name": str,                       # The name of the output tensor.
            "tinygrad": np.ndarray | None,     # The output value from tinygrad.
            "onnxruntime": np.ndarray | None,  # The output value from onnxruntime.
          },
          ...
        ]
      },
      ...
    ]
  """
  import onnx_graphsurgeon as gs
  import onnx
  import tempfile
  # rewrite the model to output all the node outputs
  # `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
  inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
  model = gs.import_onnx(inferred_model)
  model_nodes = model.nodes
  node_outputs = [n.outputs for n in model.nodes]
  model.outputs = [
      each_output for outputs in node_outputs for each_output in outputs
      if not (each_output.dtype is None and each_output.shape is None)  # output with None dtype and None shape is likely a `None` value
  ]
  rewritten_model = gs.export_onnx(model)
  # TODO: remove this once ORT supports 1.18.0
  if getattr(rewritten_model, "ir_version", 0) > 10:
    rewritten_model.ir_version = 10
  with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
    onnx.save(rewritten_model, f.name)
    rewritten_model_path = f.name
    tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
  report = []
  for node in model_nodes:
    outputs = []
    for each_output in node.outputs:
      if each_output.dtype is None and each_output.shape is None:
        continue
      name = each_output.name
      tinygrad_output = tinygrad_out[name]
      ort_output = ort_out[name]
      outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
    report.append({"node": node.name, "op": node.op, "outputs": outputs})
  return report
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
  ret = [float(x) for x in ret]
  if use_sts:
-    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
+    my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
    assert len(my_sts) < MAX_BUFS
    sts_len = 3 + 5*MAX_DIMS
    for s in my_sts:
@@ -1,8 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <database xmlns="http://nouveau.freedesktop.org/"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
+xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 <import file="freedreno_copyright.xml"/>
 <import file="adreno/adreno_common.xml"/>
 <enum name="vgt_event_type" varset="chip">
@@ -21,9 +20,9 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="HLSQ_FLUSH" value="7" variants="A3XX-A4XX"/>
 	<value name="VIZQUERY_END" value="8" variants="A2XX"/>
 	<value name="SC_WAIT_WC" value="9" variants="A2XX"/>
-	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX-"/>
+	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX"/>
-	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX-"/>
+	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX"/>
-	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX-"/>
+	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX"/>
 	<!-- Not sure that these 4 events don't have the same meaning as on A5XX+ -->
 	<value name="RST_PIX_CNT" value="13" variants="A2XX-A4XX"/>
 	<value name="RST_VTX_CNT" value="14" variants="A2XX-A4XX"/>
@@ -31,8 +30,8 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="STAT_EVENT" value="16" variants="A2XX-A4XX"/>
 	<value name="CACHE_FLUSH_AND_INV_TS_EVENT" value="20" variants="A2XX-A4XX"/>
 	<doc>
-		If A6XX_RB_SAMPLE_COUNTER_CNTL.copy is true, writes OQ Z passed
+		If A6XX_RB_SAMPLE_COUNT_CONTROL.copy is true, writes OQ Z passed
-		sample counts to RB_SAMPLE_COUNTER_BASE.  This writes to main
+		sample counts to RB_SAMPLE_COUNT_ADDR.  This writes to main
 		memory, skipping UCHE.
 	</doc>
 	<value name="ZPASS_DONE" value="21"/>
@@ -97,13 +96,6 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	</doc>
 	<value name="BLIT" value="30" variants="A5XX-"/>
 	<doc>
 	Flip between the primary and secondary LRZ buffers. This is used
 	for concurrent binning, so that BV can write to one buffer while
 	BR reads from the other.
 	</doc>
 	<value name="LRZ_FLIP_BUFFER" value="36" variants="A7XX-"/>
 	<doc>
 		Clears based on GRAS_LRZ_CNTL configuration, could clear
 		fast-clear buffer or LRZ direction.
@@ -120,12 +112,11 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="LRZ_FLUSH" value="38" variants="A5XX-"/>
 	<value name="BLIT_OP_FILL_2D" value="39" variants="A5XX-"/>
 	<value name="BLIT_OP_COPY_2D" value="40" variants="A5XX-A6XX"/>
-	<value name="LRZ_CACHE_INVALIDATE" value="40" variants="A7XX-"/>
+	<value name="UNK_40" value="40" variants="A7XX"/>
 	<value name="LRZ_Q_CACHE_INVALIDATE" value="41" variants="A7XX-"/>
 	<value name="BLIT_OP_SCALE_2D" value="42" variants="A5XX-"/>
 	<value name="CONTEXT_DONE_2D" value="43" variants="A5XX-"/>
-	<value name="VSC_BINNING_START" value="44" variants="A5XX-"/>
+	<value name="UNK_2C" value="44" variants="A5XX-"/>
-	<value name="VSC_BINNING_END" value="45" variants="A5XX-"/>
+	<value name="UNK_2D" value="45" variants="A5XX-"/>
 	<!-- a6xx events -->
 	<doc>
@@ -138,22 +129,21 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<!-- note, some of these are the same as a6xx, just named differently -->
 	<doc> Doesn't seem to do anything </doc>
-	<value name="DUMMY_EVENT" value="1" variants="A7XX-"/>
+	<value name="DUMMY_EVENT" value="1" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX"/>
-	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX-"/>
+	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX"/>
-	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX-"/>
+	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX"/>
-	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX-"/>
+	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX"/>
-	<value name="CCU_RESOLVE" value="30" variants="A7XX-"/>
+	<value name="CCU_RESOLVE" value="30" variants="A7XX"/>
-	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX-"/>
+	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX"/>
-	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX-"/>
+	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX"/>
-	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX-"/>
+	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX"/>
-	<value name="CACHE_RESET" value="48" variants="A7XX-"/>
+	<value name="CACHE_RESET" value="48" variants="A7XX"/>
-	<value name="CACHE_CLEAN" value="49" variants="A7XX-"/>
+	<value name="CACHE_CLEAN" value="49" variants="A7XX"/>
 	<!-- TODO: deal with name conflicts with other gens -->
-	<value name="CACHE_FLUSH7" value="50" variants="A7XX-"/>
+	<value name="CACHE_FLUSH7" value="50" variants="A7XX"/>
-	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX-"/>
+	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX"/>
 	<value name="DEPTH_BUFFER_FLIP" value="0x3d" variants="A8XX-"/>
 </enum>
 <enum name="pc_di_primtype">
@@ -334,7 +324,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<doc>fetch state sub-blocks and initiate shader code DMAs</doc>
 	<value name="CP_SET_STATE" value="0x25"/>
 	<doc>load constant into chip and to memory</doc>
-	<value name="CP_SET_CONSTANT" value="0x2d" variants="A2XX"/>
+	<value name="CP_SET_CONSTANT" value="0x2d"/>
 	<doc>load sequencer instruction memory (pointer-based)</doc>
 	<value name="CP_IM_LOAD" value="0x27"/>
 	<doc>load sequencer instruction memory (code embedded in packet)</doc>
@@ -381,7 +371,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_LOAD_STATE" value="0x30" variants="A3XX"/>
 	<value name="CP_LOAD_STATE4" value="0x30" variants="A4XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch enabled</doc>
-	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a" variants="A3XX-A5XX"/>
+	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a"/>
 	<doc>Conditionally load a IB based on a flag, prefetch disabled</doc>
 	<value name="CP_COND_INDIRECT_BUFFER_PFD" value="0x32" variants="A3XX"/>
 	<doc>Load a buffer with pre-fetch enabled</doc>
@@ -524,7 +514,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<!--
 	Seems to set the mode flags which control which CP_SET_DRAW_STATE
 	packets are executed, based on their ENABLE_MASK values
-
+	
 	CP_SET_MODE w/ payload of 0x1 seems to cause CP_SET_DRAW_STATE
 	packets w/ ENABLE_MASK & 0x6 to execute immediately
 	 -->
@@ -547,7 +537,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_LOAD_STATE6_GEOM" value="0x32" variants="A6XX-"/>
 	<value name="CP_LOAD_STATE6_FRAG" value="0x34" variants="A6XX-"/>
 	<!--
-	Note: For UAV state (Image/SSBOs) which have shared state across
+	Note: For IBO state (Image/SSBOs) which have shared state across
 	shader stages, for 3d pipeline CP_LOAD_STATE6 is used.  But for
 	compute shaders, CP_LOAD_STATE6_FRAG is used.  Possibly they are
 	interchangable.
@@ -576,21 +566,20 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="IN_PREEMPT" value="0x0f" variants="A6XX-"/>
 	<!-- TODO do these exist on A5xx? -->
-	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX-"/>
+	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX"/>
 	<value name="CP_REG_TO_MEM_OFFSET_MEM" value="0x74" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_REG" value="0x72" variants="A6XX-"/>
 	<value name="CP_WAIT_MEM_GTE" value="0x14" variants="A6XX"/>
 	<value name="CP_WAIT_TWO_REGS" value="0x70" variants="A6XX"/>
 	<value name="CP_MEMCPY" value="0x75" variants="A6XX-"/>
 	<value name="CP_SET_BIN_DATA5_OFFSET" value="0x2e" variants="A6XX-"/>
 	<!-- A750+, set in place of CP_SET_BIN_DATA5_OFFSET but has different values -->
 	<value name="CP_SET_UNK_BIN_DATA" value="0x2d" variants="A7XX-"/>
 	<doc>
                Write CP_CONTEXT_SWITCH_*_INFO from CP to the following dwords,
                and forcibly switch to the indicated context.
 	</doc>
 	<value name="CP_CONTEXT_SWITCH" value="0x54" variants="A6XX"/>
-	<value name="CP_SET_AMBLE" value="0x55" variants="A6XX-"/>
+	<!-- Note, kgsl calls this CP_SET_AMBLE: -->
 	<value name="CP_SET_CTXSWITCH_IB" value="0x55" variants="A6XX-"/>
 	<!--
 	Seems to always have the payload:
@@ -641,7 +630,8 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_BV_BR_COUNT_OPS" value="0x1b" variants="A7XX-"/>
 	<doc> Clears, adds to local, or adds to global timestamp </doc>
 	<value name="CP_MODIFY_TIMESTAMP" value="0x1c" variants="A7XX-"/>
-	<value name="CP_NON_CONTEXT_REG_BUNCH" value="0x5d" variants="A7XX-"/>
+	<!-- similar to CP_CONTEXT_REG_BUNCH, but discards first two dwords?? -->
 	<value name="CP_CONTEXT_REG_BUNCH2" value="0x5d" variants="A7XX-"/>
 	<doc>
 		Write to a scratch memory that is read by CP_REG_TEST with
 		SOURCE_SCRATCH_MEM set. It's not the same scratch as scratch registers.
@@ -658,11 +648,6 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<doc>Reset various on-chip state used for synchronization</doc>
 	<value name="CP_RESET_CONTEXT_STATE" value="0x1f" variants="A7XX-"/>
 	<doc>Invalidates the "CCHE" introduced on a740</doc>
 	<value name="CP_CCHE_INVALIDATE" value="0x3a" variants="A7XX-"/>
 	<value name="CP_SCOPE_CNTL" value="0x6c" variants="A7XX-"/>
 </enum>
@@ -805,14 +790,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="SB6_GS_SHADER" value="0xb"/>
 		<value name="SB6_FS_SHADER" value="0xc"/>
 		<value name="SB6_CS_SHADER" value="0xd"/>
-		<value name="SB6_UAV"       value="0xe"/>
+		<value name="SB6_IBO"       value="0xe"/>
-		<value name="SB6_CS_UAV"    value="0xf"/>
+		<value name="SB6_CS_IBO"    value="0xf"/>
 	</enum>
 	<enum name="a6xx_state_type">
 		<value name="ST6_SHADER" value="0"/>
 		<value name="ST6_CONSTANTS" value="1"/>
 		<value name="ST6_UBO" value="2"/>
-		<value name="ST6_UAV" value="3"/>
+		<value name="ST6_IBO" value="3"/>
 	</enum>
 	<enum name="a6xx_state_src">
 		<value name="SS6_DIRECT" value="0"/>
@@ -918,6 +903,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 	<stripe varset="chip" variants="A5XX-">
 		<reg32 offset="4" name="4">
 			<bitfield name="INDX_BASE_LO" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="5" name="5">
 			<bitfield name="INDX_BASE_HI" low="0" high="31"/>
 		</reg32>
 		<reg64 offset="4" name="INDX_BASE" type="address"/>
 		<reg32 offset="6" name="6">
 			<!-- max # of elements in index buffer -->
@@ -1093,10 +1084,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			<bitfield name="BINNING" pos="20" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="GMEM" pos="21" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="SYSMEM" pos="22" varset="chip" variants="A6XX-" type="boolean"/>
-			<!-- high bit is 28 until a750: -->
+			<bitfield name="GROUP_ID" low="24" high="28" type="uint"/>
 			<bitfield name="GROUP_ID" low="24" high="29" type="uint"/>
 		</reg32>
 		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_LO" low="0" high="31" type="hex"/>
 		</reg32>
@@ -1130,63 +1119,39 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>
 <enum name="a7xx_abs_mask_mode">
 	<value name="ABS_MASK" value="0x1"/>
 	<value name="NO_ABS_MASK" value="0x0"/>
 </enum>
 <domain name="CP_SET_BIN_DATA5" width="32">
 	<reg32 offset="0" name="0">
 		<bitfield name="VSC_MASK" low="0" high="15" type="hex">
 			<doc>
 				A mask of bins, starting at VSC_N, whose
 				visibility is OR'd together. A value of 0 is
 				interpreted as 1 (i.e. just use VSC_N for
 				visbility) for backwards compatibility. Only
 				exists on a7xx.
 			</doc>
 		</bitfield>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
 		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes">
 			<doc>
 				If this field is 1, VSC_MASK and VSC_N are
 				ignored and instead a new ordinal immediately
 				after specifies the full 32-bit mask of bins
 				to use. The mask is "absolute" instead of
 				relative to VSC_N.
 			</doc>
 		</bitfield>
 	</reg32>
-	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+	<reg32 offset="1" name="1">
-		<reg64 offset="1" name="BIN_DATA_ADDR" type="address"/>
+		<bitfield name="BIN_DATA_ADDR_LO" low="0" high="31" type="hex"/>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+	</reg32>
-		<reg64 offset="3" name="BIN_SIZE_ADDR" type="address"/>
+	<reg32 offset="2" name="2">
-		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<bitfield name="BIN_DATA_ADDR_HI" low="0" high="31" type="hex"/>
-		<reg64 offset="5" name="BIN_PRIM_STRM" type="address"/>
+	</reg32>
-		<!--
+	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-			a7xx adds a few more addresses to the end of the pkt
+	<reg32 offset="3" name="3">
-		 -->
+		<bitfield name="BIN_SIZE_ADDRESS_LO" low="0" high="31"/>
-		<reg64 offset="7" name="7"/>
+	</reg32>
-		<reg64 offset="9" name="9"/>
+	<reg32 offset="4" name="4">
-	</stripe>
+		<bitfield name="BIN_SIZE_ADDRESS_HI" low="0" high="31"/>
-	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
+	</reg32>
-		<reg32 offset="1" name="ABS_MASK"/>
+	<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+	<reg32 offset="5" name="5">
-		<reg64 offset="2" name="BIN_DATA_ADDR" type="address"/>
+		<bitfield name="BIN_PRIM_STRM_LO" low="0" high="31"/>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+	</reg32>
-		<reg64 offset="4" name="BIN_SIZE_ADDR" type="address"/>
+	<reg32 offset="6" name="6">
-		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<bitfield name="BIN_PRIM_STRM_HI" low="0" high="31"/>
-		<reg64 offset="6" name="BIN_PRIM_STRM" type="address"/>
+	</reg32>
-		<!--
+	<!--
-			a7xx adds a few more addresses to the end of the pkt
+		a7xx adds a few more addresses to the end of the pkt
-		 -->
+	 -->
-		<reg64 offset="8" name="8"/>
+	<reg64 offset="7" name="7"/>
-		<reg64 offset="10" name="10"/>
+	<reg64 offset="9" name="9"/>
 	</stripe>
 </domain>
 <domain name="CP_SET_BIN_DATA5_OFFSET" width="32">
@@ -1197,42 +1162,23 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
                stream is recorded.
 	</doc>
 	<reg32 offset="0" name="0">
 		<bitfield name="VSC_MASK" low="0" high="15" type="hex"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
 		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes"/>
 	</reg32>
-	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+	<reg32 offset="1" name="1">
-		<reg32 offset="1" name="1">
+		<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
-			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
-		</reg32>
+	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+	<reg32 offset="2" name="2">
-		<reg32 offset="2" name="2">
+		<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
-			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
-		</reg32>
+	<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
-		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+	<reg32 offset="3" name="3">
-		<reg32 offset="3" name="3">
+		<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
-			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
 		</reg32>
 	</stripe>
 	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
 		<reg32 offset="1" name="ABS_MASK"/>
 		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
 		<reg32 offset="2" name="2">
 			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
 		</reg32>
 		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
 		<reg32 offset="3" name="3">
 			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
 		</reg32>
 		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
 		<reg32 offset="4" name="4">
 			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
 		</reg32>
 	</stripe>
 </domain>
 <domain name="CP_REG_RMW" width="32">
@@ -1250,9 +1196,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</doc>
 	<reg32 offset="0" name="0">
 		<bitfield name="DST_REG" low="0" high="17" type="hex"/>
 		<bitfield name="DST_SCRATCH" pos="19" type="boolean" varset="chip" variants="A7XX-"/>
 		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
 		<bitfield name="SKIP_WAIT_FOR_ME" pos="23" type="boolean" varset="chip" variants="A7XX-"/>
 		<bitfield name="ROTATE" low="24" high="28" type="uint"/>
 		<bitfield name="SRC1_ADD" pos="29" type="boolean"/>
 		<bitfield name="SRC1_IS_REG" pos="30" type="boolean"/>
@@ -1266,7 +1209,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>
-<domain name="CP_REG_TO_MEM" width="32" prefix="chip">
+<domain name="CP_REG_TO_MEM" width="32">
 	<reg32 offset="0" name="0">
 		<bitfield name="REG" low="0" high="17" type="hex"/>
 		<!-- number of registers/dwords copied is max(CNT, 1). -->
@@ -1274,12 +1217,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<stripe varset="chip" variants="A2XX-A4XX">
+	<reg32 offset="1" name="1">
-		<reg32 offset="1" name="DEST" type="address"/>
+		<bitfield name="DEST" low="0" high="31"/>
-	</stripe>
+	</reg32>
-	<stripe varset="chip" variants="A5XX-">
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="DEST" type="address"/>
+		<bitfield name="DEST_HI" low="0" high="31"/>
-	</stripe>
+	</reg32>
 </domain>
 <domain name="CP_REG_TO_MEM_OFFSET_REG" width="32">
@@ -1295,7 +1238,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="DEST" type="waddress"/>
+	<reg32 offset="1" name="1">
 		<bitfield name="DEST" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
 		<bitfield name="DEST_HI" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="OFFSET0" low="0" high="17" type="hex"/>
 		<bitfield name="OFFSET0_SCRATCH" pos="19" type="boolean"/>
@@ -1315,8 +1263,18 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="DEST" type="waddress"/>
+	<reg32 offset="1" name="1">
-	<reg64 offset="3" name="OFFSET" type="waddress"/>
+		<bitfield name="DEST" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
 		<bitfield name="DEST_HI" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="OFFSET_LO" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="4" name="4">
 		<bitfield name="OFFSET_HI" low="0" high="31" type="hex"/>
 	</reg32>
 </domain>
 <domain name="CP_MEM_TO_REG" width="32">
@@ -1329,12 +1287,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- does the same thing as CP_MEM_TO_MEM::UNK31 -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
-	<stripe varset="chip" variants="A2XX-A4XX">
+	<reg32 offset="1" name="1">
-		<reg32 offset="1" name="SRC" type="address"/>
+		<bitfield name="SRC" low="0" high="31"/>
-	</stripe>
+	</reg32>
-	<stripe varset="chip" variants="A5XX-">
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="SRC" type="address"/>
+		<bitfield name="SRC_HI" low="0" high="31"/>
-	</stripe>
+	</reg32>
 </domain>
 <domain name="CP_MEM_TO_MEM" width="32">
@@ -1354,10 +1312,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- some other kind of wait -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
 	<reg64 offset="1" name="DST" type="waddress"/>
 	<reg64 offset="3" name="SRC_A" type="address"/>
 	<reg64 offset="5" name="SRC_B" type="address"/>
 	<reg64 offset="7" name="SRC_C" type="address"/>
 	<!--
 	followed by sequence of addresses.. the first is the
 	destination and the rest are N src addresses which are
@@ -1392,8 +1346,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="SCRATCH" low="20" high="22" type="uint"/>
 		<!-- number of registers/dwords copied is CNT + 1. -->
 		<bitfield name="CNT" low="24" high="26" type="uint"/>
 		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
 		<bitfield name="SKIP_WAIT_FOR_ME" pos="27" type="boolean" varset="chip" variants="A7XX-"/>
 	</reg32>
 </domain>
@@ -1416,12 +1368,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 </domain>
 <domain name="CP_MEM_WRITE" width="32">
-	<stripe varset="chip" variants="A2XX-A4XX">
+	<reg32 offset="0" name="0">
-		<reg32 offset="0" name="ADDR" type="address"/>
+		<bitfield name="ADDR_LO" low="0" high="31"/>
-	</stripe>
+	</reg32>
-	<stripe varset="chip" variants="A5XX-">
+	<reg32 offset="1" name="1">
-		<reg64 offset="0" name="ADDR" type="address"/>
+		<bitfield name="ADDR_HI" low="0" high="31"/>
-	</stripe>
+	</reg32>
 	<!-- followed by the DWORDs to write -->
 </domain>
@@ -1473,14 +1425,24 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
 		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="4" name="4">
 		<bitfield name="MASK" low="0" high="31"/>
 	</reg32>
-	<reg64 offset="5" name="WRITE_ADDR" type="waddress"/>
+	<reg32 offset="5" name="5">
 		<bitfield name="WRITE_ADDR_LO" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="6" name="6">
 		<bitfield name="WRITE_ADDR_HI" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="7" name="7">
 		<bitfield name="WRITE_DATA" low="0" high="31"/>
 	</reg32>
@@ -1495,7 +1457,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- Reserved for flags, presumably? Unused in FW -->
 		<bitfield name="RESERVED" low="0" high="31" type="hex"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
 		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1513,7 +1480,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
 		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1647,7 +1619,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	TODO what is gpuaddr for, seems to be all 0's.. maybe needed for
 	context switch?
 	 -->
-	<reg64 offset="1" name="ADDR" type="waddress"/>
+	<reg32 offset="1" name="1">
 		<bitfield name="ADDR_0_LO" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="ADDR_0_HI" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="3" name="3">
 		<!-- ??? -->
 	</reg32>
@@ -1676,8 +1653,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_SAMPLE_COUNT" pos="12" type="boolean"/>
 		<!-- Write sample count at (iova + 16) -->
 		<bitfield name="SAMPLE_COUNT_END_OFFSET" pos="13" type="boolean"/>
-		<!-- *(iova + 8) += *(iova + 16) - *iova -->
+		<!-- *(iova + 8) = *(iova + 16) - *iova -->
-		<bitfield name="WRITE_ACCUM_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>
+		<bitfield name="WRITE_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>
 		<!-- Next 4 flags are valid to set only when concurrent binning is enabled -->
 		<!-- Increment 16b BV counter. Valid only in BV pipe -->
@@ -1691,11 +1668,15 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_DST" pos="24" type="event_write_dst" addvariant="yes"/>
 		<!-- Writes into WRITE_DST from WRITE_SRC. RB_DONE_TS requires WRITE_ENABLED. -->
 		<bitfield name="WRITE_ENABLED" pos="27" type="boolean"/>
 		<bitfield name="IRQ" pos="31" type="boolean"/>
 	</reg32>
 	<stripe varset="event_write_dst" variants="EV_DST_RAM">
-		<reg64 offset="1" name="1" type="waddress"/>
+		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_0_LO" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="2" name="2">
 			<bitfield name="ADDR_0_HI" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="3" name="3">
 			<bitfield name="PAYLOAD_0" low="0" high="31"/>
 		</reg32>
@@ -1762,7 +1743,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 	</reg32>
 	<stripe varset="chip" variants="A4XX">
-		<reg32 offset="1" name="ADDR" type="address"/>
+		<reg32 offset="1" name="1">
 			<bitfield name="ADDR" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="2" name="2">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1771,7 +1754,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		</reg32>
 	</stripe>
 	<stripe varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="ADDR" type="address"/>
+		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_LO" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="2" name="2">
 			<bitfield name="ADDR_HI" low="0" high="31"/>
 		</reg32>
 		<reg32 offset="3" name="3">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1783,88 +1771,40 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 <domain name="CP_SET_MARKER" width="32" varset="chip" prefix="chip" variants="A6XX-">
 	<doc>Tell CP the current operation mode, indicates save and restore procedure</doc>
 	<enum name="set_marker_mode">
 		<value value="0" name="SET_RENDER_MODE"/>
 		<!-- IFPC - inter-frame power collapse -->
 		<value value="1" name="SET_IFPC_MODE"/>
 	</enum>
 	<enum name="a6xx_ifpc_mode">
 		<value value="0" name="IFPC_ENABLE"/>
 		<value value="1" name="IFPC_DISABLE"/>
 	</enum>
 	<enum name="a6xx_marker">
-		<value value="1" name="RM6_DIRECT_RENDER"/>
+		<value value="1" name="RM6_BYPASS"/>
-		<value value="2" name="RM6_BIN_VISIBILITY"/>
+		<value value="2" name="RM6_BINNING"/>
-		<value value="3" name="RM6_BIN_DIRECT"/>
+		<value value="4" name="RM6_GMEM"/>
-		<value value="4" name="RM6_BIN_RENDER_START"/>
+		<value value="5" name="RM6_ENDVIS"/>
-		<value value="5" name="RM6_BIN_END_OF_DRAWS"/>
+		<value value="6" name="RM6_RESOLVE"/>
-		<value value="6" name="RM6_BIN_RESOLVE"/>
+		<value value="7" name="RM6_YIELD"/>
 		<value value="7" name="RM6_BIN_RENDER_END"/>
 		<value value="8" name="RM6_COMPUTE"/>
-		<value value="12" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->
+		<value value="0xc" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->
 		<!--
 			These values come from a6xx_set_marker() in the
 			downstream kernel, and they can only be set by the kernel
 		-->
-		<value value="13" name="RM6_IB1LIST_START"/>
+		<value value="0xd" name="RM6_IB1LIST_START"/>
-		<value value="14" name="RM6_IB1LIST_END"/>
+		<value value="0xe" name="RM6_IB1LIST_END"/>
-		<value value="15" name="RM7_BIN_VISIBILITY_END"/>
+		<!-- IFPC - inter-frame power collapse -->
-
+		<value value="0x100" name="RM6_IFPC_ENABLE"/>
-		<!-- new in a8xx: -->
+		<value value="0x101" name="RM6_IFPC_DISABLE"/>
 		<value value="32" name="RM8_DEPTH_PASS_START"/>
 		<value value="33" name="RM8_DEPTH_PASS_END"/>
 	</enum>
-	<stripe varset="chip" variants="A6XX-A7XX">
+	<reg32 offset="0" name="0">
-		<reg32 offset="0" name="0">
+		<!--
-			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
+			NOTE: blob driver and some versions of freedreno/turnip set
-			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
+			b4, which is unused (at least by current sqe fw), but interferes
-
+			with parsing if we extend the size of the bitfield to include
-
+			b8 (only sent by kernel mode driver).  Really, the way the
-			<bitfield name="MODE" low="0" high="3" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			parsing works in the firmware, only b0-b3 are considered, but
-			<!-- used by preemption to determine if GMEM needs to be saved or not -->
+			if b8 is set, the low bits are interpreted differently.  To
-			<bitfield name="USES_GMEM" pos="4" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			model this, without getting confused by spurious b4, this is
-
+			described as two overlapping bitfields:
-
+		 -->
-			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
+		<bitfield name="MODE" low="0" high="8" type="a6xx_marker"/>
-
+		<bitfield name="MARKER" low="0" high="3" type="a6xx_marker"/>
-
+	</reg32>
 			<!--
 				CP_SET_MARKER is used with these bits to create a
 				critical section around a workaround for ray tracing.
 				The workaround happens after BVH building, and appears
 				to invalidate the RTU's BVH node cache. It makes sure
 				that only one of BR/BV/LPAC is executing the
 				workaround at a time, and no draws using RT on BV/LPAC
 				are executing while the workaround is executed on BR (or
 				vice versa, that no draws on BV/BR using RT are executed
 				while the workaround executes on LPAC), by
 				hooking subsequent CP_EVENT_WRITE/CP_DRAW_*/CP_EXEC_CS.
 				The blob usage is:
 				CP_SET_MARKER(RT_WA_START)
 				... workaround here ...
 				CP_SET_MARKER(RT_WA_END)
 				...
 				CP_SET_MARKER(SHADER_USES_RT)
 				CP_DRAW_INDX(...) or CP_EXEC_CS(...)
 			-->
 			<bitfield name="SHADER_USES_RT" pos="9" type="boolean" variants="A7XX-"/>
 			<bitfield name="RT_WA_START" pos="10" type="boolean" variants="A7XX-"/>
 			<bitfield name="RT_WA_END" pos="11" type="boolean" variants="A7XX-"/>
 		</reg32>
 	</stripe>
 	<stripe varset="chip" variants="A8XX-">
 		<reg32 offset="0" name="0">
 			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
 			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
 			<bitfield name="USES_GMEM" pos="7" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
 			<bitfield name="MODE" low="0" high="6" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
 			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
 			<!-- idk if the RT w/a fields apply to a8xx as well -->
 		</reg32>
 	</stripe>
 </domain>
 <domain name="CP_SET_PSEUDO_REG" width="32" varset="chip" prefix="chip" variants="A6XX-">
@@ -1890,9 +1830,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			If concurrent binning is disabled then BR also does binning so it will also
 			write the "real" registers in BR.
 		-->
-		<value value="8" name="VSC_PIPE_DATA_DRAW_BASE"/>
+		<value value="8" name="DRAW_STRM_ADDRESS"/>
-		<value value="9" name="VSC_SIZE_BASE"/>
+		<value value="9" name="DRAW_STRM_SIZE_ADDRESS"/>
-		<value value="10" name="VSC_PIPE_DATA_PRIM_BASE"/>
+		<value value="10" name="PRIM_STRM_ADDRESS"/>
 		<value value="11" name="UNK_STRM_ADDRESS"/>
 		<value value="12" name="UNK_STRM_SIZE_ADDRESS"/>
@@ -1993,11 +1933,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			a bitmask of which modes pass the test.
 		-->
-		<!-- RM6_BIN_VISIBILITY -->
+		<!-- RM6_BINNING -->
 		<bitfield name="BINNING" pos="25" variants="RENDER_MODE" type="boolean"/>
 		<!-- all others -->
 		<bitfield name="GMEM" pos="26" variants="RENDER_MODE" type="boolean"/>
-		<!-- RM6_DIRECT_RENDER -->
+		<!-- RM6_BYPASS -->
 		<bitfield name="SYSMEM" pos="27" variants="RENDER_MODE" type="boolean"/>
 		<bitfield name="BV" pos="25" variants="THREAD_MODE" type="boolean"/>
@@ -2070,45 +2010,54 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>
-<domain name="CP_SET_AMBLE" width="32">
+<domain name="CP_SET_CTXSWITCH_IB" width="32">
 	<doc>
-		Used by the userspace and kernel drivers to set various IB's
+                Used by the userspace driver to set various IB's which are
-		which are executed during context save/restore for handling
+                executed during context save/restore for handling
-		state that isn't restored by the context switch routine itself.
+                state that isn't restored by the
                context switch routine itself.
 	</doc>
-	<enum name="amble_type">
+	<enum name="ctxswitch_ib">
-		<value name="PREAMBLE_AMBLE_TYPE" value="0">
+		<value name="RESTORE_IB" value="0">
 			<doc>Executed unconditionally when switching back to the context.</doc>
 		</value>
-		<value name="BIN_PREAMBLE_AMBLE_TYPE" value="1">
+		<value name="YIELD_RESTORE_IB" value="1">
                        <doc>
 				Executed when switching back after switching
 				away during execution of
-				a CP_SET_MARKER packet with RM6_BIN_RENDER_END as the
+				a CP_SET_MARKER packet with RM6_YIELD as the
-				payload *and* skipsaverestore is set. This is
+				payload *and* the normal save routine was
-				expected to restore static register values not
+				bypassed for a shorter one. I think this is
-				saved when skipsaverestore is set.
+				connected to the "skipsaverestore" bit set by
 				the kernel when preempting.
 			</doc>
 		</value>
-		<value name="POSTAMBLE_AMBLE_TYPE" value="2">
+		<value name="SAVE_IB" value="2">
                        <doc>
 				Executed when switching away from the context,
 				except for context switches initiated via
 				CP_YIELD.
                        </doc>
 		</value>
-		<value name="KMD_AMBLE_TYPE" value="3">
+		<value name="RB_SAVE_IB" value="3">
 			<doc>
 				This can only be set by the RB (i.e. the kernel)
 				and executes with protected mode off, but
-				is otherwise similar to POSTAMBLE_AMBLE_TYPE.
+				is otherwise similar to SAVE_IB.
 				Note, kgsl calls this CP_KMD_AMBLE_TYPE
 			</doc>
 		</value>
 	</enum>
-	<reg64 offset="0" name="ADDR" type="address"/>
+	<reg32 offset="0" name="0">
 		<bitfield name="ADDR_LO" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="1" name="1">
 		<bitfield name="ADDR_HI" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="DWORDS" low="0" high="19" type="uint"/>
-		<bitfield name="TYPE" low="20" high="21" type="amble_type"/>
+		<bitfield name="TYPE" low="20" high="21" type="ctxswitch_ib"/>
 	</reg32>
 </domain>
@@ -2140,12 +2089,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="UNK_EVENT_WRITE" value="0x4"/>
 		<doc>
 			Tracks GRAS_LRZ_CNTL::GREATER, GRAS_LRZ_CNTL::DIR, and
-			GRAS_LRZ_VIEW_INFO with previous values, and if one of
+			GRAS_LRZ_DEPTH_VIEW with previous values, and if one of
 			the following is true:
 			- GRAS_LRZ_CNTL::GREATER has changed
 			- GRAS_LRZ_CNTL::DIR has changed, the old value is not
 			  CUR_DIR_GE, and the new value is not CUR_DIR_DISABLED
-			- GRAS_LRZ_VIEW_INFO has changed
+			- GRAS_LRZ_DEPTH_VIEW has changed
 			then it does a LRZ_FLUSH with GRAS_LRZ_CNTL::ENABLE
 			forced to 1.
 			Only exists in a650_sqe.fw.
@@ -2260,7 +2209,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 <domain name="CP_MEM_TO_SCRATCH_MEM" width="32">
 	<doc>
-		Best guess is that it is a faster way to fetch all the VSC_CHANNEL_VISIBILITY registers
+		Best guess is that it is a faster way to fetch all the VSC_STATE registers
 		and keep them in a local scratch memory instead of fetching every time
 		when skipping IBs.
 	</doc>
@@ -2308,25 +2257,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 		<bitfield name="CLEAR_ON_CHIP_TS" pos="0" type="boolean"/>
 		<bitfield name="CLEAR_RESOURCE_TABLE" pos="1" type="boolean"/>
-		<bitfield name="CLEAR_BV_BR_COUNTER" pos="2" type="boolean"/>
+		<bitfield name="CLEAR_GLOBAL_LOCAL_TS" pos="2" type="boolean"/>
 		<bitfield name="RESET_GLOBAL_LOCAL_TS" pos="3" type="boolean"/>
 	</reg32>
 </domain>
 <domain name="CP_SCOPE_CNTL" width="32">
 	<enum name="cp_scope">
 		<value value="0" name="INTERRUPTS"/>
 	</enum>
 	<reg32 offset="0" name="0">
 		<bitfield name="DISABLE_PREEMPTION" pos="0" type="boolean"/>
 		<bitfield low="28" high="31" name="SCOPE" type="cp_scope"/>
 	</reg32>
 </domain>
 <domain name="CP_INDIRECT_BUFFER" width="32" varset="chip" prefix="chip" variants="A5XX-">
 	<reg64 offset="0" name="IB_BASE" type="address"/>
 	<reg32 offset="2" name="2">
 		<bitfield name="IB_SIZE" low="0" high="19"/>
 	</reg32>
 </domain>
@@ -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 = {}
@@ -97,7 +97,7 @@ def parse_cmd_buf(dat):
        if state_block == SB6_CS_SHADER:
          from extra.disassemblers.adreno import disasm_raw
-          if state_type == ST6_SHADER and IOCTL > 3:
+          if state_type == ST6_SHADER and IOCTL > 2:
            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
@@ -106,30 +106,25 @@ def parse_cmd_buf(dat):
              print('constants')
              hexdump(x)
          if state_type == ST6_IBO:
-            if state_src == 0x1:
+            ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
              ibos_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
            else: ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
            CAPTURED_STATE['ibos'] = ibos_bytes[:]
            if IOCTL > 1:
              print('texture ibos')
              hexdump(ibos_bytes)
        elif state_block == SB6_CS_TEX:
          if state_type == ST6_SHADER:
-            if state_src == 0x1:
+            samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
              samplers_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
            else: samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
            CAPTURED_STATE['samplers'] = samplers_bytes[:]
            if IOCTL > 1:
              print('texture samplers')
              hexdump(samplers_bytes)
          if state_type == ST6_CONSTANTS:
-            if state_src == 0x1:
+            descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
              descriptors_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
            else: descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
            CAPTURED_STATE['descriptors'] = descriptors_bytes[:]
            if IOCTL > 1:
              print('texture descriptors')
              hexdump(descriptors_bytes)
      elif ops[opcode] == "CP_REG_TO_MEM":
        reg, cnt, b64, accum = vals[0] & 0x3FFFF, (vals[0] >> 18) & 0xFFF, (vals[0] >> 30) & 0x1, (vals[0] >> 31) & 0x1
        dest = vals[1] | (vals[2] << 32)
@@ -157,10 +152,6 @@ def parse_cmd_buf(dat):
        if IOCTL > 0:
          print(f'THREADSIZE-{(vals[0] >> 20)&0x1}\nEARLYPREAMBLE-{(vals[0] >> 23) & 0x1}\nMERGEDREGS-{(vals[0] >> 3) & 0x1}\nTHREADMODE-{vals[0] & 0x1}\nHALFREGFOOTPRINT-{(vals[0] >> 1) & 0x3f}\nFULLREGFOOTPRINT-{(vals[0] >> 7) & 0x3f}\nBRANCHSTACK-{(vals[0] >> 14) & 0x3f}\n')
          print(f'SP_CS_UNKNOWN_A9B1-{vals[1]}\nSP_CS_BRANCH_COND-{vals[2]}\nSP_CS_OBJ_FIRST_EXEC_OFFSET-{vals[3]}\nSP_CS_OBJ_START-{vals[4] | (vals[5] << 32)}\nSP_CS_PVT_MEM_PARAM-{vals[6]}\nSP_CS_PVT_MEM_ADDR-{vals[7] | (vals[8] << 32)}\nSP_CS_PVT_MEM_SIZE-{vals[9]}')
      if offset == 0xa9e8:
        CAPTURED_STATE['bindless_base'] = (vals[0] | (vals[1] << 32)) & ~0b11
        # print(hex(CAPTURED_STATE['bindless_base']))
        # hexdump(get_mem(CAPTURED_STATE['bindless_base'], 0x200))
      if offset == 0xb180:
        if IOCTL > 0:
          print('border color offset', hex(vals[1] << 32 | vals[0]))
@@ -180,8 +171,8 @@ def ioctl(fd, request, argp):
    name, stype = nrs[nr]
    s = get_struct(argp, stype)
    if IOCTL > 0: print(f"{ret:2d} = {name:40s}", ' '.join(format_struct(s)))
-    if name == "IOCTL_KGSL_GPUOBJ_INFO":
+    if name == "IOCTL_KGSL_GPUOBJ_INFO": pass
-      mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
+      # mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
    if name == "IOCTL_KGSL_GPU_COMMAND":
      for i in range(s.numcmds):
        cmd = get_struct(s.cmdlist+ctypes.sizeof(msm_kgsl.struct_kgsl_command_object)*i, msm_kgsl.struct_kgsl_command_object)
@@ -882,11 +882,6 @@ impl<'a> Thread<'a> {
                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i32;
                            (s0 * s1) as u32
                        }
                        10 => {
                            let s0 = sign_ext((s0 & 0xffffff) as u64, 24) as i64;
                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i64;
                            ((s0 * s1) >> 32) as u32
                        }
                        17 | 18 | 26 => {
                            let (s0, s1) = (s0 as i32, s1 as i32);
                            (match op {
@@ -935,7 +930,7 @@ impl<'a> Thread<'a> {
            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 | 770 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -1001,10 +996,6 @@ impl<'a> Thread<'a> {
                                    let ret = s0.wrapping_sub(s1);
                                    (ret as u32, s1 > s0)
                                }
                                770 => {
                                    let ret = s1.wrapping_sub(s0);
                                    (ret as u32, s0 > s1)
                                }
                                _ => todo_instr!(instruction)?,
                            };
                            if self.exec.read() {
@@ -1,32 +1,98 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from tinygrad import Tensor, dtypes, Device, UOp
+from typing import cast
-from tinygrad.helpers import getenv
+from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
-from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
+from tinygrad import Tensor, dtypes, Device
-from tinygrad.renderer import ProgramSpec
+from tinygrad.helpers import diskcache, OSX, getenv
 from tinygrad.engine.realize import CompiledRunner
-def get_output(asm:str, n_threads:int=1):
+@diskcache
-  input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
+def assemble(code:str) -> bytes:
-                         for ln in asm.strip().splitlines() if ln.strip()])
+  try:
-  src = f"""
+    LLVM_MC = "llvm-mc" if OSX else "/opt/rocm/llvm/bin/llvm-mc"
-  typedef long unsigned int size_t;
+    return subprocess.run([LLVM_MC, "--arch=amdgcn", "--mcpu=gfx1100", "--triple=amdgcn-amd-amdhsa", "-filetype=obj", "-o", "-"],
-  extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+                           input=code.encode("utf-8"), stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
-  extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
+  except subprocess.CalledProcessError as e:
-    int l = __ockl_get_local_id(0);
+    print("stderr:")
-    unsigned a = 0, b = 0, c = 0;
+    print(e.stderr.decode())
-    {input_asm}
+    raise
-    unsigned res;
+
-    asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
+# copied from extra/rdna
-    *(data0_1+l) = res;
+def get_prg(code:str, v_cnt:int, s_cnt:int):
-  }}"""
+  function_name = "test"
-  t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
+  metadata = f"""
-  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
+amdhsa.kernels:
-  car = CompiledRunner(prg)
+- .args:
-  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
+  - .address_space: global
-  car([t.uop.buffer], {}, wait=True)
+    .name: buf_0
-  return t.numpy()
+    .offset: 0
    .size: 8
    .type_name: unsigned int*
    .value_kind: global_buffer
  .group_segment_fixed_size: 0
  .kernarg_segment_align: 8
  .kernarg_segment_size: 8
  .language: OpenCL C
  .language_version:
  - 1
  - 2
  .max_flat_workgroup_size: 256
  .name: test
  .private_segment_fixed_size: 0
  .sgpr_count: {s_cnt}
  .sgpr_spill_count: 0
  .symbol: test.kd
  .uses_dynamic_stack: false
  .vgpr_count: {v_cnt}
  .vgpr_spill_count: 0
  .wavefront_size: 32
 amdhsa.target: amdgcn-amd-amdhsa--gfx1100
 amdhsa.version:
 - 1
 - 2
  """
  boilerplate_start = f"""
  .rodata
  .global {function_name}.kd
  .type {function_name}.kd,STT_OBJECT
  .align 0x10
  .amdhsa_kernel {function_name}"""
  kernel_desc = {
    '.amdhsa_group_segment_fixed_size': 0, '.amdhsa_private_segment_fixed_size': 0, '.amdhsa_kernarg_size': 0,
    '.amdhsa_next_free_vgpr': v_cnt,   # this matters!
    '.amdhsa_reserve_vcc': 0, '.amdhsa_reserve_xnack_mask': 0,
    '.amdhsa_next_free_sgpr': s_cnt,
    '.amdhsa_float_round_mode_32': 0, '.amdhsa_float_round_mode_16_64': 0, '.amdhsa_float_denorm_mode_32': 3, '.amdhsa_float_denorm_mode_16_64': 3,
    '.amdhsa_dx10_clamp': 1, '.amdhsa_ieee_mode': 1, '.amdhsa_fp16_overflow': 0,
    '.amdhsa_workgroup_processor_mode': 1, '.amdhsa_memory_ordered': 1, '.amdhsa_forward_progress': 0, '.amdhsa_enable_private_segment': 0,
    '.amdhsa_system_sgpr_workgroup_id_x': 1, '.amdhsa_system_sgpr_workgroup_id_y': 1, '.amdhsa_system_sgpr_workgroup_id_z': 1,
    '.amdhsa_system_sgpr_workgroup_info': 0, '.amdhsa_system_vgpr_workitem_id': 2, # is amdhsa_system_vgpr_workitem_id real?
    '.amdhsa_exception_fp_ieee_invalid_op': 0, '.amdhsa_exception_fp_denorm_src': 0,
    '.amdhsa_exception_fp_ieee_div_zero': 0, '.amdhsa_exception_fp_ieee_overflow': 0, '.amdhsa_exception_fp_ieee_underflow': 0,
    '.amdhsa_exception_fp_ieee_inexact': 0, '.amdhsa_exception_int_div_zero': 0,
    '.amdhsa_user_sgpr_dispatch_ptr': 0, '.amdhsa_user_sgpr_queue_ptr': 0, '.amdhsa_user_sgpr_kernarg_segment_ptr': 1,
    '.amdhsa_user_sgpr_dispatch_id': 0, '.amdhsa_user_sgpr_private_segment_size': 0, '.amdhsa_wavefront_size32': 1, '.amdhsa_uses_dynamic_stack': 0}
  code_start = f""".end_amdhsa_kernel
  .text
  .global {function_name}
  .type {function_name},@function
  .p2align 8
  {function_name}:
  """
  ret = ".amdgpu_metadata\n" + metadata + ".end_amdgpu_metadata" + boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) \
      + "\n" + code_start + code + f"\n.size {function_name}, .-{function_name}"
  return AMDProgram(cast(AMDDevice, Device["AMD"]), function_name, assemble(ret))
 def get_output(s:str, n_threads:int=1):
  assert n_threads <= 32
  code = "\n".join(["s_load_b64 s[0:1], s[0:1], null", "v_lshlrev_b32_e32 v0, 2, v0", s,
    "s_waitcnt 0",
    "global_store_b32 v0, v1, s[0:1]",
    "s_nop 0", "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)", "s_endpgm"])
  test = Tensor.zeros((n_threads,), dtype=dtypes.uint32).contiguous().realize().uop.buffer
  prg = get_prg(code, 32, 32)
  prg(test._buf, global_size=(1, 1, 1), local_size=(n_threads, 1, 1), wait=True)
  return test.numpy()
 def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
 def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -39,57 +105,54 @@ class TestHW(unittest.TestCase):
  def test_simple(self):
    out = get_output("""
-    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 %2 %1
+    v_mov_b32_e32 v1 v10
-    """)[0]
+    """, n_threads=2)
    np.testing.assert_equal(out, 42)
  def test_exec_mov(self):
    out = get_output("""
-    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 v10 42
    s_mov_b32_e32 exec_lo 0b10
-    v_mov_b32_e32 %1 10
+    v_mov_b32_e32 v10 10
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 %1
+    v_mov_b32_e32 v1 v10
    """, n_threads=2)
    np.testing.assert_equal(out, [42, 10])
  def test_exec_cmp_vopc(self):
    out = get_output("""
    s_mov_b32 vcc_lo 0 // reset vcc
-    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 %2 10
+    v_mov_b32_e32 v11 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmp_ne_u32 %1 %2
+    v_cmp_ne_u32 v10 v11
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 vcc_lo
+    v_mov_b32_e32 v1 vcc_lo
    """, n_threads=2)
    np.testing.assert_equal(out, 0b01)
  def test_exec_cmpx_vop3(self):
    out = get_output("""
-    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 v11 10
    v_mov_b32_e32 %2 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmpx_ne_u32 %1 %2
+    v_cmpx_ne_u32 v10 v11
    s_mov_b32_e32 s10 exec_lo
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 s10
+    v_mov_b32_e32 v1 s10
-    """, n_threads=2)[0]
+    """, n_threads=2)
-    np.testing.assert_equal(out & 0b11, 0b01)
+    np.testing.assert_equal(out, 0b01)
  def test_fmac_vop3_modifier(self):
    init_state = f"""
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
+    v_mov_b32_e32 v10 {f16_to_bits(4.0)}
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
+    v_mov_b32_e32 v11 {f16_to_bits(3.0)}
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
+    v_mov_b32_e32 v1 {f16_to_bits(2.0)}
    """
-    mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 v11 v10"), f16_to_bits(14.))
-    def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 v10"), f16_to_bits(-10.))
-    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 -v10"), f16_to_bits(14.))
    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
    self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))
  def test_s_abs_i32(self):
    def s_abs_i32(x, y, dst="s10", scc=0):
@@ -97,7 +160,7 @@ class TestHW(unittest.TestCase):
        self.assertEqual(get_output(f"""
        s_mov_b32_e32 {dst} {x}
        s_abs_i32 {dst} {dst}
-        v_mov_b32_e32 %2 {reg}
+        v_mov_b32_e32 v1 {reg}
        """)[0], val)
    s_abs_i32(0x00000001, 0x00000001, scc=1)
    s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
@@ -110,8 +173,8 @@ class TestHW(unittest.TestCase):
  def test_v_rcp_f32_neg_vop3(self):
    def v_neg_rcp_f32(x:float, y:float):
      out = get_output(f"""
-      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      v_rcp_f32_e64 %2, -%2
+      v_rcp_f32_e64 v1, -v1
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg_rcp_f32(math.inf, -0.0)
@@ -123,11 +186,10 @@ class TestHW(unittest.TestCase):
  def test_v_cndmask_b32_neg(self):
    def v_neg(x:int|float, y:float):
      # always pick -v1
      out = get_output(f"""
-      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      s_mov_b32_e32 s10 1
+      s_mov_b32_e32 s10 1 // always pick -v1
-      v_cndmask_b32 %2, %2, -%2 s10
+      v_cndmask_b32 v1, v1, -v1 s10
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg(-0.0, 0.0)
@@ -136,12 +198,5 @@ class TestHW(unittest.TestCase):
    v_neg(math.inf, -math.inf)
    v_neg(-math.inf, math.inf)
  def test_v_subrev_wrap(self):
    out = get_output("""
    v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
    v_subrev_co_u32 %2, vcc_lo, %2, %1
    """)[0]
    self.assertEqual(out, 0xffff_ffff)
 if __name__ == "__main__":
  unittest.main()
@@ -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:])
--- a/Show More
+++ b/Show More