Merge branch 'master' into algebraic_upat

more mi350x matmul work (#13138 )
* more mi350x matmul work * broken compute
2025-11-13 09:21:05 -08:00 · 2025-11-13 09:09:28 -08:00 · 2025-11-13 09:07:54 -08:00 · 2025-11-13 09:07:11 -08:00 · 2025-11-13 08:52:27 -08:00 · 2025-11-14 00:42:21 +08:00
87 changed files with 6918 additions and 7115 deletions
@@ -0,0 +1,3 @@
+[run]
+source = tinygrad
+branch = True
@@ -61,7 +61,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/pyproject.toml') }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

@@ -70,13 +70,13 @@ runs:
      uses: actions/cache@v4
      with:
        path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
+        key: downloads-cache-${{ inputs.key }}-${{ env.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.DOWNLOAD_CACHE_VERSION }}
+        key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}

    # **** Python deps ****

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

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

 jobs:
@@ -71,13 +68,10 @@ jobs:
        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
    - name: Verify Linux autogen
      run: |
-        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        ./autogen_stubs.sh libc
        ./autogen_stubs.sh io_uring
        ./autogen_stubs.sh ib
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
    - name: Verify WebGPU autogen
@@ -95,3 +89,21 @@ jobs:
        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
        ./autogen_stubs.sh mesa
        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
+  autogen-ng:
+    name: In-tree Autogen
+    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:
+        pydeps: 'clang>=20'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev
+    - name: Verify Linux autogen
+      run: |
+        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
+        python3 -c "from tinygrad.runtime.autogen import libc"
+        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
@@ -54,7 +54,7 @@ jobs:
    - name: Print macOS version
      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 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=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
@@ -199,7 +199,7 @@ jobs:
    - name: Test speed vs torch
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      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
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -320,19 +320,20 @@ jobs:
    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+    # TODO: too slow
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    # - name: Run 10 CIFAR training steps w HALF
+    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    # - name: Run 10 CIFAR training steps w BF16
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    # - name: Run full CIFAR training steps w 6 GPUS
+    #   run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
@@ -409,7 +410,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 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      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
    - name: Test tensor cores
      run: |
        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
@@ -524,17 +525,18 @@ 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=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=350 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    # - name: Run 10 CIFAR training steps w HALF
+    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    # - name: Run 10 CIFAR training steps w BF16
    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    #- name: Run full CIFAR training steps w 6 GPUS
    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
@@ -623,24 +625,26 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.9.9 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.9.9 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.9.9 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    # - name: openpilot compile3 0.9.9 driving_vision
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    # - name: openpilot compile3 0.9.9 driving_policy
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    # - name: openpilot compile3 0.9.9 dmonitoring
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      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
    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
      # TODO: ASSERT_MIN_STEP_TIME=17
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      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
    - name: openpilot compile3 0.10.1 dmonitoring
      # TODO: ASSERT_MIN_STEP_TIME=10
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_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/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -704,8 +708,9 @@ 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
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    # TODO: too slow
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
    # 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
@@ -767,8 +772,9 @@ 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
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    # TODO: too slow
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
    #- name: Run 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 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 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 DISABLE_COMPILER_CACHE=1 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 CCACHE=0 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
@@ -20,11 +20,11 @@ jobs:
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel build twine
    - name: Build and publish
      env:
        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
      run: |
-        python setup.py sdist bdist_wheel
+        python -m build
        twine upload dist/*
@@ -1,10 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
+  CACHE_VERSION: '13'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -233,7 +230,7 @@ jobs:
        python-version: '3.11'
        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 .
+      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y . --ignore-paths='tinygrad/runtime/autogen'
    - name: Lint with ruff
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
@@ -290,8 +287,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 < 18000 lines
-      run: MAX_LINE_COUNT=18000 python sz.py
+    - name: Repo line count < 18500 lines
+      run: MAX_LINE_COUNT=18500 python sz.py

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

@@ -392,7 +390,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          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
      - name: Test openpilot CL compile fp16
        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
@@ -63,3 +63,5 @@ profile_stats
 *.log
 target
 .mypy_cache
+mutants
+.mutmut-cache
@@ -21,17 +21,38 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

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

-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.
+- **Tensor library** with autograd
+- **IR and compiler** that fuse and lower kernels
+- **JIT + graph execution**
+- **nn / optim / datasets** for real training

-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.
+It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.

-## Features
+---

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

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

 ### Laziness

@@ -188,6 +188,7 @@ nv_status_codes = {}
    extra/nv_gpu_driver/g_rpc-message-header.h \
    extra/nv_gpu_driver/gsp_static_config.h \
    extra/nv_gpu_driver/vbios.h \
+    extra/nv_gpu_driver/pci_exp_table.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

@@ -253,23 +254,6 @@ generate_ib() {
  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 \
@@ -549,12 +533,10 @@ 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
@@ -563,6 +545,6 @@ elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
 elif [ "$1" == "mesa" ]; then generate_mesa
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
+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_am; generate_webgpu; generate_mesa
 else echo "usage: $0 <type>"
 fi
@@ -1,11 +1,10 @@
 import os, sys, pickle, time, re
 import numpy as np
+if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
-
-import onnx
 from tinygrad.nn.onnx import OnnxRunner

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

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

 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+  import onnx
  import onnxruntime as ort
  
  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

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

@@ -244,15 +244,16 @@ 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-len(start_tokens))*2):
-      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range(nsample):
+      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
      next_tokens[ctx[:, -1] == eot] = eot
      ctx = np.concatenate((ctx, next_tokens), axis=1)
      pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all(): break
+      if (next_tokens == eot).all() or pos == nctx: 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):]
@@ -0,0 +1,141 @@
+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"
@@ -51,11 +51,15 @@ 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}")
@@ -89,6 +89,20 @@ 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)

@@ -0,0 +1,134 @@
+/*
+ * 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
+
+
@@ -2,7 +2,8 @@ import os, pathlib, argparse
 from examples.llama3 import Tokenizer
 from tabulate import tabulate
 from tinygrad import fetch
-from tinygrad.helpers import flatten
+from tinygrad.helpers import flatten, getenv
+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())
@@ -10,19 +11,15 @@ 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.")
@@ -32,10 +29,11 @@ if __name__ == "__main__":

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

-  code_str = '\x00'.join(flatten(ret))
+  banner = "#"*40
+  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")
@@ -44,5 +42,5 @@ if __name__ == "__main__":
  print(f"code has {len(encoded)} tokens")

  if args.output:
-    write_code_to_file(args.output, ret)
-    print(f"Combined code written to {args.output}")
+    with open(args.output, 'w') as f: f.write(code_str)
+    print(f"Combined code written to {args.output}")
@@ -4,7 +4,7 @@ import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
 import tinygrad.runtime.autogen.sqtt as sqtt
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from tinygrad.helpers import round_up, flatten, all_same
+from tinygrad.helpers import round_up, flatten, all_same, temp
 from dataclasses import dataclass

 CHUNK_CLASSES = {
@@ -210,7 +210,7 @@ class RGP:
        flags=0,
        trace_shader_core_clock=0x93f05080,
        trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
        device_revision_id=0xc8,
        vgprs_per_simd=1536,
        sgprs_per_simd=128*16,
@@ -324,7 +324,7 @@ class RGP:
 if __name__ == '__main__':
  parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
  parser.add_argument('command')
-  parser.add_argument('input')
+  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
  parser.add_argument('-d', '--device')
  parser.add_argument('-o', '--output')
  args = parser.parse_args()
@@ -346,3 +346,4 @@ if __name__ == '__main__':

  if args.output is not None:
    with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
+    print(f"Saved to {args.output}")
@@ -28,18 +28,6 @@ def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
    cur_off += instr_sz
  return addr_table

-@dataclasses.dataclass
-class InstInfo:
-  typ:str=""
-  inst:str=""
-  hit:int=0
-  lat:int=0
-  stall:int=0
-  def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
-
-  def on_ev(self, ev):
-    self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
-
@dataclasses.dataclass(frozen=True)
 class InstExec:
  typ:str
@@ -49,19 +37,19 @@ class InstExec:
  time:int

@dataclasses.dataclass(frozen=True)
-class PrgExec:
-  name:str
-  wave:int
+class WaveExec:
+  wave_id:int
  cu:int
  simd:int
-  def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
+  begin_time:int
+  end_time:int
+  insts:list[InstExec]

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

    for prog in prog_evs:
      arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
@@ -72,7 +60,8 @@ class _ROCParseCtx:
    x = next(self.sqtt_evs, None)
    self.active_kern = x.kern if x is not None else None
    self.active_se = x.se if x is not None else None
-    return x
+    self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
+    return self.active_blob

  def on_occupancy_ev(self, ev):
    if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
@@ -80,19 +69,15 @@ class _ROCParseCtx:
  def on_wave_ev(self, ev):
    if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)

-    asm:dict[int, InstInfo] = {}
    inst_execs:list[InstExec] = []
    for j in range(ev.instructions_size):
      inst_ev = ev.instructions_array[j]
      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      inst_disasm = self.disasms[(self.active_kern, inst_ev.pc.address)][0]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=inst_disasm))
-      asm[inst_ev.pc.address].on_ev(inst_ev)
+      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
      inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))

    if ev.instructions_size > 0:
-      self.wave_events[key:=PrgExec(self.active_kern, ev.wave_id, ev.cu, ev.simd)] = asm
-      self.inst_execs[key] = inst_execs
+      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, ev.begin_time, ev.end_time, inst_execs))

 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
  dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -107,10 +92,10 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:

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

  @rocprof.rocprof_trace_decoder_trace_callback_t
  def trace_cb(record_type, events_ptr, n, data_ptr):
@@ -125,7 +110,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:

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

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

  try:
    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
  return ROCParseCtx

 if __name__ == "__main__":
@@ -150,7 +135,7 @@ if __name__ == "__main__":

  with args.profile.open("rb") as f: profile = pickle.load(f)
  rctx = decode(profile)
-  print('SQTT:', rctx.wave_events.keys())
+  print('SQTT:', rctx.inst_execs.keys())

  for ev in profile:
    if not isinstance(ev, ProfilePMCEvent): continue
@@ -1,7 +1,7 @@
 import os
 os.environ["PYTHONPATH"] = "."
 os.environ["SQTT"] = "1"
-os.environ["AMD"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
 os.environ["VIZ"] = "1"
 os.environ["AMD_LLVM"] = "0"

@@ -10,23 +10,25 @@ import sys, contextlib
 from tinygrad import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.device import Device, ProfileDeviceEvent

 from extra.sqtt.roc import decode, InstExec, PrgExec

-dev = Device["AMD"]
+dev = Device[os.environ["DEV"]]
+
+def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)

 def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
  name = sys._getframe(1).f_code.co_name
  def fxn(_):
    L = UOp.special(l, "lidx0")
    G = UOp.special(g, "gidx0")
-    ops:list[str] = [UOp(Ops.CUSTOM, arg="asm volatile (")]
-    for inst in instrs: ops.append(UOp(Ops.CUSTOM, src=(ops[-1],), arg=f'  "{inst}\\n\\t"'))
-    ops.append(UOp(Ops.CUSTOM, src=(ops[-1],), arg=");"))
-    return UOp.sink(*ops, L, G, arg=KernelInfo(name=name))
+    op = custom("asm volatile (")
+    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
+    op = custom(");", op)
+    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
  k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
  return k

@@ -37,9 +39,10 @@ def save_sqtt():
  sqtt:dict[PrgExec, list[InstExec]] = {}
  yield sqtt
  # decode sqtt
-  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
-  assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  sqtt.update(rctx.inst_execs)
+  if os.environ["DEV"] == "AMD":
+    rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
+    assert len(rctx.inst_execs) > 0, "empty sqtt output"
+    sqtt.update(rctx.inst_execs)

 class TestTiming(unittest.TestCase):
  def test_v_add(self):
@@ -60,18 +63,24 @@ class TestTiming(unittest.TestCase):
    assert all(s.stall == 0 for s in wave)

  def test_multi_cycle_inst(self):
+    def custom_vrcp(A, B):
+      op = custom("float a = 0.0;")
+      op = custom("float b = (*(data1_1+0));", op)
+      #op = custom('asm volatile("v_mul_f32_e32 %2 %2 %1" : "+v"(a) : "v"(b));', op)
+      op = custom('asm volatile("v_rcp_f32_e32 %2 %1" : "+v"(a) : "v"(b));', op)
+      op = custom('asm volatile("v_add_f32_e64 %1 %1 1.0" : "+v"(a));', op)
+      op = custom("*(data0_1+0) = a;", op)
+      return UOp.sink(op, A, B, arg=KernelInfo(name="custom_vrcp"))
+    out = Tensor([0.]).realize()
+    inp = Tensor([-2.0]).realize()
    with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_mov_b32_e32 v4 0x3f800000",
-        "v_rcp_f32_e32 v5 v4",
-        "v_mul_f32_e32 v6 v5 v4",
-      ]).realize()
-    w = list(sqtt.values())[0]
-    rcp, mul = w[1], w[2]
-    self.assertGreater(rcp.dur, 1) # 4 cycles on gfx11
-    self.assertEqual(mul.dur, 1)
-    # mul depends on v5, how can it run before rcp is done?
-    self.assertGreaterEqual(mul.time, rcp.time+rcp.dur)
+      Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
+
+    wave = list(sqtt.values())[0][0]
+    for i in range(len(wave.insts)):
+      if wave.insts[i].inst.startswith("global_store"):
+        print(f"store diff {wave.insts[i].time-(wave.insts[i-1].time)}")
+    self.assertEqual(out.item(), 0.5)

  def test_wmma(self):
    with save_sqtt() as sqtt:
@@ -87,12 +96,11 @@ class TestTiming(unittest.TestCase):
    n = 1
    def sleep_kernel(data0):
      assert data0.dtype.base == dtypes.ulong
-      ops:list[UOp] = []
-      ops.append(UOp(Ops.CUSTOM, arg="unsigned long long t0 = __builtin_readcyclecounter();"))
-      ops.append(UOp(Ops.CUSTOM, arg=f"__builtin_amdgcn_s_sleep({n});", src=(ops[-1],)))
-      ops.append(UOp(Ops.CUSTOM, arg="unsigned long long t1 = __builtin_readcyclecounter();", src=(ops[-1],)))
-      ops.append(UOp(Ops.CUSTOM, arg=f"data0_{data0.size}[0] = t1 - t0;", src=(ops[-1],)))
-      return UOp.sink(data0, *ops, arg=KernelInfo(name=f"sleep_{n}"))
+      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
+      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
+      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
+      op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
+      return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
    diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
    diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
    with save_sqtt() as sqtt:
@@ -0,0 +1 @@
+WARP_THREADS = 32
@@ -0,0 +1,287 @@
+import math, functools
+from typing import cast, Callable
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import AxisType, UOp, KernelInfo, Ops
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.dtype import AddrSpace, PtrDType
+from tinygrad.helpers import getenv, prod
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.tiles import RT
+
+class Group:
+  def __init__(self, warps:int, ker):
+    self.warps = warps
+    self.group_threads = warps * WARP_THREADS
+    self.threadIdx_x = ker.threadIdx_x
+    self.ker = ker
+
+  # helpers
+  @property
+  def laneid(self): return self.threadIdx_x % self.group_threads
+  @property
+  def warpid(self): return self.laneid // WARP_THREADS
+  @property
+  def groupid(self): return self.threadIdx_x // self.group_threads
+
+  # ops that only work on a single warp
+
+  clear_rid = 1000
+  def clear(self, reg:UOp, value:float=0):
+    assert self.warps == 1
+
+    i = UOp.range(reg.size, Group.clear_rid)
+    Group.clear_rid += 1
+
+    reg_store = reg.reshape((reg.size,))[i].store(value).end(i)
+
+    self.ker.push_store(reg_store, reg)
+    return reg.after(reg_store).reshape(reg.shape)
+
+  def zero(self, reg:UOp): return self.clear(reg, 0)
+  def neg_inf(self, reg:UOp): return self.clear(reg, -math.inf)
+
+  copy_rid = 300
+  def copy(self, dst:UOp, src:UOp):
+    assert self.warps == 1
+
+    assert dst.shape == src.shape
+    assert cast(PtrDType, dst.dtype).addrspace == AddrSpace.REG
+    assert cast(PtrDType, src.dtype).addrspace == AddrSpace.REG
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
+    Group.copy_rid += len(dst.shape)
+
+    dst_store = dst[*rngs_for_shape].store(src[*rngs_for_shape].cast(dst.dtype.base)).end(*rngs_for_shape)
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store).reshape(dst.shape)
+
+  mma_rid = 600
+  def mma_AB(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_inner, mma_i_width, 2+i] for i in range(2)] + [b[mma_i_inner, mma_i_width, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  def mma_ABt(self, c:UOp, a:UOp, b:UOp, after=True):
+    assert self.warps == 1
+
+    mma_i_height = UOp.range(c.shape[-3], Group.mma_rid)
+    mma_i_width = UOp.range(c.shape[-2], Group.mma_rid+1)
+    mma_i_inner = UOp.range(a.shape[-2], Group.mma_rid+2, AxisType.REDUCE)
+    Group.mma_rid += 3
+
+    wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+
+    a_in = UOp.vectorize(*[a[mma_i_height, mma_i_inner, i] for i in range(8)])
+    b_in1 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 4+i] for i in range(2)]))
+    c_out1 = UOp.vectorize(*[c[mma_i_height, mma_i_width, i] for i in range(4)])
+    b_in2 = UOp.vectorize(*([b[mma_i_width, mma_i_inner, 2+i] for i in range(2)] + [b[mma_i_width, mma_i_inner, 6+i] for i in range(2)]))
+    c_out2 = UOp.vectorize(*[c[mma_i_height, mma_i_width, 4+i] for i in range(4)])
+
+    out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
+    out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
+    c_i = [c[mma_i_height, mma_i_width, i].store(out1.gep(i)) for i in range(4)] + [c[mma_i_height, mma_i_width, 4+i].store(out2.gep(i)) for i in range(4)]
+    c_store = UOp.group(*c_i).end(mma_i_height, mma_i_width, mma_i_inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape) if after else c_store
+
+  map_rid = 400
+  def map(self, a:UOp, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
+    assert self.warps == 1
+
+    rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
+    Group.map_rid += len(a.shape)
+
+    if op.__code__.co_argcount == 1:
+      to_store = op(a[*rngs_for_shape])
+    else:
+      to_store = op(a[*rngs_for_shape], rngs_for_shape)
+
+    a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
+
+    self.ker.push_store(a_store, a)
+    return a.after(a_store).reshape(a.shape)
+
+  def row_reduce(self, vec:UOp, src:UOp, op:Callable[[UOp, UOp], UOp]):
+    assert self.warps == 1
+
+    red_local = self.ker.alloc((self.group_threads, 2), src.dtype.base, AddrSpace.LOCAL)
+    red_reg = self.ker.alloc((2,), src.dtype.base, AddrSpace.REG)
+
+    for height in self.ker.range(src.shape[-3], track=False):
+      i = UOp.range(red_reg.size, Group.clear_rid)
+      Group.clear_rid += 1
+      red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
+      reg_store = red_reg.flatten()[i].store(0.).end(i)
+      red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      for i_outer in self.ker.range(2, track=False):
+        for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
+          for i_inner in self.ker.range(4, AxisType.REDUCE, track=False):
+            elem_index = i_inner + 2 * (i_inner // 2) + i_outer * 2
+            reg_store = red_reg[i_outer].store(op(red_reg[i_outer], src[height, width, elem_index])).end(i_inner, width, i_outer)
+            red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      # store to shared memory
+      for i_outer in self.ker.range(2, track=False):
+        red_local_store = red_local[self.laneid, i_outer].store(red_reg[i_outer]).end(i_outer)
+        red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+
+      # reduce from shared memory
+      for i_outer in self.ker.range(2, track=False):
+        for i_inner in self.ker.range(3, AxisType.REDUCE, track=False):
+          offset = (self.laneid // 4) * 4 + ((self.laneid + i_inner + 1) % 4)
+          reg_store = red_reg[i_outer].store(op(red_reg[i_outer], red_local[offset, i_outer])).end(i_inner, i_outer)
+          red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      # reduce with vec
+      for i_outer in self.ker.range(2, track=False):
+        vec_store = vec[height, 0, i_outer].store(op(vec[height, 0, i_outer], red_reg[i_outer])).end(i_outer, height)
+
+    self.ker.push_store(vec_store, vec)
+    return vec.after(vec_store).reshape(vec.shape)
+
+  # ops that can work across multiple warps
+
+  LOAD_INNER = 8
+  load_rid = 100
+  def load(self, dst:UOp, src:UOp, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
+      srcf = src.flatten(-2)
+
+      load_i_height = UOp.range(dst.shape[-3], Group.load_rid)
+      load_i_width = UOp.range(dst.shape[-2], Group.load_rid+1)
+      load_i_inner = UOp.range(RT.BASE_TILE_NEPT, Group.load_rid+2)
+      Group.load_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      if not transpose:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * RT.TILE_ROW_DIM + (warp_laneid // 4)
+        col = load_i_width * RT.TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+        row_offset = ((load_i_inner % 4) // 2) * 8
+        col_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+      else:
+        row = (local_warpid * dst.shape[-3] + load_i_height) * RT.TILE_ROW_DIM + 2 * (warp_laneid % 4)
+        col = load_i_width * RT.TILE_COL_DIM + (warp_laneid // 4)
+
+        row_offset = (load_i_inner % 2) + (load_i_inner // 4) * 8
+        col_offset = ((load_i_inner % 4) // 2) * 8
+
+      src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
+
+      dst_store = dst[*dst_idxs, load_i_height, load_i_width, load_i_inner].store(srcf[*idxs[:-2], src_i_last])
+      dst_store = dst_store.end(load_i_height, load_i_width, load_i_inner)
+    elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten(-2)
+
+      srcf = src.flatten()
+      row_stride = prod(src.shape[axis+1:])
+
+      idxs = tuple(idx * dst.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * dst.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
+
+      memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
+      total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
+
+      load_i_outer = UOp.range(total_calls, Group.load_rid)
+      load_i_inner = UOp.range(Group.LOAD_INNER, Group.load_rid+1)
+      Group.load_rid += 2
+
+      load_idx = load_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
+
+      dst_i = row * dst.shape[-1] + col + load_i_inner
+      src_i += row * row_stride + col + load_i_inner
+
+      dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(load_i_outer, load_i_inner)
+    else:
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    return dst.after(dst_store.barrier()).reshape(dst.shape)
+
+  STORE_INNER = 8
+  store_rid = 200
+  def store(self, dst:UOp, src:UOp, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis=0, after=True):
+    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
+      dstf = dst.flatten(-2)
+
+      store_i_height = UOp.range(src.shape[-3], Group.store_rid)
+      store_i_width = UOp.range(src.shape[-2], Group.store_rid+1)
+      store_i_inner = UOp.range(RT.BASE_TILE_NEPT, Group.store_rid+2)
+      Group.store_rid += 3
+
+      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
+      else: local_warpid = self.warpid
+      warp_laneid = self.threadIdx_x % WARP_THREADS
+
+      row = (local_warpid * src.shape[-3] + store_i_height) * RT.TILE_ROW_DIM + (warp_laneid // 4)
+      col = store_i_width * RT.TILE_COL_DIM + 2 * (warp_laneid % 4)
+
+      row_offset = ((store_i_inner % 4) // 2) * 8
+      col_offset = (store_i_inner % 2) + (store_i_inner // 4) * 8
+
+      dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
+
+      dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, store_i_height, store_i_width, store_i_inner])
+      dst_store = dst_store.end(store_i_height, store_i_width, store_i_inner)
+    elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
+      dstf = dst.flatten()
+      row_stride = prod(dst.shape[axis+1:])
+
+      idxs = tuple(idx * src.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * src.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
+
+      srcf = src.flatten(-2)
+
+      memcpy_per_row = src.shape[-1] // Group.STORE_INNER
+      total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
+
+      store_i_outer = UOp.range(total_calls, Group.store_rid)
+      store_i_inner = UOp.range(Group.STORE_INNER, Group.store_rid+1)
+      Group.store_rid += 2
+
+      load_idx = store_i_outer * self.group_threads + self.laneid
+      row = load_idx // memcpy_per_row
+      col = (load_idx * Group.STORE_INNER) % src.shape[-1]
+
+      src_i = row * src.shape[-1] + col + store_i_inner
+      dst_i += row * row_stride + col + store_i_inner
+
+      dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(store_i_outer, store_i_inner)
+    else:
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store.barrier()).reshape(dst.shape) if after else dst_store
@@ -0,0 +1,88 @@
+from contextlib import AbstractContextManager
+from tinygrad.uop.ops import UOp, KernelInfo, AxisType, AddrSpace
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.group import Group
+from extra.thunder.tiny.tk.tiles import GL, ST, RT, RV
+
+class _tk_range:
+  user_rid = 0
+  def __init__(self, end:int, axis_type:AxisType): self.end, self.axis_type, self.done = end, axis_type, False
+  def __iter__(self): return self
+  def __next__(self):
+    if not self.done:
+      self.done = True
+      _tk_range.user_rid += 1
+      self._rng = UOp.range(self.end, _tk_range.user_rid-1, axis_type=self.axis_type)
+      return self._rng
+    raise StopIteration
+
+class Kernel(AbstractContextManager):
+  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+    self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
+    self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
+    self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
+    self.threadIdx_x = UOp.special(block_size, "lidx0")
+
+    self.range_stack = []
+    self.store_stack = []
+
+    self.global_slot = 0
+    self.shared_slot = 0
+    self.register_slot = 0
+    self.allocs = {}
+
+  @property
+  def warpid(self): return self.threadIdx_x // WARP_THREADS
+
+  def __enter__(self): return self
+  def __exit__(self, exc_type, exc_value, traceback): pass
+
+  def group(self, size:int): return Group(size, self)
+  @property
+  def warp(self): return self.group(1)
+  @property
+  def warpgroup(self): return self.group(4)
+
+  def range(self, end:int, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+    rng = _tk_range(end, axis_type)
+    if track: self.range_stack.append(rng)
+    return rng
+
+  def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
+    match addrspace:
+      case AddrSpace.GLOBAL:
+        slot = self.global_slot
+        self.global_slot += 1
+      case AddrSpace.LOCAL:
+        slot = self.shared_slot
+        self.shared_slot += 1
+      case AddrSpace.REG:
+        slot = self.register_slot
+        self.register_slot += 1
+
+    uop = UOp.placeholder(shape, dtype, slot=slot, addrspace=addrspace)
+
+    if name:
+      if (name, shape) in self.allocs: return self.allocs[(name, shape)]
+      self.allocs[(name, shape)] = uop
+
+    return uop
+
+  def gl(self, shape, dtype): return GL(shape, dtype, self)._uop
+  def st(self, shape, dtype): return ST(shape, dtype, self)._uop
+  def rt(self, shape, dtype): return RT(shape, dtype, self)._uop
+  def rv(self, length, dtype, layout="naive"): return RV(length, dtype, layout, self)._uop
+
+  def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
+
+  def finish(self):
+    # end all ranges
+    rngs = []
+    while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
+
+    return self.store_stack.pop()[0].end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+
+  def endrange(self):
+    last_store = self.store_stack.pop()
+    last_range = self.range_stack.pop()
+    return last_store[1].after(last_store[0].barrier().end(last_range._rng)).reshape(last_store[1].shape)
@@ -0,0 +1,45 @@
+from tinygrad.dtype import AddrSpace
+
+from extra.thunder.tiny.tk import WARP_THREADS
+
+class GL:
+  def __init__(self, shape, dtype, ker):
+    self.shape, self.dtype = shape, dtype
+    self._uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
+
+class ST:
+  def __init__(self, shape, dtype, ker):
+    self.shape, self.dtype = shape, dtype
+    self._uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
+
+class RT:
+  TILE_ROW_DIM, TILE_COL_DIM = 16, 16
+  BASE_TILE_NE = TILE_ROW_DIM * TILE_COL_DIM
+  BASE_TILE_NEPT = BASE_TILE_NE // WARP_THREADS
+
+  def __init__(self, shape, dtype, ker):
+    assert len(shape) == 2
+    assert shape[0] % RT.TILE_ROW_DIM == 0
+    assert shape[1] % RT.TILE_COL_DIM == 0
+
+    height = shape[0] // RT.TILE_ROW_DIM
+    width = shape[1] // RT.TILE_COL_DIM
+
+    self.shape, self.dtype = (height, width, self.BASE_TILE_NEPT), dtype
+    self._uop = ker.alloc(self.shape, dtype, AddrSpace.REG)
+
+class RV:
+  def __init__(self, length, dtype, layout, ker):
+    tiles = length // RT.TILE_ROW_DIM
+
+    match layout:
+      case "naive":
+        inner_dim = 1
+        outer_dim = (tiles + 1) // 2
+      case "ortho":
+        inner_dim = 1
+        outer_dim = tiles
+      case _: raise NotImplementedError(f"rv layout {layout} not implemented")
+
+    self.shape, self.dtype = (outer_dim, inner_dim, 2), dtype
+    self._uop = ker.alloc(self.shape, dtype, AddrSpace.REG)
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import sys, os, zlib, struct, hashlib
-from hexdump import hexdump
 from tinygrad.helpers import DEBUG, getenv, fetch
 from tinygrad.runtime.support.usb import USB3

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

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

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

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

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

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

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

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

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

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

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

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

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

-NAMES = ["chenyu","George Hotz","nimlgen","qazal","Sieds Lykles","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
@@ -40,4 +40,4 @@ for d in days:
 print("** Commits by day (last 7) **")
 print("```")
 print("\n".join([header, rule] + rows))
-print("```")
+print("```")
@@ -1,10 +0,0 @@
-[mypy]
-warn_unused_configs = True
-files = tinygrad
-ignore_missing_imports = True
-check_untyped_defs = True
-explicit_package_bases = True
-warn_unreachable = True
-warn_redundant_casts = True
-# NOTE: had to comment this out to make mypy pass on both CI and OSX
-#warn_unused_ignores = True
@@ -0,0 +1,233 @@
+[project]
+name = "tinygrad"
+version = "0.11.0"
+description = "You like pytorch? You like micrograd? You love tinygrad! <3"
+authors = [{ name = "George Hotz" }]
+
+classifiers = ["Programming Language :: Python :: 3"]
+
+license = 'MIT'
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = []
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+include-package-data = true
+packages = [
+  'tinygrad',
+  'tinygrad.apps',
+  'tinygrad.codegen',
+  'tinygrad.codegen.opt',
+  'tinygrad.codegen.late',
+  'tinygrad.engine',
+  'tinygrad.mixin',
+  'tinygrad.nn',
+  'tinygrad.renderer',
+  'tinygrad.runtime',
+  'tinygrad.runtime.autogen',
+  'tinygrad.runtime.autogen.am',
+  'tinygrad.runtime.autogen.nv',
+  'tinygrad.runtime.graph',
+  'tinygrad.runtime.support',
+  'tinygrad.runtime.support.am',
+  'tinygrad.runtime.support.nv',
+  'tinygrad.schedule',
+  'tinygrad.uop',
+  'tinygrad.viz',
+]
+
+[tool.setuptools.package-data]
+tinygrad = ["py.typed"]
+"tinygrad.viz" = ["index.html", "assets/**/*", "js/*"]
+
+
+[project.optional-dependencies]
+arm = ["unicorn"]
+triton = ["triton-nightly>=2.1.0.dev20231014192330"]
+linting = [
+  "pylint",
+  "mypy==1.18.1",
+  "typing-extensions",
+  "pre-commit",
+  "ruff",
+  "numpy",
+  "typeguard",
+]
+# mlperf = [
+#   "mlperf-logging @ git+https://github.com/mlperf/[email protected]",
+# ]
+testing_minimal = [
+  "numpy",
+  "torch==2.9.0",
+  "pytest",
+  "pytest-xdist",
+  "pytest-timeout",
+  "pytest-split",
+  "hypothesis",
+  "z3-solver",
+]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate"]
+testing = [
+  "tinygrad[testing_minimal]",
+  "pillow",
+  "onnx==1.18.0",
+  "onnx2torch",
+  "onnxruntime",
+  "opencv-python",
+  "tabulate",
+  "tqdm",
+  "safetensors",
+  "transformers",
+  "sentencepiece",
+  "tiktoken",
+  "blobfile",
+  "librosa",
+  # librosa needs numba but uv ignores python upper bounds and some numba versions require <python3.10
+  "numba>=0.55",
+  "networkx",
+  "nibabel",
+  "bottle",
+  "ggml-python",
+  "capstone",
+  "pycocotools",
+  "boto3",
+  "pandas",
+  "influxdb3-python",
+]
+docs = [
+  "mkdocs",
+  "mkdocs-material",
+  "mkdocstrings[python]",
+  "markdown-callouts",
+  "markdown-exec[ansi]",
+  "black",
+  "numpy",
+]
+
+
+[tool.mutmut]
+paths_to_mutate = ["tinygrad/"]
+do_not_mutate = [
+  "tinygrad/apps/*",
+  "tinygrad/codegen/*",
+  "tinygrad/engine/*",
+  "tinygrad/nn/*",
+  "tinygrad/renderer/*",
+  "tinygrad/runtime/*",
+  "tinygrad/schedule/*",
+  "tinygrad/uop/*",
+  "tinygrad/viz/*",
+  "tinygrad/device.py",
+  "tinygrad/dtype.py",
+  "tinygrad/gradient.py",
+  "tinygrad/helpers.py",
+  "tinygrad/tensor.py",
+]
+tests_dir = ["test/test_tiny.py", "test/test_ops.py"]
+debug = true
+
+
+[tool.mypy]
+warn_unused_configs = true
+files = ["tinygrad"]
+ignore_missing_imports = true
+check_untyped_defs = true
+explicit_package_bases = true
+warn_unreachable = true
+warn_redundant_casts = true
+# NOTE: had to comment this out to make mypy pass on both CI and OSX
+#warn_unused_ignores = true
+
+[tool.pytest.ini_options]
+norecursedirs = [
+  "extra",
+  ".hypothesis",
+  ".git",
+]
+timeout = 300
+timeout_method = "thread"
+timeout_func_only = true
+testpaths = ["test"]
+
+[tool.ruff]
+preview = true
+target-version = "py311"
+line-length = 150
+indent-width = 2
+exclude = [
+  ".git/",
+  "docs/",
+  "extra/",
+  "tinygrad/runtime/autogen",
+  "test/external/mlperf_resnet",
+  "test/external/mlperf_unet3d",
+]
+
+[tool.ruff.lint]
+select = [
+  "F",  # Pyflakes
+  "W6",
+  "E71",
+  "E72",
+  "E112",   # no-indented-block
+  "E113",   # unexpected-indentation
+  # "E124",
+  "E203",   # whitespace-before-punctuation
+  "E272",   # multiple-spaces-before-keyword
+  "E275",   # missing-whitespace-after-keyword
+  "E303",   # too-many-blank-lines
+  "E304",   # blank-line-after-decorator
+  "E501",   # line-too-long
+  # "E502",
+  "E702",   # multiple-statements-on-one-line-semicolon
+  "E703",   # useless-semicolon
+  "E731",   # lambda-assignment
+  "W191",   # tab-indentation
+  "W291",   # trailing-whitespace
+  "W293",   # blank-line-with-whitespace
+  "UP039",  # unnecessary-class-parentheses
+  "C416",   # unnecessary-comprehension
+  "RET506", # superfluous-else-raise
+  "RET507", # superfluous-else-continue
+  "A",      # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing
+  "FURB110",# if-exp-instead-of-or-operator
+  "RUF018", # assignment-in-assert
+]
+
+# detect unused imports in examples
+[tool.ruff.lint.per-file-ignores]
+"examples/**/*.py" = [
+  "W6",
+  "E71",
+  "E72",
+  "E112",
+  "E113",
+  "E203",
+  "E272",
+  "E275",
+  "E303",
+  "E304",
+  "E501",
+  "E702",
+  "E703",
+  "E731",
+  "W191",
+  "W291",
+  "W293",
+  "UP039",
+  "C416",
+  "RET506",
+  "RET507",
+  "A",
+  "FURB110",
+  "RUF018",
+  "F541",
+  "F841",
+]
+
+[tool.ruff.format]
+exclude = ["*"]
@@ -1,9 +0,0 @@
-[pytest]
-norecursedirs =
-    extra
-    .hypothesis
-    .git
-timeout = 300
-timeout_method = thread
-timeout_func_only = true
-testpaths = test
@@ -1,56 +0,0 @@
-indent-width = 2
-preview = true
-target-version = "py310"
-
-lint.select = [
-  "F",  # Pyflakes
-  "W6",
-  "E71",
-  "E72",
-  "E112",   # no-indented-block
-  "E113",   # unexpected-indentation
-  # "E124",
-  "E203",   # whitespace-before-punctuation
-  "E272",   # multiple-spaces-before-keyword
-  "E275",   # missing-whitespace-after-keyword
-  "E303",   # too-many-blank-lines
-  "E304",   # blank-line-after-decorator
-  "E501",   # line-too-long
-  # "E502",
-  "E702",   # multiple-statements-on-one-line-semicolon
-  "E703",   # useless-semicolon
-  "E731",   # lambda-assignment
-  "W191",   # tab-indentation
-  "W291",   # trailing-whitespace
-  "W293",   # blank-line-with-whitespace
-  "UP039",  # unnecessary-class-parentheses
-  "C416",   # unnecessary-comprehension
-  "RET506", # superfluous-else-raise
-  "RET507", # superfluous-else-continue
-  "A",      # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing
-  "FURB110",# if-exp-instead-of-or-operator
-  "RUF018", # assignment-in-assert
-]
-
-line-length = 150
-
-exclude = [
-  ".git/",
-  "docs/",
-  "extra/",
-  "tinygrad/runtime/autogen",
-  "test/external/mlperf_resnet",
-  "test/external/mlperf_unet3d",
-]
-
-# detect unused imports in examples
-[lint.per-file-ignores]
-"examples/**/*.py" = [
-  "W6", "E71", "E72", "E112", "E113", "E203", "E272", "E275",
-  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
-  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
-  "FURB110", "RUF018", "F541", "F841"
-]
-
-[format]
-exclude = ["*"]
@@ -1,111 +0,0 @@
-#!/usr/bin/env python3
-
-from pathlib import Path
-from setuptools import setup
-
-directory = Path(__file__).resolve().parent
-with open(directory / 'README.md', encoding='utf-8') as f:
-  long_description = f.read()
-
-testing_minimal = [
-  "numpy",
-  "torch==2.9.0",
-  "pytest",
-  "pytest-xdist",
-  "pytest-timeout",
-  "pytest-split",
-  "hypothesis",
-  "z3-solver",
-]
-
-setup(name='tinygrad',
-      version='0.11.0',
-      description='You like pytorch? You like micrograd? You love tinygrad! <3',
-      author='George Hotz',
-      license='MIT',
-      long_description=long_description,
-      long_description_content_type='text/markdown',
-      packages = [
-        'tinygrad',
-        'tinygrad.apps',
-        'tinygrad.codegen',
-        'tinygrad.codegen.opt',
-        'tinygrad.codegen.late',
-        'tinygrad.engine',
-        'tinygrad.mixin',
-        'tinygrad.nn',
-        'tinygrad.renderer',
-        'tinygrad.runtime',
-        'tinygrad.runtime.autogen',
-        'tinygrad.runtime.autogen.am',
-        'tinygrad.runtime.autogen.nv',
-        'tinygrad.runtime.graph',
-        'tinygrad.runtime.support',
-        'tinygrad.runtime.support.am',
-        'tinygrad.runtime.support.nv',
-        'tinygrad.schedule',
-        'tinygrad.uop',
-        'tinygrad.viz',
-      ],
-      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
-      classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: MIT License"
-      ],
-      install_requires=[],
-      python_requires='>=3.11',
-      extras_require={
-        'arm': ["unicorn"],
-        'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
-        'linting': [
-            "pylint",
-            "mypy==1.18.1",
-            "typing-extensions",
-            "pre-commit",
-            "ruff",
-            "numpy",
-            "typeguard",
-        ],
-        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/[email protected]"],
-        'testing_minimal': testing_minimal,
-        'testing_unit': testing_minimal + [
-            "tqdm",
-            "safetensors",
-            "tabulate",  # for sz.py
-        ],
-        'testing': testing_minimal + [
-            "pillow",
-            "onnx==1.18.0",
-            "onnx2torch",
-            "onnxruntime",
-            "opencv-python",
-            "tabulate",
-            "tqdm",
-            "safetensors",
-            "transformers",
-            "sentencepiece",
-            "tiktoken",
-            "blobfile",
-            "librosa",
-            "numba>=0.55",  # librosa needs numba but uv ignores python upper bounds and some numba versions require <python3.10
-            "networkx",
-            "nibabel",
-            "bottle",
-            "ggml-python",
-            "capstone",
-            "pycocotools",
-            "boto3",
-            "pandas",
-            "influxdb3-python"
-        ],
-        'docs': [
-            "mkdocs",
-            "mkdocs-material",
-            "mkdocstrings[python]",
-            "markdown-callouts",
-            "markdown-exec[ansi]",
-            "black",
-            "numpy",
-        ],
-      },
-      include_package_data=True)
@@ -1,237 +1,90 @@
-# ruff: noqa: E501
+# ruff: noqa: E501 E712 F401
 from tinygrad import dtypes, Device
-from tinygrad.uop.ops import UOp, AxisType, Ops
+from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen import full_rewrite
+from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
 from tinygrad.renderer import ProgramSpec
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import dedup
+from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
-from tinygrad.dtype import ImageDType
+from tinygrad.dtype import ImageDType, Invalid

-# PYTHONPATH="." DEBUG=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-# kernel 672
-# faster on d59d4cd, 50% slower with the new linearizer
+# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

-""" d59d4cd
-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=0, src=())
-c1 = UOp.range(UOp.const(dtypes.index, 64), 3, AxisType.LOOP)
-c2 = UOp.range(UOp.const(dtypes.index, 64), 4, AxisType.LOOP)
-c3 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
-c4 = (((c1*UOp.const(dtypes.index, 64))+c2)+(c3*UOp.const(dtypes.index, 4096)))
-c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), arg=1, src=())
-c6 = c5.index(c4).load()
-c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), arg=2, src=())
-c8 = UOp.range(UOp.const(dtypes.index, 48), 0, AxisType.REDUCE)
-c9 = UOp.range(UOp.const(dtypes.index, 4), 1, AxisType.REDUCE)
-c10 = c7.index(((((c8*UOp.const(dtypes.index, 4))+c9)+(c1*UOp.const(dtypes.index, 192)))+(c3*UOp.const(dtypes.index, 12288)))).load()
-c11 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), arg=3, src=())
-c12 = c11.index(((((c9*UOp.const(dtypes.index, 4))+(c2%UOp.const(dtypes.index, 4)))+(c8*UOp.const(dtypes.index, 16)))+((c2//UOp.const(dtypes.index, 4))*UOp.const(dtypes.index, 768)))).load()
-c13 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=4, src=())
-c14 = c13.index(c2).load()
-c15 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), arg=5, src=())
-c16 = c15.index(c2).load()
-c17 = (c6+(((c10*c12.cast(dtypes.float)).cast(dtypes.float).reduce(c8, c9, arg=Ops.ADD)+c14.cast(dtypes.float))*c16.cast(dtypes.float)))
-c18 = c0.index(c4).store(c17, c3, c1, c2)
-ast = c18.sink()
-more upcast axis : [(3, 320, 0, 4)]
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
-const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
-  float acc0[16];
-  int idx0 = get_global_id(0); /* 16 */
-  int idx1 = get_global_id(1); /* 512 */
-  int alu0 = (idx1>>4);
-  *(acc0+0) = 0.0f;
-  *(acc0+1) = 0.0f;
-  *(acc0+2) = 0.0f;
-  *(acc0+3) = 0.0f;
-  *(acc0+4) = 0.0f;
-  *(acc0+5) = 0.0f;
-  *(acc0+6) = 0.0f;
-  *(acc0+7) = 0.0f;
-  *(acc0+8) = 0.0f;
-  *(acc0+9) = 0.0f;
-  *(acc0+10) = 0.0f;
-  *(acc0+11) = 0.0f;
-  *(acc0+12) = 0.0f;
-  *(acc0+13) = 0.0f;
-  *(acc0+14) = 0.0f;
-  *(acc0+15) = 0.0f;
-  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
-    int alu17 = ((idx1*192)+Ridx0);
-    int alu18 = (alu17+48);
-    int alu19 = (alu17+96);
-    int alu20 = (alu17+144);
-    int alu21 = (Ridx0<<2);
-    float4 val0 = read_imagef(data3_12288, smp, (int2)(alu21,idx0));
-    float4 val1 = read_imagef(data3_12288, smp, (int2)((alu21+1),idx0));
-    float4 val2 = read_imagef(data3_12288, smp, (int2)((alu21+2),idx0));
-    float4 val3 = read_imagef(data3_12288, smp, (int2)((alu21+3),idx0));
-    float4 val4 = read_imagef(data2_393216, smp, (int2)((alu18-(3072*(((alu18>>10)*43)>>7))),alu0));
-    float4 val5 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu0));
-    float4 val6 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu0));
-    float4 val7 = read_imagef(data2_393216, smp, (int2)((alu17-(3072*(((alu17>>10)*43)>>7))),alu0));
-    *(acc0+1) = ((*(acc0+1))+(val4.x*val0.x)+(val4.y*val1.x)+(val4.z*val2.x)+(val4.w*val3.x));
-    *(acc0+5) = ((*(acc0+5))+(val4.x*val0.y)+(val4.y*val1.y)+(val4.z*val2.y)+(val4.w*val3.y));
-    *(acc0+9) = ((*(acc0+9))+(val4.x*val0.z)+(val4.y*val1.z)+(val4.z*val2.z)+(val4.w*val3.z));
-    *(acc0+13) = ((*(acc0+13))+(val4.x*val0.w)+(val4.y*val1.w)+(val4.z*val2.w)+(val4.w*val3.w));
-    *(acc0+2) = ((*(acc0+2))+(val5.x*val0.x)+(val5.y*val1.x)+(val5.z*val2.x)+(val5.w*val3.x));
-    *(acc0+6) = ((*(acc0+6))+(val5.x*val0.y)+(val5.y*val1.y)+(val5.z*val2.y)+(val5.w*val3.y));
-    *(acc0+10) = ((*(acc0+10))+(val5.x*val0.z)+(val5.y*val1.z)+(val5.z*val2.z)+(val5.w*val3.z));
-    *(acc0+14) = ((*(acc0+14))+(val5.x*val0.w)+(val5.y*val1.w)+(val5.z*val2.w)+(val5.w*val3.w));
-    *(acc0+3) = ((*(acc0+3))+(val6.x*val0.x)+(val6.y*val1.x)+(val6.z*val2.x)+(val6.w*val3.x));
-    *(acc0+7) = ((*(acc0+7))+(val6.x*val0.y)+(val6.y*val1.y)+(val6.z*val2.y)+(val6.w*val3.y));
-    *(acc0+11) = ((*(acc0+11))+(val6.x*val0.z)+(val6.y*val1.z)+(val6.z*val2.z)+(val6.w*val3.z));
-    *(acc0+15) = ((*(acc0+15))+(val6.x*val0.w)+(val6.y*val1.w)+(val6.z*val2.w)+(val6.w*val3.w));
-    *(acc0+0) = ((*(acc0+0))+(val7.x*val0.x)+(val7.y*val1.x)+(val7.z*val2.x)+(val7.w*val3.x));
-    *(acc0+4) = ((*(acc0+4))+(val7.x*val0.y)+(val7.y*val1.y)+(val7.z*val2.y)+(val7.w*val3.y));
-    *(acc0+8) = ((*(acc0+8))+(val7.x*val0.z)+(val7.y*val1.z)+(val7.z*val2.z)+(val7.w*val3.z));
-    *(acc0+12) = ((*(acc0+12))+(val7.x*val0.w)+(val7.y*val1.w)+(val7.z*val2.w)+(val7.w*val3.w));
-  }
-  int alu39 = (idx0<<2);
-  half4 val8 = (*((__global half4*)((data4_64+alu39))));
-  half4 val9 = (*((__global half4*)((data5_64+alu39))));
-  int alu40 = (idx0+(idx1<<6));
-  int2 cast0 = (int2)((alu40&1023),alu0);
-  float4 val10 = read_imagef(data1_131072, smp, cast0);
-  int2 cast1 = (int2)(((alu40+16)&1023),alu0);
-  float4 val11 = read_imagef(data1_131072, smp, cast1);
-  int2 cast2 = (int2)(((alu40+32)&1023),alu0);
-  float4 val12 = read_imagef(data1_131072, smp, cast2);
-  int2 cast3 = (int2)(((alu40+48)&1023),alu0);
-  float4 val13 = read_imagef(data1_131072, smp, cast3);
-  float cast4 = ((float)(val8.x));
-  float cast5 = ((float)(val9.x));
-  float cast6 = ((float)(val8.y));
-  float cast7 = ((float)(val9.y));
-  float cast8 = ((float)(val8.z));
-  float cast9 = ((float)(val9.z));
-  float cast10 = ((float)(val8.w));
-  float cast11 = ((float)(val9.w));
-  write_imagef(data0_131072, cast0, (float4)((val10.x+(((*(acc0+0))+cast4)*cast5)),(val10.y+(((*(acc0+4))+cast6)*cast7)),(val10.z+(((*(acc0+8))+cast8)*cast9)),(val10.w+(((*(acc0+12))+cast10)*cast11))));
-  write_imagef(data0_131072, cast1, (float4)((val11.x+(((*(acc0+1))+cast4)*cast5)),(val11.y+(((*(acc0+5))+cast6)*cast7)),(val11.z+(((*(acc0+9))+cast8)*cast9)),(val11.w+(((*(acc0+13))+cast10)*cast11))));
-  write_imagef(data0_131072, cast2, (float4)((val12.x+(((*(acc0+2))+cast4)*cast5)),(val12.y+(((*(acc0+6))+cast6)*cast7)),(val12.z+(((*(acc0+10))+cast8)*cast9)),(val12.w+(((*(acc0+14))+cast10)*cast11))));
-  write_imagef(data0_131072, cast3, (float4)((val13.x+(((*(acc0+3))+cast4)*cast5)),(val13.y+(((*(acc0+7))+cast6)*cast7)),(val13.z+(((*(acc0+11))+cast8)*cast9)),(val13.w+(((*(acc0+15))+cast10)*cast11))));
-}
-*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    322.55us/    77.83ms (    157 GFLOPS    4|160    GB/s) ['mul', '__add__', 'conv2d']
-"""
+def vision_conv_143():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 0)
+  c2 = UOp.range(32, 3, AxisType.LOOP)
+  c5 = UOp.range(128, 4, AxisType.LOOP)
+  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<35)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<67)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((64, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(128), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)

-""" master 99e76f33a0f4ec84c79c1271dbc955fe6b5a7778
-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
-c2 = UOp.range(64, 3, AxisType.LOOP)
-c4 = UOp.range(64, 4, AxisType.LOOP)
-c7 = UOp.range(32, 2, AxisType.LOOP)
-c10 = (((c2*64)+c4)+(c7*4096))
-c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
-c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
-c16 = UOp.range(48, 0, AxisType.REDUCE)
-c19 = UOp.range(4, 1, AxisType.REDUCE)
-c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
-c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
-c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
-c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
-c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
-c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
-ast = c52.sink()
-more upcast axis : [(3, 320, 0, 4)]
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void r_512_16_4_4_48_4(write_only image2d_t data0_131072, read_only image2d_t data1_131072, read_only image2d_t data2_393216, read_only image2d_t data3_12288, __global half* data4_64, __global half* data5_64) {
-const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
-  float acc0[16];
-  int idx0 = get_global_id(0); /* 16 */
-  int idx1 = get_global_id(1); /* 512 */
-  *(acc0+0) = 0.0f;
-  *(acc0+1) = 0.0f;
-  *(acc0+2) = 0.0f;
-  *(acc0+3) = 0.0f;
-  *(acc0+4) = 0.0f;
-  *(acc0+5) = 0.0f;
-  *(acc0+6) = 0.0f;
-  *(acc0+7) = 0.0f;
-  *(acc0+8) = 0.0f;
-  *(acc0+9) = 0.0f;
-  *(acc0+10) = 0.0f;
-  *(acc0+11) = 0.0f;
-  *(acc0+12) = 0.0f;
-  *(acc0+13) = 0.0f;
-  *(acc0+14) = 0.0f;
-  *(acc0+15) = 0.0f;
-  int alu16 = (idx0<<2);
-  half4 val0 = (*((__global half4*)((data4_64+alu16))));
-  half4 val1 = (*((__global half4*)((data5_64+alu16))));
-  int alu17 = (idx0+(idx1<<6));
-  int alu18 = (idx1>>4);
-  int2 cast0 = (int2)((alu17&1023),alu18);
-  float4 val2 = read_imagef(data1_131072, smp, cast0);
-  int2 cast1 = (int2)(((alu17+16)&1023),alu18);
-  float4 val3 = read_imagef(data1_131072, smp, cast1);
-  int2 cast2 = (int2)(((alu17+32)&1023),alu18);
-  float4 val4 = read_imagef(data1_131072, smp, cast2);
-  int2 cast3 = (int2)(((alu17+48)&1023),alu18);
-  float4 val5 = read_imagef(data1_131072, smp, cast3);
-  for (int Ridx0 = 0; Ridx0 < 48; Ridx0++) {
-    int alu19 = ((idx1*192)+Ridx0);
-    int alu20 = (alu19+48);
-    int alu21 = (alu19+96);
-    int alu22 = (alu19+144);
-    int alu23 = (Ridx0<<2);
-    float4 val6 = read_imagef(data3_12288, smp, (int2)(alu23,idx0));
-    float4 val7 = read_imagef(data3_12288, smp, (int2)((alu23+1),idx0));
-    float4 val8 = read_imagef(data3_12288, smp, (int2)((alu23+2),idx0));
-    float4 val9 = read_imagef(data3_12288, smp, (int2)((alu23+3),idx0));
-    float4 val10 = read_imagef(data2_393216, smp, (int2)((alu20-(3072*(((alu20>>10)*43)>>7))),alu18));
-    *(acc0+1) = ((*(acc0+1))+(val10.x*val6.x)+(val10.y*val7.x)+(val10.z*val8.x)+(val10.w*val9.x));
-    *(acc0+5) = ((*(acc0+5))+(val10.x*val6.y)+(val10.y*val7.y)+(val10.z*val8.y)+(val10.w*val9.y));
-    *(acc0+9) = ((*(acc0+9))+(val10.x*val6.z)+(val10.y*val7.z)+(val10.z*val8.z)+(val10.w*val9.z));
-    *(acc0+13) = ((*(acc0+13))+(val10.x*val6.w)+(val10.y*val7.w)+(val10.z*val8.w)+(val10.w*val9.w));
-    float4 val11 = read_imagef(data2_393216, smp, (int2)((alu21-(3072*(((alu21>>10)*43)>>7))),alu18));
-    *(acc0+2) = ((*(acc0+2))+(val11.x*val6.x)+(val11.y*val7.x)+(val11.z*val8.x)+(val11.w*val9.x));
-    *(acc0+6) = ((*(acc0+6))+(val11.x*val6.y)+(val11.y*val7.y)+(val11.z*val8.y)+(val11.w*val9.y));
-    *(acc0+10) = ((*(acc0+10))+(val11.x*val6.z)+(val11.y*val7.z)+(val11.z*val8.z)+(val11.w*val9.z));
-    *(acc0+14) = ((*(acc0+14))+(val11.x*val6.w)+(val11.y*val7.w)+(val11.z*val8.w)+(val11.w*val9.w));
-    float4 val12 = read_imagef(data2_393216, smp, (int2)((alu22-(3072*(((alu22>>10)*43)>>7))),alu18));
-    *(acc0+3) = ((*(acc0+3))+(val12.x*val6.x)+(val12.y*val7.x)+(val12.z*val8.x)+(val12.w*val9.x));
-    *(acc0+7) = ((*(acc0+7))+(val12.x*val6.y)+(val12.y*val7.y)+(val12.z*val8.y)+(val12.w*val9.y));
-    *(acc0+11) = ((*(acc0+11))+(val12.x*val6.z)+(val12.y*val7.z)+(val12.z*val8.z)+(val12.w*val9.z));
-    *(acc0+15) = ((*(acc0+15))+(val12.x*val6.w)+(val12.y*val7.w)+(val12.z*val8.w)+(val12.w*val9.w));
-    float4 val13 = read_imagef(data2_393216, smp, (int2)((alu19-(3072*(((alu19>>10)*43)>>7))),alu18));
-    *(acc0+0) = ((*(acc0+0))+(val13.x*val6.x)+(val13.y*val7.x)+(val13.z*val8.x)+(val13.w*val9.x));
-    *(acc0+4) = ((*(acc0+4))+(val13.x*val6.y)+(val13.y*val7.y)+(val13.z*val8.y)+(val13.w*val9.y));
-    *(acc0+8) = ((*(acc0+8))+(val13.x*val6.z)+(val13.y*val7.z)+(val13.z*val8.z)+(val13.w*val9.z));
-    *(acc0+12) = ((*(acc0+12))+(val13.x*val6.w)+(val13.y*val7.w)+(val13.z*val8.w)+(val13.w*val9.w));
-  }
-  float cast4 = ((float)(val0.x));
-  float cast5 = ((float)(val1.x));
-  float cast6 = ((float)(val0.y));
-  float cast7 = ((float)(val1.y));
-  float cast8 = ((float)(val0.z));
-  float cast9 = ((float)(val1.z));
-  float cast10 = ((float)(val0.w));
-  float cast11 = ((float)(val1.w));
-  write_imagef(data0_131072, cast0, (float4)((val2.x+(((*(acc0+0))+cast4)*cast5)),(val2.y+(((*(acc0+4))+cast6)*cast7)),(val2.z+(((*(acc0+8))+cast8)*cast9)),(val2.w+(((*(acc0+12))+cast10)*cast11))));
-  write_imagef(data0_131072, cast1, (float4)((val3.x+(((*(acc0+1))+cast4)*cast5)),(val3.y+(((*(acc0+5))+cast6)*cast7)),(val3.z+(((*(acc0+9))+cast8)*cast9)),(val3.w+(((*(acc0+13))+cast10)*cast11))));
-  write_imagef(data0_131072, cast2, (float4)((val4.x+(((*(acc0+2))+cast4)*cast5)),(val4.y+(((*(acc0+6))+cast6)*cast7)),(val4.z+(((*(acc0+10))+cast8)*cast9)),(val4.w+(((*(acc0+14))+cast10)*cast11))));
-  write_imagef(data0_131072, cast3, (float4)((val5.x+(((*(acc0+3))+cast4)*cast5)),(val5.y+(((*(acc0+7))+cast6)*cast7)),(val5.z+(((*(acc0+11))+cast8)*cast9)),(val5.w+(((*(acc0+15))+cast10)*cast11))));
-}
-*** QCOM     672 r_512_16_4_4_48_4                              arg  6 mem   0.10 GB tm    527.97us/    78.94ms (     96 GFLOPS    3|98     GB/s) ['conv2d', 'mul', '__add__']
-"""
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

-c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 0)
-c2 = UOp.range(64, 3, AxisType.LOOP)
-c4 = UOp.range(64, 4, AxisType.LOOP)
-c7 = UOp.range(32, 2, AxisType.LOOP)
-c10 = (((c2*64)+c4)+(c7*4096))
-c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 1024, 4)), (), 1)
-c14 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((32, 3072, 4)), (), 2)
-c16 = UOp.range(48, 0, AxisType.REDUCE)
-c19 = UOp.range(4, 1, AxisType.REDUCE)
-c28 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 192, 4)), (), 3)
-c40 = (c14.index(((((c16*4)+c19)+(c2*192))+(c7*12288)))*c28.index(((((c19*4)+(c4%4))+(c16*16))+((c4//4)*768))))
-c42 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 4)
-c46 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(64), (), 5)
-c50 = (c12.index(c10)+((c40.reduce(c16, c19, arg=Ops.ADD)+c42.index(c4).cast(dtypes.float))*c46.index(c4).cast(dtypes.float)))
-c52 = c0.index(c10, ptr=True).store(c50).end(c7, c2, c4)
-ast = c52.sink()
+def vision_conv_153():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 1024, 4)), (), 0)
+  c2 = UOp.range(16, 3, AxisType.LOOP)
+  c5 = UOp.range(256, 4, AxisType.LOOP)
+  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c16 = UOp.range(7, 0, AxisType.REDUCE)
+  c17 = c8*2+c16
+  c24 = ((c17<3)!=True)&(c17<19)
+  c26 = UOp.range(7, 1, AxisType.REDUCE)
+  c27 = c2*2+c26
+  c32 = ((c27<3)!=True)&(c27<35)
+  c34 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((16, 1024, 4)), (), 1)
+  c38 = c5//2
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
+  c49 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((128, 49, 4)), (), 2)
+  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
+  c63 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(256), (), 3)
+  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
+  c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
+
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
+  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+def dm_conv_172():
+  c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((1, 240, 4)), (), 0)
+  c2 = UOp.range(960, 4, AxisType.LOOP)
+  c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((8, 384, 4)), (), 1)
+  c7 = UOp.range(32, 0, AxisType.REDUCE)
+  c10 = UOp.range(4, 1, AxisType.REDUCE)
+  c13 = UOp.range(12, 3, AxisType.REDUCE)
+  c18 = UOp.range(8, 2, AxisType.REDUCE)
+  c23 = UOp(Ops.DEFINE_GLOBAL, dtypes.imageh((240, 128, 4)), (), 2)
+  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
+  c37 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(960), (), 3)
+  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
+  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
+  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
+  c55 = c0.index(c2, ptr=True).store(c53).end(c2)
+
+  opts = None
+  # JITBEAM=2
+  # (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
+  return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
+
+ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()

 compiler = Device.default.compiler
 renderer = Device.default.renderer
@@ -240,20 +93,13 @@ allocator = Device.default.allocator
 uops = full_rewrite(ast, renderer)
 src = renderer.render(uops)

-# NOLOCALS=1 IMAGE=2 DEV=CL
 lib = compiler.compile(src)
-# r_64_8_16_4_4_48_4
-# NOLOCALS: r_512_16_4_4_48_4
-ps = ProgramSpec("r_512_16_4_4_48_4", src, Device.DEFAULT, ast, uops)
-print(ps.src)
-print(ps.applied_opts)
-# (Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None))
+ps = ProgramSpec("conv", src, Device.DEFAULT, ast, uops)
 cr = CompiledRunner(ps, precompiled=lib)

 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
-print(len(gs))
-print([g.dtype for g in gs])
-
+# print(len(gs))
+# print([g.dtype for g in gs])
 bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]

 t = cr(bufs, wait=True)
@@ -112,7 +112,6 @@ backend_test.exclude('test_dequantizelinear_e5m2_cpu')
 backend_test.exclude('test_dequantizelinear_float4e2m1_cpu')

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

 # no support for int pow
 backend_test.exclude('test_pow_types_int32_int32_cpu')
@@ -9,6 +9,11 @@ def custom_arange_kernel(C:UOp) -> UOp:
  i = UOp.range(C.size, 0)
  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))

+def custom_eye_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.shape[0], 0)
+  j = UOp.range(C.shape[1], 1)
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
+
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
  A,B = A.flatten(), B.flatten()
  assert B.size == A.size
@@ -125,6 +130,12 @@ class TestCustomKernel(unittest.TestCase):
    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
    self.assertTrue((ref == tst).all().item())

+  def test_eye(self):
+    ref = Tensor.eye(1024).contiguous().realize()
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
  def test_flip_contract(self):
    a = Tensor.randn(10,4)
    b = Tensor.empty_like(a)
@@ -4,7 +4,7 @@ from dataclasses import replace

 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.gpudims import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops, GroupOp
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, PatternMatcher, graph_rewrite, UPat
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
@@ -38,6 +38,22 @@ class TestLinearizer(unittest.TestCase):
    np.testing.assert_equal(a.numpy(), ta)
    np.testing.assert_equal(b.numpy(), tb)

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

+  @unittest.expectedFailure # TODO: investigate
  def test_two_nested_range_alt_indexing(self):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
@@ -262,6 +279,8 @@ class TestLinearizer(unittest.TestCase):
    _assert_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
    #                              2             ->             3
    _assert_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
+    #                              2             ->             2
+    _assert_grouped_dims("gidx", (65536,2), (65535,65535,65535), False, [32768,4], False)
    # test when the only divisor is the square root of dim
    _assert_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)

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

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

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
@@ -300,6 +300,12 @@ class TestRangeify(unittest.TestCase):
    w2 = Tensor.empty(12, 8, 3, 3)
    x.conv2d(w1).conv2d(w2).realize()

+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    x.conv2d(w1).conv2d(w2).realize()
+
  def test_xception_conv2d(self):
    # NOTE: this fusion is bad, it's recomputing the inner many times
    x = Tensor.empty(1, 4, 32, 32)
@@ -1573,6 +1573,13 @@ class TestSchedule(unittest.TestCase):
  def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)
  def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4)

+  def test_resnet_conv2d(self):
+    x = Tensor.empty(1, 8, 32, 32)
+    w1 = Tensor.empty(8, 8, 3, 3)
+    w2 = Tensor.empty(8, 8, 1, 1)
+    out = x.conv2d(w1).conv2d(w2)
+    check_schedule(out, 2)
+
  @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
  def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half)
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -3,7 +3,7 @@ import torch
 import unittest, copy, mmap, random, math, array
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _METADATA
-from tinygrad.helpers import getenv, temp, mv_address
+from tinygrad.helpers import Context, getenv, temp, mv_address
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
 from tinygrad.device import is_dtype_supported
@@ -846,6 +846,16 @@ class TestTensorMetadata(unittest.TestCase):
    #self.assertEqual(len(bw), 1)
    #self.assertEqual(bw[0].name, "sigmoid")

+  def test_tracemeta_0(self):
+    with Context(TRACEMETA=0):
+      x = Tensor.rand(3, requires_grad=True)
+      y = Tensor.rand(3, requires_grad=True)
+      out = (x.relu() * y.sigmoid()).sum()
+      self.assertIsNone(out.uop.metadata)
+      self.assertIsNone(out.uop.src[0].metadata)
+      si = out.schedule()[-1]
+      self.assertEqual(si.metadata, ())
+
 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
    if ast.op is op: return ast
@@ -0,0 +1,405 @@
+import unittest, math
+
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.helpers import CI
+from tinygrad.renderer.ptx import PTXRenderer
+import numpy as np
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.kernel import Kernel
+
+@unittest.skipUnless(Device.DEFAULT in ["CUDA", "NV"], "only cuda")
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
+class TestTK(unittest.TestCase):
+  @unittest.skipIf(CI, "no wmma in ci")
+  def test_simple_matmul(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.bfloat16)
+      b = ker.gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem, transpose=True)
+
+        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b, dtype=dtypes.float32).float()
+
+    np.testing.assert_allclose(c.numpy(), ref.numpy())
+
+  @unittest.skipIf(CI, "no wmma in ci")
+  def test_simple_matmul_transposed(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      c = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.bfloat16)
+      b = ker.gl((1, 1, N, N), dtypes.bfloat16)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      c_reg = warp.zero(c_reg)
+      for tile in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
+        b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
+
+        a_reg = warp.load(a_reg, a_smem)
+        b_reg = warp.load(b_reg, b_smem)
+
+        c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg)
+      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (c, a, b)])
+    for _ in range(5): ei.run(wait=True)
+    c = c.float()
+
+    ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float()
+
+    np.testing.assert_allclose(c.numpy(), ref.numpy())
+
+  def test_load_store(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
+      a_reg = warp.load(a_reg, a_smem)
+      b_reg = warp.copy(b_reg, a_reg)
+      b_smem = warp.store(b_smem, b_reg)
+      b = warp.store(b, b_smem, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float()
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+
+  def test_max(self):
+    N = 16
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      max_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        max_reg = warp.neg_inf(max_reg.after(tile_row))
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        max_reg = ker.endrange()
+
+        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+
+  def test_max_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, M), dtypes.float32)
+      a = ker.gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      max_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_N):
+        max_reg = warp.neg_inf(max_reg.after(tile_row))
+
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+        max_reg = ker.endrange()
+
+        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+
+  def test_sum(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_SIZE):
+        sum_reg = warp.zero(sum_reg.after(tile_row))
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
+
+  def test_sum_nonsquare(self):
+    N, M = 16, 64
+    BLOCK_N, BLOCK_M = 16, 64
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, M), dtypes.float32)
+      a = ker.gl((1, 1, N, M), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+
+      sum_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
+
+      for tile_row in ker.range(N // BLOCK_N):
+        sum_reg = warp.zero(sum_reg.after(tile_row))
+
+        for tile_col in ker.range(M // BLOCK_M):
+          a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
+          a_reg = warp.load(a_reg, a_smem)
+          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+        sum_reg = ker.endrange()
+
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
+        b_smem = warp.store(b_smem, b_reg)
+
+        for tile_col in ker.range(M // BLOCK_M):
+          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, M, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
+
+  @unittest.skip("fake range not ended")
+  def test_softmax(self):
+    N = 32
+    BLOCK_SIZE = 16
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
+      a = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+      max_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+      norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+
+      max_vec = warp.neg_inf(max_vec)
+      norm_vec = warp.zero(norm_vec)
+
+      for tile_col in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg = warp.load(a_reg, a_smem)
+
+        a_reg = warp.map(a_reg, lambda x: x * (1.0 / math.log(2)))
+
+        max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
+        max_vec = warp.row_reduce(max_vec, a_reg, lambda a, b: a.maximum(b))
+        a_reg = warp.map(a_reg, lambda x, idx: (x - max_vec[idx[0], 0, (idx[2]%4)//2]).exp2())
+        max_vec_last = warp.map(max_vec_last, lambda x, idx: (x - max_vec[*idx]).exp2())
+        norm_vec = warp.map(norm_vec, lambda x, idx: x * max_vec_last[*idx])
+        norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
+      norm_vec = ker.endrange()
+
+      for tile_col in ker.range(N // BLOCK_SIZE):
+        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg = warp.load(a_reg, a_smem)
+
+        a_reg = warp.map(a_reg, lambda x: x * (1.0 / math.log(2)))
+        a_reg = warp.map(a_reg, lambda x, idx: (x - max_vec[idx[0], 0, (idx[2]%4)//2]).exp2())
+        a_reg = warp.map(a_reg, lambda x, idx: x / norm_vec[idx[0], 0, (idx[2]%4)//2])
+
+        a_smem = warp.store(a_smem, a_reg)
+        b = warp.store(b, a_smem, (0, 0, 0, tile_col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, BLOCK_SIZE, N, dtype="float32")
+      b = Tensor.empty(1, 1, BLOCK_SIZE, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().softmax(axis=3)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -0,0 +1,47 @@
+import ctypes, subprocess, tempfile, unittest
+from tinygrad.helpers import WIN
+from tinygrad.runtime.support.c import Struct
+
+class TestAutogen(unittest.TestCase):
+  def test_packed_struct_sizeof(self):
+    layout = [('a', ctypes.c_char), ('b', ctypes.c_int, 5), ('c', ctypes.c_char)]
+    class X(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
+    class Y(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
+    class Z(Struct): _packed_, _fields_ = True, layout
+    self.assertNotEqual(ctypes.sizeof(X), 4) # ctypes bug! gcc-13.3.0 says this should have size 4
+    self.assertEqual(ctypes.sizeof(Y), 6)
+    self.assertEqual(ctypes.sizeof(Z), 3)
+    layout = [('a', ctypes.c_int, 31), ('b', ctypes.c_int, 31), ('c', ctypes.c_int, 1), ('d', ctypes.c_int, 1)]
+    class Foo(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
+    class Bar(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
+    class Baz(Struct): _fields_, _packed_ = layout, True
+    self.assertEqual(ctypes.sizeof(Foo), 12)
+    self.assertEqual(ctypes.sizeof(Bar), 12)
+    self.assertEqual(ctypes.sizeof(Baz), 8)
+
+  @unittest.skipIf(WIN, "doesn't compile on windows")
+  def test_packed_struct_interop(self):
+    class Baz(Struct): pass
+    Baz._packed_ = True
+    Baz._fields_ = [('a', ctypes.c_int, 30), ('b', ctypes.c_int, 30), ('c', ctypes.c_int, 2), ('d', ctypes.c_int, 2)]
+    src = '''
+      struct __attribute__((packed)) baz {
+        int a:30;
+        int b:30;
+        int c:2;
+        int d:2;
+      };
+
+      int test(struct baz x) {
+        return x.a + x.b + x.c + x.d;
+      }
+    '''
+    args = ('-x', 'c', '-fPIC', '-shared')
+    with tempfile.NamedTemporaryFile(suffix=".so") as f:
+      subprocess.check_output(('clang',) + args + ('-', '-o', f.name), input=src.encode('utf-8'))
+      b = Baz(0xAA000, 0x00BB0, 0, 1)
+      test = ctypes.CDLL(f.name).test
+      test.argtypes = [Baz]
+      self.assertEqual(test(b), b.a + b.b + b.c + b.d)
+
+if __name__ == "__main__": unittest.main()
@@ -62,6 +62,7 @@ class TestConv(unittest.TestCase):
    np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
    np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)

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

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

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

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

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

@@ -211,5 +211,39 @@ class TestPatternMatcher(unittest.TestCase):
      return u.src[0]
    for a,b in zip(simple_src(a), simple_src(b)): self._assert_eq_upat(a, b)

+class TestAlgebraic(unittest.TestCase):
+  def test_plus_0(self):
+    pm = PatternMatcher([
+      (UPat.var("x") + 0, UPat.var("x")), # x+0 -> x
+    ])
+    expr = UOp.const(dtypes.int, 4)+0
+    print(expr)
+    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4))
+
+  def test_div_mul(self):
+    pm = PatternMatcher([
+      ((UPat.var("x") * UPat.var("x2")) / UPat.var("x2"), UPat.var("x")), # (x*x2)/x2 -> x
+    ])
+    expr = UOp.const(dtypes.float, 4)/2*2
+    print(expr)
+    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4))
+
+  def test_mul_is_and(self):
+    pm = PatternMatcher([
+      (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), UPat.var('x') & UPat.var('y')),
+    ])
+    expr = UOp.const(dtypes.bool, True)*UOp.const(dtypes.bool, True)
+    print(expr)
+    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.bool, True)&UOp.const(dtypes.bool, True))
+
+  def test_div_neg_1(self):
+    pm = PatternMatcher([
+      (UPat.var("x") // -1, UPat.var("x") * -1), # x//-1 -> x * -1
+    ])
+    expr = UOp.const(dtypes.float, 4)//-1
+    print(expr)
+    self.assertEqual(pm.rewrite(expr), UOp.const(dtypes.int, 4) * -1)
+
+
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -643,6 +643,10 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
    self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")

+  def test_div_mod_recombine_partial(self):
+    gidx = Variable("gidx", 0, 15)
+    self.helper_test_variable((gidx//2)%4+(gidx//8)*4, 0, 7, "gidx//2")
+
  def test_div_mod_recombine_folded_mod(self):
    a = Variable("a", 0, 2)
    b = Variable("b", 0, 100)
@@ -366,8 +366,8 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIBB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

-class TestVizProfiler(unittest.TestCase):
-  def test_perfetto_node(self):
+class TestVizProfiler(BaseTestViz):
+  def test_node(self):
    prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=False),
            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]

@@ -381,7 +381,7 @@ class TestVizProfiler(unittest.TestCase):
    self.assertEqual(event['dur'], 10)
    assert event['ref'] is None

-  def test_perfetto_copy_node(self):
+  def test_copy_node(self):
    prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
@@ -399,7 +399,7 @@ class TestVizProfiler(unittest.TestCase):

    self.assertEqual(j["dur"], (event2["st"]+event2["dur"])-event["st"])

-  def test_perfetto_graph(self):
+  def test_graph(self):
    prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
            ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
            ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV', name='E_25_4n2', st_id=0, en_id=1, is_copy=False),
@@ -436,6 +436,12 @@ class TestVizProfiler(unittest.TestCase):
    sz = len(get_profile(prof))
    self.assertLessEqual(sz/n_events, 26)

+  def test_calltrace(self):
+    def fxn(): return Tensor.empty(10).mul(2).realize()
+    fxn()
+    trace = get_viz_list()[0]["steps"][0]["trace"]
+    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in trace), str(trace)
+
  # can pack up to 1hr 11 min of trace events
  def test_trace_duration(self):
    dur_mins = 72
@@ -136,7 +136,7 @@ def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
  """

  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
-  assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
  lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
  if SPEC: type_verify(lst, program_spec)
  return lst
@@ -26,28 +26,35 @@ def _split_dims(dims, max_sizes):
  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)

 def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
-  if reverse: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-  # try to split up dims: (a,) -> (b, c)
-  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
+  if max_sizes is None: limited = dims
+  else:
+    # try to group first: (a, b, c, d) -> (ab, c, d)
+    limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
+    # check if grouping failed
+    if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+    # try to split up dims: (a,) -> (b, c)
+    if limited == dims: limited = _split_dims(dims, max_sizes)
+  raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
  if len(limited) < len(dims):
    ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
+    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
    for idx, contraction_group in zip(raw_idxs, contraction):
      for c in contraction_group[:-1]:
        ret.append(idx % dims[c])
        idx //= dims[c]
      ret.append(idx)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
-  return ret[::-1] if reverse else ret
+    return ret
+  elif (a:=len(limited)) > (b:=len(dims)):
+    if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
+    if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
+    if a == 3 and b == 2: return [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+  elif limited != dims:
+    # Convert to 1D
+    flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(dims) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
+    # Get back original indices from 1D
+    return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
+  return raw_idxs

 def add_gpudims(ctx:Renderer, s:UOp):
  if s.arg is None: return None
@@ -1,48 +1,59 @@
 import heapq
+from typing import Any
 from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
-from tinygrad.helpers import prod
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
+from tinygrad.helpers import prod, getenv, TUPLE_ORDER

-def linearize(u:UOp) -> list[UOp]:
+def linearize(sink:UOp) -> list[UOp]:
  # this is a toposort with priority
-  lst = list(u.toposort())
+  lst = list(sink.toposort())
  consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
  in_degree:dict[UOp, int] = {}
-  priorities:dict[UOp, tuple[int, int]] = {}
+  out_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, tuple[int, int, Any]] = {}

  # get consumers and assign priorities
  # NOTE: this requires the lst be locally toposorted
  for u in reversed(lst):
    for s in u.src: consumers[s].append(u)
    in_degree[u] = len(u.src)
+    out_degree[u] = len(consumers[u])

    # we place UOps with higher run_counts later
-    # this will cause ranges to be placed late and ends to be placed early
    run_count = prod([int(r.vmax)+1 for r in u.ranges])

-    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
-    priority = [0] + [priorities[x][1] for x in consumers[u]]
-    if u.op is Ops.LOAD: priority.append(-1000)
-    if u.op is Ops.BARRIER: priority.append(-1500)
-    # ranges are scheduled as late as possible so anything that can be outside is
-    # if u.op is Ops.RANGE: priority = [2000]
-    if u.op is Ops.END: priority = [-1000]
-    # move defines and consts to the top
-    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
-    priorities[u] = (run_count, min(priority))
+    # simple priority override. this is all bottom up now, smaller numbers will be closer to the top
+    extra = None
+    match u.op:
+      # the order and placement of these defines is important
+      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
+      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
+      case Ops.DEFINE_LOCAL: priority = -18
+      case Ops.DEFINE_REG: priority = -17
+      case Ops.CONST: priority = -10  # early consts
+      case Ops.LOAD: priority = -1    # place loads early
+      case Ops.STORE: priority = 1    # place stores late
+      case Ops.RANGE: priority = 5    # placing RANGE is good
+      case Ops.END: priority = -5     # placing END is bad
+      case _: priority = 0            # everything else has priority 0
+    priorities[u] = (run_count, priority, extra)

  # number the uops in "ideal" order
-  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}

-  # then force then to be toposorted in as close to the ideal order as possible
-  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
+  # then force them to be toposorted in as close to the ideal order as possible
+  heap = [(-nkey[sink], sink)]
  newlst = []
  while heap:
    newlst.append(u:=heapq.heappop(heap)[1])
-    for v in consumers[u]:
-      in_degree[v] -= 1
-      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
-  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
+    for v in u.src:
+      out_degree[v] -= 1
+      if out_degree[v] == 0: heapq.heappush(heap, (-nkey[v],v))
+  newlst = newlst[::-1]
+
+  if getenv("DEBUG_LINEARIZE"):
+    for i,u in enumerate(newlst):
+      print(f"{i:4d} {str(u.op):20s} {multirange_str(u.ranges, color=True, pad=10)} {priorities[u]}")
  return newlst

 class CFGContext:
@@ -81,10 +81,10 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
            return k

  # are we grouping? (requires local shape support)
-  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
-    for sz in [16]:
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
+    for axis, sz in itertools.product((0, 1, 2), (16,)):
      try:
-        k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
+        k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
        break
      except KernelOptError: pass

@@ -107,7 +107,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # potentially do more upcasts of non reduce axes based on a heuristic
  is_dsp = k.ren is not None and k.ren.device == "DSP"
  upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
    xb_choices = []
    # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
    for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
@@ -135,8 +135,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # if last reduce dim is small(ish), loop unroll the reduce
  # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
  try:
-    upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
-    if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (upcast_size < 64):
+    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
        k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
        # if it's small, upcast a second reduce dimension too
@@ -105,6 +105,8 @@ class Scheduler:
  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]

+  def upcast_size(self) -> int: return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+
  # copied from kernel.py
  @property
  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
@@ -217,8 +219,7 @@ class Scheduler:
    return ret

  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
-    reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
-    if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
    reduceop = reduceops[0]
    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
      mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
@@ -312,9 +313,10 @@ class Scheduler:

  # helpers for hand_coded_optimizations
  @property
+  def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
+  @property
  def reduceop(self) -> UOp|None:
-    red = [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
-    if not len(red): return None
+    if not (red := self.reduceops): return None
    return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
  @property
  def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
@@ -34,7 +34,6 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
        # check if it simplifies
        if count_divmod(nidx) <= count_divmod(u):
          u = nidx
-          continue
  return u

 pm_simplify_ranges = PatternMatcher([
@@ -4,8 +4,8 @@ from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
-from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM
+from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer

@@ -266,7 +266,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
 class CompileError(Exception): pass

 class Compiler:
-  def __init__(self, cachekey:str|None=None): self.cachekey = None if DISABLE_COMPILER_CACHE else cachekey
+  def __init__(self, cachekey:str|None=None): self.cachekey = cachekey if CCACHE else None
  def compile(self, src:str) -> bytes: return src.encode()   # NOTE: empty compiler is the default
  def compile_cached(self, src:str) -> bytes:
    if self.cachekey is None or (lib := diskcache_get(self.cachekey, src)) is None:
@@ -291,8 +291,8 @@ class Compiled:
    if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
    for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)

-    try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers))
-    except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc
+    self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers,
+                                                       f"No compiler for {self.device} is available")

    if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")

@@ -300,10 +300,6 @@ class Compiled:
    compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
    return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"

-  def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]:
-    for renderer, compiler in compilers:
-      with contextlib.suppress(Exception): yield renderer(), compiler()
-
  def synchronize(self):
    """
    Synchronize all pending operations on the device.
@@ -343,6 +339,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
  # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
  if dtype == dtypes.half:
    if device == "CL": return not CI and not OSX
+    if device == "QCOM": return False # QCOM compiler is flaky with half
    if device in ["CUDA", "NV"]: return not CI
    if device == "CPU" and CPU_LLVM: return OSX
    if device == "PYTHON": return sys.version_info >= (3, 12)
@@ -3,7 +3,7 @@ import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.helpers import unwrap
+from tinygrad.helpers import unwrap, disable_gc
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -13,6 +13,7 @@ from tinygrad.codegen.opt import Opt

 # **************** Program Creation ****************

+@disable_gc()
@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
 def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
  """
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
+import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
@@ -114,6 +114,13 @@ def suppress_finalizing(func):
      if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
  return wrapper

+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T:
+  excs = []
+  for typ in candidates:
+    try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+    except Exception as e: excs.append(e)
+  raise ExceptionGroup(err_msg, excs)
+
 def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t

 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
@@ -166,19 +173,22 @@ TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
 PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
-DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
-CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 1)
+CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
 DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
+# set to 1, this uses tuplize in the linearizer sort order
+TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
+# set to 0 to disable the compiler cache
+CCACHE = ContextVar("CCACHE", 1)

@dataclass(frozen=True)
 class Metadata:
@@ -231,11 +241,29 @@ class Profiling(contextlib.ContextDecorator):

 def perf_counter_us() -> decimal.Decimal: return decimal.Decimal(time.perf_counter_ns())/1000

+@functools.cache
+def lines(fn) -> list[str]:
+  try:
+    with open(fn, encoding="utf-8") as f: return f.readlines()
+  except (FileNotFoundError, OSError): return []
+
+def printable(loc:tuple[str, int]) -> str:
+  try: return lines(loc[0])[loc[1]-1].strip()
+  except IndexError: return "<missing>"
+
+def get_stacktrace(frm, max_frames=30) -> tuple[tuple, ...]:
+  ret:list[tuple] = []
+  for i in range(max_frames):
+    if (frm:=frm.f_back) is None: break
+    ret.append(((fc:=frm.f_code).co_filename, frm.f_lineno, fc.co_name, printable((fc.co_filename, frm.f_lineno))))
+  return tuple(ret)
+
@dataclass(frozen=True)
 class TracingKey:
  display_name:str                       # display name of this trace event
  keys:tuple[Any, ...]=()                # optional keys to search for related traces
  ret:Any=None
+  tb:tuple[tuple, ...]|None=field(default_factory=lambda: get_stacktrace(sys._getframe(1)) if VIZ else None)

 class ProfileEvent: pass

@@ -352,10 +380,12 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip

 # *** Exec helpers

+def system(cmd, **kwargs): return subprocess.check_output(cmd.split(), **kwargs).decode().strip()
+
 def cpu_objdump(lib, objdump_tool='objdump'):
  with tempfile.NamedTemporaryFile(delete=True) as f:
    pathlib.Path(f.name).write_bytes(lib)
-    print(subprocess.check_output([objdump_tool, '-d', f.name]).decode('utf-8'))
+    print(system(f"{objdump_tool} -d {f.name}"))

 def capstone_flatdump(lib: bytes):
  try: import capstone
@@ -433,6 +463,13 @@ class tqdm(Generic[T]):
 class trange(tqdm):
  def __init__(self, n:int, **kwargs): super().__init__(iterable=range(n), total=n, **kwargs)

+class disable_gc(contextlib.ContextDecorator):
+  def __enter__(self):
+    self._was_enabled = gc.isenabled()
+    if self._was_enabled: gc.disable()
+  def __exit__(self, *exc):
+    if self._was_enabled: gc.enable()
+
 # *** universal support for code object pickling

 def _reconstruct_code(*args): return types.CodeType(*args)
@@ -2,16 +2,21 @@
 import functools
 from typing import TypeAlias, TYPE_CHECKING, Self
 from tinygrad.uop import Ops
-from tinygrad.helpers import prod, argfix, flatten, dedup
-if TYPE_CHECKING:
-  from tinygrad.uop.ops import UOp
-  sint:TypeAlias = UOp|int
+from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
+from tinygrad.uop.ops import resolve, smax
+if TYPE_CHECKING: from tinygrad.uop.ops import UOp
+sint: TypeAlias = "UOp | int"
+
+def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
+  # unsqueeze left to make every shape same length
+  max_dim = max(len(shape) for shape in shapes)
+  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)

 class MovementMixin:
  # required to implement
  def _mop(self, op:Ops, arg) -> Self: raise NotImplementedError
  @property
-  def shape(self) -> tuple["sint", ...]: raise NotImplementedError
+  def shape(self) -> tuple[sint, ...]: raise NotImplementedError

  # great functions you get!
  @property
@@ -26,7 +31,7 @@ class MovementMixin:
    """
    return len(self.shape)

-  def numel(self) -> "sint":
+  def numel(self) -> sint:
    """
    Returns the total number of elements in the tensor.

@@ -42,6 +47,33 @@ class MovementMixin:
    if not -max(1, total) <= dim <= max(1, total)-1: raise IndexError(f"{dim=} out of range {[-max(1, total), max(1, total)-1]}")
    return dim + total if dim < 0 else dim

+  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Self:
+    if self.shape == new_shape: return self
+    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
+    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
+    shape, _ = _align_left(self.shape, new_shape)
+    # for each dimension, check either dim is 1, or it does not change
+    if not all(s == ns or s == 1 for s,ns in zip(shape, new_shape)):
+      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
+    reshaped = self.reshape(shape)
+    ret = reshaped._mop(Ops.EXPAND, arg=new_shape)
+    return reshaped if ret.shape == reshaped.shape else ret
+
+  def expand(self, shape, *args) -> Self:
+    """
+    Returns a tensor that is expanded to the shape that is specified.
+    Expand can also increase the number of dimensions that a tensor has.
+
+    Passing a `-1` or `None` to a dimension means that its size will not be changed.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.expand(4, -1).numpy())
+    ```
+    """
+    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
+    return self._broadcast_to(new_shape)
+
  def reshape(self, shape, *args) -> Self:
    """
    Returns a tensor with the same data as the original tensor but with a different shape.
@@ -61,7 +93,7 @@ class MovementMixin:
    ret = self._mop(Ops.RESHAPE, arg=new_shape)
    return self if ret.shape == self.shape else ret

-  def shrink(self, arg:tuple[tuple["sint", "sint"]|None, ...]) -> Self:
+  def shrink(self, arg:tuple[tuple[sint, sint]|None, ...]) -> Self:
    """
    Returns a tensor that shrinks the each axis based on input arg.
    `arg` must have the same length as `self.ndim`.
@@ -124,6 +156,9 @@ class MovementMixin:

  # **** high level ****

+  def shrink_to(self, shape, *args) -> Self:
+    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
+
  def view(self, shape, *args) -> Self:
    """`.view` is an alias for `.reshape`."""
    return self.reshape(shape, *args)
@@ -257,3 +292,59 @@ class MovementMixin:
    for i, name in enumerate(lhs): assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
    t = t.permute([lhs.index(name) for name in rhs])
    return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0]<dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
+
+  # *** movement ops with expand ***
+
+  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Self:
+    """
+    Repeats elements of a tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat_interleave(2).numpy())
+    ```
+    """
+    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
+    shp = x.shape
+    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
+
+  def repeat(self, repeats, *args) -> Self:
+    """
+    Repeats tensor number of times along each dimension specified by `repeats`.
+    `repeats` can be passed as a tuple or as separate arguments.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([1, 2, 3])
+    print(t.repeat(4, 2).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(t.repeat(4, 2, 1).shape)
+    ```
+    """
+    repeats = argfix(repeats, *args)
+    base_shape = _align_left(self.shape, repeats)[0]
+    unsqueezed_shape = flatten([[s] if r == 1 else [1, s] for r,s in zip(repeats, base_shape)])
+    expanded_shape = flatten([[s] if r == 1 else [r, s] for r,s in zip(repeats, base_shape)])
+    final_shape = [r*s for r,s in zip(repeats, base_shape)]
+    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)
+
+  # **** pool level ****
+
+  def _pool(self, k_:tuple[sint, ...], stride:int|tuple[int, ...]=1, dilation:int|tuple[int, ...]=1) -> Self:
+    assert len(self.shape) >= len(k_), f"can't pool {self.shape} with {k_}"
+    s_, d_ = make_tuple(stride, len(k_)), make_tuple(dilation, len(k_))
+    assert len(k_) == len(s_) == len(d_), f"stride/dilation mismatch kernel:{k_} stride:{s_} dilation:{d_}"
+    noop, i_ = [None] * (self.ndim-len(k_)), self.shape[-len(k_):]
+    assert all(resolve(d*(k-1)+1 <= i) for k,d,i in zip(k_,d_,i_)), "kernel size cannot be greater than actual input size"
+    o_ = [ceildiv(i-d*(k-1), s) for i,d,k,s in zip(i_,d_,k_,s_)]
+    # input size scaling factor to make sure shrink for stride is possible
+    f_ = [smax(1, ceildiv(o*s - d, i)) for o,s,i,d in zip(o_,s_,i_,d_)]
+    # repeats such that we don't need padding
+    x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
+    # handle dilation
+    x = x.shrink_to(noop + [k*(i*f+d) for k,i,d,f in zip(k_,i_,d_,f_)]).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
+    # handle stride
+    x = x.shrink_to(noop + flatten((k,o*s) for k,o,s in zip(k_,o_,s_))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
+    x = x.shrink_to(noop + flatten((k,o,1) for k,o in zip(k_,o_))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
+    # permute to move reduce to the end
+    return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
@@ -1124,6 +1124,16 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)

  # ***** Indexing Ops *****
+  def NonZero(x:Tensor):
+    mask = (x!=0).flatten()
+    flat_idx = Tensor.arange(mask.numel(), dtype=dtypes.int64, device=x.device).masked_select(mask)
+    if flat_idx.ndim == 0: flat_idx = flat_idx.reshape(1)
+    if x.ndim == 0:
+      return Tensor.zeros((0, flat_idx.shape[0]), dtype=dtypes.int64, device=x.device, requires_grad=False)
+    strides = [prod(int(s) for s in x.shape[i+1:]) if i+1 < x.ndim else 1 for i in range(x.ndim)]
+    coords = [((flat_idx // stride) % int(dim)) for stride, dim in zip(strides, x.shape)]
+    return Tensor.stack(*coords, dim=0)
+
  def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]

  def Gather(x:Tensor, indices:Tensor, axis:int=0):
@@ -450,16 +450,9 @@ class AMDRenderer(CStyleLanguage):
      ]) + base_rewrite
  def __reduce__(self): return self.__class__, (self.arch,)

-  # language options
-  ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
-  ocml = [(f"__ocml_{name}_f{n}", f"{dt}, {dt}" if "fmax" == name else dt, dt, atr)
-            for dt, n in [(dtype.name, dtype.itemsize * 8) for dtype in [dtypes.float, dtypes.double, dtypes.half]]
-            for name, atr in [("fmax", "const"), ("exp2", "pure"), ("log2", "pure"), ("sqrt", "const"), ("sin", ""), ("trunc", "")]]
-
-  kernel_typedef = "\n".join(f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml)
  # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
  # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters
-  kernel_typedef += '\nextern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {launch_bounds})))'
+  kernel_typedef = 'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {launch_bounds})))'
  code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})",
                       "i": lambda x: f"(__ockl_get_group_id({x})*__ockl_get_local_size({x})+__ockl_get_local_id({x}))"}
  code_for_op = { **CStyleLanguage.code_for_op,
@@ -490,15 +483,25 @@ class AMDRenderer(CStyleLanguage):
           f"{vec} make_{vec}({', '.join([f'{scal} {x}' for x in _nms[:dtype.count]])}) {{ return {{ {', '.join(_nms[:dtype.count])} }}; }}"

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
-    prefix = ["#define INFINITY (__builtin_inff())","#define NAN (__builtin_nanf(\"\"))","typedef long unsigned int size_t;","#define half _Float16"]
+    prefix, ockl = [], []
    type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
    used_dtypes = uops_to_dtypes(uops)
+    if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops):
+      prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
+    if any(u.op is Ops.SPECIAL for u in uops):
+      prefix.append("typedef long unsigned int size_t;")
+      ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
+    ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")}
+    ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.itemsize * 8}", dt.name, dt.name, ocml_ops[op][1])
+      for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
    if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes): prefix.append("typedef unsigned short hip_bfloat16;")
+    if any(dt.scalar() == dtypes.half for dt in used_dtypes): prefix.append("#define half _Float16")
    if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes):
      prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"]
      prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
  v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
  return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")
+    prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml]
    prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if dt.count > 1]

    for name, (N, M, K), dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
@@ -0,0 +1,17 @@
+import importlib, pathlib
+from tinygrad.helpers import system
+
+root = (here:=pathlib.Path(__file__).parent).parents[2]
+
+def load(name, dll, files, **kwargs):
+  if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists():
+    files = files() if callable(files) else files
+    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(dll, files, **kwargs))
+  return importlib.import_module(f"{path}.{name.replace('/', '.')}")
+
+def __getattr__(nm):
+  match nm:
+    case "libc": return load("libc", ["find_library('c')"], lambda: (
+      [i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
+      ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), use_errno=True)
+    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -7357,6 +7357,113 @@ class struct_c__SA_FWSECLIC_FRTS_CMD(Structure):
     ]

 FWSECLIC_FRTS_CMD = struct_c__SA_FWSECLIC_FRTS_CMD
+PCIEXPTBL_H = True # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE = 0x00 # macro
+NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT = 0xE0 # macro
+PCI_EXP_ROM_SIGNATURE = 0xaa55 # macro
+PCI_EXP_ROM_SIGNATURE_NV = 0x4e56 # macro
+PCI_EXP_ROM_SIGNATURE_NV2 = 0xbb77 # macro
+def IS_VALID_PCI_ROM_SIG(sig):  # macro
+   return ((sig==0xaa55) or (sig==0x4e56) or (sig==0xbb77))
+OFFSETOF_PCI_EXP_ROM_SIG = 0x0 # macro
+OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET = 0x16 # macro
+OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR = 0x18 # macro
+PCI_DATA_STRUCT_SIGNATURE = 0x52494350 # macro
+PCI_DATA_STRUCT_SIGNATURE_NV = 0x5344504E # macro
+PCI_DATA_STRUCT_SIGNATURE_NV2 = 0x53494752 # macro
+def IS_VALID_PCI_DATA_SIG(sig):  # macro
+   return ((sig==0x52494350) or (sig==0x5344504E) or (sig==0x53494752))
+# PCI_LAST_IMAGE = NVBIT ( 7 ) # macro
+PCI_ROM_IMAGE_BLOCK_SIZE = 512 # macro
+OFFSETOF_PCI_DATA_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID = 0x4 # macro
+OFFSETOF_PCI_DATA_STRUCT_LEN = 0xa # macro
+OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE = 0xd # macro
+OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE = 0x14 # macro
+OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN = 0x10 # macro
+OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE = 0x15 # macro
+NV_PCI_DATA_EXT_SIG = 0x4544504E # macro
+NV_PCI_DATA_EXT_REV_10 = 0x100 # macro
+NV_PCI_DATA_EXT_REV_11 = 0x101 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SIG = 0x0 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LEN = 0x6 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_REV = 0x4 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN = 0x8 # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE = 0xa # macro
+OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS = 0xb # macro
+PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED = 0x04 # macro
+class struct__PCI_EXP_ROM_STANDARD(Structure):
+    pass
+
+struct__PCI_EXP_ROM_STANDARD._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_STANDARD._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 22),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_STANDARD = struct__PCI_EXP_ROM_STANDARD
+PPCI_EXP_ROM_STANDARD = ctypes.POINTER(struct__PCI_EXP_ROM_STANDARD)
+class struct__PCI_EXP_ROM_NBSI(Structure):
+    pass
+
+struct__PCI_EXP_ROM_NBSI._pack_ = 1 # source:False
+struct__PCI_EXP_ROM_NBSI._fields_ = [
+    ('sig', ctypes.c_uint16),
+    ('reserved', ctypes.c_ubyte * 20),
+    ('nbsiDataOffset', ctypes.c_uint16),
+    ('pciDataStrucPtr', ctypes.c_uint16),
+    ('sizeOfBlock', ctypes.c_uint32),
+]
+
+PCI_EXP_ROM_NBSI = struct__PCI_EXP_ROM_NBSI
+PPCI_EXP_ROM_NBSI = ctypes.POINTER(struct__PCI_EXP_ROM_NBSI)
+class union__PCI_EXP_ROM(Union):
+    _pack_ = 1 # source:False
+    _fields_ = [
+    ('standard', PCI_EXP_ROM_STANDARD),
+    ('nbsi', PCI_EXP_ROM_NBSI),
+     ]
+
+PCI_EXP_ROM = union__PCI_EXP_ROM
+PPCI_EXP_ROM = ctypes.POINTER(union__PCI_EXP_ROM)
+class struct__PCI_DATA_STRUCT(Structure):
+    pass
+
+struct__PCI_DATA_STRUCT._pack_ = 1 # source:False
+struct__PCI_DATA_STRUCT._fields_ = [
+    ('sig', ctypes.c_uint32),
+    ('vendorID', ctypes.c_uint16),
+    ('deviceID', ctypes.c_uint16),
+    ('deviceListPtr', ctypes.c_uint16),
+    ('pciDataStructLen', ctypes.c_uint16),
+    ('pciDataStructRev', ctypes.c_ubyte),
+    ('classCode', ctypes.c_ubyte * 3),
+    ('imageLen', ctypes.c_uint16),
+    ('vendorRomRev', ctypes.c_uint16),
+    ('codeType', ctypes.c_ubyte),
+    ('lastImage', ctypes.c_ubyte),
+    ('maxRunTimeImageLen', ctypes.c_uint16),
+]
+
+PCI_DATA_STRUCT = struct__PCI_DATA_STRUCT
+PPCI_DATA_STRUCT = ctypes.POINTER(struct__PCI_DATA_STRUCT)
+class struct__NV_PCI_DATA_EXT_STRUCT(Structure):
+    pass
+
+struct__NV_PCI_DATA_EXT_STRUCT._pack_ = 1 # source:False
+struct__NV_PCI_DATA_EXT_STRUCT._fields_ = [
+    ('signature', ctypes.c_uint32),
+    ('nvPciDataExtRev', ctypes.c_uint16),
+    ('nvPciDataExtLen', ctypes.c_uint16),
+    ('subimageLen', ctypes.c_uint16),
+    ('privLastImage', ctypes.c_ubyte),
+    ('flags', ctypes.c_ubyte),
+]
+
+NV_PCI_DATA_EXT_STRUCT = struct__NV_PCI_DATA_EXT_STRUCT
+PNV_PCI_DATA_EXT_STRUCT = ctypes.POINTER(struct__NV_PCI_DATA_EXT_STRUCT)
 __all__ = \
    ['ACPI_DATA', 'ACPI_DSM_CACHE', 'ACPI_DSM_FUNCTION_COUNT',
    'ACPI_DSM_FUNCTION_CURRENT', 'ACPI_DSM_FUNCTION_GPS',
@@ -7480,12 +7587,16 @@ __all__ = \
    'NVDM_TYPE_UEFI_XTL_DEBUG_INTR',
    'NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX', 'NVGPU_ENGINE_CAPS_MASK_BITS',
    'NV_ACPI_GENERIC_FUNC_COUNT',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE',
+    'NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_AVAILABLE',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_UNAVAILABLE',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V1',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V2',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V3',
    'NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V4',
+    'NV_PCI_DATA_EXT_REV_10', 'NV_PCI_DATA_EXT_REV_11',
+    'NV_PCI_DATA_EXT_SIG', 'NV_PCI_DATA_EXT_STRUCT',
    'NV_RPC_UPDATE_PDE_BAR_1', 'NV_RPC_UPDATE_PDE_BAR_2',
    'NV_RPC_UPDATE_PDE_BAR_INVALID', 'NV_RPC_UPDATE_PDE_BAR_TYPE',
    'NV_RPC_UPDATE_PDE_BAR_TYPE__enumvalues',
@@ -7775,8 +7886,31 @@ __all__ = \
    'NV_VGPU_PTE_64_INDEX_SHIFT', 'NV_VGPU_PTE_64_PAGE_SIZE',
    'NV_VGPU_PTE_64_SIZE', 'NV_VGPU_PTE_INDEX_MASK',
    'NV_VGPU_PTE_INDEX_SHIFT', 'NV_VGPU_PTE_PAGE_SIZE',
-    'NV_VGPU_PTE_SIZE', 'PACKED_REGISTRY_ENTRY',
-    'PACKED_REGISTRY_TABLE', 'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
+    'NV_VGPU_PTE_SIZE', 'OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_LEN',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_REV',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE',
+    'OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE',
+    'OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN',
+    'OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE',
+    'OFFSETOF_PCI_DATA_STRUCT_LEN', 'OFFSETOF_PCI_DATA_STRUCT_SIG',
+    'OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID',
+    'OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET',
+    'OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR',
+    'OFFSETOF_PCI_EXP_ROM_SIG', 'PACKED_REGISTRY_ENTRY',
+    'PACKED_REGISTRY_TABLE', 'PCIEXPTBL_H',
+    'PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED', 'PCI_DATA_STRUCT',
+    'PCI_DATA_STRUCT_SIGNATURE', 'PCI_DATA_STRUCT_SIGNATURE_NV',
+    'PCI_DATA_STRUCT_SIGNATURE_NV2', 'PCI_EXP_ROM',
+    'PCI_EXP_ROM_NBSI', 'PCI_EXP_ROM_SIGNATURE',
+    'PCI_EXP_ROM_SIGNATURE_NV', 'PCI_EXP_ROM_SIGNATURE_NV2',
+    'PCI_EXP_ROM_STANDARD', 'PCI_ROM_IMAGE_BLOCK_SIZE',
+    'PNV_PCI_DATA_EXT_STRUCT', 'PPCI_DATA_STRUCT', 'PPCI_EXP_ROM',
+    'PPCI_EXP_ROM_NBSI', 'PPCI_EXP_ROM_STANDARD',
+    'REGISTRY_TABLE_ENTRY_TYPE_BINARY',
    'REGISTRY_TABLE_ENTRY_TYPE_DWORD',
    'REGISTRY_TABLE_ENTRY_TYPE_STRING',
    'REGISTRY_TABLE_ENTRY_TYPE_UNKNOWN', 'RM_ENGINE_TYPE',
@@ -8343,6 +8477,8 @@ __all__ = \
    'struct_UpdateBarPde_v15_00',
    'struct_VIRTUAL_DISPLAY_GET_MAX_RESOLUTION_PARAMS',
    'struct_VIRTUAL_DISPLAY_GET_NUM_HEADS_PARAMS',
+    'struct__NV_PCI_DATA_EXT_STRUCT', 'struct__PCI_DATA_STRUCT',
+    'struct__PCI_EXP_ROM_NBSI', 'struct__PCI_EXP_ROM_STANDARD',
    'struct_alloc_object_FERMI_CONTEXT_SHARE_A_v04_00',
    'struct_alloc_object_FERMI_VASPACE_A_v03_00',
    'struct_alloc_object_GF100_DISP_SW_v03_00',
@@ -8595,7 +8731,7 @@ __all__ = \
    'union_NV2080_CTRL_FB_FS_INFO_QUERY_DATA_v26_04',
    'union_NV2080_CTRL_GRMGR_GR_FS_INFO_QUERY_DATA_v1A_1D',
    'union_NV2080_CTRL_INTERNAL_PFM_REQ_HNDLR_STATE_SYNC_DATA_type_v21_04',
-    'union_alloc_object_params_v25_08',
+    'union__PCI_EXP_ROM', 'union_alloc_object_params_v25_08',
    'union_alloc_object_params_v26_00',
    'union_alloc_object_params_v27_00',
    'union_alloc_object_params_v29_06', 'union_c__SA_GspFwWprMeta_0',
@@ -594,7 +594,8 @@ class AMDProgram(HCQProgram):
      self.dev.synchronize()

      for se, buf in enumerate(self.dev.sqtt_buffers):
-        wptr = ((self.dev.sqtt_wptrs.cpu_view().view(fmt='I')[se]&0x1FFFFFFF)-(((buf.va_addr//32)&0x1FFFFFFF) if self.dev.target[0] == 11 else 0))*32
+        wptr = (self.dev.sqtt_wptrs.cpu_view().view(fmt='I')[se] & 0x1FFFFFFF) * 32
+        if self.dev.target[:2] == (11, 0): wptr -= ((buf.va_addr // 32) & 0x1FFFFFFF) * 32

        if DEBUG >= 5: print(f'\t{self.dev.device}: SE {se} blob size {wptr:#x}')
        assert wptr >= 0 and wptr <= buf.size, f"{wptr} > {buf.size}, should never happen"
@@ -785,7 +786,7 @@ class PCIIface(PCIIfaceBase):
  gpus:ClassVar[list[str]] = []

  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550, 0x7590], bars=[0, 2, 5], vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x744c, 0x7480, 0x7550, 0x7590])], bars=[0, 2, 5], vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
    self._setup_adev(self.pci_dev)
    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
@@ -911,8 +912,8 @@ class AMDDevice(HCQCompiled):
    max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
    self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))

-    compilers:list[CompilerPairT] = [(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch)),
-                                     (functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch))]
+    compilers:list[CompilerPairT] = [(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
+                                     (functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch))]

    super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                     functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
@@ -5,7 +5,7 @@ from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler
 from tinygrad.runtime.ops_cpu import CPUAllocator
 from tinygrad.dtype import dtypes, DType, PtrDType
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, DEBUG
+from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.runtime.autogen import libc, qcom_dsp
 if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import
@@ -123,10 +123,9 @@ class ClangCompiler(Compiler):

  def compile(self, src:str) -> bytes:
    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
-    with tempfile.NamedTemporaryFile(delete=True) as output_file:
-      subprocess.check_output([getenv("CC", 'clang'), *self.args, '-O2', '-Wall', '-Werror', '-x', 'c', '-fPIC', '-ffreestanding', '-nostdlib',
-                               '-', '-o', str(output_file.name)], input=src.encode('utf-8'))
-      return pathlib.Path(output_file.name).read_bytes()
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+      system(f"{getenv('CC','clang')} {' '.join(self.args)} -O2 -Wall -Werror -x c -fPIC -ffreestanding -nostdlib - -o {f.name}", input=src.encode())
+      return pathlib.Path(f.name).read_bytes()

  def disassemble(self, lib:bytes): return cpu_objdump(lib, self.objdump_tool)

@@ -389,12 +389,12 @@ class NVKIface:

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
    # Uncached memory is "system". Use huge pages only for gpu memory.
-    page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (mmap.PAGESIZE if MOCKGPU else 4 << 10))
    size = round_up(size, page_size)
-    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access)
+    va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access) if (alloced:=cpu_addr is None) else cpu_addr

    if host:
-      va_addr = cpu_addr or FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
+      if alloced: va_addr = FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)

      flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \
            | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30)
@@ -456,8 +456,8 @@ class PCIIface(PCIIfaceBase):
  gpus:ClassVar[list[str]] = []

  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x10de, devices=[0x2204, 0x2684, 0x2b85], bars=[0, 1], vram_bar=1,
-      va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
+    super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
+      bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
    if not OSX: System.reserve_hugepages(64)

    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
@@ -471,7 +471,7 @@ class PCIIface(PCIIfaceBase):
  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
    # Force use of huge pages for large allocations. NVDev will attempt to use huge pages in any case,
    # but if the size is not aligned, the tail will be allocated with 4KB pages, increasing TLB pressure.
-    page_size = (2 << 20) if size >= (8 << 20) and not uncached and not host else (4 << 10)
+    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << 10))
    return super().alloc(round_up(size, page_size), host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, **kwargs)

  def setup_usermode(self): return 0xce000000, self.pci_dev.map_bar(bar=0, fmt='I', off=0xbb0000, size=0x10000)
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import os, ctypes, functools, mmap, struct, array, math, sys, weakref
+import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
 from types import SimpleNamespace
 from typing import Any, cast
@@ -9,16 +9,18 @@ from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, adreno
 from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile
+from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE
+from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import

 BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2

 #Parse C-style defines: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:
 # qreg.<regname>(<field_x>=..., <field_y>=..., ..., <field_n>=...)
-def _qreg_exec(reg, __val=0, **kwargs):
+def _qreg_exec(__reg, __val=0, **kwargs):
  for k, v in kwargs.items():
-    __val |= (getattr(adreno, f'{reg[4:]}_{k.upper()}') if v else 0) if type(v) is bool else (v << getattr(adreno, f'{reg[4:]}_{k.upper()}__SHIFT'))
+    reg_name = f"{__reg[4:]}_{k.removeprefix('_').upper()}"
+    __val |= (getattr(adreno, reg_name) if v else 0) if type(v) is bool else (v << getattr(adreno, f'{reg_name}__SHIFT'))
  return __val
 qreg: Any = type("QREG", (object,), {name[4:].lower(): functools.partial(_qreg_exec, name) for name in adreno.__dict__.keys() if name[:4] == 'REG_'})

@@ -45,6 +47,10 @@ class QCOMSignal(HCQSignal):
      kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.owner.fd, context_id=self.owner.ctx, timestamp=self.owner.last_cmd, timeout=0xffffffff)

 class QCOMComputeQueue(HWQueue):
+  def __init__(self, dev:QCOMDevice):
+    self.dev = dev
+    super().__init__()
+
  def __del__(self):
    if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))

@@ -54,7 +60,7 @@ class QCOMComputeQueue(HWQueue):

  def _cache_flush(self, write_back=True, invalidate=False, sync=True, memsync=False):
    # TODO: 7xx support.
-    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(QCOMDevice.dummy_addr), 0) # dirty cache write-back.
+    if write_back: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_FLUSH_TS, *data64_le(self.dev.dummy_addr), 0) # dirty cache write-back.
    if invalidate: self.cmd(adreno.CP_EVENT_WRITE, adreno.CACHE_INVALIDATE) # invalidate cache lines (following reads from RAM).
    if memsync: self.cmd(adreno.CP_WAIT_MEM_WRITES)
    if sync: self.cmd(adreno.CP_WAIT_FOR_IDLE)
@@ -63,18 +69,20 @@ class QCOMComputeQueue(HWQueue):
    self._cache_flush(write_back=True, invalidate=True, sync=True, memsync=True)
    return self

-  def signal(self, signal:QCOMSignal, value=0, ts=False):
+  def signal(self, signal:QCOMSignal, value=0):
    self.cmd(adreno.CP_WAIT_FOR_IDLE)
-    if QCOMDevice.gpu_id < 700:
-      self.cmd(adreno.CP_EVENT_WRITE, qreg.cp_event_write_0(event=adreno.CACHE_FLUSH_TS, timestamp=ts),
-               *data64_le(signal.timestamp_addr if ts else signal.value_addr), qreg.cp_event_write_3(value & 0xFFFFFFFF))
+    if self.dev.gpu_id[:2] < (7, 3):
+      self.cmd(adreno.CP_EVENT_WRITE, qreg.cp_event_write_0(event=adreno.CACHE_FLUSH_TS), *data64_le(signal.value_addr), lo32(value))
      self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
    else:
      # TODO: support devices starting with 8 Gen 1. Also, 700th series have convenient CP_GLOBAL_TIMESTAMP and CP_LOCAL_TIMESTAMP
      raise RuntimeError('CP_EVENT_WRITE7 is not supported')
    return self

-  def timestamp(self, signal:QCOMSignal): return self.signal(signal, 0, ts=True)
+  def timestamp(self, signal:QCOMSignal):
+    self.cmd(adreno.CP_WAIT_FOR_IDLE)
+    self.cmd(adreno.CP_REG_TO_MEM, qreg.cp_reg_to_mem_0(reg=adreno.REG_A6XX_CP_ALWAYS_ON_COUNTER, cnt=2, _64b=True),*data64_le(signal.timestamp_addr))
+    return self

  def wait(self, signal:QCOMSignal, value=0):
    self.cmd(adreno.CP_WAIT_REG_MEM, qreg.cp_wait_reg_mem_0(function=adreno.WRITE_GE, poll=adreno.POLL_MEMORY),*data64_le(signal.value_addr),
@@ -314,12 +322,9 @@ class QCOMAllocator(HCQAllocatorBase):
    self.dev._gpu_free(opaque)

 class QCOMDevice(HCQCompiled):
-  gpu_id: int = 0
-  dummy_addr: int = 0
-
  def __init__(self, device:str=""):
    self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
-    QCOMDevice.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)
+    self.dummy_addr = cast(int, self._gpu_alloc(0x1000).va_addr)

    flags = kgsl.KGSL_CONTEXT_PREAMBLE | kgsl.KGSL_CONTEXT_PWR_CONSTRAINT | kgsl.KGSL_CONTEXT_NO_FAULT_TOLERANCE | kgsl.KGSL_CONTEXT_NO_GMEM_ALLOC \
              | kgsl.KGSL_CONTEXT_PRIORITY(getenv("QCOM_PRIORITY", 8)) | kgsl.KGSL_CONTEXT_PREEMPT_STYLE(kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
@@ -339,11 +344,17 @@ class QCOMDevice(HCQCompiled):
    # Load info about qcom device
    info = kgsl.struct_kgsl_devinfo()
    kgsl.IOCTL_KGSL_DEVICE_GETPROPERTY(self.fd, type=kgsl.KGSL_PROP_DEVICE_INFO, value=ctypes.addressof(info), sizebytes=ctypes.sizeof(info))
-    QCOMDevice.gpu_id = ((info.chip_id >> 24) & 0xFF) * 100 + ((info.chip_id >> 16) & 0xFF) * 10 + ((info.chip_id >>  8) & 0xFF)
-    if QCOMDevice.gpu_id >= 700: raise RuntimeError(f"Unsupported GPU: {QCOMDevice.gpu_id}")
+    self.gpu_id = (info.chip_id >> 24, (info.chip_id >> 16) & 0xFF, (info.chip_id >> 8) & 0xFF)
+
+    # a7xx start with 730x or 'Cxxx', a8xx starts 'Exxx'
+    if self.gpu_id[:2] >= (7, 3): raise RuntimeError(f"Unsupported GPU: chip_id={info.chip_id:#x}")
+
+    if PROFILE and self.gpu_id[:2] < (7, 3):
+      System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")

    compilers = [(QCOMRenderer, functools.partial(QCOMCompiler, device))]
-    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal, QCOMComputeQueue, None)
+    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
+                     functools.partial(QCOMComputeQueue, self), None)

  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
    flags |= kgsl.KGSL_MEMALIGN(alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
@@ -365,3 +376,7 @@ class QCOMDevice(HCQCompiled):
      self.synchronize()
      self._gpu_free(self._stack)
      self._stack = self._gpu_alloc(sz)
+
+  def _at_profile_finalize(self):
+    super()._at_profile_finalize()
+    with contextlib.suppress(RuntimeError): System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", "10", "Failed to reenable suspend mode")
@@ -0,0 +1,126 @@
+import ctypes.util, importlib.metadata, itertools, re, functools, os
+from tinygrad.helpers import flatten, unwrap
+from clang.cindex import Config, Index, CursorKind as CK, TranslationUnit as TU, LinkageKind as LK, TokenKind as ToK, TypeKind as TK
+from clang.cindex import PrintingPolicy as PP, PrintingPolicyProperty as PPP, SourceRange
+
+assert importlib.metadata.version('clang')[:2] == "20"
+if not Config.loaded: Config.set_library_file(os.getenv("LIBCLANG_PATH", ctypes.util.find_library("clang-20")))
+
+def fst(c): return next(c.get_children())
+def last(c): return list(c.get_children())[-1]
+def readext(f, fst, snd=None):
+  with open(f, "r") as f:
+    f.seek(start:=(fst.start.offset if isinstance(fst, SourceRange) else fst))
+    return f.read((fst.end.offset if isinstance(fst, SourceRange) else snd)-start)
+def attrs(c): return list(filter(lambda k: (v:=k.value) >= 400 and v < 500, map(lambda c: c.kind, c.get_children())))
+
+base_rules = [(r'\s*\\\n\s*', ' '), (r'\s*\n\s*', ' '), (r'//.*', ''), (r'/\*.*?\*/', ''), (r'\b(0[xX][0-9a-fA-F]+|\d+)[uUlL]+\b', r'\1'),
+              (r'\b0+(?=\d)', ''), (r'\s*&&\s*', r' and '), (r'\s*\|\|\s*', r' or '), (r'\s*!\s*', ' not '),
+              (r'(struct|union|enum)\s*([a-zA-Z_][a-zA-Z0-9_]*\b)', r'\1_\2'),
+              (r'\((unsigned )?(char|uint64_t)\)', ''), (r'^.*\d+:\d+.*$', ''), (r'^.*\w##\w.*$', '')]
+
+ints = (TK.INT, TK.UINT, TK.LONG, TK.ULONG, TK.LONGLONG, TK.ULONGLONG)
+
+def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_errno=False, anon_names={}, types={}, parse_macros=True):
+  macros, lines, anoncnt, types = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}
+  def tname(t, suggested_name=None, typedef=None) -> str:
+    suggested_name = anon_names.get(f"{(decl:=t.get_declaration()).location.file}:{decl.location.line}", suggested_name)
+    nonlocal lines, types, anoncnt
+    tmap = {TK.VOID:"None", TK.CHAR_U:"ctypes.c_ubyte", TK.UCHAR:"ctypes.c_ubyte", TK.CHAR_S:"ctypes.c_char", TK.SCHAR:"ctypes.c_char",
+            **{getattr(TK, k):f"ctypes.c_{k.lower()}" for k in ["BOOL", "WCHAR", "FLOAT", "DOUBLE", "LONGDOUBLE"]},
+            **{getattr(TK, k):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
+               [(16, "USHORT"), (16, "SHORT"), (32, "UINT"), (32, "INT"), (64, "ULONG"), (64, "LONG"), (64, "ULONGLONG"), (64, "LONGLONG")]}}
+
+    if t.kind in tmap: return tmap[t.kind]
+    if t.spelling in types and types[t.spelling][1]: return types[t.spelling][0]
+    if ((f:=t).kind in (fks:=(TK.FUNCTIONPROTO, TK.FUNCTIONNOPROTO))) or (t.kind == TK.POINTER and (f:=t.get_pointee()).kind in fks):
+      return f"ctypes.CFUNCTYPE({tname(f.get_result())}{(', '+', '.join(map(tname, f.argument_types()))) if f.kind==TK.FUNCTIONPROTO else ''})"
+    match t.kind:
+      case TK.POINTER: return "ctypes.c_void_p" if (ptr:=t.get_pointee()).kind == TK.VOID else f"ctypes.POINTER({tname(ptr)})"
+      case TK.ELABORATED: return tname(t.get_named_type(), suggested_name)
+      case TK.TYPEDEF if t.spelling == t.get_canonical().spelling: return tname(t.get_canonical())
+      case TK.TYPEDEF:
+        defined, nm = (canon:=t.get_canonical()).spelling in types, tname(canon, typedef=t.spelling.replace('::', '_'))
+        types[t.spelling] = nm if t.spelling.startswith("__") else t.spelling.replace('::', '_'), True
+        # RECORDs need to handle typedefs specially to allow for self-reference
+        if canon.kind != TK.RECORD or defined: lines.append(f"{t.spelling.replace('::', '_')} = {nm}")
+        return types[t.spelling][0]
+      case TK.RECORD:
+        # TODO: packed unions
+        # TODO: pragma pack support
+        # check for forward declaration
+        if t.spelling in types: types[t.spelling] = (nm:=types[t.spelling][0]), len(list(t.get_fields())) != 0
+        else:
+          if decl.is_anonymous():
+            types[t.spelling] = (nm:=(suggested_name or (f"_anon{'struct' if decl.kind == CK.STRUCT_DECL else 'union'}{anoncnt()}")), True)
+          else: types[t.spelling] = (nm:=t.spelling.replace(' ', '_').replace('::', '_')), len(list(t.get_fields())) != 0
+          lines.append(f"class {nm}({'Struct' if decl.kind==CK.STRUCT_DECL else 'ctypes.Union'}): pass")
+          if typedef: lines.append(f"{typedef} = {nm}")
+        acnt = itertools.count().__next__
+        ll=["  ("+((fn:=f"'_{acnt()}'")+f", {tname(f.type, nm+fn[1:-1])}" if f.is_anonymous_record_decl() else f"'{f.spelling}', "+
+            tname(f.type, f'{nm}_{f.spelling}'))+(f',{f.get_bitfield_width()}' if f.is_bitfield() else '')+")," for f in t.get_fields()]
+        lines.extend(([f"{nm}._anonymous_ = ["+", ".join(f"'_{i}'" for i in range(n))+"]"] if (n:=acnt()) else [])+
+                     ([f"{nm}._packed_ = True"] * (CK.PACKED_ATTR in attrs(decl)))+([f"{nm}._fields_ = [",*ll,"]"] if ll else []))
+        return nm
+      case TK.ENUM:
+        # TODO: C++ and GNU C have forward declared enums
+        if decl.is_anonymous(): types[t.spelling] = suggested_name or f"_anonenum{anoncnt()}", True
+        else: types[t.spelling] = t.spelling.replace(' ', '_').replace('::', '_'), True
+        lines.append(f"{types[t.spelling][0]} = CEnum({tname(decl.enum_type)})\n" +
+                     "\n".join(f"{e.spelling} = {types[t.spelling][0]}.define('{e.spelling}', {e.enum_value})" for e in decl.get_children()
+                     if e.kind == CK.ENUM_CONSTANT_DECL) + "\n")
+        return types[t.spelling][0]
+      case TK.CONSTANTARRAY:
+        return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * {t.get_array_size()})"
+      case TK.INCOMPLETEARRAY: return f"({tname(t.get_array_element_type(), suggested_name.rstrip('s') if suggested_name else None)} * 0)"
+      case _: raise NotImplementedError(f"unsupported type {t.kind}")
+
+  for f in files:
+    tu = Index.create().parse(f, args, options=TU.PARSE_DETAILED_PROCESSING_RECORD)
+    (pp:=PP.create(tu.cursor)).set_property(PPP.TerseOutput, 1)
+    for c in tu.cursor.walk_preorder():
+      if str(c.location.file) != str(f) and (not recsym or c.kind not in (CK.FUNCTION_DECL,)): continue
+      rollback = lines, types
+      try:
+        match c.kind:
+          case CK.FUNCTION_DECL if c.linkage == LK.EXTERNAL and dll:
+            # TODO: we could support name-mangling
+            lines.append(f"# {c.pretty_printed(pp)}\ntry: ({c.spelling}:=dll.{c.spelling}).restype, {c.spelling}.argtypes = "
+              f"{tname(c.result_type)}, [{', '.join(tname(arg.type) for arg in c.get_arguments())}]\nexcept AttributeError: pass\n")
+          case CK.STRUCT_DECL | CK.UNION_DECL | CK.TYPEDEF_DECL | CK.ENUM_DECL: tname(c.type)
+          case CK.MACRO_DEFINITION if parse_macros and len(toks:=list(c.get_tokens())) > 1:
+            if toks[1].spelling == '(' and toks[0].extent.end.column == toks[1].extent.start.column:
+              it = iter(toks[1:])
+              _args = [t.spelling for t in itertools.takewhile(lambda t:t.spelling!=')', it) if t.kind == ToK.IDENTIFIER]
+              if len(body:=list(it)) == 0: continue
+              macros += [f"{c.spelling} = lambda {','.join(_args)}: {readext(f, body[0].location.offset, toks[-1].extent.end.offset)}"]
+            else: macros += [f"{c.spelling} = {readext(f, toks[1].location.offset, toks[-1].extent.end.offset)}"]
+          case CK.VAR_DECL if c.linkage == LK.INTERNAL:
+            if (c.type.kind == TK.CONSTANTARRAY and c.type.get_array_element_type().get_canonical().kind in ints and
+                (init:=last(c)).kind == CK.INIT_LIST_EXPR and all(re.match(r"\[.*\].*=", readext(f, c.extent)) for c in init.get_children())):
+              cs = init.get_children()
+              macros += [f"{c.spelling} = {{{','.join(f'{readext(f,next(it:=c.get_children()).extent)}:{readext(f,next(it).extent)}' for c in cs)}}}"]
+            elif c.type.get_canonical().kind in ints: macros += [f"{c.spelling} = {readext(f, last(c).extent)}"]
+            else: macros += [f"{c.spelling} = {tname(c.type)}({readext(f, last(c).extent)})"]
+          case CK.VAR_DECL if c.linkage == LK.EXTERNAL and dll:
+            lines.append(f"try: {c.spelling} = {tname(c.type)}.in_dll(dll, '{c.spelling}')\nexcept (ValueError,AttributeError): pass")
+      except NotImplementedError as e:
+        print(f"skipping {c.spelling}: {e}")
+        lines, types = rollback
+  main = (f"# mypy: ignore-errors\nimport ctypes{', os' if any('os' in s for s in dll) else ''}\n"
+    "from tinygrad.helpers import unwrap\nfrom tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR\n" + '\n'.join([*prolog,
+      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)),
+      *(["def dll():",*flatten([[f"  try: return ctypes.CDLL(unwrap({d}){', use_errno=True' if use_errno else ''})",'  except: pass'] for d in dll]),
+         "  return None", "dll = dll()\n"]*bool(dll)), *lines]) + '\n')
+  macros = [r for m in macros if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]
+  while True:
+    try:
+      exec(main + '\n'.join(macros), {})
+      break
+    except (SyntaxError, NameError, TypeError) as e:
+      macrono = unwrap(e.lineno if isinstance(e, SyntaxError) else unwrap(unwrap(e.__traceback__).tb_next).tb_lineno) - main.count('\n') - 1
+      assert macrono >= 0 and macrono < len(macros), f"error outside macro range: {e}"
+      print(f"skipping {macros[macrono]}: {e}")
+      del macros[macrono]
+    except Exception as e: raise Exception("parsing failed") from e
+  return main + '\n'.join(macros + epilog)
@@ -0,0 +1,72 @@
+import ctypes, functools, sys
+from typing import TYPE_CHECKING
+
+def _do_ioctl(__idir, __base, __nr, __struct, __fd, **kwargs):
+  import tinygrad.runtime.support.hcq as hcq, fcntl
+  ioctl = __fd.ioctl if isinstance(__fd, hcq.FileIOInterface) else functools.partial(fcntl.ioctl, __fd)
+  if (rc:=ioctl((__idir<<30)|(ctypes.sizeof(out:=__struct(**kwargs))<<16)|(__base<<8)|__nr, out)): raise RuntimeError(f"ioctl returned {rc}")
+  return out
+
+def _IO(base, nr): return functools.partial(_do_ioctl, 0, ord(base) if isinstance(base, str) else base, nr, None)
+def _IOW(base, nr, typ): return functools.partial(_do_ioctl, 1, ord(base) if isinstance(base, str) else base, nr, typ)
+def _IOR(base, nr, typ): return functools.partial(_do_ioctl, 2, ord(base) if isinstance(base, str) else base, nr, typ)
+def _IOWR(base, nr, typ): return functools.partial(_do_ioctl, 3, ord(base) if isinstance(base, str) else base, nr, typ)
+
+def CEnum(typ: type[ctypes._SimpleCData]):
+  class _CEnum(typ): # type: ignore
+    _val_to_name_: dict[int,str] = {}
+
+    @classmethod
+    def from_param(cls, val): return val if isinstance(val, cls) else cls(val)
+    @classmethod
+    def get(cls, val, default="unknown"): return cls._val_to_name_.get(val.value if isinstance(val, cls) else val, default)
+    @classmethod
+    def items(cls): return cls._val_to_name_.items()
+    @classmethod
+    def define(cls, name, val):
+      cls._val_to_name_[val] = name
+      return val
+
+    def __eq__(self, other): return self.value == other
+    def __repr__(self): return self.get(self) if self.value in self.__class__._val_to_name_ else str(self.value)
+
+  return _CEnum
+
+# supports gcc (C11) __attribute__((packed))
+if TYPE_CHECKING: Struct = ctypes.Structure
+else:
+  class MetaStruct(type(ctypes.Structure)):
+    def __new__(mcs, name, bases, dct):
+      fields = dct.pop("_fields_", None)
+      cls = super().__new__(mcs, name, bases, dct)
+      if dct.get("_packed_", False) and fields is not None: mcs._build(cls, fields)
+      return cls
+
+    def __setattr__(cls, k, v):
+      # NB: _fields_ must be set after _packed_ because PyCStructType_setattro marks _fields_ as final.
+      if k == "_fields_" and getattr(cls, "_packed_", False): type(cls)._build(cls, v)
+      elif k == "_packed_" and hasattr(cls, "_fields_"): type(cls)._build(cls, cls._fields_)
+      else: super().__setattr__(k, v)
+
+    @staticmethod
+    def _build(cls, fields):
+      o = 0
+      for n,t,b in [(f[0], f[1], f[2] if len(f) == 3 else 0) for f in fields]:
+        if b == 0: o = (o + 7) & ~7
+        m = (1 << (sz:=ctypes.sizeof(t)*8 if b == 0 else b)) - 1
+        def _s(self,v,m,s,b): self._data[:] = ((int.from_bytes(self._data,sys.byteorder)&~(m<<s))|((v&m)<<s)).to_bytes(len(self._data), sys.byteorder)
+        setattr(cls, n, property(functools.partial(lambda self,m,s:(int.from_bytes(self._data,sys.byteorder)>>s)&m,m=m,s=o),
+                                 functools.partial(_s,m=m,s=o,b=b)))
+        o += sz
+
+      type(ctypes.Structure).__setattr__(cls, '_fields_', [('_data', ctypes.c_ubyte * ((o + 7) // 8))])
+      type(ctypes.Structure).__setattr__(cls, '_packed_', True)
+      setattr(cls, '_packed_fields_', fields)
+
+  class Struct(ctypes.Structure, metaclass=MetaStruct):
+    def __init__(self, *args, **kwargs):
+      if hasattr(self, '_packed_fields_'):
+        for f,v in zip(self._packed_fields_, args): setattr(self, f[0], v)
+        for k,v in kwargs.items(): setattr(self, k, v)
+      else: super().__init__(*args, **kwargs)
+
@@ -1,4 +1,5 @@
-import ctypes, subprocess
+import ctypes
+from tinygrad.helpers import system
 import tinygrad.runtime.autogen.comgr as comgr
 assert comgr.AMD_COMGR_LANGUAGE_HIP == 4
 try:
@@ -13,7 +14,7 @@ from tinygrad.runtime.support.compiler_cpu import LLVMCompiler
 from tinygrad.helpers import OSX, to_char_p_p

 def amdgpu_disassemble(lib:bytes):
-  asm = subprocess.check_output(["llvm-objdump" if OSX else "/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib).decode("utf-8").splitlines()
+  asm = system(f"{'llvm-objdump' if OSX else '/opt/rocm/llvm/bin/llvm-objdump'} -d -", input=lib).splitlines()
  while asm and ("s_nop 0" in asm[-1] or "s_code_end" in asm[-1]): asm.pop()
  print("\n".join(asm))

@@ -1,6 +1,6 @@
 import subprocess, hashlib, tempfile, ctypes, re, pathlib
 from typing import Callable
-from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
+from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv, system
 import tinygrad.runtime.autogen.nvrtc as nvrtc
 from tinygrad.device import Compiler, CompileError

@@ -37,7 +37,7 @@ def cuda_disassemble(lib:bytes, arch:str):
    fn = (pathlib.Path(tempfile.gettempdir()) / f"tinycuda_{hashlib.md5(lib).hexdigest()}").as_posix()
    with open(fn, "wb") as f: f.write(lib)
    subprocess.run(["ptxas", f"-arch={arch}", "-o", fn, fn], check=False, stderr=subprocess.DEVNULL) # optional ptx -> sass step for CUDA=1
-    print(subprocess.check_output(['nvdisasm', fn]).decode('utf-8'))
+    print(system(f'nvdisasm {fn}'))
  except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains ptxas/nvdisasm binary of compatible version.")

 class CUDACompiler(Compiler):
@@ -1,6 +1,6 @@
-import base64, ctypes, pathlib, tempfile, hashlib, subprocess
+import base64, ctypes, pathlib, tempfile, hashlib
 from tinygrad.device import Compiler
-from tinygrad.helpers import cpu_objdump
+from tinygrad.helpers import cpu_objdump, system
 import tinygrad.runtime.autogen.mesa as mesa
 from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, expect, cerr
 try: import tinygrad.runtime.autogen.llvm as llvm
@@ -82,5 +82,5 @@ class NAKCompiler(NIRCompiler):
    try:
      fn = (pathlib.Path(tempfile.gettempdir()) / f"tinynak_{hashlib.md5(lib).hexdigest()}").as_posix()
      with open(fn, "wb") as f: f.write(lib[ctypes.sizeof(mesa.struct_nak_shader_info):])
-      print(subprocess.check_output(['nvdisasm', "-b", f"SM{self.arch[3:]}", fn]).decode('utf-8'))
+      print(system(f"nvdisasm -b SM{self.arch[3:]} {fn}"))
    except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains nvdisasm binary of compatible version.")
@@ -1,7 +1,7 @@
 import struct, ctypes, ctypes.util
 from dataclasses import dataclass
 from tinygrad.helpers import getbits, i2u, unwrap
-import tinygrad.runtime.autogen.libc as libc
+from tinygrad.runtime.autogen import libc

@dataclass(frozen=True)
 class ElfSection: name:str; header:libc.Elf64_Shdr; content:bytes # noqa: E702
@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, traceback, collections
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent
+from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent, select_first_inited
 from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
@@ -437,19 +437,10 @@ class HCQCompiled(Compiled, Generic[SignalType]):
    except MemoryError: buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
    return buf, realloced

-  def _make_no_iface_error(self, errs:str, err_short:str) -> RuntimeError:
-    # Keep it in a separate function to avoid creating a traceback <-> locals ref cycle
-    e = RuntimeError(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")
-    if hasattr(e, "add_note"): e.add_note(errs + err_short)
-    return e
-
  def _select_iface(self, *ifaces:Type):
-    errs, err_short = "", ""
    if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper()))
-    for iface_t in ifaces:
-      try: return iface_t(self, self.device_id)
-      except Exception as e: errs, err_short = errs + f"\n{iface_t.__name__}: {traceback.format_exc()}", err_short + f"\n{iface_t.__name__}: {e}."
-    raise self._make_no_iface_error(errs, err_short)
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in ifaces],
+                               f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available")

  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"

@@ -1,5 +1,5 @@
-import ctypes.util, os, sys, subprocess
-from tinygrad.helpers import DEBUG, OSX, getenv
+import ctypes.util, os, sys
+from tinygrad.helpers import DEBUG, OSX, getenv, system

 if sys.platform == 'win32':
  # Windows llvm distribution doesn't seem to add itself to PATH or anywhere else where it can be easily retrieved from.
@@ -10,7 +10,7 @@ if sys.platform == 'win32':
 elif OSX:
  # Will raise FileNotFoundError if brew is not installed
  # `brew --prefix` will return even if formula is not installed
-  if not os.path.exists(brew_prefix:=subprocess.check_output(['brew', '--prefix', 'llvm@20']).decode().strip()):
+  if not os.path.exists(brew_prefix:=system("brew --prefix llvm@20")):
    raise FileNotFoundError('LLVM not found, you can install it with `brew install llvm@20`')
  LLVM_PATH: str|None = os.path.join(brew_prefix, 'lib', 'libLLVM.dylib')
 else:
@@ -49,8 +49,9 @@ class NVRpcQueue:
    self.seq += 1
    self.gsp.nvdev.NV_PGSP_QUEUE_HEAD[0].write(0x0)

-  def wait_resp(self, cmd:int) -> memoryview:
-    while True:
+  def wait_resp(self, cmd:int, timeout=10000) -> memoryview:
+    start_time = int(time.perf_counter() * 1000)
+    while (int(time.perf_counter() * 1000) - start_time) < timeout:
      System.memory_barrier()
      if self.rx.readPtr == self.tx.writePtr: continue

@@ -73,6 +74,7 @@ class NVRpcQueue:

      if hdr.rpc_result != 0: raise RuntimeError(f"RPC call {hdr.function} failed with result {hdr.rpc_result}")
      if hdr.function == cmd: return msg
+    raise RuntimeError(f"Timeout waiting for RPC response for command {cmd}")

 class NV_FLCN(NV_IP):
  def init_sw(self):
@@ -89,10 +91,18 @@ class NV_FLCN(NV_IP):
    self.prep_booter()

  def prep_ucode(self):
-    expansion_rom_off, bit_addr = {"GA": 0x16600, "AD": 0x14e00}[self.nvdev.chip_name[:2]], 0x1b0
-    vbios_bytes = bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x98e00)//4]))
+    vbios_bytes, vbios_off = memoryview(bytes(array.array('I', self.nvdev.mmio[0x00300000//4:(0x00300000+0x100000)//4]))), 0
+    while True:
+      pci_blck = vbios_bytes[vbios_off + nv.OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR:].cast('H')[0]
+      imglen = vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN:].cast('H')[0] * nv.PCI_ROM_IMAGE_BLOCK_SIZE
+      match vbios_bytes[vbios_off + pci_blck + nv.OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE]:
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE: block_size = imglen
+        case nv.NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT:
+          expansion_rom_off = vbios_off - block_size
+          break
+      vbios_off += imglen

-    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[bit_addr:bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
+    bit_header = nv.BIT_HEADER_V1_00.from_buffer_copy(vbios_bytes[(bit_addr:=0x1b0):bit_addr + ctypes.sizeof(nv.BIT_HEADER_V1_00)])
    assert bit_header.Signature == 0x00544942, f"Invalid BIT header signature {hex(bit_header.Signature)}"

    for i in range(bit_header.TokenEntries):
@@ -12,6 +12,11 @@ MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000,
 class PCIBarInfo: addr:int; size:int # noqa: E702

 class _System:
+  def write_sysfs(self, path:str, value:str, msg:str, expected:str|None=None):
+    if FileIOInterface(path, os.O_RDONLY).read().splitlines()[0] != (expected or value):
+      os.system(cmd:=f"sudo sh -c 'echo {value} > {path}'")
+      if FileIOInterface(path, os.O_RDONLY).read().splitlines()[0] != (expected or value): raise RuntimeError(f"{msg}. Please run {cmd} manually.")
+
  @functools.cached_property
  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None

@@ -26,9 +31,7 @@ class _System:

  @functools.cached_property
  def pagemap(self) -> FileIOInterface:
-    if FileIOInterface(reloc_sysfs:="/proc/sys/vm/compact_unevictable_allowed", os.O_RDONLY).read()[0] != "0":
-      os.system(cmd:=f"sudo sh -c 'echo 0 > {reloc_sysfs}'")
-      assert FileIOInterface(reloc_sysfs, os.O_RDONLY).read()[0] == "0", f"Failed to disable migration of locked pages. Please run {cmd} manually."
+    self.write_sysfs("/proc/sys/vm/compact_unevictable_allowed", "0", "Failed to disable migration of locked pages")
    return FileIOInterface("/proc/self/pagemap", os.O_RDONLY)

  @functools.cached_property
@@ -90,12 +93,12 @@ class _System:
    if data is not None: sysmem_view[:len(data)] = data
    return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

-  def pci_scan_bus(self, target_vendor:int, target_devices:list[int]) -> list[str]:
+  def pci_scan_bus(self, target_vendor:int, target_devices:list[tuple[int, list[int]]]) -> list[str]:
    result = []
    for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
      vendor = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/vendor").read(), 16)
      device = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/device").read(), 16)
-      if vendor == target_vendor and device in target_devices: result.append(pcibus)
+      if vendor == target_vendor and any((device & mask) in devlist for mask, devlist in target_devices): result.append(pcibus)
    return sorted(result)

  def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, PCIBarInfo]:
@@ -244,7 +247,7 @@ class LNXPCIIfaceBase:
  dev_impl:PCIDevImplBase
  gpus:ClassVar[list[str]] = []

-  def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
+  def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size):
    if len((cls:=type(self)).gpus) == 0:
      cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices))

@@ -285,6 +288,7 @@ class LNXPCIIfaceBase:
 class APLPCIIfaceBase(LNXPCIIfaceBase):
  def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
    self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(dev.__class__.__name__[:2], pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
+    assert (read_vendor:=self.pci_dev.read_config(0x00, 2)) == vendor, f"Vendor ID mismatch: expected {vendor:#x}, got {read_vendor:#x}"
  def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")

 PCIIfaceBase:type = APLPCIIfaceBase if OSX else LNXPCIIfaceBase
@@ -1,10 +1,10 @@
-import ctypes.util, os, subprocess, platform, sysconfig
-from tinygrad.helpers import OSX
+import ctypes.util, os, platform, sysconfig
+from tinygrad.helpers import system, OSX

 WEBGPU_PATH: str | None

 if OSX:
-  if not os.path.exists(brew_prefix:=subprocess.check_output(['brew', '--prefix', 'dawn']).decode().strip()):
+  if not os.path.exists(brew_prefix:=system("brew --prefix dawn")):
    raise FileNotFoundError('dawn library not found. Install it with `brew tap wpmed92/dawn && brew install dawn`')
  WEBGPU_PATH = os.path.join(brew_prefix, 'lib', 'libwebgpu_dawn.dylib')
 elif platform.system() == "Windows":
@@ -239,7 +239,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
    # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it
    if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape):
-      ending_ranges[x] = list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
+      ending_ranges[x] += list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())

    # REDUCE_AXIS creates ranges for the axes it is reducing
    if x.op is Ops.REDUCE_AXIS:
@@ -247,15 +247,23 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:

    if debug:
      realized_ranges = rctx.realize_map.get(x, None)
-      disp = []
-      for i, (ri, ro) in enumerate(zip([r.render() for r in rngs], [r.render() for r in out_rngs])):
-        rng = f"{ri}" if ri == ro else f"{ri} -> {ro}"
-        if realized_ranges is not None and i in realized_ranges: rng = colored(rng, "yellow")
-        disp.append("["+rng+"]")
-      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}", ''.join(disp))
+      if x.op is Ops.RESHAPE or len(rngs) != len(out_rngs):
+        disp = render_ranges(rngs, realized=realized_ranges) + " -> " + render_ranges(out_rngs, realized=realized_ranges)
+      else:
+        disp = render_ranges(rngs, out_rngs, realized=realized_ranges)
+      print("***" if x in rctx.realize_map else "   ",
+            f"{len(consumer_map[x]):2d} {str(x.op):20s} {str(x._shape):35s} {len(ending_ranges[x]):2d}", disp)

    # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
    rctx.range_map[x] = (rngs, out_rngs)

  tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
  return tsink, rctx
+
+def render_ranges(*rngs_list, realized) -> str:
+  disp = []
+  for i, rs in enumerate(zip(*[[r.render() for r in rngs] for rngs in rngs_list])):
+    rng = rs[0] if all_same(rs) else " -> ".join(rs)
+    if realized is not None and i in realized: rng = colored(rng, "yellow")
+    disp.append("["+rng+"]")
+  return ''.join(disp)
@@ -5,7 +5,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
-from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap
+from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap, disable_gc
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -525,6 +525,7 @@ replace_contiguous = PatternMatcher([
  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
 ])

+@disable_gc()
@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
@@ -538,6 +539,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:

  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
  tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
@@ -10,6 +10,7 @@ from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, p
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin
+from tinygrad.mixin.movement import _align_left
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Device, Buffer
@@ -79,10 +80,6 @@ def _apply_winograd_matrix(mat, t:Tensor, dims:int) -> Tensor:
  assert isinstance(ret, Tensor), "sum didn't return a Tensor"
  return ret

-def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
-  # unsqueeze left to make every shape same length
-  max_dim = max(len(shape) for shape in shapes)
-  return tuple((1,) * (max_dim - len(shape)) + shape for shape in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
  return tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*_align_left(*shapes)))

@@ -175,7 +172,7 @@ class Tensor(OpMixin):

  def _apply_uop(self, fxn:Callable, *x:Tensor, extra_args=(), **kwargs) -> Tensor:
    new_uop: UOp = fxn(*[t.uop for t in (self,)+x], *extra_args, **kwargs)
-    if (metadata:=_METADATA.get()) is not None: all_metadata[new_uop] = (metadata,)
+    if (metadata:=_METADATA.get()) is not None and TRACEMETA >= 1: all_metadata[new_uop] = (metadata,)
    needs_input_grad = [t.requires_grad for t in (self,)+x]
    return Tensor(new_uop, device=new_uop.device, requires_grad=True if any(needs_input_grad) else None if None in needs_input_grad else False)

@@ -1040,21 +1037,6 @@ class Tensor(OpMixin):

  def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)

-  def expand(self, shape, *args) -> Tensor:
-    """
-    Returns a tensor that is expanded to the shape that is specified.
-    Expand can also increase the number of dimensions that a tensor has.
-
-    Passing a `-1` or `None` to a dimension means that its size will not be changed.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.expand(4, -1).numpy())
-    ```
-    """
-    new_shape = tuple(from_ if to == -1 or to is None else to for from_, to in zip(*(_align_left(self.shape, argfix(shape, *args)))))
-    return self._broadcast_to(new_shape)
-
  def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Tensor:
    """
    Returns a tensor with padding applied based on the input `padding`.
@@ -1122,8 +1104,6 @@ class Tensor(OpMixin):
  def pad_to(self, shape, *args):
    if len(new_shape := argfix(shape, *args)) != self.ndim: raise ValueError(f"dim mismatch, cannot pad {self.shape} to {new_shape}")
    return self.pad(tuple([None if ns is None else (0, ns-s) for s,ns in zip(self.shape, new_shape)]))
-  def shrink_to(self, shape, *args):
-    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))

  # ***** movement high level ops *****

@@ -1343,39 +1323,6 @@ class Tensor(OpMixin):
    # checks for shapes and number of dimensions delegated to cat
    return Tensor.cat(*[t.unsqueeze(dim) for t in argfix(self, *args)], dim=dim)

-  def repeat_interleave(self, repeats:int, dim:int|None=None) -> Tensor:
-    """
-    Repeats elements of a tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat_interleave(2).numpy())
-    ```
-    """
-    x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
-    shp = x.shape
-    return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])
-
-  def repeat(self, repeats, *args) -> Tensor:
-    """
-    Repeats tensor number of times along each dimension specified by `repeats`.
-    `repeats` can be passed as a tuple or as separate arguments.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([1, 2, 3])
-    print(t.repeat(4, 2).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(t.repeat(4, 2, 1).shape)
-    ```
-    """
-    repeats = argfix(repeats, *args)
-    base_shape = _align_left(self.shape, repeats)[0]
-    unsqueezed_shape = flatten([[1, s] for s in base_shape])
-    expanded_shape = flatten([[r, s] for r,s in zip(repeats, base_shape)])
-    final_shape = [r*s for r,s in zip(repeats, base_shape)]
-    return self.reshape(unsqueezed_shape).expand(expanded_shape).reshape(final_shape)
-
  def split(self, sizes:int|Sequence[int], dim:int=0) -> tuple[Tensor, ...]:
    """
    Splits the tensor into chunks along the dimension specified by `dim`.
@@ -2146,31 +2093,6 @@ class Tensor(OpMixin):

  # ***** processing ops *****

-  def _pool(self, k_:tuple[sint, ...], stride:int|tuple[int, ...]=1, dilation:int|tuple[int, ...]=1) -> Tensor:
-    assert len(self.shape) >= len(k_), f"can't pool {self.shape} with {k_}"
-    s_, d_ = make_tuple(stride, len(k_)), make_tuple(dilation, len(k_))
-    assert len(k_) == len(s_) == len(d_), f"stride/dilation mismatch kernel:{k_} stride:{s_} dilation:{d_}"
-    noop, i_ = [None] * (self.ndim-len(k_)), self.shape[-len(k_):]
-    assert all(resolve(d*(k-1)+1 <= i) for k,d,i in zip(k_,d_,i_)), "kernel size cannot be greater than actual input size"
-    o_ = [ceildiv(i-d*(k-1), s) for i,d,k,s in zip(i_,d_,k_,s_)]
-    if any(resolve(k > s) for k,s in zip(k_,s_)) or any(d != 1 for d in d_):
-      # input size scaling factor to make sure shrink for stride is possible
-      f_ = [1 + int(resolve(o*s > (i - d*(k-1)))) for o,s,i,d,k in zip(o_,s_,i_,d_,k_)]
-      # # repeats such that we don't need padding
-      x = self.repeat([1]*len(noop) + [ceildiv(k*(i*f+d),i) for k,i,d,f in zip(k_,i_,d_,f_)])
-      # handle dilation
-      x = x.shrink(tuple(noop + [(0,k*(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)])).reshape(noop + flatten((k,(i*f+d)) for k,i,d,f in zip(k_,i_,d_,f_)))
-      # handle stride
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o*s)) for k,o,s in zip(k_,o_,s_)))).reshape(noop + flatten((k,o,s) for k,o,s in zip(k_,o_,s_)))
-      x = x.shrink(tuple(noop + flatten(((0,k), (0,o), (0,1)) for k,o in zip(k_,o_)))).reshape(noop + flatten((k,o) for k,o in zip(k_,o_)))
-      # permute to move reduce to the end
-      return x.permute(*range(len(noop)), *[len(noop)+i*2+1 for i in range(len(i_))], *[len(noop)+i*2 for i in range(len(i_))])
-    # TODO: once the shapetracker can optimize well, remove this alternative implementation
-    x = self.pad(tuple(noop + [(0, max(0,o*s-i)) for i,o,s in zip(i_,o_,s_)])).shrink(tuple(noop + [(0,o*s) for o,s in zip(o_,s_)]))
-    x = x.reshape(noop + flatten(((o,s) for o,s in zip(o_,s_))))
-    x = x.shrink(tuple(noop + flatten(((0,o), (0,k)) for o,k in zip(o_,k_))))
-    return x.permute(*range(len(noop)), *[len(noop)+i*2 for i in range(len(i_))], *[len(noop)+i*2+1 for i in range(len(i_))])
-
  def _resolve_pool_pads(self, padding:int|Sequence[int], dims:int) -> Sequence[int]:
    if not isinstance(padding, int) and not (len(padding) == 2*dims or len(padding) == dims):
      raise ValueError(f"Padding must be an int or a sequence of length {dims} or {2*dims}, but got {padding=} for {self.shape=} with {dims=}.")
@@ -3407,18 +3329,8 @@ class Tensor(OpMixin):
    return self / (1 + self.abs())

  # ***** broadcasted elementwise ops *****
-  def _broadcast_to(self, new_shape:tuple[sint, ...]) -> Tensor:
-    if self.shape == new_shape: return self
-    if self.ndim > len(new_shape): raise ValueError(f"cannot broadcast tensor to fewer dimensions. shape={self.shape} to {new_shape=}")
-    # first unsqueeze left with 1s https://data-apis.org/array-api/latest/API_specification/broadcasting.html
-    shape, _ = _align_left(self.shape, new_shape)
-    # for each dimension, check either dim is 1, or it does not change
-    if not all(resolve(s == ns) or resolve(s == 1) for s,ns in zip(shape, new_shape)):
-      raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
-    # NOTE: this cast is no-op in forward and uses sum_acc_dtype in the backward sum
-    return self.reshape(shape).cast(sum_acc_dtype(self.dtype))._apply_uop(UOp.expand, arg=new_shape).cast(self.dtype)

-  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True) -> tuple[Tensor, Tensor]:
+  def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True, backward_cast:bool=True) -> tuple[Tensor, Tensor]:
    x: Tensor = self
    if not isinstance(y, Tensor):
      # make y a Tensor
@@ -3434,8 +3346,13 @@ class Tensor(OpMixin):

    if reverse: x, y = y, x

+    # compute the output shape
+    out_shape = _broadcast_shape(x.shape, y.shape)
+
    # broadcast
-    return x._broadcast_to(out_shape:=_broadcast_shape(x.shape, y.shape)), y._broadcast_to(out_shape)
+    # NOTE: the backward cast is no-op in forward and uses sum_acc_dtype in the backward sum
+    return x.cast(sum_acc_dtype(x.dtype) if backward_cast else x.dtype)._broadcast_to(out_shape).cast(x.dtype), \
+           y.cast(sum_acc_dtype(y.dtype) if backward_cast else y.dtype)._broadcast_to(out_shape).cast(y.dtype)

  def sub(self, x:Tensor|ConstType, reverse=False) -> Tensor:
    """
@@ -4175,18 +4092,17 @@ class Tensor(OpMixin):

  def image_dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor:
    # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1)
-    x, dx, dw = self, self.ndim, w.ndim
-    if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D")
-    if x.shape[-1] != w.shape[-min(w.ndim, 2)]: raise RuntimeError(f"cannot image_dot {x.shape} and {w.shape}")
+    if not (self.ndim > 0 and w.ndim > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {self.ndim=}, {w.ndim=}")
+    if self.shape[-1] != w.shape[-min(w.ndim, 2)]: raise RuntimeError(f"cannot image_dot {self.shape} and {w.shape}")

    bs, groups, cin, cout = prod(self.shape[0:-2]), prod(w.shape[0:-2]), w.shape[-2], w.shape[-1]
-    out_shape_t = self.shape[0:-2] + (cout,-1) if len(self.shape) > 1 else (cout, )
+    out_shape_t = self.shape[0:-2] + (cout,-1) if len(self.shape) > 1 else (cout,)

    # NOTE: with NHWC we can remove the transposes
    # bs x groups*cin x H x W
-    cx = self.transpose(self.ndim-1, self.ndim-2).reshape((bs//groups, groups*cin, -1, 1))
+    cx = self.transpose(self.ndim-1, self.ndim-2).reshape(bs//groups, groups*cin, -1, 1)
    # groups*cout x cin x H, W
-    cw = w.transpose(w.ndim-1, w.ndim-2).reshape((groups*cout, cin, 1, 1))
+    cw = w.transpose(w.ndim-1, w.ndim-2).reshape(groups*cout, cin, 1, 1)
    return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2)

  def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor:
@@ -4199,10 +4115,9 @@ class Tensor(OpMixin):
    if cin % 4 != 0 and not (cin == 1 and groups%4 == 0):
      x = x.reshape(bs, groups, cin, iy, ix)   # do this always?
      added_input_channels = 4 - (cin % 4)
-      w = w.pad(tuple((0, added_input_channels) if i == 2 else None for i in range(w.ndim)))
-      x = x.pad(tuple((0, added_input_channels) if i == 2 else None for i in range(x.ndim)))
      cin = cin + added_input_channels
-      x = x.reshape(bs, groups*cin, iy, ix)
+      w = w.pad_to(None, None, cin, None, None)
+      x = x.pad_to(None, None, cin, None, None).reshape(bs, groups*cin, iy, ix)

    # hack for non multiples of 4 on rcout
    added_output_channels = 0
@@ -4210,7 +4125,7 @@ class Tensor(OpMixin):
      added_output_channels = 4 - (rcout % 4)
      rcout += added_output_channels
      cout = groups * rcout
-      w = w.pad(tuple((0, added_output_channels) if i == 1 else None for i in range(w.ndim)))
+      w = w.pad_to(None, rcout, None, None, None)

    # packed (note: flipping bs and iy would make the auto-padding work)
    x = x.permute(0,2,3,1)
@@ -4224,18 +4139,18 @@ class Tensor(OpMixin):
    x, w = x.contiguous(), w.contiguous()

    # expand out
-    rcin_hi, rcin_lo = cin//4 if cin >= 4 else 1, 4 if cin >= 4 else 1
-    cout_expand = [groups//4 if cin == 1 else groups, 4 if cin == 1 else 1, rcout//4 if rcout >= 4 else 1, 4 if rcout >= 4 else 1]
+    rcin_hi, rcin_lo = (cin//4, 4) if cin >= 4 else (1, 1)
+    group_shape, rcout_expand = (groups//4, 4) if cin == 1 else (groups, 1), (rcout//4, 4) if rcout >= 4 else (1, 1)
    x = x.reshape(bs, iy, ix, groups, rcin_hi, rcin_lo)
    if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo)
    else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)

    # prepare input
    x = x.permute(0,3,4,5,1,2).pad(self._resolve_pool_pads(padding,2))._pool((H,W), stride, dilation)# -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
-    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *cout_expand[0:2], 1, 1, rcin_hi, rcin_lo, H, W)
+    x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *group_shape, 1, 1, rcin_hi, rcin_lo, H, W)

    # prepare weights
-    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *cout_expand, rcin_hi, rcin_lo, H, W))
+    w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *group_shape, *rcout_expand, rcin_hi, rcin_lo, H, W))

    # the conv!
    ret = (x*w).cast(base_image_type((bs*oy, ox*cout//4, 4)) if IMAGE >= 2 else dtypes.float32).sum((-4, -3, -2, -1), dtype=dtype)
@@ -4263,7 +4178,7 @@ _METADATA: _ContextVar[Metadata|None] = _ContextVar(default=None)

 def _metadata_wrapper(fn: Callable[P, T]) -> Callable[P, T]:
  def _wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
-    if _METADATA.get() is not None: return fn(*args, **kwargs)
+    if TRACEMETA < 1 or _METADATA.get() is not None: return fn(*args, **kwargs)

    if TRACEMETA >= 2:
      caller_frame = sys._getframe(frame := 1)
@@ -1,3 +1,5 @@
+# flake8: noqa: E702
+# allow semicolons to put multiple ops on one line
 from enum import auto, IntEnum, Enum

 # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
@@ -9,9 +11,21 @@ class FastEnum(IntEnum):

 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
+  # ** 1 -- defines/special **
+
+  # TODO: unify these ops into the levels of the memory hierarchy
+  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto()
+
+  # this is for symbolic shapes
+  DEFINE_VAR = auto(); BIND = auto()
+
+  # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
+  SPECIAL = auto()
+
+  # ** 2 -- non op uops **
+
  # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702
-  SENTINEL = auto()
+  NOOP = auto(); SINK = auto(); PRECAST = auto()

  # AFTER passes src[0] through and promises in the toposort that any consumers of the AFTER run after src[1:]
  AFTER = auto()
@@ -19,64 +33,70 @@ class Ops(FastEnum):
  # GROUP is a NOOP that just merges things together
  GROUP = auto()

-  # buffer ops
-  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
+  # vector creation / item selection
+  GEP = auto(); VECTORIZE = auto()

-  # create buffer
-  BUFFERIZE = auto()
-
-  # ops that adjust the behavior of the scheduler
-  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto() # noqa: E702
-
-  # movement ops! these only exist in the tensor graph
-  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto() # noqa: E702
-  MULTI = auto()  # MULTI is really a movement op
-
-  # TODO: unify these ops into the levels of the memory hierarchy. depends on ASSIGN is STORE
-  DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto() # noqa: E702
-
-  # this is for symbolic shapes
-  DEFINE_VAR = auto(); BIND = auto() # noqa: E702
-
-  # this is a RANGE for GPU dimensions, similar to symbolic shapes but not exactly
-  SPECIAL = auto()
-
-  # reduce
-  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702
-
-  # optimization helper ops
-  UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702
-
-  # UnaryOps
-  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto() # noqa: E702
-
-  # load/store before math
-  LOAD = auto(); STORE = auto() # noqa: E702
-  ASSIGN = auto()  # TODO: ASSIGN is STORE, remove ASSIGN
-
-  # tensor core math op, not elementwise
-  WMMA = auto()
+  # ** 3 -- load/store **

  # INDEX is a BinaryOp similar to ADD, but it operates on pointers
  INDEX = auto()

+  # load/store before math
+  LOAD = auto(); STORE = auto()
+
+  # ** 4 -- math **
+
+  # tensor core math op, not elementwise
+  WMMA = auto()
+
+  # UnaryOps
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()
+  SQRT = auto(); RECIPROCAL = auto(); NEG = auto(); TRUNC = auto()
+
  # BinaryOps
-  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto() # noqa: E702
-  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto() # noqa: E702
-  XOR = auto(); OR = auto(); AND = auto() # noqa: E702
-  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto() # noqa: E702
+  ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto()
+  CMPLT = auto(); CMPNE = auto(); CMPEQ = auto()
+  XOR = auto(); OR = auto(); AND = auto()
+  THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto()

  # TernaryOps
-  WHERE = auto(); MULACC = auto() # noqa: E702
+  WHERE = auto(); MULACC = auto()
+
+  # ** 5 -- control flow / consts / custom **

  # control flow ops
-  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto() # noqa: E702
+  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto()

  # consts. VCONST is a vectorized const
-  VCONST = auto(); CONST = auto() # noqa: E702
+  VCONST = auto(); CONST = auto()

  # CUSTOM/CUSTOMI are used to output strings into codegen. the I makes the string inline
-  CUSTOM = auto(); CUSTOMI = auto() # noqa: E702
+  CUSTOM = auto(); CUSTOMI = auto()
+
+  # ** 6 -- ops that don't exist in programs **
+
+  # tensor graph ops
+  UNIQUE = auto(); DEVICE = auto(); KERNEL = auto()
+  ASSIGN = auto()
+
+  # buffer ops
+  BUFFERIZE = auto(); COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto()
+
+  # ops that adjust the behavior of the scheduler
+  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
+
+  # movement ops! these only exist in the tensor graph
+  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
+  MULTI = auto()  # MULTI is really a movement op
+
+  # reduce
+  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto()
+
+  # errors/placeholders
+  REWRITE_ERROR = auto(); SENTINEL = auto()
+
+  # expander ops
+  UNROLL = auto(); CONTRACT = auto(); CAT = auto(); PTRCAT = auto()

 class GroupOp:
  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC}
@@ -4,11 +4,10 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
 from dataclasses import dataclass
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
-from tinygrad.mixin import OpMixin
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType, AddrSpace
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CI
-from tinygrad.helpers import strip_parens, colored, ansilen
+from tinygrad.helpers import strip_parens, colored, ansilen, printable
 if TYPE_CHECKING:
  from tinygrad.device import Buffer, MultiBuffer

@@ -48,6 +47,11 @@ def range_str(u:UOp, color=False) -> str:
  ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
  return colored(ret, axis_colors[u.arg[-1]]) if color else ret

+def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
+  ret = ','.join([range_str(x, color=color) for x in sorted(rngs, key=lambda x: x.arg)])
+  if pad is not None: ret += " " * (pad-ansilen(ret))
+  return ret
+
 def consumer_map_from_toposort(lst:Iterable[UOp]):
  ret: dict[UOp, dict[UOp, None]] = {}
  for u in lst:
@@ -102,6 +106,9 @@ class recursive_property(property):
        s.__dict__[self.nm] = val = self.fxn(s)
    return val

+# we import this late so we can use resolve/smax in mixins
+from tinygrad.mixin import OpMixin
+
 # NOTE: this should be frozen, but frozen is slower
@dataclass(eq=False, slots=True)
 class UOp(OpMixin, metaclass=UOpMetaClass):
@@ -212,10 +219,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

      # ops with custom handling
      case Ops.KERNEL: return self.arg.ast._shape
-      case Ops.STORE:
-        if isinstance(self.dtype, PtrDType): return (self.ptrdtype.size,)
-        if self.dtype is not dtypes.void: return self.src[0].src[0].shape
-        return None

      # TODO: disallow shape changing bitcast
      case Ops.BITCAST:
@@ -267,7 +270,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
          return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))

    # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE}):
+    if self.op in (GroupOp.Elementwise-{Ops.BITCAST}).union({Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
      # TODO: remove this hack for 3 op assign
      input_shapes = [x._shape for x in (self.src[:2] if self.op is Ops.ASSIGN else self.src) if x._shape is not None]
      if len(input_shapes) == 0: return None
@@ -554,7 +557,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  # in these four, if the shape doesn't change we can return self
  def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
  #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
-  def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
+  #def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
  #def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
  def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)

@@ -784,8 +787,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

  # *** uop high level syntactic sugar ***

-  def shrink_to(self, arg:tuple[sint, ...]): return self.shrink(tuple([(0,x) for x in arg]))
-
  @staticmethod
  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
    lookup = {AddrSpace.GLOBAL: Ops.DEFINE_GLOBAL, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
@@ -851,10 +852,10 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
 # ***** uop helpers *****

 def print_uops(uops:list[UOp]):
+  uops_index = {u:i for i,u in enumerate(uops)}
  for i,u in enumerate(uops):
-    formatted_srcs = [(uops.index(x) if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
-    formatted_range = ','.join([range_str(r, color=True) for r in sorted(u.ranges, key=lambda x: x.arg)])
-    print(f"{i:4d} {str(u.op):20s}: {(formatted_range)+' '*(10-ansilen(formatted_range))} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
+    formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src]
+    print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")

 # ***** pattern matcher *****

@@ -866,14 +867,6 @@ def get_location() -> tuple[str, int]:
    frm = frm.f_back
  return frm.f_code.co_filename, frm.f_lineno

-@functools.cache
-def lines(fn) -> list[str]:
-  with open(fn) as f: return f.readlines()
-
-def printable(loc:tuple[str, int]) -> str:
-  try: return lines(loc[0])[loc[1]-1].strip()
-  except FileNotFoundError: return "<missing>"
-
 class UPat(OpMixin):
  __slots__ = ("op", "dtype", "arg", "name", "src")
  def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|None=None,
@@ -993,11 +986,18 @@ def upat_interpret(p:UPat, fxn:Callable) -> Callable:
      return None
  return universal_match

+def fixup_pm_function(fxn) -> Callable:
+  if isinstance(fxn, UPat):
+    # TODO: write this
+    raise NotImplementedError("rhs UPat is not supported")
+  if isinstance(fxn, tuple): return types.FunctionType(*fxn)
+  return fxn
+
 class PatternMatcher:
  def __init__(self, patterns:Sequence[tuple[UPat, Callable|tuple]], compiled=bool(getenv("UPAT_COMPILE", 1))):
    if compiled: from tinygrad.uop.upat import upat_compile
    # if this comes from a pickle, we reconstruct the lambda functions here
-    self.patterns:list[tuple[UPat, Callable]] = [(p,types.FunctionType(*fxn) if isinstance(fxn, tuple) else fxn) for p,fxn in patterns]
+    self.patterns:list[tuple[UPat, Callable]] = [(p,fixup_pm_function(fxn)) for p,fxn in patterns]
    # NOTE: use of DefaultDict here is very dangerous! all keys will live for the lifetime of the PatternMatcher!
    self.pdict: dict[Ops, list[tuple[UPat, Callable, set]]] = {}
    # uop is required, arg is optional
@@ -134,10 +134,6 @@ shared_codegen_spec = PatternMatcher([
  # WMMA has a <a, b, acc>
  (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),

-  # UNROLL/CONTRACT is used here for WMMA
-  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
-
  # VECTORIZE/GEP
  (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
  (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
@@ -166,6 +162,10 @@ kernel_spec = PatternMatcher([
  # index is allowed here
  (UPat(GroupOp.Elementwise|{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR}, dtype=dtypes.index), lambda: True),

+  # UNROLL/CONTRACT is used here for WMMA
+  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
+  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
+
  # END can end multiple axes here
  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),

@@ -48,8 +48,10 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
  (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
  ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base),  # (x%y)%y = -> x%y (rewritten with base for speed)
-  # 4 variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
+  # variations of (x%c)+(x//c)*c = x   TODO: add sorting to remove some variations
  (UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
+  ((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
+    lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
  ((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
    lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
  ((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),
@@ -428,10 +430,10 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
        newuops = [uop.substitute({X:newX}) for X,newX in candidate]
        if any(u is uop for u in newuops): continue  # if any branch doesnt appear in uop, skip
        newuops = [u.simplify().substitute({newX:X}).simplify(full_symbolic=False) for (X,newX),u in zip(candidate,newuops)]
-        if uop.op is Ops.VECTORIZE and len(uop.src) == 2:
+        if all_same(newuops): uop = newuops[0]
+        elif uop.op is Ops.VECTORIZE and len(uop.src) == 2:
          if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1]))
          if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1]))
-        elif all_same(newuops): uop = newuops[0]

  # try all the valids together (but only the whole expressions)
  if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop:
@@ -491,7 +493,8 @@ pm_move_where_on_load = PatternMatcher([
 pm_simplify_valid = PatternMatcher([
  # simplify valid
  (UPat(Ops.AND, name="valid"), simplify_valid),
-  (UPat.var("c").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda c,x,i: c.where(uop_given_valid(c, x, try_simplex=False), i)),
+  # TODO: this regressed openpilot, not having this regressed cifar
+  # (UPat.var("c").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda c,x,i: c.where(uop_given_valid(c, x, try_simplex=False), i)),
 ])

 # this is symbolic 2.0
@@ -292,6 +292,7 @@
  }
  .raw-text > pre {
    display: inline-block;
+    min-width: 100%;
  }
  .raw-text code {
    max-height: none !important;
@@ -25,7 +25,7 @@ const colored = n => d3.create("span").call(s => s.selectAll("span").data(typeof
 const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();

 let timeout = null;
-const updateProgress = ({ start }) => {
+const updateProgress = ({ start, err }) => {
  clearTimeout(timeout);
  const msg = document.getElementById("progress-message");
  msg.style.display = "none";
@@ -33,6 +33,11 @@ const updateProgress = ({ start }) => {
    msg.innerText = "Rendering new graph...";
    timeout = setTimeout(() => { msg.style.display = "block"; }, 2000);
  }
+  d3.select("#custom").html("");
+  if (err) {
+    displaySelection("#custom");
+    d3.select("#custom").append(() => d3.create("div").classed("raw-text", true).call(s => s.append(() => codeBlock(err, "txt"))).node());
+  }
 }

 function intersectRect(r1, r2) {
@@ -136,6 +141,10 @@ function renderDag(graph, additions, recenter, layoutOpts) {
    }).attr("class", e => e.value.label.type).attr("id", e => `${e.v}-${e.w}`).datum(e => e.value.label.text));
    if (recenter) document.getElementById("zoom-to-fit-btn").click();
  };
+  worker.onerror = (e) => {
+    e.preventDefault();
+    updateProgress({ err:"Error in graph layout:\n"+e.message });
+  }
 }

 // ** profiler graph
@@ -259,7 +268,10 @@ async function renderProfiler() {
        html.append(() => tabulate([["Name", colored(e.name)], ["Duration", formatTime(e.dur)], ["Start Time", formatTime(e.st)]]).node());
        html.append("div").classed("args", true);
        if (e.info != null) html.append("p").style("white-space", "pre-wrap").text(e.info);
-        if (shapeRef != null) html.append("a").text("View codegen rewrite").on("click", () => switchCtx(shapeRef.ctx, shapeRef.step));
+        if (shapeRef != null) {
+          html.append("a").text("View codegen rewrite").on("click", () => switchCtx(shapeRef.ctx, shapeRef.step));
+          html.append("a").text("View program").on("click", () => switchCtx(shapeRef.ctx, ctxs[shapeRef.ctx+1].steps.findIndex(s => s.name==="View Program")));
+        }
        // tiny device events go straight to the rewrite rule
        const key = k.startsWith("TINY") ? null : `${k}-${j}`;
        if (key != null) shapeMetadata.set(key, html.node());
@@ -488,12 +500,15 @@ async function renderProfiler() {
    }
  }

-  canvas.addEventListener("click", e => {
+  const clickShape = (e) => {
    e.preventDefault();
    const foundRect = findRectAtPosition(e.clientX, e.clientY);
-    if (foundRect?.step != null && foundRect?.key == null) { return switchCtx(foundRect.ctx, foundRect.step); }
+    if (foundRect?.step != null && (foundRect?.key == null || e.type == "dblclick")) { return switchCtx(foundRect.ctx, foundRect.step); }
    if (foundRect?.key != focusedShape) { focusShape(foundRect); }
-  });
+  }
+  canvas.addEventListener("click", clickShape);
+
+  canvas.addEventListener("dblclick", clickShape);

  canvas.addEventListener("mousemove", e => {
    const foundRect = findRectAtPosition(e.clientX, e.clientY);
@@ -536,6 +551,7 @@ document.getElementById("zoom-to-fit-btn").addEventListener("click", () => {

 // **** main VIZ interfacae

+const pathLink = (fp, lineno) => d3.create("a").attr("href", "vscode://file/"+fp+":"+lineno).text(`${fp.split("/").at(-1)}:${lineno}`);
 function codeBlock(st, language, { loc, wrap }={}) {
  const code = document.createElement("code");
  // plaintext renders like a terminal print, otherwise render with syntax highlighting
@@ -544,11 +560,7 @@ function codeBlock(st, language, { loc, wrap }={}) {
  code.className = "hljs";
  const ret = document.createElement("pre");
  if (wrap) ret.className = "wrap";
-  if (loc != null) {
-    const link = ret.appendChild(document.createElement("a"));
-    link.href = "vscode://file/"+loc.join(":");
-    link.textContent = `${loc[0].split("/").at(-1)}:${loc[1]}`+"\n\n";
-  }
+  if (loc != null) ret.appendChild(pathLink(loc[0], loc[1]).style("margin-bottom", "4px").node());
  ret.appendChild(code);
  return ret;
 }
@@ -650,7 +662,7 @@ async function main() {
        u.li = list.appendChild(document.createElement("ul"));
        u.li.id = `step-${i}-${j}`;
        const p = u.li.appendChild(document.createElement("p"));
-        p.innerText = `${u.name}`+(u.match_count ? ` - ${u.match_count}` : '');
+        p.appendChild(colored(`${u.name}`+(u.match_count ? ` - ${u.match_count}` : '')));
        p.onclick = (e) => {
          e.stopPropagation();
          const subrewrites = getSubrewrites(e.currentTarget.parentElement);
@@ -661,7 +673,7 @@ async function main() {
      }
      for (const l of ul.querySelectorAll("ul > ul > p")) {
        const subrewrites = getSubrewrites(l.parentElement);
-        if (subrewrites.length > 0) { l.innerText += ` (${subrewrites.length})`; l.parentElement.classList.add("has-children"); }
+        if (subrewrites.length > 0) { l.appendChild(d3.create("span").text(` (${subrewrites.length})`).node()); l.parentElement.classList.add("has-children"); }
      }
    }
    return setState({ currentCtx:-1 });
@@ -716,8 +728,8 @@ async function main() {
        }
      }
      metadata.appendChild(tabulate(ret.summary.map(s => {
-        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
-        return [s.label.trim(), div.node()];
+        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "100%").style("height", "100%");
+        return [s.label.trim(), div.text(s.value.toLocaleString()).node()];
      })).node());
    } else root.appendChild(codeBlock(ret.src, ret.lang || "txt"));
    return document.querySelector("#custom").replaceChildren(root);
@@ -748,6 +760,15 @@ async function main() {
  // ** right sidebar code blocks
  const codeElement = codeBlock(ret[currentRewrite].uop, "python", { wrap:false });
  metadata.replaceChildren(toggleLabel, codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeElement);
+  if (step.trace) {
+    const trace = d3.create("pre").append("code").classed("hljs", true);
+    for (let i=step.trace.length-1; i>=0; i--) {
+      const [fp, lineno, fn, code] = step.trace[i];
+      trace.append("div").style("margin-bottom", "2px").style("display","flex").text(fn+" ").append(() => pathLink(fp, lineno).node());
+      trace.append("div").html(hljs.highlight(code, { language: "python" }).value).style("margin-bottom", "1ex");
+    }
+    metadata.insertBefore(trace.node().parentNode, codeElement);
+  }
  // ** rewrite steps
  if (step.match_count >= 1) {
    const rewriteList = metadata.appendChild(document.createElement("div"));
@@ -1,14 +1,15 @@
 #!/usr/bin/env python3
 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
 import subprocess, ctypes, pathlib, traceback
-from contextlib import redirect_stdout
+from contextlib import redirect_stdout, redirect_stderr
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, TypeVar, Generator, Callable
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA, ProfileEvent, ProfileRangeEvent, TracingKey, ProfilePointEvent, temp
-from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, printable, GroupOp, srender, sint, sym_infer, range_str, pyrender
-from tinygrad.uop.ops import print_uops, range_start
+from tinygrad.helpers import printable
+from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, pyrender
+from tinygrad.uop.ops import print_uops, range_start, multirange_str
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
@@ -30,7 +31,7 @@ ref_map:dict[Any, int] = {}
 def get_rewrites(t:RewriteTrace) -> list[dict]:
  ret = []
  for i,(k,v) in enumerate(zip(t.keys, t.rewrites)):
-    steps = [{"name":s.name, "loc":s.loc, "match_count":len(s.matches), "code_line":printable(s.loc),
+    steps = [{"name":s.name, "loc":s.loc, "match_count":len(s.matches), "code_line":printable(s.loc), "trace":k.tb if j == 0 else None,
              "query":f"/ctxs?ctx={i}&idx={j}", "depth":s.depth} for j,s in enumerate(v)]
    if isinstance(k.ret, ProgramSpec):
      steps.append({"name":"View UOp List", "query":f"/render?ctx={i}&fmt=uops", "depth":0})
@@ -78,11 +79,14 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
        label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
    try:
      if len(rngs:=u.ranges):
-        label += f"\n({','.join([range_str(x, color=True) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
-      if u.op not in {Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u._shape is not None:
+        label += f"\n({multirange_str(rngs, color=True)})"
+      if u._shape is not None:
        label += f"\n{shape_to_str(u.shape)}"
      if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
        label += f"\n{u.render()}"
+        ranges: list[UOp] = []
+        for us in u.src[1:]: ranges += [s for s in us.toposort() if s.op in {Ops.RANGE, Ops.SPECIAL}]
+        if ranges: label += "\n"+' '.join([f"{s.render()}={s.vmax+1}" for s in ranges])
      if u.op in {Ops.END, Ops.REDUCE} and len(trngs:=list(UOp.sink(*u.src[range_start[u.op]:]).ranges)):
        label += "\n"+' '.join([f"{range_str(s, color=True)}({s.vmax+1})" for s in trngs])
    except Exception:
@@ -147,17 +151,22 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
  for st,et,dur,e in dev_events:
    if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.arg["name"]] = e
    if dur == 0: continue
-    name, info, key = e.name, None, None
+    name, fmt, key = e.name, [], None
    if (ref:=ref_map.get(name)) is not None:
      name = ctxs[ref]["name"]
      if isinstance(p:=trace.keys[ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
-        info = f"{sym_infer(p.estimates.ops, ei.arg['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei.arg['var_vals'])/t:4.1f}"+ \
-               f"|{sym_infer(p.estimates.lds,ei.arg['var_vals'])/t:.1f} GB/s\n{[str(m) for m in (ei.arg['metadata'] or ())]}"
+        flops = sym_infer(p.estimates.ops, var_vals:=ei.arg['var_vals'])/(t:=dur*1e-6)
+        membw, ldsbw = sym_infer(p.estimates.mem, var_vals)/t, sym_infer(p.estimates.lds, var_vals)/t
+        fmt = [f"{flops*1e-9:.0f} GFLOPS" if flops < 1e14 else f"{flops*1e-12:.0f} TFLOPS",
+              (f"{membw*1e-9:.0f} GB/s" if membw < 1e13 else f"{membw*1e-12:.0f} TB/s")+" mem",
+              (f"{ldsbw*1e-9:.0f} GB/s" if ldsbw < 1e15 else f"{ldsbw*1e-12:.0f} TB/s")+" lds"]
+        if (metadata_str:=",".join([str(m) for m in (ei.arg['metadata'] or ())])): fmt.append(metadata_str)
+        if isinstance(e, ProfileGraphEntry): fmt.append("(batched)")
        key = ei.key
    elif isinstance(e.name, TracingKey):
      name = e.name.display_name
      ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), st-start_ts, dur, enum_str(info or "", scache)))
+    events.append(struct.pack("<IIIIfI", enum_str(name, scache), option(ref), option(key), st-start_ts, dur, enum_str("\n".join(fmt), scache)))
  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None

 def encode_mem_free(key:int, ts:int, execs:list[ProfilePointEvent], scache:dict) -> bytes:
@@ -201,14 +210,31 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
    return ctxs.append({"name":"Counters", "steps":[step]})
  try: from extra.sqtt.roc import decode
  except Exception: return err("DECODER IMPORT ISSUE")
-  try:
-    rctx = decode(profile)
-    steps = [{"name":str(x[0]), "depth":0, "data":{"rows":[(e.inst, e.time, e.time-x[1][i-1].time if i else 0, e.dur, e.stall,
-                                                            str(e.typ).split("_")[-1]) for i,e in enumerate(x[1])],
-                                              "cols":["Instruction", "Clk", "Wait", "Duration", "Stall", "Type"], "summary":[]},
-              "query":f"/render?ctx={len(ctxs)}&step={i}&fmt=counters"} for i,x in enumerate(rctx.inst_execs.items())]
-    if not steps: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
+  try: rctx = decode(profile)
  except Exception: return err("DECODER ERROR")
+  if not rctx.inst_execs: return err("EMPTY SQTT OUTPUT", f"{len(sqtt_events)} SQTT events recorded, none got decoded")
+  steps:list[dict] = []
+  for name,waves in rctx.inst_execs.items():
+    if (r:=ref_map.get(name)): name = ctxs[r]["name"]
+    steps.append({"name":name, "depth":0, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                  "data":{"src":trace.keys[r].ret.src if r else name, "lang":"cpp"}})
+
+    # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
+    #           The idle time can be caused by:
+    #             * Arbiter loss
+    #             * Source or destination register dependency
+    #             * Instruction cache miss
+    # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
+    # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
+    for w in waves:
+      rows, prev_instr = [], w.begin_time
+      for i,e in enumerate(w.insts):
+        rows.append((e.inst, e.time, max(0, e.time-prev_instr), e.dur, e.stall, str(e.typ).split("_")[-1]))
+        prev_instr = max(prev_instr, e.time + e.dur)
+      summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"CU", "value":w.cu},
+                 {"label":"SIMD", "value":w.simd}]
+      steps.append({"name":f"Wave {w.wave_id}", "depth":1, "query":f"/render?ctx={len(ctxs)}&step={len(steps)}&fmt=counters",
+                    "data":{"rows":rows, "cols":["Instruction", "Clk", "Idle", "Duration", "Stall", "Type"], "summary":summary}})
  ctxs.append({"name":"Counters", "steps":steps})

 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
@@ -268,8 +294,11 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
  for i,usage in instr_usage.items(): rows[i].append([[k, v, (v/max_usage)*100] for k,v in usage.items()])
  return {"rows":rows, "cols":["Opcode", "Latency", {"title":"HW Resources", "labels":resource_labels}], "summary":summary}

-def get_stdout(f:Callable) -> str:
-  with redirect_stdout(buf:=io.StringIO()): f()
+def get_stdout(f: Callable) -> str:
+  buf = io.StringIO()
+  try:
+    with redirect_stdout(buf), redirect_stderr(buf): f()
+  except Exception: traceback.print_exc(file=buf)
  return buf.getvalue()

 def get_render(i:int, j:int, fmt:str) -> dict|None:
@@ -277,8 +306,8 @@ def get_render(i:int, j:int, fmt:str) -> dict|None:
  if not isinstance(prg:=trace.keys[i].ret, ProgramSpec): return None
  if fmt == "uops": return {"src":get_stdout(lambda: print_uops(prg.uops or [])), "lang":"txt"}
  if fmt == "src": return {"src":prg.src, "lang":"cpp"}
-  lib = (compiler:=Device[prg.device].compiler).compile(prg.src)
-  disasm_str = get_stdout(lambda: compiler.disassemble(lib))
+  compiler = Device[prg.device].compiler
+  disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(prg.src)))
  from tinygrad.runtime.support.compiler_cpu import llvm, LLVMCompiler
  if isinstance(compiler, LLVMCompiler):
    mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode()
Author	SHA1	Message	Date
George HotzandGitHub	86d1e42ed8	Merge branch 'master' into algebraic_upat	2025-11-13 09:21:05 -08:00
George HotzandGitHub	faf68c03a8	more mi350x matmul work (#13138 ) * more mi350x matmul work * broken compute	2025-11-13 09:09:28 -08:00
George HotzandGitHub	5c13504bc1	Merge branch 'master' into algebraic_upat	2025-11-13 09:07:54 -08:00
Ayman Jabr GitHub George Hotz	256f81bb02	Fix tracemeta 0 (#13049 ) * chore: tclesius branch resolved * fix: indentation --------- Co-authored-by: George Hotz <[email protected]>	2025-11-13 09:07:11 -08:00
alpharushandGitHub	7e0aaadecd	feat: add repro command to summary (#10930 )	2025-11-13 08:52:27 -08:00
nimlgenandGitHub	6be86dde17	nv: add timeout when repsonding to rpc (#13260 )	2025-11-14 00:42:21 +08:00
nimlgenandGitHub	f9b7586e08	roc: fix blob gc (#13256 )	2025-11-13 23:38:35 +08:00
George HotzandGitHub	263b724143	one cache and bump it (#13258 )	2025-11-13 07:33:31 -08:00
George HotzandGitHub	5efa727b83	move _pool to MovementMixins (#13257 )	2025-11-13 07:28:52 -08:00
geohot	bcdfc109b5	hotfix: disable flaky test	2025-11-13 06:19:28 -08:00
qazalandGitHub	006dea4c3e	roc: only save instruction execs (#13254 )	2025-11-13 21:28:40 +08:00
nimlgenandGitHub	f9586b38ba	system: pci mask and val (#13251 )	2025-11-13 20:44:58 +08:00
George HotzandGitHub	7316da3253	new readme (#13250 ) * new readme * update	2025-11-13 00:48:28 -08:00
geohot	17aa3379e9	hotfix: improve self_tokenize	2025-11-13 00:18:57 -08:00
chenyuandGitHub	4e5a9132e7	JIT_BATCH_SIZE=0 in compile3 (#13245 ) fixed some enqueue time	2025-11-12 23:12:45 -05:00
wozeparrotandGitHub	759557f633	feat: move tk tests to testextra (#13242 )	2025-11-12 17:06:53 -08:00
chenyuandGitHub	3f939f3d3c	update pm_simplify_valid (#13241 ) * update pm_simplify_valid fixed openpilot conv regression * IMAGE training is broken	2025-11-12 19:40:02 -05:00
chenyuandGitHub	f9851a852f	minor update to uop_given_valid [pr] (#13243 ) split from #13241	2025-11-12 19:03:18 -05:00
qazalandGitHub	fe2876a6d8	hotfix: second GB/s in viz (#13240 )	2025-11-13 07:14:27 +08:00
George HotzandGitHub	a23dea202b	actually make AMD_LLVM not default (#13238 )	2025-11-12 15:07:23 -08:00
George HotzandGitHub	ab9fa964d8	DISABLE_COMPILER_CACHE -> CCACHE (#13234 ) * DISABLE_COMPILER_CACHE -> CCACHE * Fix cachekey assignment in Compiler constructor	2025-11-12 15:07:09 -08:00
qazalandGitHub	be2e24cb25	roc: requires sudo to install (#13237 )	2025-11-12 16:59:22 -05:00
geohot	8f1f195b6d	hotfix: no hexdump for usbgpu patch.py	2025-11-12 12:05:37 -08:00
nimlgenandGitHub	9a53fcbde4	amd: sqtt on rdna3.5 (#13233 )	2025-11-13 03:30:42 +08:00
George HotzandGitHub	13f10a31dc	AMD_LLVM default off (#13232 )	2025-11-12 11:06:33 -08:00
qazalandGitHub	8b26cf2b3d	sqtt: update rcp timing test (#13231 ) * sqtt: assert correct output in timing test * found why	2025-11-13 02:01:54 +08:00
Jan AkhremchikandGitHub	bc8e537423	Add NONZERO op to onnx backend (#13211 )	2025-11-12 08:55:51 -08:00
nimlgenandGitHub	af17e07251	viz: sqtt touchups (#13228 ) * viz: sqtt touchups * revert * matches	2025-11-12 22:40:37 +08:00
qazalandGitHub	7a6853fa40	viz: show python callstack in the first graph (#13218 )	2025-11-12 20:52:28 +08:00
nimlgenandGitHub	82eb63d3ad	qcom: auto switch idle timer when profiling (#13230 ) * qcom: auto switch idle timer when profiling * fi	2025-11-12 20:31:24 +08:00
nimlgenandGitHub	fcd8d0751a	test_timing for hip (#13229 )	2025-11-12 20:28:58 +08:00
qazalandGitHub	74b9d33acb	viz: direct link to program source (#13227 )	2025-11-12 16:27:13 +08:00
wozeparrotandGitHub	371c1f2355	tk: move tiles to class (#13224 )	2025-11-11 21:53:46 -08:00
sirhcmandGitHub	41a098a82d	In-tree autogen: libc.py (#13217 ) * checkout changes from autogen branch * parents * pylint happy * move sys to system in helpers.py * typo * typo	2025-11-11 19:13:48 -08:00
wozeparrotandGitHub	222bb12ddf	tk softmax (#13205 )	2025-11-11 15:13:16 -08:00
wozeparrotandGitHub	787f0070ed	feat: don't use output reg as local reduce reg (#13203 )	2025-11-11 14:35:16 -08:00
chenyuandGitHub	ece1415def	clean up image_dot and image_conv2d (#13222 ) * clean up image_dot and image_conv2d * those are fine * interesting	2025-11-11 15:53:03 -05:00
nimlgenandGitHub	2f0ea29b34	qcom: 48bit timestamps (#13214 ) * qcom: 48bit timestamps * f * lol * fix	2025-11-12 04:14:33 +08:00
qazalandGitHub	bc55bc4849	cleanup test_viz profiler tests (#13221 )	2025-11-12 03:46:48 +08:00
chenyuandGitHub	23b90945c3	add a benchmark for openpilot vision with DEBUG=2 (#13219 ) see per kernel speed, also disable the jobs for 0.9.9	2025-11-11 14:41:52 -05:00
George HotzandGitHub	c2075f3613	gc disable during big rewrites (#13215 ) * gc disable during big rewrites * cleaner with helper	2025-11-11 10:30:47 -08:00
Roelof van DijkandGitHub	e59313da08	migrate pytest and ruff (#13216 )	2025-11-11 13:27:51 -05:00
Gaétan LepageandGitHub	6fd7ce3832	migrate to pyproject.toml (#13189 ) * migrate to pyproject.toml * move mypy config to pyproject.toml	2025-11-11 09:09:27 -08:00
qazalandGitHub	8002921a04	viz: improve the program run tooltip (#13212 ) * add tflops to tooltip format * show if the run was batched	2025-11-12 00:56:03 +08:00
qazalandGitHub	f91e366a17	viz: display the graph layout recursion error (#13194 ) * viz: display the graph layout recursion error * share styles * +min-width * same thing * inline the append	2025-11-11 15:25:12 +08:00
wozeparrotandGitHub	73497af4c0	clean: use np for allclose (#13204 )	2025-11-10 23:02:43 -08:00
George HotzandGitHub	a6360fd94d	store can have shape (#13202 ) * store can have shape * _shape	2025-11-10 22:16:47 -08:00
b1tg GitHub chenyu	f3692b7406	clean up hip renderer (#13063 ) * clean up hip renderer * ocml --------- Co-authored-by: chenyu <[email protected]>	2025-11-11 00:44:24 -05:00
chenyuandGitHub	22b8579234	one last regressed dm kernel (#13201 )	2025-11-10 23:30:52 -05:00
chenyuandGitHub	58b7e4fab3	GROUPTOP heuristic on more axes (#13206 ) fixed dm speed	2025-11-10 23:30:37 -05:00
chenyuandGitHub	829cdafccc	update openpilot slow conv uop ast (#13197 ) the two remaining slow ones	2025-11-10 17:03:20 -05:00
George HotzandGitHub	0c978d45e6	stub attention (#13196 ) * stub attention * name the kernels	2025-11-10 13:48:38 -08:00
chenyuandGitHub	58c30fc7ce	minor image_conv2d cleanup (#13193 )	2025-11-10 16:05:40 -05:00
chenyuandGitHub	60e55d9a2d	line count 18500 (#13191 )	2025-11-10 13:52:13 -05:00
nimlgenandGitHub	09a59c2203	qcom: support new chip versioning (#13185 ) * qcom: support new chip versioning * ops * nit * fix * f	2025-11-10 23:57:29 +08:00
qazalandGitHub	50934050bc	sqtt: append all wave execs (#13190 )	2025-11-10 23:50:08 +08:00
qazalandGitHub	38a24731a1	cleanup sqtt tooling (#13188 ) * cleanup viz/serve.py * use latest profile in rgptool.py * unwrap nullable in roc.py, fix disasms typing	2025-11-10 20:52:57 +08:00
qazalandGitHub	845a24dcc6	viz: group sqtt waves by program (#13187 ) * viz: group sqtt waves by program * color the names	2025-11-10 19:25:23 +08:00
George HotzandGitHub	fd6803000e	mutmut cfg (#13184 ) * mutmut cfg * coveragerc	2025-11-09 23:29:29 -08:00
wozeparrotandGitHub	6252831ceb	feat: initial tk library (#13160 )	2025-11-09 22:54:29 -08:00
George HotzandGitHub	925231aec1	repeat does less reshape for 1s (#13183 )	2025-11-09 19:43:02 -08:00
geohot	d7369de048	hotfix: update weekly commits table	2025-11-09 19:37:06 -08:00
chenyuandGitHub	6c48c87e51	improved ASSERT_MIN_STEP_TIME (#13182 ) * improved ASSERT_MIN_STEP_TIME getting close, current time +1ms then round up * relax	2025-11-09 16:41:12 -05:00
nimlgenandGitHub	17715688c7	system: validate vendor for APLPCIIfaceBase (#13181 )	2025-11-10 02:49:21 +08:00
nimlgenandGitHub	614783693e	nv: remove hardcoded expansion_rom_off (#13180 ) * nv: remove hardcoded expansion_rom_off * to max size	2025-11-09 21:43:19 +08:00
chenyuandGitHub	e1d46de8f8	update GROUPTOP heuristic more (#13178 ) reverts #13176	2025-11-09 02:31:12 -05:00
chenyuandGitHub	41e45c20ff	minor stuff reading the printed code [pr] (#13177 )	2025-11-09 00:58:51 -05:00
chenyuandGitHub	8e868dced8	only GROUPTOP one reduce kernel (#13176 ) * only GROUPTOP one reduce kernel * ALLOWED_GATED_READ_IMAGE=148	2025-11-08 22:38:44 -05:00
chenyuandGitHub	834067d91f	move onnx import in compile3 (#13172 ) only used in test_vs_onnx	2025-11-08 09:44:34 -08:00
nimlgenandGitHub	7f3240dbfe	nv: cleanup alloc (#13170 ) * nv: cleanup alloc * okay okay	2025-11-09 00:14:46 +08:00
qazalandGitHub	7250fc0354	viz: double click on kernel run goes to codegen (#13147 )	2025-11-08 23:40:50 +08:00
qazalandGitHub	8a7fa9e7b4	sqtt: show total cycles of kernel in viz (#13169 )	2025-11-08 21:00:40 +08:00
chenyuandGitHub	2ba8b4946f	external_benchmark_op_cat.py (#13168 ) * external_benchmark_op_cat.py cat kernel that's 1ms on master and 50us with no GROUP and with NOLOCALS * fix	2025-11-08 01:54:10 -05:00
chenyuandGitHub	a62496cb3d	clean up get_grouped_dims [pr] (#13159 )	2025-11-08 01:53:54 -05:00
wozeparrotandGitHub	eb0192b0bb	feat: print ranges that aren't ended (#13167 )	2025-11-07 22:01:29 -08:00
George HotzandGitHub	b41541bc44	bounty: Remove Tensor._pool alternative implementation and verify kernels remain the same (#13164 )	2025-11-07 16:59:48 -08:00
George HotzandGitHub	ffb9e8396f	fix indexing bug with convs * minimal difference for ONE_POOL=1 * fix indexing bug * improve indexing debugger * more debugger improvements * always for reshape	2025-11-07 16:45:19 -08:00
chenyuandGitHub	6a509da7f3	Scheduler.reduceops helper [pr] (#13162 )	2025-11-07 18:59:46 -05:00
George HotzandGitHub	2413311289	make _pool simpler (#13161 ) * make _pool simpler * just syntax * more correct and smaller * try this now * Revert "try this now" This reverts commit `607cdc2164`. * ONE_POOL	2025-11-07 15:58:44 -08:00
George HotzandGitHub	70054cdb14	move backward cast to broadcasted, expand to mixins (#13156 ) * shrink_to mixin * move backward cast into _broadcasted * expand to movement mixin * move a few more * fix spec issue	2025-11-07 15:07:47 -08:00
George HotzandGitHub	f2519ea0ba	shrink_to mixin (#13155 )	2025-11-07 11:46:24 -08:00
C TandGitHub	0f9d7f650d	whisper: fix oob, explicit dtype (#13144 ) * fix dtype depending on numpy version numpy v2 np.array returns int64 which Tensor passed through for the first decode call, swallowing the <\|notimestamps\|> token and corrupting the sequence * fix whisper OOB global limit on whisper's context length * enforce whisper max_tokens_to_sample (match openai) local limit on max tokens decoded	2025-11-07 12:55:01 -05:00
Ahmed Harmouche GitHub chenyu	3ecff3a8da	Fix dim splitting bug for len(dim) == len(limited) case (#13142 ) * Fix gpudims bug on webgpu * Fix split dim bug * Remove webgpu_bug from examples * Add test for shape correctness * Fix 3D indexing --------- Co-authored-by: chenyu <[email protected]>	2025-11-07 12:31:06 -05:00
nimlgenandGitHub	b8e48effcb	device: no compilers message with reasons (#13146 ) * device: no compilers message with reasons * typings * mypy	2025-11-07 23:01:45 +08:00
nimlgenandGitHub	35e461ef69	hcq: use exception group (#12616 ) * hcq: use exception group * fix	2025-11-07 21:23:12 +08:00
nimlgenandGitHub	10dc8335d2	tinygpu: fix teardown crash (#13143 ) * tinygpu: fix crash * um? * double relase * restore	2025-11-07 19:52:54 +08:00
qazalandGitHub	d4a216d7d9	viz: display compiler errors (#13141 )	2025-11-07 18:09:50 +08:00
qazalandGitHub	7e94369464	add helper for test_timing custom ops (#13140 )	2025-11-07 17:13:55 +08:00
nimlgenandGitHub	95620426d5	tinygpu: unmap dma when client closed (#13129 ) * tinygpu: unmap dma when client closed * syn * tiny fixes	2025-11-07 16:08:43 +08:00
wozeparrotandGitHub	500d7661fa	feat: show range len on index in viz (#13139 )	2025-11-06 23:21:27 -08:00
George HotzandGitHub	bb6364d7c7	tuplize from linearizer behind flag (#13136 ) * remove tuplize from linearizer * optional tuplize	2025-11-06 20:15:03 -08:00
chenyuandGitHub	bb8cf948f2	variation of (x%c)+(x//c)*c = x (#13135 ) when x is in the form of y//b, the idiv term might have combined	2025-11-06 18:53:28 -05:00
George HotzandGitHub	42b34cf83d	bottom up linearizer (#13133 ) * bottom up linearizer * late stores * more complete * remove broken heuristic * upcast size * opt * more conservative * it needs that * disable opencl half on QCOM * fix * make that a real test * cpu test okay * ptx skip * end is after the range	2025-11-06 15:30:32 -08:00
geohot	e0d828dba8	little cleanups	2025-11-06 13:58:19 -08:00
chenyuandGitHub	bfb0c0391f	test custom eye function (#13134 ) this version is also faster with NOOPT	2025-11-06 14:51:55 -05:00
George HotzandGitHub	290441dd44	do loads early (#13131 ) * do loads early * local and reg	2025-11-06 09:57:09 -08:00
George HotzandGitHub	097264853d	very simple priority (#13130 ) * very simple priority * still simple	2025-11-06 09:25:28 -08:00
George HotzandGitHub	07b415e831	fixup op order (#13128 ) * fixup op order * more order * move a few more * more * DEBUG_LINEARIZE	2025-11-06 08:50:04 -08:00
geohot	6538935441	write tests for algebraic UPat	2025-10-05 08:21:38 +08:00