ONE_POOL

Revert "try this now"
This reverts commit 607cdc2164.
2025-11-07 15:56:11 -08:00 · 2025-11-07 15:52:33 -08:00 · 2025-11-07 15:50:08 -08:00 · 2025-11-07 15:44:19 -08:00 · 2025-11-07 15:40:38 -08:00 · 2025-11-07 15:33:55 -08:00
202 changed files with 6167 additions and 4236 deletions
@@ -302,4 +302,4 @@ runs:
    - name: Install mesa (macOS)
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa
+      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -2,7 +2,7 @@ name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
+  PYTHON_CACHE_VERSION: '4'
  APT_CACHE_VERSION: '1'
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
@@ -131,7 +131,7 @@ jobs:
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -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 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      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
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -211,6 +211,7 @@ jobs:
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
+        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
    - name: Run Tensor Core GEMM (PTX)
      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
@@ -238,6 +239,8 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+    - name: Run quantized LLaMA3
+      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -271,6 +274,7 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
+          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -319,9 +323,9 @@ jobs:
    - 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=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=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=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=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
@@ -405,7 +409,7 @@ jobs:
    #    python3 -c "import torch; print(torch.__version__)"
    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      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
    - name: Test tensor cores
      run: |
        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
@@ -523,7 +527,7 @@ jobs:
    - 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=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=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
@@ -619,22 +623,24 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: benchmark openpilot 0.9.9 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    - name: benchmark openpilot 0.9.9 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    - name: benchmark openpilot 0.9.9 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision 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: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy 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: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 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 Space Lab policy + vision
-      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
+      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
+    - 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
+    - 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
+    - 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
+    - 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
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -645,16 +651,6 @@ jobs:
        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (comma)
-        path: |
-          openpilot_compile_0_9_4.txt
-          openpilot_compile_0_9_7.txt
-          openpilot_0_9_4.txt
-          openpilot_0_9_7.txt
-          openpilot_image_0_9_4.txt
-          openpilot_image_0_9_7.txt

  testreddriverbenchmark:
    name: AM Benchmark
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -2,7 +2,7 @@ name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
+  PYTHON_CACHE_VERSION: '4'
  APT_CACHE_VERSION: '1'
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
@@ -204,7 +204,7 @@ jobs:
        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
@@ -230,7 +230,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: linting-only
-        python-version: '3.10'
+        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 .
@@ -243,8 +243,9 @@ jobs:
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+    # broken because of UPatAny
+    #- name: Run TYPED=1
+    #  run: TYPED=1 python -c "import tinygrad"

  unittest:
    name: Unit Tests
@@ -264,8 +265,6 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Check SPEC=1
-      run: SPEC=1 python3 test/test_tiny.py
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
    # TODO: too slow
@@ -294,6 +293,25 @@ jobs:
    - name: Repo line count < 18000 lines
      run: MAX_LINE_COUNT=18000 python sz.py

+  spec:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: SPEC=2 (${{ matrix.group }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: spec-unit
+        deps: testing_unit
+    - name: Test SPEC=2
+      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+
  fuzzing:
    name: Fuzzing
    runs-on: ubuntu-latest
@@ -351,7 +369,7 @@ jobs:
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -374,15 +392,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot alt model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot fastvits model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      # - name: Test openpilot simple_plan vision model correctness (float32)
-      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
-      - name: Test openpilot LLVM compile
-        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      - name: Test openpilot CL compile fp16
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+      - name: Test openpilot CL compile fp32 (test correctness)
+        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+      - name: Test openpilot LLVM compile fp16
+        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -524,11 +540,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py

  testdsp:
    name: Linux (DSP)
@@ -629,6 +645,7 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
+          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
-        pass_filenames: false
+        pass_filenames: false
@@ -31,7 +31,9 @@ $(for p in "$@"; do echo "  $p,"; done)
 ]
 def _try_dlopen_$name():
  library = ctypes.util.find_library("$name")
-  if library: return ctypes.CDLL(library)
+  if library:
+    try: return ctypes.CDLL(library)
+    except OSError: pass
  for candidate in PATHS_TO_TRY:
    try: return ctypes.CDLL(candidate)
    except OSError: pass
@@ -432,11 +434,13 @@ generate_sqtt() {
    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
    $ROCPROF_SRC/include/trace_decoder_instrument.h \
    $ROCPROF_SRC/include/trace_decoder_types.h \
-    -o extra/sqtt/rocprof/rocprof.py
-  fixup extra/sqtt/rocprof/rocprof.py
-  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
+    -o $BASE/rocprof.py
+  fixup $BASE/rocprof.py
+  sed -i '1s/^/# pylint: skip-file\n/' $BASE/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" $BASE/rocprof.py
+  patch_dlopen $BASE/rocprof.py rocprof-trace-decoder "'/usr/local/lib/librocprof-trace-decoder.so'" "'/usr/local/lib/librocprof-trace-decoder.dylib'"
+  sed -i "s/def _try_dlopen_rocprof-trace-decoder():/def _try_dlopen_rocprof_trace_decoder():/g" $BASE/rocprof.py
+  sed -i "s|FunctionFactoryStub()|_try_dlopen_rocprof_trace_decoder()|g" $BASE/rocprof.py
 }

 generate_webgpu() {
@@ -520,23 +524,18 @@ generate_mesa() {
  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)

  fixup $BASE/mesa.py
-  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "f'{brew_prefix()}/lib/libtinymesa_cpu.dylib'"
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
-  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
-def brew_prefix():
-  try: return subprocess.check_output(['brew', '--prefix', 'tinymesa']).decode().strip()
-  except Exception: return ''
-EOF
-  sed -i "/in_dll/s/.*/try: &\nexcept AttributeError: pass/" $BASE/mesa.py
+  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
-  echo "def __getattr__(nm): raise AttributeError() if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
+  echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
  # bitfield bug in clang2py
  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
-  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
+  sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
  python3 -c "import tinygrad.runtime.autogen.mesa"
 }

@@ -53,9 +53,7 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 idx = UOp.const(dtypes.index, 0)
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
-ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
-alu = ld_1 + ld_2
+alu = buf_1.index(idx) + buf_2.index(idx)
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
 st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))
@@ -1,109 +0,0 @@
-# Kernel Creation
-
-Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
-
- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
-
-`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
-
-## Kernelize flow
-
-Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
-
-```py
-# initialize 3 input buffers on the device
-a = Tensor([1]).realize()
-b = Tensor([2]).realize()
-c = Tensor([3]).realize()
-
-# create the Tensor graph
-mul = a*b
-out = mul+c
-
-print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
-print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
-
-out.kernelize()
-
-print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
-print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
-```
-
-The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
-
-```py
-print(out.uop)
-```
-
-The first source is the output BUFFER:
-
-```
-UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
-  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
-```
-
-And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
-
-```
-UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
-    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
-  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
-     x1,
-    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
-```
-
-KERNEL describes the compute AST, metadata and memory dependencies.
-
-BUFFER holds a reference to the device memory where the output will be stored.
-
-Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
-
-```py
-child = out+2
-child.kernelize()
-print(child.uop.src[1].arg.ast)
-```
-
-```
-UOp(Ops.SINK, dtypes.void, arg=None, src=(
-  UOp(Ops.STORE, dtypes.void, arg=None, src=(
-    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
-    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
-    UOp(Ops.ADD, dtypes.int, arg=None, src=(
-      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
-         x2,)),
-      UOp(Ops.CONST, dtypes.int, arg=2, src=(
-         x2,)),)),)),))
-```
-
-`Tensor.realize` will execute the kernels and write outputs to memory:
-
-```py
-Tensor.realize(out)
-print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
-print(out.item()) # 5
-```
-
-<hr />
-
-**Summary**
-
- The large Tensor graph is built from a mix of data, compute and movement Ops.
-
- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
-
- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
-
- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -41,7 +41,7 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -145,6 +145,41 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear

+def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
+  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
+  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
+  scale = fp8_max / x.abs().max()
+  x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
+  return x_scl_sat.cast(dtype), scale.float().reciprocal()
+
+class FP8Linear:
+  def __init__(self, in_features, out_features, bias=True):
+    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
+    self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
+    self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
+
+  def __call__(self, x:Tensor):
+    y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
+    if self.bias is not None: y = y + self.bias.cast(y.dtype)
+    return y.cast(x.dtype)
+
+  @staticmethod
+  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
+    assert not quantize_embeds
+    new_tensors = {}
+    for name,v in tensors.items():
+      if "feed_forward" in name or "attention.w" in name:
+        assert "weight" in name, name
+        fp8_weight, scale = quantize_to_fp8(v)
+        new_tensors[name] = fp8_weight
+        new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
+        if isinstance(device, tuple):
+          new_tensors[name].shard_(device, axis=-1)
+          new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
+      else:
+        new_tensors[name] = v
+    return new_tensors
+
 MODEL_PARAMS = {
  "1B": {
    "args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
@@ -167,6 +202,7 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  # build model
  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
+  elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

@@ -242,7 +278,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -15,7 +15,7 @@ export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

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

 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
@@ -1,9 +1,5 @@
 import os, sys, pickle, time, re
 import numpy as np
-if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
-if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
-if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
-if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
@@ -21,11 +17,14 @@ def compile(onnx_file):

  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
+
  # Float inputs and outputs to tinyjits for openpilot are always float32
+  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
+  if not getenv("NPY_IMG"):
+    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")

  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -33,8 +32,6 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
-              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
@@ -69,14 +66,9 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return test_val
+  return inputs, test_val

-def test_vs_compile(run, new_inputs, test_val=None):
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
-
-  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
-  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
-            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
+def test_vs_compile(run, inputs, test_val=None):

  # run 20 times
  step_times = []
@@ -93,68 +85,57 @@ def test_vs_compile(run, new_inputs, test_val=None):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

-  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

  # test that changing the numpy changes the model outputs
-  if any([x.device == 'NPY' for x in inputs.values()]):
-    for v in new_inputs_numpy.values(): v *= 2
-    out = run(**inputs)
-    changed_val = out.numpy()
-    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  inputs_2x = {k: Tensor(v.numpy()*2, device=v.device) for k,v in inputs.items()}
+  out = run(**inputs_2x)
+  changed_val = out.numpy()
+  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val

-def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
-  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
+  import onnxruntime as ort
+  
+  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

-  timings = []
-  if ort:
-    # test with onnxruntime
-    import onnxruntime as ort
-    onnx_session = ort.InferenceSession(onnx_file)
-    for _ in range(1 if test_val is not None else 5):
-      st = time.perf_counter()
-      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
-      timings.append(time.perf_counter() - st)
-    new_torch_out = onnx_output[0]
-  else:
-    # test with torch
-    import torch
-    from onnx2torch import convert
-    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
-    torch_model = convert(onnx_model).float()
-    with torch.no_grad():
-      for _ in range(1 if test_val is not None else 5):
-        st = time.perf_counter()
-        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
-        timings.append(time.perf_counter() - st)
-      new_torch_out = torch_out.numpy()
+  ORT_TO_NP_DTYPES: dict[str, np.dtype] = {
+    'tensor(float)': np.dtype('float32'),
+    'tensor(float16)': np.dtype('float16'),
+    'tensor(uint8)': np.dtype('uint8'),
+  }

-  if test_val is not None:
-    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
-    print("test vs onnx passed")
+  timings = []
+  onnx_session = ort.InferenceSession(onnx_file)
+  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
+  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+
+  for _ in range(1 if test_val is not None else 5):
+    st = time.perf_counter()
+    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
+    timings.append(time.perf_counter() - st)
+
+  np.testing.assert_allclose(onnx_output[0].reshape(test_val.shape), test_val, atol=tol, rtol=tol)
+  print("test vs onnx passed")
  return timings

+def bench(run, inputs):
+  from extra.bench_log import WallTimeEvent, BenchEvent
+  for _ in range(10):
+    with WallTimeEvent(BenchEvent.STEP):
+      run(**inputs).numpy()
+
 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
-  test_val = compile(onnx_file) if not getenv("RUN") else None
+  inputs, outputs = compile(onnx_file)

  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

-  # same randomness as compile
-  Tensor.manual_seed(100)
-  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
-                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
-
-  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
-  if getenv("BENCHMARK"):
-    for be in ["torch", "ort"]:
-      try:
-        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
-        print(f"timing {be}: {min(timings)*1000:.2f} ms")
-      except Exception as e:
-        print(f"{be} fail with {e}")
-  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))
+  test_vs_compile(pickle_loaded, inputs, outputs)
+  if getenv("SELFTEST"):
+    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)

+  if getenv("BENCHMARK_LOG", ""):
+    bench(pickle_loaded, inputs)
@@ -99,6 +99,7 @@ if __name__ == "__main__":
  parser.add_argument('--timing',      action='store_true', help="Print timing per step")
  parser.add_argument('--noshow',      action='store_true', help="Don't show the image")
  parser.add_argument('--fp16',        action='store_true', help="Cast the weights to float16")
+  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

  N = 1
@@ -112,19 +113,22 @@ if __name__ == "__main__":

  model = StableDiffusionV2(**params)

-  default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
-  weights_fn = args.weights_fn
-  if not weights_fn:
-    weights_url = args.weights_url if args.weights_url else default_weights_url
-    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
-
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, safe_load(weights_fn), strict=False)
+    if not args.fakeweights:
+      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+      weights_fn = args.weights_fn
+      if not weights_fn:
+        weights_url = args.weights_url if args.weights_url else default_weights_url
+        weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
+
+        load_state_dict(model, safe_load(weights_fn), strict=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16).realize())
+          v.replace(v.cast(dtypes.float16))
+
+    Tensor.realize(*get_state_dict(model).values())

  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -263,14 +263,16 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
+  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

  model = StableDiffusion()

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
-    load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
+    if not args.fakeweights:
+      model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
+      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -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,4 @@
+# source extra/cl_android.sh
+export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
+export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
+
@@ -1,354 +1,168 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
+from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv, colored, prod, unwrap
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
-from tinygrad.codegen.opt.swizzler import merge_views, view_left
-
-def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
+from tinygrad.helpers import getenv

 N = 4096
+M = K = N
 run_count = 5

-BN = 128
-BM = 128
-BK = 8
+# ---------------------------
+# launch/config constants
+# ---------------------------

-TN = 4
-TM = 4
+WARP_SIZE = 32

-# NOTE: this is from testgrad
-# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
-# src->r->view  -->   src->view->r
-def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
-  if r.tag is not None: return None
-  # confirm the input is in order
-  # TODO: replace this with a UOp that allows for nothing else then remove this
-  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
-  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"
+# Threadblock tile sizes (block-level tile of C that a block computes)
+BLOCK_N = 128   # columns of C (N-dim) per block
+BLOCK_M = 128   # rows of C (M-dim) per block
+BLOCK_K = 8     # K-slice per block iteration

-  # append the reduce shape to each of the views
-  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
-  rstrides = strides_for_shape(rshape)
-  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
-                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
+# Register tile sizes (per-thread accumulator tile of C)
+TN = 4     # columns per thread
+TM = 4     # rows per thread

-  # no reshape required with shrinking REDUCE_AXIS
-  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
-             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
+assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
+assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
+assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0

-pm = PatternMatcher([
-  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
-])
+WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
+WAVE_TILE_N = 128 if is_kernel5 else 64
+WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
+assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
+assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
+WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
+WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
+assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"

-def rangeify_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  sink = c.schedule()[-1].ast
-  #print(sink)
+LANES_PER_WAVE_X = 8
+LANES_PER_WAVE_Y = 4
+ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
+ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
+assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
+assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"

-  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
-  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
-  opts += [Opt(OptOps.UNROLL, 0, 8)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+  assert dest.shape == src.shape
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  return dest.after(copy) if set else copy

-  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+def hand_spec_kernel3():
+  # ---------------------------
+  # block indices & placeholders
+  # ---------------------------
+  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
+  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")

-def top_spec_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  sink = c.schedule()[-1].ast
-  L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
-  sink = graph_rewrite(sink, view_left+pm)
-  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
-  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
+  a = UOp.placeholder((N, N), dtypes.float, slot=1)
+  b = UOp.placeholder((N, N), dtypes.float, slot=2)
+  c = UOp.placeholder((N, N), dtypes.float, slot=0)

-def hl_spec_kernel3():
-  nbIterWaveM = 2
-  nbIterWaveN = 2
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]

-  # define buffers
-  # TODO: remove these views once the defines have a shape
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
+  # open the main reduction range
+  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
+  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]

-  # shape buffers. TODO: permutes
-  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
-  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
-  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
-  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
-  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
-  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
-  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
-  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)
+  # globals are no longer used, they are already in the indexes
+  del blockIdx_y, blockIdx_x

-  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
-  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
-  assert len(expanded_shape) == 20
-  permute_a = list(range(len(expanded_shape)))
-  permute_b = permute_a[:]
+  # ---------------------------
+  # GLOBAL -> LOCAL (As, Bs)
+  # ---------------------------
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")

-  # this makes all the global loads match
-  # this can also be more simply done by rebinding the RANGEs
-  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
-  permute_a[17:20] = [11,12,13]
-  permute_a[11:14] = [17,18,19]
-  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
-  permute_a[2:7] = [3,4,5,6,2]
+  # A: read BM x BK tiles (permute on store into locals)
+  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)

-  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
-  permute_b[17:20] = [5,6,7]
+  # B: read BK x BN tiles
+  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)

-  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
-  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(As_store, Bs_store)
+  As, Bs = As.after(barrier), Bs.after(barrier)

-  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
-  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  # open inner k range
+  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)

-  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
-  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
-  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))
+  # ---------------------------
+  # LOCAL -> REG (per-wave tiles)
+  # ---------------------------
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y

-  axis_types = (
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
-    AxisType.REDUCE, AxisType.REDUCE)
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y

-  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
-  sink = graph_rewrite(sink, merge_views)
-  return sink
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)

-def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
-  BLOCK_SIZE = 128 if kernel5 else 256
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)

-  nbWaves = BLOCK_SIZE // 32
-  WN = 128 if kernel5 else 64
-  WM = BN * BM // nbWaves // WN
+  # ---------------------------
+  # FMA: c_regs += A_col * B_row
+  # ---------------------------
+  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  i = UOp.range(c_regs.size, 16)
+  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))

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

-  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
-  waveIndex = threadIdx_x // 32
-  waveIdx = waveIndex % nbWaveX
-  waveIdy = waveIndex // nbWaveX
-  indexInWave = threadIdx_x % 32
+  # Close k, sync, and close K tiles
+  sink = sink.end(k).barrier().end(k_tile_range)

-  nbThreadXPerWave = 8
-  nbThreadYPerWave = 4
+  # ---------------------------
+  # REG -> GLOBAL (epilogue)
+  # ---------------------------
+  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
+                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c[waveIdy, :, laneIdy, :,
+        waveIdx, :, laneIdx, :]
+  sink = copy(c, c_regs.after(sink), rng=600)

-  idxInWave = indexInWave % nbThreadXPerWave
-  idyInWave = indexInWave // nbThreadXPerWave
+  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()

-  nbIterWaveN = WN // (nbThreadXPerWave * TN)
-  nbIterWaveM = WM // (nbThreadYPerWave * TM)
+def test_matmul(sink:UOp, N=N):
+  with Context(DEBUG=0):
+    a = Tensor.randn(N, N)
+    b = Tensor.randn(N, N)
+    hc = Tensor.empty(N, N)
+    Tensor.realize(a, b, hc)

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

-  # Thread mapping to read BKxBN block from A
-  rAIdx = threadIdx_x % BK
-  rAIdy = threadIdx_x // BK
-  # Thread mapping to read BNxBK block from B
-  rBIdx = threadIdx_x % BN
-  rBIdy = threadIdx_x // BN
+  GlobalCounters.reset()
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(run_count):
+      ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")

-  strideReadB = BLOCK_SIZE // BN
-  strideReadA = BLOCK_SIZE // BK
-  nbReadsB = BN * BK // BLOCK_SIZE
-  nbReadsA = BM * BK // BLOCK_SIZE
-
-  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
-  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
-
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
-
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
-
-  BM_As_stride = (BM+4) if kernel5 else BM
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
-
-  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
-
-  i = UOp.range(c_regs.dtype.size, 16)
-  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
-
-  if kernel4:
-    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
-    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
-
-    # initial load from globals into locals (0)
-    kId = 0
-
-    # load from globals into locals
-    i = UOp.range(nbReadsB, 0)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 1)
-    index_x = rAIdx + kId
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
-
-    # iterate over the middle chunk
-    kId_range = UOp.range(N//BK-1, 2)
-    kId = kId_range*BK
-
-    barrier = UOp.barrier(As_store, Bs_store)
-
-    # load from globals into registers (next round)
-    i = UOp.range(nbReadsB, 3)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId + BK
-    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 4)
-    index_x = rAIdx + kId + BK
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
-
-    def inner_loop(first_range, inp_dep=()):
-      # inner unroll
-      k = UOp.range(BK, first_range+0)
-
-      # load from locals into registers
-      iterWave = UOp.range(nbIterWaveN, first_range+1)
-      i = UOp.range(TN, first_range+2)
-      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
-
-      iterWave = UOp.range(nbIterWaveM, first_range+3)
-      i = UOp.range(TM, first_range+4)
-      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
-
-      # do the GEMM math
-      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
-      yt = UOp.range(TM, first_range+6)
-      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
-      xt = UOp.range(TN, first_range+8)
-      x = iterWaveN * TN + xt
-      y = iterWaveM * TM + yt
-      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-      # sketchy, this should end the kId_range but it doesn't
-      sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                              iterWaveM, iterWaveN, yt, xt, k)
-      return sink
-
-    # TODO: kId_range should endrange after a barrier
-    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
-
-    # load from registers into locals
-    i = UOp.range(nbReadsB, 14)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId + BK
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
-
-    i = UOp.range(nbReadsA, 15)
-    index_x = rAIdx + kId + BK
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
-
-    # final iteration without the copy
-    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
-  else:
-    kId_range = UOp.range(N//BK, 0)
-    kId = kId_range*BK
-
-    # load from globals into locals
-    i = UOp.range(nbReadsB, 1)
-    index_x = BN * blockIdx_x + rBIdx
-    index_y = rBIdy + i * strideReadB + kId
-    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
-
-    i = UOp.range(nbReadsA, 2)
-    index_x = rAIdx + kId
-    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
-
-    barrier = UOp.barrier(As_store, Bs_store)
-
-    k = UOp.range(BK, 3)
-
-    # load from locals into registers
-    iterWave = UOp.range(nbIterWaveN, 4)
-    i = UOp.range(TN, 5)
-    index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-    B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
-
-    iterWave = UOp.range(nbIterWaveM, 6)
-    i = UOp.range(TM, 7)
-    index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-    A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
-
-    # do the GEMM math
-    iterWaveM = UOp.range(nbIterWaveM, 8)
-    yt = UOp.range(TM, 9)
-    iterWaveN = UOp.range(nbIterWaveN, 10)
-    xt = UOp.range(TN, 12)
-    x = iterWaveN * TN + xt
-    y = iterWaveM * TM + yt
-    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-    sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                            iterWaveM, iterWaveN, yt, xt, k, kId_range)
-
-  # store c_regs into c
-  iterWaveM = UOp.range(nbIterWaveM, 1000)
-  yt = UOp.range(TM, 1001)
-  iterWaveN = UOp.range(nbIterWaveN, 1002)
-  xt = UOp.range(TN, 1003)
-  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
-  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
-  indexC = N * (yOut + yt) + xOut + xt
-  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink),
-                         iterWaveM, iterWaveN, yt, xt)
-
-  return sink.sink(arg=KernelInfo(name="tinygemm"))
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    tc = (a @ b).realize()
+  with Context(DEBUG=0):
+    err = (hc - tc).square().mean().item()
+  print(f"mean squared error {err}")
+  if err > 1e-06:
+    raise RuntimeError("matmul is wrong!")

 if __name__ == "__main__":
-  HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
-  elif HL == 1: hprg = hl_spec_kernel3()
-  else: hprg = hand_spec_kernel3()
-  if HL == 3:
-    with Context(BLOCK_REORDER=0):
-      prg = get_program(hprg, Device.default.renderer)
-  else:
-    prg = get_program(hprg, Device.default.renderer)
-  print(prg.src)
-  if getenv("SRC"): exit(0)
-  hrunner = CompiledRunner(prg)
-
-  a = Tensor.randn(N, N).realize()
-  b = Tensor.randn(N, N).realize()
-  hc = Tensor.zeros(N, N).contiguous().realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    for _ in range(run_count): tc = (a@b).realize()
-
-  GlobalCounters.reset()
-  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
-  ei = ExecItem(hrunner, buffers)
-  with Context(DEBUG=2):
-    for _ in range(run_count): ei.run(wait=True)
-  err = (hc-tc).square().mean().item()
-  print(f"hrunner {err}")
-  if err > 1e-06: raise RuntimeError("matmul is wrong!")
+  test_matmul(hand_spec_kernel3(), N=N)
@@ -1,17 +1,11 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
-from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program

 # for copied uops
-from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
-from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.codegen.opt.search import Opt, OptOps
-from tinygrad import Device, dtypes, Tensor
-from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
+from tinygrad import dtypes
+from tinygrad.dtype import DTYPES_DICT

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

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

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

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

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

-dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
-acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("FP8E4M3") else dtypes.fp8e5m2 if getenv("FP8E5M2") else dtypes.float)
+acc_dtype = (dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else
+            dtypes.fp8e4m3 if getenv("ACC_FP8E4M3") else dtypes.fp8e5m2 if getenv("ACC_FP8E5M2") else None)
 if getenv("INT"):  dtype_in, acc_dtype = dtypes.int8, dtypes.int32
 if getenv("UINT"): dtype_in, acc_dtype = dtypes.uint8, dtypes.int32

@@ -14,8 +16,10 @@ N = getenv("N", 4096)
 M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
-ATOL = getenv("ATOL", 1e-4)
-RTOL = getenv("RTOL", 3e-2)
+
+atol, rtol = {dtypes.half:{1e-3, 1e-2}, dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
+ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
+
 INT_LOW = getenv("INT_LOW", 0)
 INT_HIGH = getenv("INT_HIGH", 10)

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

-for dtype in [torch.float32, torch.float16]:
+for dtype in [torch.float32, torch.float16, torch.bfloat16]:
  for N in [256, 512, 1024, 2048, 4096]:
    FLOPS = N*N*N*2

    b = torch.rand((N,N), dtype=dtype)
    c = torch.rand((N,N), dtype=dtype)
    if CUDA: b,c = b.cuda(),c.cuda()
+    if MPS: b,c = b.to('mps'),c.to('mps')

    def torch_prog(b, c):
      st = time.perf_counter()
      a = b@c
      if CUDA: torch.cuda.synchronize()
+      if MPS: torch.mps.synchronize()
      return time.perf_counter() - st
    tm = min([torch_prog(b, c) for _ in range(20)])
    print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS {N:4d}x{N:4d}x{N:4d} matmul in {dtype}")
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3

-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
+import argparse, glob, os, time, subprocess, sys

 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
@@ -12,7 +11,7 @@ def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
  return devs

-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
@@ -7,31 +7,34 @@ import os
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
 NUM_WAVES = 2
-FLOPS_PER_MATMUL =  16*16*16*2
-INTERNAL_LOOP  = 1_000_000
-INSTRUCTIONS_PER_LOOP = 1_000
+FLOPS_PER_MATMUL = 16*16*16*2
+INTERNAL_LOOP = 1_000_00
+INSTRUCTIONS_PER_LOOP = 200
+DIRECTIVE = ".amdhsa_wavefront_size32 1"

 assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()

-def launchBenchmark(instruction, vgprIndices, dense = True):
-  if dense:
+def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
+  if accum:
+    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
+                                                                vgprIndices[1], vgprIndices[2],
+                                                                vgprIndices[1], vgprIndices[2], extra)
+  elif dense:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2]) * INSTRUCTIONS_PER_LOOP
+                                                                vgprIndices[1], vgprIndices[2])
  else:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[3], vgprIndices[4],
-                                                                vgprIndices[5]) * INSTRUCTIONS_PER_LOOP
-  src = assemblyTemplate.replace("INSTRUCTION", instructions)
+                                                                  vgprIndices[1], vgprIndices[2],
+                                                                  vgprIndices[3], vgprIndices[4],
+                                                                  vgprIndices[5])
+  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
+  src = src.replace("DIRECTIVE", DIRECTIVE)
  lib = COMPILER.compile(src)
  fxn = AMDProgram(DEV, "matmul", lib)
-  start = time.perf_counter()
-  fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) #For some reason the returned time is very small after the first kernel execution
-  end = time.perf_counter()
-  elapsed = end-start
+  elapsed = fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-  print("{:<29} : {} T(FL)OPS".format(instruction, round(FLOPs/elapsed/10**12, 2)))
+  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

 if __name__=="__main__":
  DEVICENUM = os.getenv("DEVICENUM", "0")
@@ -40,18 +43,17 @@ if __name__=="__main__":
  except:
    raise RuntimeError("Error while initiating AMD device")

-  if (ARCH := DEV.arch) not in ['gfx1100', 'gfx1201']:
-    raise RuntimeError("only gfx1100 and gfx1201 supported")
-  COMPILER = HIPCompiler(ARCH)
-
-  if ARCH == 'gfx1100':
+  COMPILER = HIPCompiler(DEV.arch)
+  if DEV.arch in {'gfx1100', 'gfx1103'}:
+    if DEV.arch == 'gfx1103':
+      NUM_WORKGROUPS = 8
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  if ARCH == 'gfx1201':
+  elif DEV.arch == 'gfx1201':
    NUM_WORKGROUPS = 64
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
@@ -76,4 +78,20 @@ if __name__=="__main__":
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
    FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
+    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
+  elif DEV.arch == 'gfx950':
+    DIRECTIVE = ".amdhsa_accum_offset 4"
+    NUM_WORKGROUPS = 256
+    WAVE_SIZE = 64
+    NUM_WAVES = 4
+    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
+    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
+    FLOPS_PER_MATMUL = 16*16*32*2
+    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
+    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
+    FLOPS_PER_MATMUL = 16*16*128*2
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
+    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
+  else:
+    raise RuntimeError(f"arch {DEV.arch} not supported.")
@@ -1,9 +1,9 @@
    .text
    .globl matmul
-    .p2align 8 
+    .p2align 8
    .type matmul,@function
 matmul:
-    s_mov_b32 s1, 1000000
+    s_mov_b32 s1, INTERNAL_LOOP
    s_mov_b32 s2, 0
    inner_loop:
        INSTRUCTION
@@ -17,7 +17,7 @@ matmul:
 .amdhsa_kernel matmul
  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
+  DIRECTIVE
 .end_amdhsa_kernel

 .amdgpu_metadata
@@ -28,7 +28,7 @@ amdhsa.version:
 amdhsa.kernels:
  - .name: matmul
    .symbol: matmul.kd
-    .kernarg_segment_size:  0
+    .kernarg_segment_size: 0
    .group_segment_fixed_size: 0
    .private_segment_fixed_size: 0
    .kernarg_segment_align: 4
@@ -36,6 +36,5 @@ amdhsa.kernels:
    .sgpr_count: 8
    .vgpr_count: 32
    .max_flat_workgroup_size: 1024
-    .args:
 ...
 .end_amdgpu_metadata
@@ -1,7 +1,8 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <database xmlns="http://nouveau.freedesktop.org/"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
+xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
+<import file="freedreno_copyright.xml"/>
 <import file="adreno/adreno_common.xml"/>

 <enum name="vgt_event_type" varset="chip">
@@ -20,9 +21,9 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="HLSQ_FLUSH" value="7" variants="A3XX-A4XX"/>
 	<value name="VIZQUERY_END" value="8" variants="A2XX"/>
 	<value name="SC_WAIT_WC" value="9" variants="A2XX"/>
-	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX"/>
-	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX"/>
-	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX"/>
+	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX-"/>
+	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX-"/>
+	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX-"/>
 	<!-- Not sure that these 4 events don't have the same meaning as on A5XX+ -->
 	<value name="RST_PIX_CNT" value="13" variants="A2XX-A4XX"/>
 	<value name="RST_VTX_CNT" value="14" variants="A2XX-A4XX"/>
@@ -30,8 +31,8 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="STAT_EVENT" value="16" variants="A2XX-A4XX"/>
 	<value name="CACHE_FLUSH_AND_INV_TS_EVENT" value="20" variants="A2XX-A4XX"/>
 	<doc>
-		If A6XX_RB_SAMPLE_COUNT_CONTROL.copy is true, writes OQ Z passed
-		sample counts to RB_SAMPLE_COUNT_ADDR.  This writes to main
+		If A6XX_RB_SAMPLE_COUNTER_CNTL.copy is true, writes OQ Z passed
+		sample counts to RB_SAMPLE_COUNTER_BASE.  This writes to main
 		memory, skipping UCHE.
 	</doc>
 	<value name="ZPASS_DONE" value="21"/>
@@ -96,6 +97,13 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	</doc>
 	<value name="BLIT" value="30" variants="A5XX-"/>

+	<doc>
+	Flip between the primary and secondary LRZ buffers. This is used
+	for concurrent binning, so that BV can write to one buffer while
+	BR reads from the other.
+	</doc>
+	<value name="LRZ_FLIP_BUFFER" value="36" variants="A7XX-"/>
+
 	<doc>
 		Clears based on GRAS_LRZ_CNTL configuration, could clear
 		fast-clear buffer or LRZ direction.
@@ -112,11 +120,12 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="LRZ_FLUSH" value="38" variants="A5XX-"/>
 	<value name="BLIT_OP_FILL_2D" value="39" variants="A5XX-"/>
 	<value name="BLIT_OP_COPY_2D" value="40" variants="A5XX-A6XX"/>
-	<value name="UNK_40" value="40" variants="A7XX"/>
+	<value name="LRZ_CACHE_INVALIDATE" value="40" variants="A7XX-"/>
+	<value name="LRZ_Q_CACHE_INVALIDATE" value="41" variants="A7XX-"/>
 	<value name="BLIT_OP_SCALE_2D" value="42" variants="A5XX-"/>
 	<value name="CONTEXT_DONE_2D" value="43" variants="A5XX-"/>
-	<value name="UNK_2C" value="44" variants="A5XX-"/>
-	<value name="UNK_2D" value="45" variants="A5XX-"/>
+	<value name="VSC_BINNING_START" value="44" variants="A5XX-"/>
+	<value name="VSC_BINNING_END" value="45" variants="A5XX-"/>

 	<!-- a6xx events -->
 	<doc>
@@ -129,21 +138,22 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<!-- note, some of these are the same as a6xx, just named differently -->

 	<doc> Doesn't seem to do anything </doc>
-	<value name="DUMMY_EVENT" value="1" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX"/>
-	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX"/>
-	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX"/>
-	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX"/>
-	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX"/>
-	<value name="CCU_RESOLVE" value="30" variants="A7XX"/>
-	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX"/>
-	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX"/>
-	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX"/>
-	<value name="CACHE_RESET" value="48" variants="A7XX"/>
-	<value name="CACHE_CLEAN" value="49" variants="A7XX"/>
+	<value name="DUMMY_EVENT" value="1" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX-"/>
+	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX-"/>
+	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX-"/>
+	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX-"/>
+	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX-"/>
+	<value name="CCU_RESOLVE" value="30" variants="A7XX-"/>
+	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX-"/>
+	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX-"/>
+	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX-"/>
+	<value name="CACHE_RESET" value="48" variants="A7XX-"/>
+	<value name="CACHE_CLEAN" value="49" variants="A7XX-"/>
 	<!-- TODO: deal with name conflicts with other gens -->
-	<value name="CACHE_FLUSH7" value="50" variants="A7XX"/>
-	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX"/>
+	<value name="CACHE_FLUSH7" value="50" variants="A7XX-"/>
+	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX-"/>
+	<value name="DEPTH_BUFFER_FLIP" value="0x3d" variants="A8XX-"/>
 </enum>

 <enum name="pc_di_primtype">
@@ -324,7 +334,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<doc>fetch state sub-blocks and initiate shader code DMAs</doc>
 	<value name="CP_SET_STATE" value="0x25"/>
 	<doc>load constant into chip and to memory</doc>
-	<value name="CP_SET_CONSTANT" value="0x2d"/>
+	<value name="CP_SET_CONSTANT" value="0x2d" variants="A2XX"/>
 	<doc>load sequencer instruction memory (pointer-based)</doc>
 	<value name="CP_IM_LOAD" value="0x27"/>
 	<doc>load sequencer instruction memory (code embedded in packet)</doc>
@@ -371,7 +381,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_LOAD_STATE" value="0x30" variants="A3XX"/>
 	<value name="CP_LOAD_STATE4" value="0x30" variants="A4XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch enabled</doc>
-	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a"/>
+	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a" variants="A3XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch disabled</doc>
 	<value name="CP_COND_INDIRECT_BUFFER_PFD" value="0x32" variants="A3XX"/>
 	<doc>Load a buffer with pre-fetch enabled</doc>
@@ -514,7 +524,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<!--
 	Seems to set the mode flags which control which CP_SET_DRAW_STATE
 	packets are executed, based on their ENABLE_MASK values
-	
+
 	CP_SET_MODE w/ payload of 0x1 seems to cause CP_SET_DRAW_STATE
 	packets w/ ENABLE_MASK & 0x6 to execute immediately
 	 -->
@@ -537,7 +547,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_LOAD_STATE6_GEOM" value="0x32" variants="A6XX-"/>
 	<value name="CP_LOAD_STATE6_FRAG" value="0x34" variants="A6XX-"/>
 	<!--
-	Note: For IBO state (Image/SSBOs) which have shared state across
+	Note: For UAV state (Image/SSBOs) which have shared state across
 	shader stages, for 3d pipeline CP_LOAD_STATE6 is used.  But for
 	compute shaders, CP_LOAD_STATE6_FRAG is used.  Possibly they are
 	interchangable.
@@ -566,20 +576,21 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="IN_PREEMPT" value="0x0f" variants="A6XX-"/>

 	<!-- TODO do these exist on A5xx? -->
-	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX"/>
+	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_MEM" value="0x74" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_REG" value="0x72" variants="A6XX-"/>
 	<value name="CP_WAIT_MEM_GTE" value="0x14" variants="A6XX"/>
 	<value name="CP_WAIT_TWO_REGS" value="0x70" variants="A6XX"/>
 	<value name="CP_MEMCPY" value="0x75" variants="A6XX-"/>
 	<value name="CP_SET_BIN_DATA5_OFFSET" value="0x2e" variants="A6XX-"/>
+	<!-- A750+, set in place of CP_SET_BIN_DATA5_OFFSET but has different values -->
+	<value name="CP_SET_UNK_BIN_DATA" value="0x2d" variants="A7XX-"/>
 	<doc>
                Write CP_CONTEXT_SWITCH_*_INFO from CP to the following dwords,
                and forcibly switch to the indicated context.
 	</doc>
 	<value name="CP_CONTEXT_SWITCH" value="0x54" variants="A6XX"/>
-	<!-- Note, kgsl calls this CP_SET_AMBLE: -->
-	<value name="CP_SET_CTXSWITCH_IB" value="0x55" variants="A6XX-"/>
+	<value name="CP_SET_AMBLE" value="0x55" variants="A6XX-"/>

 	<!--
 	Seems to always have the payload:
@@ -630,8 +641,7 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 	<value name="CP_BV_BR_COUNT_OPS" value="0x1b" variants="A7XX-"/>
 	<doc> Clears, adds to local, or adds to global timestamp </doc>
 	<value name="CP_MODIFY_TIMESTAMP" value="0x1c" variants="A7XX-"/>
-	<!-- similar to CP_CONTEXT_REG_BUNCH, but discards first two dwords?? -->
-	<value name="CP_CONTEXT_REG_BUNCH2" value="0x5d" variants="A7XX-"/>
+	<value name="CP_NON_CONTEXT_REG_BUNCH" value="0x5d" variants="A7XX-"/>
 	<doc>
 		Write to a scratch memory that is read by CP_REG_TEST with
 		SOURCE_SCRATCH_MEM set. It's not the same scratch as scratch registers.
@@ -648,6 +658,11 @@ xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">

 	<doc>Reset various on-chip state used for synchronization</doc>
 	<value name="CP_RESET_CONTEXT_STATE" value="0x1f" variants="A7XX-"/>
+
+	<doc>Invalidates the "CCHE" introduced on a740</doc>
+	<value name="CP_CCHE_INVALIDATE" value="0x3a" variants="A7XX-"/>
+
+	<value name="CP_SCOPE_CNTL" value="0x6c" variants="A7XX-"/>
 </enum>


@@ -790,14 +805,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="SB6_GS_SHADER" value="0xb"/>
 		<value name="SB6_FS_SHADER" value="0xc"/>
 		<value name="SB6_CS_SHADER" value="0xd"/>
-		<value name="SB6_IBO"       value="0xe"/>
-		<value name="SB6_CS_IBO"    value="0xf"/>
+		<value name="SB6_UAV"       value="0xe"/>
+		<value name="SB6_CS_UAV"    value="0xf"/>
 	</enum>
 	<enum name="a6xx_state_type">
 		<value name="ST6_SHADER" value="0"/>
 		<value name="ST6_CONSTANTS" value="1"/>
 		<value name="ST6_UBO" value="2"/>
-		<value name="ST6_IBO" value="3"/>
+		<value name="ST6_UAV" value="3"/>
 	</enum>
 	<enum name="a6xx_state_src">
 		<value name="SS6_DIRECT" value="0"/>
@@ -903,12 +918,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>

 	<stripe varset="chip" variants="A5XX-">
-		<reg32 offset="4" name="4">
-			<bitfield name="INDX_BASE_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="5" name="5">
-			<bitfield name="INDX_BASE_HI" low="0" high="31"/>
-		</reg32>
 		<reg64 offset="4" name="INDX_BASE" type="address"/>
 		<reg32 offset="6" name="6">
 			<!-- max # of elements in index buffer -->
@@ -1084,8 +1093,10 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			<bitfield name="BINNING" pos="20" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="GMEM" pos="21" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="SYSMEM" pos="22" varset="chip" variants="A6XX-" type="boolean"/>
-			<bitfield name="GROUP_ID" low="24" high="28" type="uint"/>
+			<!-- high bit is 28 until a750: -->
+			<bitfield name="GROUP_ID" low="24" high="29" type="uint"/>
 		</reg32>
+		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_LO" low="0" high="31" type="hex"/>
 		</reg32>
@@ -1119,39 +1130,63 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

+<enum name="a7xx_abs_mask_mode">
+	<value name="ABS_MASK" value="0x1"/>
+	<value name="NO_ABS_MASK" value="0x0"/>
+</enum>
+
 <domain name="CP_SET_BIN_DATA5" width="32">
 	<reg32 offset="0" name="0">
+		<bitfield name="VSC_MASK" low="0" high="15" type="hex">
+			<doc>
+				A mask of bins, starting at VSC_N, whose
+				visibility is OR'd together. A value of 0 is
+				interpreted as 1 (i.e. just use VSC_N for
+				visbility) for backwards compatibility. Only
+				exists on a7xx.
+			</doc>
+		</bitfield>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
+		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes">
+			<doc>
+				If this field is 1, VSC_MASK and VSC_N are
+				ignored and instead a new ordinal immediately
+				after specifies the full 32-bit mask of bins
+				to use. The mask is "absolute" instead of
+				relative to VSC_N.
+			</doc>
+		</bitfield>
 	</reg32>
-	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-	<reg32 offset="1" name="1">
-		<bitfield name="BIN_DATA_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="BIN_DATA_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
-	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-	<reg32 offset="3" name="3">
-		<bitfield name="BIN_SIZE_ADDRESS_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="4" name="4">
-		<bitfield name="BIN_SIZE_ADDRESS_HI" low="0" high="31"/>
-	</reg32>
-	<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
-	<reg32 offset="5" name="5">
-		<bitfield name="BIN_PRIM_STRM_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="6" name="6">
-		<bitfield name="BIN_PRIM_STRM_HI" low="0" high="31"/>
-	</reg32>
-	<!--
-		a7xx adds a few more addresses to the end of the pkt
-	 -->
-	<reg64 offset="7" name="7"/>
-	<reg64 offset="9" name="9"/>
+	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg64 offset="1" name="BIN_DATA_ADDR" type="address"/>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg64 offset="3" name="BIN_SIZE_ADDR" type="address"/>
+		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<reg64 offset="5" name="BIN_PRIM_STRM" type="address"/>
+		<!--
+			a7xx adds a few more addresses to the end of the pkt
+		 -->
+		<reg64 offset="7" name="7"/>
+		<reg64 offset="9" name="9"/>
+	</stripe>
+	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
+		<reg32 offset="1" name="ABS_MASK"/>
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg64 offset="2" name="BIN_DATA_ADDR" type="address"/>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg64 offset="4" name="BIN_SIZE_ADDR" type="address"/>
+		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+		<reg64 offset="6" name="BIN_PRIM_STRM" type="address"/>
+		<!--
+			a7xx adds a few more addresses to the end of the pkt
+		 -->
+		<reg64 offset="8" name="8"/>
+		<reg64 offset="10" name="10"/>
+	</stripe>
 </domain>

 <domain name="CP_SET_BIN_DATA5_OFFSET" width="32">
@@ -1162,23 +1197,42 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
                stream is recorded.
 	</doc>
 	<reg32 offset="0" name="0">
+		<bitfield name="VSC_MASK" low="0" high="15" type="hex"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
+		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes"/>
 	</reg32>
-	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-	<reg32 offset="1" name="1">
-		<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
-	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-	<reg32 offset="2" name="2">
-		<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
-	<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
-	<reg32 offset="3" name="3">
-		<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
-	</reg32>
+	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg32 offset="1" name="1">
+			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg32 offset="2" name="2">
+			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+		<reg32 offset="3" name="3">
+			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+	</stripe>
+	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
+		<reg32 offset="1" name="ABS_MASK"/>
+		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+		<reg32 offset="2" name="2">
+			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+		<reg32 offset="3" name="3">
+			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+		<reg32 offset="4" name="4">
+			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+		</reg32>
+	</stripe>
 </domain>

 <domain name="CP_REG_RMW" width="32">
@@ -1196,6 +1250,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</doc>
 	<reg32 offset="0" name="0">
 		<bitfield name="DST_REG" low="0" high="17" type="hex"/>
+		<bitfield name="DST_SCRATCH" pos="19" type="boolean" varset="chip" variants="A7XX-"/>
+		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
+		<bitfield name="SKIP_WAIT_FOR_ME" pos="23" type="boolean" varset="chip" variants="A7XX-"/>
 		<bitfield name="ROTATE" low="24" high="28" type="uint"/>
 		<bitfield name="SRC1_ADD" pos="29" type="boolean"/>
 		<bitfield name="SRC1_IS_REG" pos="30" type="boolean"/>
@@ -1209,7 +1266,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_REG_TO_MEM" width="32">
+<domain name="CP_REG_TO_MEM" width="32" prefix="chip">
 	<reg32 offset="0" name="0">
 		<bitfield name="REG" low="0" high="17" type="hex"/>
 		<!-- number of registers/dwords copied is max(CNT, 1). -->
@@ -1217,12 +1274,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="1" name="DEST" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="1" name="DEST" type="address"/>
+	</stripe>
 </domain>

 <domain name="CP_REG_TO_MEM_OFFSET_REG" width="32">
@@ -1238,12 +1295,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="1" name="DEST" type="waddress"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="OFFSET0" low="0" high="17" type="hex"/>
 		<bitfield name="OFFSET0_SCRATCH" pos="19" type="boolean"/>
@@ -1263,18 +1315,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="DEST" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="DEST_HI" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="3" name="3">
-		<bitfield name="OFFSET_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="4" name="4">
-		<bitfield name="OFFSET_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="DEST" type="waddress"/>
+	<reg64 offset="3" name="OFFSET" type="waddress"/>
 </domain>

 <domain name="CP_MEM_TO_REG" width="32">
@@ -1287,12 +1329,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- does the same thing as CP_MEM_TO_MEM::UNK31 -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="SRC" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
-		<bitfield name="SRC_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="1" name="SRC" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="1" name="SRC" type="address"/>
+	</stripe>
 </domain>

 <domain name="CP_MEM_TO_MEM" width="32">
@@ -1312,6 +1354,10 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- some other kind of wait -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
+	<reg64 offset="1" name="DST" type="waddress"/>
+	<reg64 offset="3" name="SRC_A" type="address"/>
+	<reg64 offset="5" name="SRC_B" type="address"/>
+	<reg64 offset="7" name="SRC_C" type="address"/>
 	<!--
 	followed by sequence of addresses.. the first is the
 	destination and the rest are N src addresses which are
@@ -1346,6 +1392,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="SCRATCH" low="20" high="22" type="uint"/>
 		<!-- number of registers/dwords copied is CNT + 1. -->
 		<bitfield name="CNT" low="24" high="26" type="uint"/>
+		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
+		<bitfield name="SKIP_WAIT_FOR_ME" pos="27" type="boolean" varset="chip" variants="A7XX-"/>
 	</reg32>
 </domain>

@@ -1368,12 +1416,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 </domain>

 <domain name="CP_MEM_WRITE" width="32">
-	<reg32 offset="0" name="0">
-		<bitfield name="ADDR_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_HI" low="0" high="31"/>
-	</reg32>
+	<stripe varset="chip" variants="A2XX-A4XX">
+		<reg32 offset="0" name="ADDR" type="address"/>
+	</stripe>
+	<stripe varset="chip" variants="A5XX-">
+		<reg64 offset="0" name="ADDR" type="address"/>
+	</stripe>
 	<!-- followed by the DWORDs to write -->
 </domain>

@@ -1425,24 +1473,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="4" name="4">
 		<bitfield name="MASK" low="0" high="31"/>
 	</reg32>
-	<reg32 offset="5" name="5">
-		<bitfield name="WRITE_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="6" name="6">
-		<bitfield name="WRITE_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="5" name="WRITE_ADDR" type="waddress"/>
 	<reg32 offset="7" name="7">
 		<bitfield name="WRITE_DATA" low="0" high="31"/>
 	</reg32>
@@ -1457,12 +1495,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- Reserved for flags, presumably? Unused in FW -->
 		<bitfield name="RESERVED" low="0" high="31" type="hex"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1480,12 +1513,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
-	</reg32>
+	<reg64 offset="1" name="POLL_ADDR" type="address"/>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1619,12 +1647,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	TODO what is gpuaddr for, seems to be all 0's.. maybe needed for
 	context switch?
 	 -->
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_0_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="2" name="2">
-		<bitfield name="ADDR_0_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="1" name="ADDR" type="waddress"/>
 	<reg32 offset="3" name="3">
 		<!-- ??? -->
 	</reg32>
@@ -1653,8 +1676,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_SAMPLE_COUNT" pos="12" type="boolean"/>
 		<!-- Write sample count at (iova + 16) -->
 		<bitfield name="SAMPLE_COUNT_END_OFFSET" pos="13" type="boolean"/>
-		<!-- *(iova + 8) = *(iova + 16) - *iova -->
-		<bitfield name="WRITE_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>
+		<!-- *(iova + 8) += *(iova + 16) - *iova -->
+		<bitfield name="WRITE_ACCUM_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>

 		<!-- Next 4 flags are valid to set only when concurrent binning is enabled -->
 		<!-- Increment 16b BV counter. Valid only in BV pipe -->
@@ -1668,15 +1691,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_DST" pos="24" type="event_write_dst" addvariant="yes"/>
 		<!-- Writes into WRITE_DST from WRITE_SRC. RB_DONE_TS requires WRITE_ENABLED. -->
 		<bitfield name="WRITE_ENABLED" pos="27" type="boolean"/>
+		<bitfield name="IRQ" pos="31" type="boolean"/>
 	</reg32>

 	<stripe varset="event_write_dst" variants="EV_DST_RAM">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR_0_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="2" name="2">
-			<bitfield name="ADDR_0_HI" low="0" high="31"/>
-		</reg32>
+		<reg64 offset="1" name="1" type="waddress"/>
 		<reg32 offset="3" name="3">
 			<bitfield name="PAYLOAD_0" low="0" high="31"/>
 		</reg32>
@@ -1743,9 +1762,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 	</reg32>
 	<stripe varset="chip" variants="A4XX">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR" low="0" high="31"/>
-		</reg32>
+		<reg32 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="2" name="2">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1754,12 +1771,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		</reg32>
 	</stripe>
 	<stripe varset="chip" variants="A5XX-">
-		<reg32 offset="1" name="1">
-			<bitfield name="ADDR_LO" low="0" high="31"/>
-		</reg32>
-		<reg32 offset="2" name="2">
-			<bitfield name="ADDR_HI" low="0" high="31"/>
-		</reg32>
+		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="3" name="3">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1771,40 +1783,88 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_SET_MARKER" width="32" varset="chip" prefix="chip" variants="A6XX-">
 	<doc>Tell CP the current operation mode, indicates save and restore procedure</doc>
+	<enum name="set_marker_mode">
+		<value value="0" name="SET_RENDER_MODE"/>
+		<!-- IFPC - inter-frame power collapse -->
+		<value value="1" name="SET_IFPC_MODE"/>
+	</enum>
+	<enum name="a6xx_ifpc_mode">
+		<value value="0" name="IFPC_ENABLE"/>
+		<value value="1" name="IFPC_DISABLE"/>
+	</enum>
 	<enum name="a6xx_marker">
-		<value value="1" name="RM6_BYPASS"/>
-		<value value="2" name="RM6_BINNING"/>
-		<value value="4" name="RM6_GMEM"/>
-		<value value="5" name="RM6_ENDVIS"/>
-		<value value="6" name="RM6_RESOLVE"/>
-		<value value="7" name="RM6_YIELD"/>
+		<value value="1" name="RM6_DIRECT_RENDER"/>
+		<value value="2" name="RM6_BIN_VISIBILITY"/>
+		<value value="3" name="RM6_BIN_DIRECT"/>
+		<value value="4" name="RM6_BIN_RENDER_START"/>
+		<value value="5" name="RM6_BIN_END_OF_DRAWS"/>
+		<value value="6" name="RM6_BIN_RESOLVE"/>
+		<value value="7" name="RM6_BIN_RENDER_END"/>
 		<value value="8" name="RM6_COMPUTE"/>
-		<value value="0xc" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->
+		<value value="12" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->

 		<!--
 			These values come from a6xx_set_marker() in the
 			downstream kernel, and they can only be set by the kernel
 		-->
-		<value value="0xd" name="RM6_IB1LIST_START"/>
-		<value value="0xe" name="RM6_IB1LIST_END"/>
-		<!-- IFPC - inter-frame power collapse -->
-		<value value="0x100" name="RM6_IFPC_ENABLE"/>
-		<value value="0x101" name="RM6_IFPC_DISABLE"/>
+		<value value="13" name="RM6_IB1LIST_START"/>
+		<value value="14" name="RM6_IB1LIST_END"/>
+		<value value="15" name="RM7_BIN_VISIBILITY_END"/>
+
+		<!-- new in a8xx: -->
+		<value value="32" name="RM8_DEPTH_PASS_START"/>
+		<value value="33" name="RM8_DEPTH_PASS_END"/>
 	</enum>
-	<reg32 offset="0" name="0">
-		<!--
-			NOTE: blob driver and some versions of freedreno/turnip set
-			b4, which is unused (at least by current sqe fw), but interferes
-			with parsing if we extend the size of the bitfield to include
-			b8 (only sent by kernel mode driver).  Really, the way the
-			parsing works in the firmware, only b0-b3 are considered, but
-			if b8 is set, the low bits are interpreted differently.  To
-			model this, without getting confused by spurious b4, this is
-			described as two overlapping bitfields:
-		 -->
-		<bitfield name="MODE" low="0" high="8" type="a6xx_marker"/>
-		<bitfield name="MARKER" low="0" high="3" type="a6xx_marker"/>
-	</reg32>
+	<stripe varset="chip" variants="A6XX-A7XX">
+		<reg32 offset="0" name="0">
+			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
+			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
+
+
+			<bitfield name="MODE" low="0" high="3" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<!-- used by preemption to determine if GMEM needs to be saved or not -->
+			<bitfield name="USES_GMEM" pos="4" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+
+
+			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
+
+
+			<!--
+				CP_SET_MARKER is used with these bits to create a
+				critical section around a workaround for ray tracing.
+				The workaround happens after BVH building, and appears
+				to invalidate the RTU's BVH node cache. It makes sure
+				that only one of BR/BV/LPAC is executing the
+				workaround at a time, and no draws using RT on BV/LPAC
+				are executing while the workaround is executed on BR (or
+				vice versa, that no draws on BV/BR using RT are executed
+				while the workaround executes on LPAC), by
+				hooking subsequent CP_EVENT_WRITE/CP_DRAW_*/CP_EXEC_CS.
+				The blob usage is:
+
+
+				CP_SET_MARKER(RT_WA_START)
+				... workaround here ...
+				CP_SET_MARKER(RT_WA_END)
+				...
+				CP_SET_MARKER(SHADER_USES_RT)
+				CP_DRAW_INDX(...) or CP_EXEC_CS(...)
+			-->
+			<bitfield name="SHADER_USES_RT" pos="9" type="boolean" variants="A7XX-"/>
+			<bitfield name="RT_WA_START" pos="10" type="boolean" variants="A7XX-"/>
+			<bitfield name="RT_WA_END" pos="11" type="boolean" variants="A7XX-"/>
+		</reg32>
+	</stripe>
+	<stripe varset="chip" variants="A8XX-">
+		<reg32 offset="0" name="0">
+			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
+			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
+			<bitfield name="USES_GMEM" pos="7" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<bitfield name="MODE" low="0" high="6" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
+			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
+			<!-- idk if the RT w/a fields apply to a8xx as well -->
+		</reg32>
+	</stripe>
 </domain>

 <domain name="CP_SET_PSEUDO_REG" width="32" varset="chip" prefix="chip" variants="A6XX-">
@@ -1830,9 +1890,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			If concurrent binning is disabled then BR also does binning so it will also
 			write the "real" registers in BR.
 		-->
-		<value value="8" name="DRAW_STRM_ADDRESS"/>
-		<value value="9" name="DRAW_STRM_SIZE_ADDRESS"/>
-		<value value="10" name="PRIM_STRM_ADDRESS"/>
+		<value value="8" name="VSC_PIPE_DATA_DRAW_BASE"/>
+		<value value="9" name="VSC_SIZE_BASE"/>
+		<value value="10" name="VSC_PIPE_DATA_PRIM_BASE"/>
 		<value value="11" name="UNK_STRM_ADDRESS"/>
 		<value value="12" name="UNK_STRM_SIZE_ADDRESS"/>

@@ -1933,11 +1993,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			a bitmask of which modes pass the test.
 		-->

-		<!-- RM6_BINNING -->
+		<!-- RM6_BIN_VISIBILITY -->
 		<bitfield name="BINNING" pos="25" variants="RENDER_MODE" type="boolean"/>
 		<!-- all others -->
 		<bitfield name="GMEM" pos="26" variants="RENDER_MODE" type="boolean"/>
-		<!-- RM6_BYPASS -->
+		<!-- RM6_DIRECT_RENDER -->
 		<bitfield name="SYSMEM" pos="27" variants="RENDER_MODE" type="boolean"/>

 		<bitfield name="BV" pos="25" variants="THREAD_MODE" type="boolean"/>
@@ -2010,54 +2070,45 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_SET_CTXSWITCH_IB" width="32">
+<domain name="CP_SET_AMBLE" width="32">
 	<doc>
-                Used by the userspace driver to set various IB's which are
-                executed during context save/restore for handling
-                state that isn't restored by the
-                context switch routine itself.
+		Used by the userspace and kernel drivers to set various IB's
+		which are executed during context save/restore for handling
+		state that isn't restored by the context switch routine itself.
 	</doc>
-	<enum name="ctxswitch_ib">
-		<value name="RESTORE_IB" value="0">
+	<enum name="amble_type">
+		<value name="PREAMBLE_AMBLE_TYPE" value="0">
 			<doc>Executed unconditionally when switching back to the context.</doc>
 		</value>
-		<value name="YIELD_RESTORE_IB" value="1">
+		<value name="BIN_PREAMBLE_AMBLE_TYPE" value="1">
                        <doc>
 				Executed when switching back after switching
 				away during execution of
-				a CP_SET_MARKER packet with RM6_YIELD as the
-				payload *and* the normal save routine was
-				bypassed for a shorter one. I think this is
-				connected to the "skipsaverestore" bit set by
-				the kernel when preempting.
+				a CP_SET_MARKER packet with RM6_BIN_RENDER_END as the
+				payload *and* skipsaverestore is set. This is
+				expected to restore static register values not
+				saved when skipsaverestore is set.
 			</doc>
 		</value>
-		<value name="SAVE_IB" value="2">
+		<value name="POSTAMBLE_AMBLE_TYPE" value="2">
                        <doc>
 				Executed when switching away from the context,
 				except for context switches initiated via
 				CP_YIELD.
                        </doc>
 		</value>
-		<value name="RB_SAVE_IB" value="3">
+		<value name="KMD_AMBLE_TYPE" value="3">
 			<doc>
 				This can only be set by the RB (i.e. the kernel)
 				and executes with protected mode off, but
-				is otherwise similar to SAVE_IB.
-
-				Note, kgsl calls this CP_KMD_AMBLE_TYPE
+				is otherwise similar to POSTAMBLE_AMBLE_TYPE.
 			</doc>
 		</value>
 	</enum>
-	<reg32 offset="0" name="0">
-		<bitfield name="ADDR_LO" low="0" high="31"/>
-	</reg32>
-	<reg32 offset="1" name="1">
-		<bitfield name="ADDR_HI" low="0" high="31"/>
-	</reg32>
+	<reg64 offset="0" name="ADDR" type="address"/>
 	<reg32 offset="2" name="2">
 		<bitfield name="DWORDS" low="0" high="19" type="uint"/>
-		<bitfield name="TYPE" low="20" high="21" type="ctxswitch_ib"/>
+		<bitfield name="TYPE" low="20" high="21" type="amble_type"/>
 	</reg32>
 </domain>

@@ -2089,12 +2140,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="UNK_EVENT_WRITE" value="0x4"/>
 		<doc>
 			Tracks GRAS_LRZ_CNTL::GREATER, GRAS_LRZ_CNTL::DIR, and
-			GRAS_LRZ_DEPTH_VIEW with previous values, and if one of
+			GRAS_LRZ_VIEW_INFO with previous values, and if one of
 			the following is true:
 			- GRAS_LRZ_CNTL::GREATER has changed
 			- GRAS_LRZ_CNTL::DIR has changed, the old value is not
 			  CUR_DIR_GE, and the new value is not CUR_DIR_DISABLED
-			- GRAS_LRZ_DEPTH_VIEW has changed
+			- GRAS_LRZ_VIEW_INFO has changed
 			then it does a LRZ_FLUSH with GRAS_LRZ_CNTL::ENABLE
 			forced to 1.
 			Only exists in a650_sqe.fw.
@@ -2209,7 +2260,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_MEM_TO_SCRATCH_MEM" width="32">
 	<doc>
-		Best guess is that it is a faster way to fetch all the VSC_STATE registers
+		Best guess is that it is a faster way to fetch all the VSC_CHANNEL_VISIBILITY registers
 		and keep them in a local scratch memory instead of fetching every time
 		when skipping IBs.
 	</doc>
@@ -2257,7 +2308,25 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 		<bitfield name="CLEAR_ON_CHIP_TS" pos="0" type="boolean"/>
 		<bitfield name="CLEAR_RESOURCE_TABLE" pos="1" type="boolean"/>
-		<bitfield name="CLEAR_GLOBAL_LOCAL_TS" pos="2" type="boolean"/>
+		<bitfield name="CLEAR_BV_BR_COUNTER" pos="2" type="boolean"/>
+		<bitfield name="RESET_GLOBAL_LOCAL_TS" pos="3" type="boolean"/>
+	</reg32>
+</domain>
+
+<domain name="CP_SCOPE_CNTL" width="32">
+	<enum name="cp_scope">
+		<value value="0" name="INTERRUPTS"/>
+	</enum>
+	<reg32 offset="0" name="0">
+		<bitfield name="DISABLE_PREEMPTION" pos="0" type="boolean"/>
+		<bitfield low="28" high="31" name="SCOPE" type="cp_scope"/>
+	</reg32>
+</domain>
+
+<domain name="CP_INDIRECT_BUFFER" width="32" varset="chip" prefix="chip" variants="A5XX-">
+	<reg64 offset="0" name="IB_BASE" type="address"/>
+	<reg32 offset="2" name="2">
+		<bitfield name="IB_SIZE" low="0" high="19"/>
 	</reg32>
 </domain>

@@ -10,6 +10,8 @@ import ctypes


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

        if state_block == SB6_CS_SHADER:
          from extra.disassemblers.adreno import disasm_raw
-          if state_type == ST6_SHADER and IOCTL > 2:
+          if state_type == ST6_SHADER and IOCTL > 3:
            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
@@ -106,25 +106,30 @@ def parse_cmd_buf(dat):
              print('constants')
              hexdump(x)
          if state_type == ST6_IBO:
-            ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
+            if state_src == 0x1:
+              ibos_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
            CAPTURED_STATE['ibos'] = ibos_bytes[:]
            if IOCTL > 1:
              print('texture ibos')
              hexdump(ibos_bytes)
        elif state_block == SB6_CS_TEX:
          if state_type == ST6_SHADER:
-            samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
+            if state_src == 0x1:
+              samplers_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
            CAPTURED_STATE['samplers'] = samplers_bytes[:]
            if IOCTL > 1:
              print('texture samplers')
              hexdump(samplers_bytes)
          if state_type == ST6_CONSTANTS:
-            descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
+            if state_src == 0x1:
+              descriptors_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
+            else: descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
            CAPTURED_STATE['descriptors'] = descriptors_bytes[:]
            if IOCTL > 1:
              print('texture descriptors')
              hexdump(descriptors_bytes)
-
      elif ops[opcode] == "CP_REG_TO_MEM":
        reg, cnt, b64, accum = vals[0] & 0x3FFFF, (vals[0] >> 18) & 0xFFF, (vals[0] >> 30) & 0x1, (vals[0] >> 31) & 0x1
        dest = vals[1] | (vals[2] << 32)
@@ -152,6 +157,10 @@ def parse_cmd_buf(dat):
        if IOCTL > 0:
          print(f'THREADSIZE-{(vals[0] >> 20)&0x1}\nEARLYPREAMBLE-{(vals[0] >> 23) & 0x1}\nMERGEDREGS-{(vals[0] >> 3) & 0x1}\nTHREADMODE-{vals[0] & 0x1}\nHALFREGFOOTPRINT-{(vals[0] >> 1) & 0x3f}\nFULLREGFOOTPRINT-{(vals[0] >> 7) & 0x3f}\nBRANCHSTACK-{(vals[0] >> 14) & 0x3f}\n')
          print(f'SP_CS_UNKNOWN_A9B1-{vals[1]}\nSP_CS_BRANCH_COND-{vals[2]}\nSP_CS_OBJ_FIRST_EXEC_OFFSET-{vals[3]}\nSP_CS_OBJ_START-{vals[4] | (vals[5] << 32)}\nSP_CS_PVT_MEM_PARAM-{vals[6]}\nSP_CS_PVT_MEM_ADDR-{vals[7] | (vals[8] << 32)}\nSP_CS_PVT_MEM_SIZE-{vals[9]}')
+      if offset == 0xa9e8:
+        CAPTURED_STATE['bindless_base'] = (vals[0] | (vals[1] << 32)) & ~0b11
+        # print(hex(CAPTURED_STATE['bindless_base']))
+        # hexdump(get_mem(CAPTURED_STATE['bindless_base'], 0x200))
      if offset == 0xb180:
        if IOCTL > 0:
          print('border color offset', hex(vals[1] << 32 | vals[0]))
@@ -171,8 +180,8 @@ def ioctl(fd, request, argp):
    name, stype = nrs[nr]
    s = get_struct(argp, stype)
    if IOCTL > 0: print(f"{ret:2d} = {name:40s}", ' '.join(format_struct(s)))
-    if name == "IOCTL_KGSL_GPUOBJ_INFO": pass
-      # mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
+    if name == "IOCTL_KGSL_GPUOBJ_INFO":
+      mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
    if name == "IOCTL_KGSL_GPU_COMMAND":
      for i in range(s.numcmds):
        cmd = get_struct(s.cmdlist+ctypes.sizeof(msm_kgsl.struct_kgsl_command_object)*i, msm_kgsl.struct_kgsl_command_object)
@@ -882,6 +882,11 @@ impl<'a> Thread<'a> {
                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i32;
                            (s0 * s1) as u32
                        }
+                        10 => {
+                            let s0 = sign_ext((s0 & 0xffffff) as u64, 24) as i64;
+                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i64;
+                            ((s0 * s1) >> 32) as u32
+                        }
                        17 | 18 | 26 => {
                            let (s0, s1) = (s0 as i32, s1 as i32);
                            (match op {
@@ -930,7 +935,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 | 770 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -996,6 +1001,10 @@ impl<'a> Thread<'a> {
                                    let ret = s0.wrapping_sub(s1);
                                    (ret as u32, s1 > s0)
                                }
+                                770 => {
+                                    let ret = s1.wrapping_sub(s0);
+                                    (ret as u32, s0 > s1)
+                                }
                                _ => todo_instr!(instruction)?,
                            };
                            if self.exec.read() {
@@ -1,98 +1,32 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from typing import cast
-from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import diskcache, OSX, getenv
+from tinygrad import Tensor, dtypes, Device, UOp
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner

-@diskcache
-def assemble(code:str) -> bytes:
-  try:
-    LLVM_MC = "llvm-mc" if OSX else "/opt/rocm/llvm/bin/llvm-mc"
-    return subprocess.run([LLVM_MC, "--arch=amdgcn", "--mcpu=gfx1100", "--triple=amdgcn-amd-amdhsa", "-filetype=obj", "-o", "-"],
-                           input=code.encode("utf-8"), stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
-  except subprocess.CalledProcessError as e:
-    print("stderr:")
-    print(e.stderr.decode())
-    raise
-
-# copied from extra/rdna
-def get_prg(code:str, v_cnt:int, s_cnt:int):
-  function_name = "test"
-  metadata = f"""
-amdhsa.kernels:
- .args:
-  - .address_space: global
-    .name: buf_0
-    .offset: 0
-    .size: 8
-    .type_name: unsigned int*
-    .value_kind: global_buffer
-  .group_segment_fixed_size: 0
-  .kernarg_segment_align: 8
-  .kernarg_segment_size: 8
-  .language: OpenCL C
-  .language_version:
-  - 1
-  - 2
-  .max_flat_workgroup_size: 256
-  .name: test
-  .private_segment_fixed_size: 0
-  .sgpr_count: {s_cnt}
-  .sgpr_spill_count: 0
-  .symbol: test.kd
-  .uses_dynamic_stack: false
-  .vgpr_count: {v_cnt}
-  .vgpr_spill_count: 0
-  .wavefront_size: 32
-amdhsa.target: amdgcn-amd-amdhsa--gfx1100
-amdhsa.version:
- 1
- 2
-  """
-  boilerplate_start = f"""
-  .rodata
-  .global {function_name}.kd
-  .type {function_name}.kd,STT_OBJECT
-  .align 0x10
-  .amdhsa_kernel {function_name}"""
-  kernel_desc = {
-    '.amdhsa_group_segment_fixed_size': 0, '.amdhsa_private_segment_fixed_size': 0, '.amdhsa_kernarg_size': 0,
-    '.amdhsa_next_free_vgpr': v_cnt,   # this matters!
-    '.amdhsa_reserve_vcc': 0, '.amdhsa_reserve_xnack_mask': 0,
-    '.amdhsa_next_free_sgpr': s_cnt,
-    '.amdhsa_float_round_mode_32': 0, '.amdhsa_float_round_mode_16_64': 0, '.amdhsa_float_denorm_mode_32': 3, '.amdhsa_float_denorm_mode_16_64': 3,
-    '.amdhsa_dx10_clamp': 1, '.amdhsa_ieee_mode': 1, '.amdhsa_fp16_overflow': 0,
-    '.amdhsa_workgroup_processor_mode': 1, '.amdhsa_memory_ordered': 1, '.amdhsa_forward_progress': 0, '.amdhsa_enable_private_segment': 0,
-    '.amdhsa_system_sgpr_workgroup_id_x': 1, '.amdhsa_system_sgpr_workgroup_id_y': 1, '.amdhsa_system_sgpr_workgroup_id_z': 1,
-    '.amdhsa_system_sgpr_workgroup_info': 0, '.amdhsa_system_vgpr_workitem_id': 2, # is amdhsa_system_vgpr_workitem_id real?
-    '.amdhsa_exception_fp_ieee_invalid_op': 0, '.amdhsa_exception_fp_denorm_src': 0,
-    '.amdhsa_exception_fp_ieee_div_zero': 0, '.amdhsa_exception_fp_ieee_overflow': 0, '.amdhsa_exception_fp_ieee_underflow': 0,
-    '.amdhsa_exception_fp_ieee_inexact': 0, '.amdhsa_exception_int_div_zero': 0,
-    '.amdhsa_user_sgpr_dispatch_ptr': 0, '.amdhsa_user_sgpr_queue_ptr': 0, '.amdhsa_user_sgpr_kernarg_segment_ptr': 1,
-    '.amdhsa_user_sgpr_dispatch_id': 0, '.amdhsa_user_sgpr_private_segment_size': 0, '.amdhsa_wavefront_size32': 1, '.amdhsa_uses_dynamic_stack': 0}
-  code_start = f""".end_amdhsa_kernel
-  .text
-  .global {function_name}
-  .type {function_name},@function
-  .p2align 8
-  {function_name}:
-  """
-  ret = ".amdgpu_metadata\n" + metadata + ".end_amdgpu_metadata" + boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) \
-      + "\n" + code_start + code + f"\n.size {function_name}, .-{function_name}"
-  return AMDProgram(cast(AMDDevice, Device["AMD"]), function_name, assemble(ret))
-
-def get_output(s:str, n_threads:int=1):
-  assert n_threads <= 32
-  code = "\n".join(["s_load_b64 s[0:1], s[0:1], null", "v_lshlrev_b32_e32 v0, 2, v0", s,
-    "s_waitcnt 0",
-    "global_store_b32 v0, v1, s[0:1]",
-    "s_nop 0", "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)", "s_endpgm"])
-  test = Tensor.zeros((n_threads,), dtype=dtypes.uint32).contiguous().realize().uop.buffer
-  prg = get_prg(code, 32, 32)
-  prg(test._buf, global_size=(1, 1, 1), local_size=(n_threads, 1, 1), wait=True)
-  return test.numpy()
+def get_output(asm:str, n_threads:int=1):
+  input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
+                         for ln in asm.strip().splitlines() if ln.strip()])
+  src = f"""
+  typedef long unsigned int size_t;
+  extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+  extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
+    int l = __ockl_get_local_id(0);
+    unsigned a = 0, b = 0, c = 0;
+    {input_asm}
+    unsigned res;
+    asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
+    *(data0_1+l) = res;
+  }}"""
+  t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
+  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
+  car = CompiledRunner(prg)
+  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
+  car([t.uop.buffer], {}, wait=True)
+  return t.numpy()

 def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
 def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -105,54 +39,57 @@ class TestHW(unittest.TestCase):

  def test_simple(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v1 v10
-    """, n_threads=2)
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 %1
+    """)[0]
    np.testing.assert_equal(out, 42)

  def test_exec_mov(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
+    v_mov_b32_e32 %1 42
    s_mov_b32_e32 exec_lo 0b10
-    v_mov_b32_e32 v10 10
+    v_mov_b32_e32 %1 10
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 v10
+    v_mov_b32_e32 %2 %1
    """, n_threads=2)
    np.testing.assert_equal(out, [42, 10])

  def test_exec_cmp_vopc(self):
    out = get_output("""
    s_mov_b32 vcc_lo 0 // reset vcc
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v11 10
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmp_ne_u32 v10 v11
+    v_cmp_ne_u32 %1 %2
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 vcc_lo
+    v_mov_b32_e32 %2 vcc_lo
    """, n_threads=2)
    np.testing.assert_equal(out, 0b01)

  def test_exec_cmpx_vop3(self):
    out = get_output("""
-    v_mov_b32_e32 v10 42
-    v_mov_b32_e32 v11 10
+    s_mov_b32_e32 exec_lo 0b11
+    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 %2 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmpx_ne_u32 v10 v11
+    v_cmpx_ne_u32 %1 %2
    s_mov_b32_e32 s10 exec_lo
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 v1 s10
-    """, n_threads=2)
-    np.testing.assert_equal(out, 0b01)
+    v_mov_b32_e32 %2 s10
+    """, n_threads=2)[0]
+    np.testing.assert_equal(out & 0b11, 0b01)

  def test_fmac_vop3_modifier(self):
    init_state = f"""
-    v_mov_b32_e32 v10 {f16_to_bits(4.0)}
-    v_mov_b32_e32 v11 {f16_to_bits(3.0)}
-    v_mov_b32_e32 v1 {f16_to_bits(2.0)}
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
+    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
    """
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 v11 v10"), f16_to_bits(14.))
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 v10"), f16_to_bits(-10.))
-    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 -v10"), f16_to_bits(14.))
+    mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
+    def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
+    self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))

  def test_s_abs_i32(self):
    def s_abs_i32(x, y, dst="s10", scc=0):
@@ -160,7 +97,7 @@ class TestHW(unittest.TestCase):
        self.assertEqual(get_output(f"""
        s_mov_b32_e32 {dst} {x}
        s_abs_i32 {dst} {dst}
-        v_mov_b32_e32 v1 {reg}
+        v_mov_b32_e32 %2 {reg}
        """)[0], val)
    s_abs_i32(0x00000001, 0x00000001, scc=1)
    s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
@@ -173,8 +110,8 @@ class TestHW(unittest.TestCase):
  def test_v_rcp_f32_neg_vop3(self):
    def v_neg_rcp_f32(x:float, y:float):
      out = get_output(f"""
-      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      v_rcp_f32_e64 v1, -v1
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      v_rcp_f32_e64 %2, -%2
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg_rcp_f32(math.inf, -0.0)
@@ -186,10 +123,11 @@ class TestHW(unittest.TestCase):

  def test_v_cndmask_b32_neg(self):
    def v_neg(x:int|float, y:float):
+      # always pick -v1
      out = get_output(f"""
-      v_mov_b32_e32 v1 {f32_to_bits(x)}
-      s_mov_b32_e32 s10 1 // always pick -v1
-      v_cndmask_b32 v1, v1, -v1 s10
+      v_mov_b32_e32 %2 {f32_to_bits(x)}
+      s_mov_b32_e32 s10 1
+      v_cndmask_b32 %2, %2, -%2 s10
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg(-0.0, 0.0)
@@ -198,5 +136,12 @@ class TestHW(unittest.TestCase):
    v_neg(math.inf, -math.inf)
    v_neg(-math.inf, math.inf)

+  def test_v_subrev_wrap(self):
+    out = get_output("""
+    v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
+    v_subrev_co_u32 %2, vcc_lo, %2, %1
+    """)[0]
+    self.assertEqual(out, 0xffff_ffff)
+
 if __name__ == "__main__":
  unittest.main()
@@ -2,8 +2,6 @@

 ## Getting SQ Thread Trace

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

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
@@ -1,68 +0,0 @@
-import ctypes
-from dataclasses import dataclass
-import tinygrad.runtime.autogen.comgr as comgr
-from tinygrad.runtime.support.compiler_amd import check
-
-@dataclass
-class InstrCtx:
-  pc:int=0
-  inst:str=""
-
-@comgr.amd_comgr_create_disassembly_info.argtypes[2]
-def instr_cb(text, user_data):
-  c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
-  c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
-  return comgr.AMD_COMGR_STATUS_SUCCESS
-
-# nop callback
-@comgr.amd_comgr_create_disassembly_info.argtypes[3]
-def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
-
-def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
-  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
-  lib_buf = ctypes.create_string_buffer(lib, len(lib))
-  check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
-  check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
-
-  @comgr.amd_comgr_create_disassembly_info.argtypes[1]
-  def memory_cb(from_addr, to, size, _):
-    base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
-    start = int(from_addr) - base
-    if start < 0 or start >= buf_len: return 0
-    ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
-    return n
-
-  info_src = comgr.amd_comgr_disassembly_info_t()
-  check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
-
-  @comgr.amd_comgr_iterate_symbols.argtypes[1]
-  def sym_callback(sym, udata):
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
-    if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
-                                                                        ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
-    check(nobits.value)
-    base = ctypes.addressof(lib_buf)
-    pc = base + offset.value
-    end = pc + size.value
-    addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
-    instr_ref = ctypes.py_object(ctx:=InstrCtx())
-    instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
-    while pc < end:
-      size_read = ctypes.c_uint64(0)
-      ctx.pc = pc
-      st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
-      if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
-        rel = (pc - base) - offset.value
-        addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
-        pc += size_read.value
-      else: # don't inf loop if comgr fails
-        b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
-        addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
-        pc += 1
-    return comgr.AMD_COMGR_STATUS_SUCCESS
-  addr_table:dict[int, tuple[str, int]] = {}
-  check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
-  return addr_table
@@ -13,6 +13,6 @@ if __name__ == "__main__":
    os.chmod(fp, 0o755)
    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
  else:
-    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
+    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
  shutil.copy2(lib, DEST)
  print(f"Installed {lib.name} to", DEST)
@@ -154,8 +154,22 @@ class RGP:
      if device not in device_events: raise RuntimeError(f"Device {device} not found in profile, devices in profile: {', '.join(device_events.keys())} ")
      device_event = device_events[device]
    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
+    device_props = device_event.props
+    # merge events per SE
+    merged_sqtt_events:dict[int, ProfileSQTTEvent] = {}
+    for ev in sqtt_events:
+      if ev.se not in merged_sqtt_events: merged_sqtt_events[ev.se] = ev
+      else:
+        merged_sqtt_events[ev.se] = ProfileSQTTEvent(
+          device=ev.device,
+          kern=ev.kern,
+          se=ev.se,
+          itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
+          blob=merged_sqtt_events[ev.se].blob + ev.blob,
+        )
+    sqtt_events = list(merged_sqtt_events.values())
+
    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
-    device_props = sqtt_events[0].props
    gfx_ver = device_props['gfx_target_version'] // 10000
    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
@@ -1,9 +1,32 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses
-from extra.sqtt.rocprof import rocprof
-from extra.sqtt.disasm import comgr_get_address_table
-from tinygrad.helpers import temp, DEBUG
-from tinygrad.device import ProfileEvent, ProfileProgramEvent
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+from tinygrad.helpers import temp, unwrap, DEBUG
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
+from tinygrad.runtime.autogen import llvm, rocprof
+from tinygrad.runtime.support.elf import elf_loader
+
+# to pass NULL to callbacks
+llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
+def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
+  llvm.LLVMInitializeAMDGPUTargetInfo()
+  llvm.LLVMInitializeAMDGPUTargetMC()
+  llvm.LLVMInitializeAMDGPUAsmParser()
+  llvm.LLVMInitializeAMDGPUDisassembler()
+  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
+
+  image, sections, relocs = elf_loader(lib)
+  text = next((sh.header for sh in sections if sh.name == ".text"), None)
+  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
+
+  addr_table:dict[int, tuple[str, int]] = {}
+  out = ctypes.create_string_buffer(128)
+  cur_off = off
+  while cur_off < sz + off:
+    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
+    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
+    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
+    cur_off += instr_sz
+  return addr_table

@dataclasses.dataclass
 class InstInfo:
@@ -17,47 +40,70 @@ class InstInfo:
  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
+  inst:str
+  stall:int
+  dur:int
+  time:int
+
+@dataclasses.dataclass(frozen=True)
+class PrgExec:
+  name:str
+  wave:int
+  cu:int
+  simd:int
+  def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
+
 class _ROCParseCtx:
-  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
-    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
-    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
+  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]] = {}

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

-  def next_sqtt(self): return next(self.sqtt_evs, None)
-  def find_program(self, addr): return self.addr2prg[addr]
+  def next_sqtt(self):
+    x = next(self.sqtt_evs, None)
+    self.active_kern = x.kern if x is not None else None
+    self.active_se = x.se if x is not None else None
+    return x

  def on_occupancy_ev(self, ev):
-    if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
+    if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)

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

-    asm = {}
+    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]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
+      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_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))

-    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
+    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

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

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

  @rocprof.rocprof_trace_decoder_se_data_callback_t
  def copy_cb(buf, buf_size, data_ptr):
@@ -74,12 +120,12 @@ if __name__ == "__main__":
      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
      case _:
-        if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
+        if DEBUG >= 5: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS

  @rocprof.rocprof_trace_decoder_isa_callback_t
  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(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
@@ -92,5 +138,27 @@ if __name__ == "__main__":

    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS

-  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  print(ROCParseCtx.wave_events.keys())
+  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
+  return ROCParseCtx
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  args = parser.parse_args()
+
+  with args.profile.open("rb") as f: profile = pickle.load(f)
+  rctx = decode(profile)
+  print('SQTT:', rctx.wave_events.keys())
+
+  for ev in profile:
+    if not isinstance(ev, ProfilePMCEvent): continue
+    print(f"PMC Event: dev={ev.device} kern={ev.kern}")
+    ptr = 0
+    for s in ev.sched:
+      view = memoryview(ev.blob).cast('Q')
+      print(f"\t{s.name}")
+      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
+        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
+        ptr += 1
@@ -0,0 +1,105 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+os.environ["AMD"] = "1"
+os.environ["VIZ"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+import unittest
+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.engine.realize import CompiledRunner
+from tinygrad.device import Device, ProfileDeviceEvent
+
+from extra.sqtt.roc import decode, InstExec, PrgExec
+
+dev = Device["AMD"]
+
+def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
+
+def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
+  name = sys._getframe(1).f_code.co_name
+  def fxn(_):
+    L = UOp.special(l, "lidx0")
+    G = UOp.special(g, "gidx0")
+    op = custom("asm volatile (")
+    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
+    op = custom(");", op)
+    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
+  k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
+  return k
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[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)
+
+class TestTiming(unittest.TestCase):
+  def test_v_add(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([f"v_add_f32 v{10+i} v{10+i+1} {10+i}" for i in range(3)]).realize()
+    wave = list(sqtt.values())[0][:-1]
+    assert all(s.dur == 1 for s in wave)
+    assert all(s.stall == 0 for s in wave)
+
+  def test_chain_v_add_1l(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([
+        "v_add_f32_e32 v1 v0 v0",
+        "v_add_f32_e32 v2 v1 v1",
+      ]).realize()
+    wave = list(sqtt.values())[0][:-1]
+    assert all(s.dur == 1 for s in wave)
+    assert all(s.stall == 0 for s in wave)
+
+  def test_multi_cycle_inst(self):
+    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)
+
+  def test_wmma(self):
+    with save_sqtt() as sqtt:
+      asm_kernel([
+        "v_wmma_f32_16x16x16_f16 v[16:23], v[0:7], v[8:15], v[16:23]",
+        "v_add_f32_e32 v0 v16 v0",
+      ], l=32*4).realize()
+    assert len(sqtt) == 2, f"expected two waves, got {len(sqtt)} {list(sqtt.keys())}"
+    wmma = list(sqtt.values())[0][0]
+    self.assertGreater(wmma.dur, 1) # rgp says 32 clocks
+
+  def test_sleep(self):
+    n = 1
+    def sleep_kernel(data0):
+      assert data0.dtype.base == dtypes.ulong
+      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
+      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
+      op = custom(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:
+      diff_hw_reg.realize()
+    diff_sqtt = list(sqtt.values())[0][2]
+    self.assertEqual(diff_sqtt.dur, diff_hw_reg.item()-1) # 1 cycle for reading the counter register
+
+if __name__ == "__main__":
+  unittest.main()
@@ -0,0 +1,106 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+constexpr int NUM_WORKERS = 4;
+constexpr int PIPE_STAGES = 3;
+
+constexpr int ATTN_B = 16;
+constexpr int ATTN_N = 1024;
+constexpr int ATTN_H = 16;
+constexpr int ATTN_D = 64;
+
+template<int D> constexpr size_t ROWS = 16*(64/D); // height of each worker tile (rows)
+template<int D, typename T=bf16, typename L=row_l> using qkvo_tile = rt<T, ROWS<D>, D, L>;
+template<int D, typename T=float> using attn_tile = rt<T, ROWS<D>, ROWS<D>>;
+template<int D> using shared_tile = st_bf<ROWS<D>, D>;
+template<int D> using global_layout = gl<bf16, -1, -1, -1, D>; // B, N, H, specified at runtime, D known at compile time for this kernel
+template<int D> struct globals { global_layout<D> Qg, Kg, Vg, Og; };
+
+__launch_bounds__(NUM_WORKERS*WARP_THREADS, 1)
+__global__ void attend_ker(bf16 *O_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr) {
+    constexpr int D = ATTN_D;
+    global_layout<D> Qg{Q_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
+    global_layout<D> Kg{K_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
+    global_layout<D> Vg{V_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
+    global_layout<D> Og{O_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
+    globals<D> g(Qg, Kg, Vg, Og);
+
+    using load_group = kittens::group<2>; // pairs of workers collaboratively load k, v tiles
+    int loadid = load_group::groupid(), workerid = kittens::warpid(); // which worker am I?
+    constexpr int LOAD_BLOCKS = NUM_WORKERS / load_group::GROUP_WARPS;
+    const int batch = blockIdx.z, head = blockIdx.y, q_seq = blockIdx.x * NUM_WORKERS + workerid;
+
+    extern __shared__ alignment_dummy __shm[];
+    shared_allocator al((int*)&__shm[0]);
+
+    shared_tile<D> (&k_smem)[LOAD_BLOCKS][PIPE_STAGES] = al.allocate<shared_tile<D>, LOAD_BLOCKS, PIPE_STAGES>();
+    shared_tile<D> (&v_smem)[LOAD_BLOCKS][PIPE_STAGES] = al.allocate<shared_tile<D>, LOAD_BLOCKS, PIPE_STAGES>();
+
+    shared_tile<D> (&qo_smem)[NUM_WORKERS] = reinterpret_cast<shared_tile<D>(&)[NUM_WORKERS]>(k_smem);
+    // Initialize all of the register tiles.
+    qkvo_tile<D, bf16> q_reg, k_reg; // Q and K are both row layout, as we use mma_ABt.
+    qkvo_tile<D, bf16, col_l> v_reg; // V is column layout, as we use mma_AB.
+    qkvo_tile<D, float> o_reg; // Output tile.
+    attn_tile<D, float> att_block; // attention tile, in float. (We want to use float wherever possible.)
+    attn_tile<D, bf16> att_block_mma; // bf16 attention tile for the second mma_AB. We cast right before that op.
+    typename attn_tile<D, float>::col_vec max_vec_last, max_vec, norm_vec; // these are column vectors for the in-place softmax.
+    // each warp loads its own Q tile of 16x64
+    if (q_seq*ROWS<D> < g.Qg.depth()) {
+      warp::load<1, false>(qo_smem[workerid], g.Qg, {batch, q_seq, head, 0});  // going through shared memory improves coalescing of dram reads.
+      __syncwarp();
+      warp::load(q_reg, qo_smem[workerid]);
+    }
+    __syncthreads();
+
+    if constexpr(D == 64) q_reg *= __float2bfloat16(0.125f * 1.44269504089f);
+    else if constexpr(D == 128) q_reg *= __float2bfloat16(0.08838834764f * 1.44269504089f);
+
+    max_vec = base_types::constants<float>::neg_infty();
+    norm_vec = 0.f;
+    o_reg = 0.f;
+    // launch the load of the first k, v tiles
+    int kv_blocks = (g.Kg.depth() + LOAD_BLOCKS*ROWS<D>-1) / (LOAD_BLOCKS*ROWS<D>), tic = 0;
+    load_group::load_async<1, false>(k_smem[loadid][0], g.Kg, {batch, loadid, head, 0});
+    load_group::load_async<1, false>(v_smem[loadid][0], g.Vg, {batch, loadid, head, 0});
+    // iterate over k, v for these q's that have been loaded
+    for(auto kv_idx = 0; kv_idx < kv_blocks; kv_idx++, tic=(tic+1)%3) {
+        int next_load_idx = (kv_idx+1)*LOAD_BLOCKS + loadid;
+        if(next_load_idx*ROWS<D> < g.Kg.depth()) {
+            int next_tic = (tic+1)%3;
+            load_group::load_async<1, false>(k_smem[loadid][next_tic], g.Kg, {batch, next_load_idx, head, 0});
+            load_group::load_async<1, false>(v_smem[loadid][next_tic], g.Vg, {batch, next_load_idx, head, 0});
+            load_async_wait<1>(); // next k, v can stay in flight.
+        }
+        else load_async_wait();
+        __syncthreads();
+
+        #pragma unroll LOAD_BLOCKS
+        for(int subtile = 0; subtile < LOAD_BLOCKS && (kv_idx*LOAD_BLOCKS + subtile)*ROWS<D> < g.Kg.depth(); subtile++) {
+          warp::load(k_reg, k_smem[subtile][tic]); // load k from shared into registers
+          att_block = 0.f; // zero 16x16 attention tile
+          warp::mma<transpose::N, transpose::T>(att_block, q_reg, k_reg, att_block); // [email protected]
+          // int first_index = (kv_idx*LOAD_BLOCKS + subtile)*ROWS<D>; // one past the last KV index of this tile
+          // int start_fill = g.Kg.depth()-first_index < ROWS<D> ? g.Kg.depth()-first_index : ROWS<D>;
+          // right_fill(att_block, att_block, start_fill, base_types::constants<float>::neg_infty());
+          max_vec_last = max_vec;
+          max_vec = warp::max<axis::COL>(att_block, max_vec); 
+          att_block = warp::exp2(att_block - max_vec); 
+          max_vec_last = warp::exp2(max_vec_last - max_vec); 
+          norm_vec *= max_vec_last; 
+          norm_vec = warp::sum<axis::COL>(att_block, norm_vec); 
+          att_block_mma = att_block; // copy to bf16 tile
+          warp::load(v_reg, v_smem[subtile][tic]);
+          o_reg *= max_vec_last;
+          warp::mma<transpose::N, transpose::N>(o_reg, att_block_mma, v_reg, o_reg);
+        }
+    }
+
+    o_reg /= norm_vec;
+    __syncthreads();
+    if (q_seq*ROWS<D> < g.Og.depth()) { // write out o.
+      warp::store(qo_smem[workerid], o_reg); // going through shared memory improves coalescing of dram writes.
+      __syncwarp();
+      warp::store<1, false>(g.Og, qo_smem[workerid], {batch, q_seq, head, 0});
+    }
+}
@@ -0,0 +1,43 @@
+import pathlib
+from tinygrad import Device, Tensor
+from tinygrad.helpers import Context
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler
+
+if __name__ == "__main__":
+  code = (pathlib.Path(__file__).parent / "fa.cu").read_text()
+  device = Device["CUDA"]
+  kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr", "-DKITTENS_4090"]
+  lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
+  kernel_name = lib.decode().split(".globl\t")[1].split("\n")[0]
+  print("kernel name", kernel_name)
+  print(pretty_ptx(lib.decode()))
+
+  prg = device.runtime(kernel_name, lib)
+  prg.smem = 16384 * 3
+
+  B, N, H, D = 16, 1024, 16, 64
+  q = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
+  k = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
+  v = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
+  out = Tensor.empty(B, N, H, D, device='CUDA', dtype="bfloat16")
+  Tensor.realize(q, k, v, out)
+
+  NUM_WORKERS = 4
+  ROWS = 16 * (64 // D)
+
+  gsz = (N // (ROWS*NUM_WORKERS), H, B)
+  for _ in range(5):
+    et = prg(out.uop.buffer.ensure_allocated()._buf, q.uop.buffer._buf, k.uop.buffer._buf, v.uop.buffer._buf,
+             global_size=gsz, local_size=(ROWS*NUM_WORKERS,1,1), wait=True)
+
+    attn_flops = 2 * B * H * N * N * D + \
+                 4 * B * H * N * N + \
+                 2 * B * H * N * N * D
+    print(f"{attn_flops/(et*1e9):2f} GFLOPS")
+
+  for _ in range(5):
+    with Context(DEBUG=2):
+      ref = q.scaled_dot_product_attention(k, v)
+
+  ref, out = ref.float(), out.float()
+  print((ref-out).mean().item(), (ref-out).max().item())
@@ -46,9 +46,9 @@ __device__ static inline void arrive(int id) {
 #include "memory/memory.cuh"
 #include "shared/shared.cuh"
 #include "register/register.cuh"
+#include "mma/mma.cuh"

 #ifdef KITTENS_HOPPER
-#include "mma/mma.cuh"

 template<int n_reg> __device__ static inline void increase_registers() {
    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
@@ -93,4 +93,4 @@ __device__ static inline void sync() {
 using warp = group<1>;      // scope used by most pre-Hopper GPUs, and also for most register operations.
 using warpgroup = group<4>; // special scope commonly used by Hopper and later.

-}
+}
@@ -65,8 +65,8 @@ template<typename _T, int _axis=-9999, bool _swizzle_flag=true> struct descripto
 namespace detail {
 template<typename... Args>
 struct descriptor_dict {
-    __host__ descriptor_dict() {}
-    template<typename T> __host__ descriptor_dict(T _, int b, int d, int r, int c) {}
+    __host__ __device__ descriptor_dict() {}
+    template<typename T> __host__ __device__ descriptor_dict(T _, int b, int d, int r, int c) {}
    __host__ __device__ descriptor_dict(const descriptor_dict &other) {}
 #ifdef KITTENS_HOPPER
    template<typename T, int U> __device__ const CUtensorMap* get() const {
@@ -85,8 +85,8 @@ struct descriptor_dict<_T, Args...> {
    using DESC = kittens::tma::descriptor<_T>; // copy or initialize with a default value
    CUtensorMap tma_desc;
    descriptor_dict<Args...> other_descs;
-    __host__ descriptor_dict() {}
-    __host__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
+    __host__ __device__ descriptor_dict() {}
+    __host__ __device__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
        kittens::detail::tma::create_tensor_map<typename DESC::T, DESC::axis, DESC::swizzle_flag>(&tma_desc, data, b, d, r, c);
    }
    __host__ __device__ inline descriptor_dict(const descriptor_dict &other) :
@@ -135,7 +135,7 @@ struct gl {

    detail::descriptor_dict<TMA_Types...> tma_descs;

-    __host__ inline gl(T *_data,
+    __host__ __device__ inline gl(T *_data,
                        ducks::gl::make_arg_t<b> _batch,
                        ducks::gl::make_arg_t<d> _depth,
                        ducks::gl::make_arg_t<r> _rows,
@@ -425,4 +425,4 @@ __host__ static inline CUtensorMap* allocate_and_create_tensor_map(const typenam

 } // namespace tma
 } // namespace detail
-} // namespace kittens
+} // namespace kittens
@@ -0,0 +1,45 @@
+// https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/matmul/educational/level_04.cu
+#include "kittens.cuh"
+using namespace kittens;
+
+constexpr int g_N = 8192;
+constexpr int BLOCK_SIZE = 32;
+#define NUM_WORKERS (1)
+
+using sub_tile = st_bf<BLOCK_SIZE,BLOCK_SIZE>;
+using tile_gl =  gl<bf16, 1, 1, g_N, g_N>;
+
+__launch_bounds__(NUM_WORKERS*WARP_THREADS, 1)
+__global__ void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
+  tile_gl g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  extern __shared__ alignment_dummy __shm[];
+  shared_allocator al((int*)&__shm[0]);
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &As = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &Bs = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> A_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> B_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE, ducks::rt_layout::col> B_reg_col;
+  rt_fl<BLOCK_SIZE,BLOCK_SIZE> C_accum;
+
+  int col = blockIdx.x;
+  int row = blockIdx.y;
+
+  warp::zero(C_accum);
+  int num_tiles = (g_N + BLOCK_SIZE - 1) / BLOCK_SIZE;
+  for (int tile = 0; tile < num_tiles; ++tile) {
+    warp::load(As, g_A, {0, 0, row, tile});
+    warp::load(Bs, g_B, {0, 0, tile, col});
+    __syncthreads();
+    warp::load(A_reg, As);
+    warp::load(B_reg, Bs);
+    warp::swap_layout(B_reg_col, B_reg);
+    __syncthreads();
+    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
+    __syncthreads();
+  }
+  warp::store(g_C, C_accum, {0, 0, row, col});
+}
@@ -0,0 +1,55 @@
+import pathlib
+from tinygrad import Device, Tensor
+from tinygrad.helpers import Context, getenv
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler
+
+if __name__ == "__main__":
+  if getenv("MATMUL2"):
+    code = (pathlib.Path(__file__).parent / "matmul2.cu").read_text()
+  else:
+    code = (pathlib.Path(__file__).parent / "matmul.cu").read_text()
+
+  device = Device["CUDA"]
+  kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr"]
+  lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
+  kernel_name = lib.decode().split(".globl\t")[1].split("\n")[0]
+  print("kernel name", kernel_name)
+  print(pretty_ptx(lib.decode()))
+
+  prg = device.runtime(kernel_name, lib)
+  if getenv("MATMUL2"):
+    prg.smem = 16384 * 2
+  else:
+    prg.smem = 10000
+
+  N = 8192
+  a = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
+  b = Tensor.randn(N, N, device='CUDA', dtype="bfloat16")
+  c = Tensor.empty(N, N, device='CUDA', dtype="bfloat16")
+  Tensor.realize(a, b, c)
+
+  WARP_THREADS = 32
+  if getenv("MATMUL2"):
+    SUPER_N = 2
+    SUPER_M = 2
+    NUM_WORKERS = SUPER_N * SUPER_M
+    BLOCK_SIZE = 32
+    gsz = (N // (BLOCK_SIZE * SUPER_N), N // (BLOCK_SIZE * SUPER_M), 1)
+  else:
+    NUM_WORKERS = 1
+    BLOCK_SIZE = 32
+    gsz = (N // (BLOCK_SIZE), N // (BLOCK_SIZE), 1)
+
+  for _ in range(5):
+    et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
+             global_size=gsz, local_size=(NUM_WORKERS*WARP_THREADS,1,1), wait=True)
+    print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
+
+  # print(c.tolist())
+
+  for _ in range(5):
+    with Context(DEBUG=2):
+      ref = (a@b).realize()
+
+  ref, c = ref.float(), c.float()
+  print((ref-c).mean().item(), (ref-c).max().item())
@@ -0,0 +1,105 @@
+#include "kittens.cuh"
+using namespace kittens;
+
+constexpr int g_N = 8192;
+
+constexpr int SUPER_N = 2;
+constexpr int SUPER_M = 2;
+constexpr int NUM_WORKERS = SUPER_N * SUPER_M;
+constexpr int LOAD_TASKS = SUPER_N + SUPER_M;
+
+constexpr int WORKER_M = 32;
+constexpr int WORKER_N = 32;
+
+constexpr int BLOCK_K = 32;
+constexpr int BLOCK_M = WORKER_M * SUPER_M;
+constexpr int BLOCK_N = WORKER_N * SUPER_N;
+
+constexpr int PIPE_STAGES = 2;
+
+using reg_tile_A = rt_bf<WORKER_M, BLOCK_K>;
+using reg_tile_B_col = rt_bf<BLOCK_K, WORKER_N, ducks::rt_layout::col>;
+using reg_tile_C = rt_fl<WORKER_M, WORKER_N>;
+
+using shared_tile_A = st_bf<WORKER_M, BLOCK_K>;
+using shared_tile_B = st_bf<BLOCK_K, WORKER_N>;
+using shared_tile_C = st_bf<WORKER_M, WORKER_N>;
+
+using gl_tile_A = gl<bf16, 1, 1, g_N, g_N, shared_tile_A>;
+using gl_tile_B = gl<bf16, 1, 1, g_N, g_N, shared_tile_B>;
+using gl_tile_C = gl<bf16, 1, 1, g_N, g_N, shared_tile_C>;
+
+__launch_bounds__(NUM_WORKERS *WARP_THREADS, 1) __global__
+    void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
+  gl_tile_C g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
+  gl_tile_A g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
+  gl_tile_B g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  extern __shared__ alignment_dummy __shm[];
+  shared_allocator al((int *)&__shm[0]);
+
+  shared_tile_A(&As)[SUPER_M][PIPE_STAGES] =
+      al.allocate<shared_tile_A, SUPER_M, PIPE_STAGES>();
+  shared_tile_B(&Bs)[SUPER_N][PIPE_STAGES] =
+      al.allocate<shared_tile_B, SUPER_N, PIPE_STAGES>();
+
+  reg_tile_A A_reg;
+  reg_tile_B_col B_reg_col;
+  reg_tile_C C_accum;
+
+  int warpid = kittens::warpid();
+  int warp_m = warpid % SUPER_M;
+  int warp_n = warpid / SUPER_M;
+
+  int load_group_id = warpgroup::groupid();
+
+  int block_row = blockIdx.y * SUPER_M;
+  int block_col = blockIdx.x * SUPER_N;
+
+  warp::zero(C_accum);
+  int num_tiles = (g_N + BLOCK_K - 1) / BLOCK_K;
+
+  for (int load_tile = 0; load_tile < (PIPE_STAGES - 1); load_tile++) {
+    if (load_tile < num_tiles) {
+      int load_smem_idx = load_tile % PIPE_STAGES;
+      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
+        if (task_id < SUPER_M) {
+          warp::load_async(As[task_id][load_smem_idx], g_A, {0, 0, block_row + task_id, load_tile});
+        } else {
+          int n_index = task_id - SUPER_M;
+          warp::load_async(Bs[n_index][load_smem_idx], g_B, {0, 0, load_tile, block_col + n_index});
+        }
+      }
+    }
+  }
+
+  for (int tile = 0; tile < num_tiles; tile++) {
+    int compute_smem_idx = tile % PIPE_STAGES;
+
+    int load_tile = tile + PIPE_STAGES - 1;
+    int load_smem_idx = load_tile % PIPE_STAGES;
+
+    if (load_tile < num_tiles) {
+      for (int task_id = warpid; task_id < LOAD_TASKS; task_id += NUM_WORKERS) {
+        if (task_id < SUPER_M) {
+          warp::load_async(As[task_id][load_smem_idx], g_A,
+                           {0, 0, block_row + task_id, load_tile});
+        } else {
+          int n_index = task_id - SUPER_M;
+          warp::load_async(Bs[n_index][load_smem_idx], g_B,
+                           {0, 0, load_tile, block_col + n_index});
+        }
+      }
+      load_async_wait<1>();
+    } else
+      load_async_wait();
+    __syncthreads();
+
+    warp::load(A_reg, As[warp_m][compute_smem_idx]);
+    warp::load(B_reg_col, Bs[warp_n][compute_smem_idx]);
+
+    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
+    __syncthreads();
+  }
+  warp::store(g_C, C_accum, {0, 0, block_row + warp_m, block_col + warp_n});
+}
@@ -3,9 +3,9 @@ import argparse

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

  Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()
@@ -1,75 +0,0 @@
-import torch
-
-#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
-#some changes: classic momentum instead of weighting gradient
-#added ns_steps, ns_params, nesterov as hyperparams
-def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
-  """
-  Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-  quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-  of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-  zero even beyond the point where the iteration no longer converges all the way to one everywhere
-  on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-  where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-  performance at all relative to UV^T, where USV^T = G is the SVD.
-  """
-  assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
-
-  a, b, c = params
-  X = G
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  # Ensure spectral norm is at most 1
-  X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
-  # Perform the NS iterations
-  for _ in range(steps):
-    A = X @ X.mT
-    B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
-    X = a * X + B @ X
-
-  if G.size(-2) > G.size(-1):
-    X = X.mT
-
-  return X
-
-def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-  if beta:
-    momentum.mul_(beta).add_(grad)
-    update = grad.add(momentum,alpha=beta) if nesterov else momentum
-  else: update = grad
-  if update.ndim == 4: # for the case of conv filters
-    update = update.view(len(update), -1)
-  update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
-  return update
-
-class SingleDeviceMuon(torch.optim.Optimizer):
-  """
-  Muon variant for usage in non-distributed settings.
-  """
-  def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
-    defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
-    super().__init__(params, defaults)
-
-  @torch.no_grad()
-  def step(self, closure=None):
-
-    loss = None
-    if closure is not None:
-      with torch.enable_grad():
-        loss = closure()
-
-    for group in self.param_groups:
-      for p in group["params"]:
-        if p.grad is None:
-          p.grad = torch.zeros_like(p)  # Force synchronization
-        state = self.state[p]
-        if len(state) == 0:
-          state["momentum_buffer"] = torch.zeros_like(p)
-        update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
-                             ns_params=group["ns_params"], nesterov=group["nesterov"])
-        p.mul_(1.0 - group["lr"] * group["weight_decay"])
-
-        p.add_(update.reshape(p.shape), alpha=-group["lr"])
-
-    return loss
@@ -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
 	}
@@ -321,7 +321,7 @@
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
 				CODE_SIGN_IDENTITY = "-";
-				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				CURRENT_PROJECT_VERSION = 1;
@@ -357,7 +357,7 @@
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
 				CODE_SIGN_IDENTITY = "-";
-				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
 				CODE_SIGN_STYLE = Automatic;
 				COMBINE_HIDPI_IMAGES = YES;
 				CURRENT_PROJECT_VERSION = 1;
@@ -502,7 +502,7 @@
 			buildSettings = {
 				AD_HOC_CODE_SIGNING_ALLOWED = YES;
 				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
-				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_IDENTITY = "-";
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 1;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
@@ -530,7 +530,7 @@
 			buildSettings = {
 				AD_HOC_CODE_SIGNING_ALLOWED = YES;
 				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
-				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_IDENTITY = "-";
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 1;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
@@ -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;
 }
@@ -25,8 +25,6 @@ nav:
    - Layout: developer/layout.md
    - Speed: developer/speed.md
    - UOp: developer/uop.md
-    - Grouper:
-      - developer/kernelize.md
    - Runtime:
      - developer/runtime.md
      - HCQ: developer/hcq.md
@@ -1,5 +1,8 @@
 [pytest]
-norecursedirs = extra
+norecursedirs =
+    extra
+    .hypothesis
+    .git
 timeout = 300
 timeout_method = thread
 timeout_func_only = true
@@ -1,6 +1,6 @@
 indent-width = 2
 preview = true
-target-version = "py310"
+target-version = "py311"

 lint.select = [
  "F",  # Pyflakes
@@ -9,10 +9,11 @@ with open(directory / 'README.md', encoding='utf-8') as f:

 testing_minimal = [
  "numpy",
-  "torch==2.8.0",
+  "torch==2.9.0",
  "pytest",
  "pytest-xdist",
  "pytest-timeout",
+  "pytest-split",
  "hypothesis",
  "z3-solver",
 ]
@@ -31,6 +32,7 @@ setup(name='tinygrad',
        'tinygrad.codegen.opt',
        'tinygrad.codegen.late',
        'tinygrad.engine',
+        'tinygrad.mixin',
        'tinygrad.nn',
        'tinygrad.renderer',
        'tinygrad.runtime',
@@ -51,7 +53,7 @@ setup(name='tinygrad',
        "License :: OSI Approved :: MIT License"
      ],
      install_requires=[],
-      python_requires='>=3.10',
+      python_requires='>=3.11',
      extras_require={
        'arm': ["unicorn"],
        'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
@@ -54,6 +54,8 @@ def gen_diff(table_old, table_new):

 def display_diff(diff): return "+"+str(diff) if diff > 0 else str(diff)

+NONCORE_DIRS = {"tinygrad/apps", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
+
 if __name__ == "__main__":
  if len(sys.argv) == 3:
    headers = ["Name", "Lines", "Diff", "Tokens/Line", "Diff"]
@@ -76,9 +78,12 @@ if __name__ == "__main__":
    else:
      print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", floatfmt=".1f")+"\n")
      groups = sorted([('/'.join(x[0].rsplit("/", 1)[0].split("/")[0:2]), x[1], x[2]) for x in table])
+      dir_sizes = {}
      for dir_name, group in itertools.groupby(groups, key=lambda x:x[0]):
-        print(f"{dir_name:30s} : {sum([x[1] for x in group]):6d}")
+        dir_sizes[dir_name] = sum([x[1] for x in group])
+        print(f"{dir_name:30s} : {dir_sizes[dir_name]:6d}")
+      print(f"\n core line count: {sum([v for k,v in dir_sizes.items() if k not in NONCORE_DIRS])}")
      total_lines = sum([x[1] for x in table])
-      print(f"\ntotal line count: {total_lines}")
+      print(f"total line count: {total_lines}")
      max_line_count = int(os.getenv("MAX_LINE_COUNT", "-1"))
      assert max_line_count == -1 or total_lines <= max_line_count, f"OVER {max_line_count} LINES"
@@ -0,0 +1,260 @@
+# ruff: noqa: E501
+from tinygrad import dtypes, Device
+from tinygrad.uop.ops import UOp, AxisType, Ops
+from tinygrad.codegen import full_rewrite
+from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import dedup
+from tinygrad.device import Buffer
+from tinygrad.dtype import ImageDType
+
+# 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
+
+""" 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']
+"""
+
+""" 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__']
+"""
+
+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()
+
+compiler = Device.default.compiler
+renderer = Device.default.renderer
+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))
+cr = CompiledRunner(ps, precompiled=lib)
+
+gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.DEFINE_GLOBAL]), key=lambda u: u.arg)
+print(len(gs))
+print([g.dtype for g in gs])
+
+bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
+
+t = cr(bufs, wait=True)
+print(f"{t*1e6:.2f} us")
@@ -1,63 +0,0 @@
-import time, sys, hashlib
-from pathlib import Path
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad import Tensor, dtypes, TinyJit
-from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange
-import numpy as np
-from extra.bench_log import BenchEvent, WallTimeEvent
-
-OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
-
-if __name__ == "__main__":
-  run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))
-
-  Tensor.manual_seed(100)
-  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
-  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}
-
-  # benchmark
-  for _ in range(5):
-    GlobalCounters.reset()
-    st = time.perf_counter_ns()
-    ret = next(iter(run_onnx(new_inputs_junk).values())).cast(dtypes.float32).numpy()
-    print(f"unjitted: {(time.perf_counter_ns() - st)*1e-6:7.4f} ms")
-
-  # NOTE: the inputs to a JIT must be first level arguments
-  run_onnx_jit = TinyJit(lambda **kwargs: run_onnx(kwargs), prune=True)
-  step_times = []
-  for _ in range(20):
-    GlobalCounters.reset()
-    st = time.perf_counter_ns()
-    with WallTimeEvent(BenchEvent.STEP):
-      # Need to cast non-image inputs from numpy, this is only realistic way to run model
-      inputs = {**{k:v for k,v in new_inputs_junk.items() if 'img' in k},
-                **{k:Tensor(v) for k,v in new_inputs_junk_numpy.items() if 'img' not in k}}
-      ret = next(iter(run_onnx_jit(**inputs).values())).cast(dtypes.float32).numpy()
-      step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
-    print(f"jitted:  {t:7.4f} ms")
-
-  suffix = ""
-  if IMAGE.value < 2: suffix += f"_image{IMAGE.value}" # image=2 has no suffix for compatibility
-  if getenv("FLOAT16") == 1: suffix += "_float16"
-  path = Path(__file__).parent / "openpilot" / f"{hashlib.md5(OPENPILOT_MODEL.encode()).hexdigest()}{suffix}.npy"
-
-  # validate if we have records
-  tinygrad_out = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
-  if getenv("SAVE_OUTPUT"):
-    np.save(path, tinygrad_out)
-    print(f"saved output to {path}!")
-  elif getenv("FUZZ") and path.exists():
-    known_good_out = np.load(path)
-    for _ in trange(1000):
-      ret = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
-      np.testing.assert_allclose(known_good_out, ret, atol=1e-2, rtol=1e-2)
-    print(colored("fuzz validated!", "green"))
-  elif path.exists():
-    known_good_out = np.load(path)
-    np.testing.assert_allclose(known_good_out, tinygrad_out, atol=1e-2, rtol=1e-2)
-    print(colored("outputs validated!", "green"))
-  else:
-    print(colored("skipping validation", "yellow"))
@@ -2,8 +2,9 @@ from extra.models.resnet import ResNet50
 from tinygrad import Tensor, nn, Device
 from tinygrad.helpers import Profiling, Timing, getenv
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.uop.spec import type_verify
+from tinygrad.codegen import full_rewrite_to_sink
+from tinygrad.codegen.late.linearizer import linearize
+from tinygrad.uop.spec import type_verify, program_spec

 if __name__ == "__main__":
  mdl = ResNet50()
@@ -28,18 +29,17 @@ if __name__ == "__main__":
        asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
        if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]

-        rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False)
        with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
          with Timing("***** model rewrite in   "):
            rewritten_uops = []
            for u in asts:
-              rewritten_uops.append(apply_rewrites(u, rewrites))
+              rewritten_uops.append(full_rewrite_to_sink(u, ren=Device.default.renderer))

        if LINEARIZE:
          with Timing("***** model linearize in "):
            uops_line = []
            for u in rewritten_uops:
-              uops_line.append(apply_rewrites(u, rewrites_for_linearizer))
+              uops_line.append(linearize(u))
          with Timing("***** model verify in    "):
-            for u in uops_line: type_verify(u.arg.lst)
-          print(sum(len(u.arg.lst) for u in uops_line))
+            for u in uops_line: type_verify(u, program_spec)
+          print(sum(len(u) for u in uops_line))
@@ -1,39 +0,0 @@
-import subprocess
-import random
-import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-def run_test(i, full_run=False):
-  print(f"\rRunning iteration {i}...", end=" ", flush=True)
-
-  p = subprocess.Popen(['python3', 'test/test_tiny.py', 'TestTiny.test_plus'],  stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-  if not full_run:
-    time.sleep(random.uniform(0, 1200) / 1000)
-    p.kill()
-    _, stderr = p.communicate()
-  else:
-    _, stderr = p.communicate()
-
-  if full_run:
-    stderr_text = stderr.decode()
-    print(stderr_text)
-    assert "Ran 1 test in" in stderr_text and "OK" in stderr_text
-
-max_workers = 4
-with ThreadPoolExecutor(max_workers=max_workers) as executor:
-  futures = []
-  for i in range(1000000):
-    if i % 100 == 0:
-      for future in as_completed(futures):
-        try: future.result()
-        except Exception as e:
-          print(f"\nError in iteration: {e}")
-      futures = []
-
-      run_test(i, True)
-    else:
-      future = executor.submit(run_test, i, False)
-      futures.append(future)
-
-    if len(futures) > max_workers * 2: futures = [f for f in futures if not f.done()]
@@ -0,0 +1,44 @@
+import subprocess
+import random
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tinygrad.helpers import getenv
+
+# checks that HCQ drivers can be killed during operation without causing issues
+
+def run_test(i, full_run=False, force_ok=False):
+  print(f"\rRunning iteration {i}...", end=" ", flush=True)
+
+  p = subprocess.Popen(["python3", "test/test_tiny.py", "TestTiny.test_plus"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+  if not full_run:
+    time.sleep(random.uniform(0, 1200) / 1000.0)
+    p.kill()
+    _, stderr = p.communicate()
+  else:
+    _, stderr = p.communicate()
+    stderr_text = stderr.decode()
+    assert ("Ran 1 test in" in stderr_text and "OK" in stderr_text) or (not force_ok and "Failed to take lock file" in stderr_text), stderr_text
+
+if __name__ == "__main__":
+  max_workers = getenv("MAX_WORKERS", 4)
+  with ProcessPoolExecutor(max_workers=max_workers) as executor:
+    futures = []
+    for i in range(1000000):
+      if i % 100 == 0:
+        # wait for everything we launched so far
+        for f in as_completed(futures):
+          try:
+            f.result()
+          except Exception as e:
+            print(f"\nError in iteration: {e}")
+        futures = []
+
+        # do a full run in the main proc
+        run_test(i, True, force_ok=True)
+      else:
+        futures.append(executor.submit(run_test, i, bool(getenv("FULL_RUN", 0))))
+
+      # keep list small
+      if len(futures) > max_workers * 2:
+        futures = [f for f in futures if not f.done()]
@@ -0,0 +1,38 @@
+from tinygrad import Tensor, nn, Context, GlobalCounters
+
+if __name__ == "__main__":
+  conv = nn.Conv2d(64, 128, 3)
+  img = Tensor.randn((1,64,128,128))
+  with Context(DEBUG=0, BEAM=0):
+    Tensor.realize(img, conv.weight, conv.bias)
+
+  tst = conv(img).permute(0,2,3,1).realize()
+  print(tst.shape)
+
+  print("NEW")
+  img_perm = img.permute(0,2,3,1).contiguous()
+  print(img_perm.shape)
+  pp = img_perm.permute(0,3,1,2)._pool((3,3)).permute(0,2,3,4,5,1)
+
+  def hwio(pp, conv):
+    pp = pp.unsqueeze(-1)
+    weight = conv.weight.permute(2,3,1,0).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-4,-3, -2])
+
+  def ohwi(pp, conv):
+    pp = pp.unsqueeze(-4)
+    weight = conv.weight.permute(0,2,3,1).contiguous()
+    print(pp.shape, weight.shape, (pp*weight).shape)
+    return (pp * weight).sum([-3,-2,-1])
+
+  for f in [hwio, ohwi]:
+    GlobalCounters.reset()
+    print("\n**************", f.__name__, "**************")
+    out = f(pp, conv)
+    out.realize()
+    print(out.shape)
+
+    with Context(DEBUG=0, BEAM=0):
+      err = (tst-out).square()
+      print(err.mean().item(), err.max().item())
@@ -0,0 +1,20 @@
+import os
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+
+import unittest, time
+from tinygrad import Device
+
+class TestOpen(unittest.TestCase):
+  def generate_test_open(n):
+    def test(self):
+      dev = Device[Device.DEFAULT]
+      for i in range(10):
+        dev.allocator.alloc(10 << 20)
+        time.sleep(0.5)
+    test.__name__ = f'test_open_{n}'
+    return test
+
+  for i in range(64): locals()[f'test_open_{i}'] = generate_test_open(i)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -272,6 +272,10 @@ class TestMainOnnxOps(TestOnnxOps):
  def test_qlinearmatmul_2D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 2)
  def test_qlinearmatmul_3D_int8_float32(self): self._run_qlinearmatmul_test(np.int8, np.float32, 3)

+  def test_reduce_l2_half(self):
+    inputs = {"data": np.random.randn(1, 1, 32, 32, 32).astype(np.half)*100}
+    self.helper_test_single_op("ReduceL2", inputs, {}, ["reduced"])
+
 class TestTrainingOnnxOps(TestOnnxOps):
  # NOTE: ORT doesn't actually support training ops on cpu so we test using functions provided by onnx
  DOMAIN = AI_ONNX_PREVIEW_TRAINING_DOMAIN
@@ -282,11 +286,11 @@ class TestTrainingOnnxOps(TestOnnxOps):
    tiny_out = runner(inps)
    onnx_out = onnx_fxn(**inps, **opts)
    for (nm, t_out), o_out in  zip(tiny_out.items(), onnx_out):
-      np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-3, atol=1e-6, err_msg=f"{nm} failed")
+      np.testing.assert_allclose(t_out.numpy(), o_out, rtol=1e-6, atol=1e-6, err_msg=f"{nm} failed")

-  def test_adagrad_t_greater_than_zero(self):
+  def test_adagrad_t(self):
    from onnx.backend.test.case.node.adagrad import apply_adagrad
-    for t in [1, 3, 100]:
+    for t in [0, 1, 3, 100]:
      inputs = {
        "r": np.array(0.01, dtype=np.float32),
        "t": np.array(t, dtype=np.int32),
@@ -298,10 +302,10 @@ class TestTrainingOnnxOps(TestOnnxOps):
      outputs = ["X_out", "H_out"]
      self._validate_training("Adagrad", apply_adagrad, inputs, attributes, outputs)

-  def test_momentum_t_greater_than_zero(self):
+  def test_momentum(self):
    from onnx.backend.test.case.node.momentum import apply_momentum, apply_nesterov
    for onnx_fxn, mode in ((apply_momentum, "standard"), (apply_nesterov, "nesterov")):
-      for t in [1, 3, 100]:
+      for t in [0, 1, 3, 100]:
        inputs = {
          "r": np.array(0.01, dtype=np.float32),
          "t": np.array(t, dtype=np.int32),
@@ -313,9 +317,9 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

-  def test_adam_t_greater_than_zero(self):
+  def test_adam(self):
    from onnx.backend.test.case.node.adam import apply_adam
-    for t in [1, 3, 100]:
+    for t in [0, 1, 3, 100]:
      inputs = {
        "r": np.array(0.01, dtype=np.float32),
        "t": np.array(t, dtype=np.int32),
@@ -422,6 +426,7 @@ class TestContribOnnxOps(TestOnnxOps):
        outputs = ["C"]
        self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate

+  def test_qlinear_add_round_half_to_even(self):
    with self.subTest(test_case="round_half_to_even"):
      inputs = {
        "A": np.array([1, 1, 1, 1], dtype=np.int8),
@@ -435,7 +440,7 @@ class TestContribOnnxOps(TestOnnxOps):
      }
      attributes = {}
      outputs = ["C"]
-      self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs)
+      self.helper_test_single_op("QLinearAdd", inputs, attributes, outputs, atol=1) # TODO: look into why this is inaccurate

  def test_qlinear_mul(self):
    for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
@@ -486,4 +491,4 @@ class TestContribOnnxOps(TestOnnxOps):
        self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)

 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -1,7 +1,7 @@
 import random
 import z3
 from tinygrad import dtypes
-from tinygrad.uop.spec import uops_to_z3, z3_cdiv
+from tinygrad.uop.validate import uops_to_z3, z3_cdiv
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)
@@ -207,7 +207,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")

        # stop if kernel uops repeat
-        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.opts).uops)
+        try: tuops = tuplize_uops(get_program(test_lin.get_optimized_ast(), test_lin.ren).uops)
        except KeyboardInterrupt: raise
        except BaseException as e:
          print(test_lin.ast)
@@ -224,7 +224,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
        if state1 is not None and validate_device is not None:
          validate_lin = test_lin.copy()
-          validate_lin.opts = validate_device.renderer
+          validate_lin.ren = validate_device.renderer
          if validate_rawbufs is None:
            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
@@ -2,7 +2,7 @@ import random, operator
 import z3
 from tinygrad import Variable, dtypes
 from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import uops_to_z3
+from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context

 seed = random.randint(0, 100)
@@ -13,7 +13,7 @@ try:
  from tinygrad.engine.realize import get_program
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
  from tinygrad.codegen.opt import Opt
-  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
+  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
  from tinygrad.device import Device
 except ImportError as e:
  print(repr(e))
@@ -51,8 +51,8 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
-  # NOTE: this always uses the opts_to_apply path
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=p.applied_opts)
+  # the ast.arg is non None if we are inside of search.py
+  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
  input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
  # if no renderer was provided, open the device to get it
  if renderer is None: renderer = Device[p.device].renderer
@@ -85,17 +85,16 @@ class TestKernelSpeed(unittest.TestCase):
    gbs = mems / tm / 1e9
    self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)

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

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

  # theoretical is nv_gbs=1008, amd_gbs=960
  def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=840, amd_gbs=750)
-  def test_gemv_4096_16384(self): self._test_matmul(4096, 16384, 1, nv_gbs=830, amd_gbs=750)
+  def test_gemv_4096_16384(self): self._test_matmul(4096, 16384, 1, nv_gbs=820, amd_gbs=750)

 if __name__ == '__main__':
  unittest.main()
@@ -85,6 +85,8 @@ class AMDDriver(VirtDriver):
      VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/gpu_id', functools.partial(TextFileDesc, text=f"{gpu_id}")),
      VirtFile(f'/sys/devices/virtual/kfd/kfd/topology/nodes/{gpu_id}/properties',
        functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id))),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/power_dpm_force_performance_level',
+               functools.partial(TextFileDesc, text='profile_standard\n')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
               functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
@@ -14,6 +14,9 @@ regSQ_THREAD_TRACE_BUF0_BASE = 0x39e8 + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_BUF0_SIZE = 0x39e9 + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_WPTR = 0x39ef + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_STATUS = 0x39f4 + amd_gpu.GC_BASE__INST0_SEG1
+regCP_PERFMON_CNTL = 0x3808 + amd_gpu.GC_BASE__INST0_SEG1
+regCPG_PERFCOUNTER1_LO = 0x3000 + amd_gpu.GC_BASE__INST0_SEG1
+regGUS_PERFCOUNTER_HI = 0x3643 + amd_gpu.GC_BASE__INST0_SEG1

 class SQTT_EVENTS:
  THREAD_TRACE_FINISH = 0x00000037
@@ -130,7 +133,7 @@ class PM4Executor(AMDQueue):
    _src_addr_hi = self._next_dword()
    dst_addr_lo = self._next_dword()
    dst_addr_hi = self._next_dword()
-    assert copy_data_flags == 0x100204, hex(copy_data_flags) # better fail than silently do the wrong thing
+    assert copy_data_flags in {0x100204, 0x000204}, hex(copy_data_flags) # better fail than silently do the wrong thing
    to_mv(dst_addr_hi<<32|dst_addr_lo, 4).cast('I')[0] = self.gpu.regs[src_addr_lo]

  def _exec_wait_reg_mem(self, n):
@@ -280,6 +283,9 @@ class AMDGPURegisters:
    self.regs: dict[tuple[int, int], int] = {}
  def __getitem__(self, addr:int) -> int:
    if addr == regGRBM_GFX_INDEX: return self.grbm_index
+    if regCPG_PERFCOUNTER1_LO < addr < regGUS_PERFCOUNTER_HI:
+      assert self.regs[(regCP_PERFMON_CNTL, 0)] == 0x401, "read mode should be enabled"
+      return addr << 16 | self.grbm_index
    return self.regs[(addr, getbits(self.grbm_index, 16, 23))]
  def __setitem__(self, addr:int, val:int):
    if addr == regGRBM_GFX_INDEX: self.grbm_index = val
@@ -112,7 +112,7 @@ class TestRealWorld(unittest.TestCase):
        loss.backward()
        optimizer.step()

-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 102)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 103)

  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
  def test_forward_cifar(self):
@@ -176,7 +176,7 @@ class TestRealWorld(unittest.TestCase):
      for v in data.values(): v.to_(Device.DEFAULT)

      helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
-          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 358)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 427)

 if __name__ == '__main__':
  unittest.main()
@@ -14,6 +14,8 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 # TODO: write a clean version of this
 from test.test_linearizer import helper_realized_ast, helper_linearizer_opt

+# NOTE: get_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
+
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                         ensure_triggered:bool=True):
  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
@@ -41,7 +43,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
  if dtype_in == dtypes.bfloat16: r = r.float()
  realized_ast, bufs = helper_realized_ast(r)
  opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, opts=opts), device=Device.DEFAULT))
+  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
  if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
  assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
  prg.exec(bufs)
@@ -68,7 +70,7 @@ class TestTensorCores(unittest.TestCase):
      n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
      r = a.matmul(b, dtype=tc.dtype_out)
-      prg = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
+      prg = get_program(r.schedule()[-1].ast, Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
      if Device.DEFAULT == "CPU" and CPU_LLVM:
        assert "0x201000" in prg.src
      elif Device.DEFAULT == "AMD" and AMD_LLVM:
@@ -154,7 +156,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.STORE

@@ -167,7 +169,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -182,7 +184,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out).relu()
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, opts=opts).uops:
+    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -1,5 +1,5 @@
 import unittest, itertools, math
-from tinygrad import Tensor, Device, dtypes
+from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
@@ -24,7 +24,6 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))

-  @unittest.expectedFailure  # no two level fold
  def test_neg_folding(self):
    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -68,12 +67,9 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_tensor_one_mul(self):
    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))

-  # TODO: these will be fixed with better folding
-  @unittest.expectedFailure
  def test_bool_tensor_mul_bool(self):
    _check_ast_count(0, Tensor([True, False]) * True)
    _check_ast_count(0, Tensor([True, False]) * False)
-  @unittest.expectedFailure
  def test_bool_mul_bool_tensor(self):
    _check_ast_count(0, True * Tensor([True, False]))
    _check_ast_count(0, False * Tensor([True, False]))
@@ -83,10 +79,8 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_div_tensor_one(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))

-  @unittest.expectedFailure  # TODO: fix
  def test_idiv_literal_one(self):
    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  @unittest.expectedFailure  # TODO: fix
  def test_idiv_tensor_one(self):
    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))

@@ -132,7 +126,8 @@ class TestBitcastConstFolding(unittest.TestCase):
    t({dtypes.int64: 4598983288165178391, dtypes.uint64: 4598983288165178391, dtypes.float64: 0.29485681936461233})

  def test_vec_bitcast(self):
-    r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
+    with Context(SPEC=0):
+      r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
    self.assertEqual(r.op, Ops.VECTORIZE)
    self.assertEqual(r.dtype, dtypes.uint32.vec(3))
    self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75))
@@ -188,7 +183,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
    np.testing.assert_equal(Tensor(4).sum().numpy(), 4)

  def test_padded_const_sum(self):
-    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).sum())
+    _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)

    # NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
@@ -0,0 +1,221 @@
+import unittest
+from tinygrad import Tensor, UOp, Context
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import KernelInfo, AxisType
+
+# **** kernels ****
+
+def custom_arange_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.size, 0)
+  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
+
+def custom_eye_kernel(C:UOp) -> UOp:
+  i = UOp.range(C.shape[0], 0)
+  j = UOp.range(C.shape[1], 1)
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
+
+def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
+  A,B = A.flatten(), B.flatten()
+  assert B.size == A.size
+  i = UOp.range(A.size, 0)
+  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.size}"))
+
+def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
+  C,A,B = C.flatten(), A.flatten(), B.flatten()
+  i = UOp.range(C.size, 0)
+  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.size}")).simplify()
+
+def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
+  C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
+  assert C.size == D.size
+  i = UOp.range(C.size, 0)
+  store_c = C[i].store(A[i]+B[i])
+  store_d = D[i].store(A[i]*B[i])
+  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
+
+def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  assert A.shape[1] == B.shape[0]
+  i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
+  C = C[i, j].set(0.0)
+  C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
+  prog = C.end(i, j)
+  return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
+
+def custom_sum(B:UOp, A:UOp) -> UOp:
+  i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
+  B = B[0].set(0.0)
+  B = B[0].set(B.after(i)[0] + A[i], end=i)
+  return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
+
+def flip_contract_kernel(dest:UOp, src:UOp):
+  i = UOp.range(dest.shape[0], 0)
+  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
+  vec = src[i, j].contract(j)
+  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
+  return store.end(i).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
+
+def slice_sum_kernel(dest:UOp, src:UOp):
+  G = UOp.range(src.shape[0], 0)
+  slice_src = src[G, :]
+  reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
+  reg = reg.after(G)[0].set(0)
+  R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
+  reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
+  ast = dest[G].set(reg[0], end=G)
+  return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
+
+def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
+  # attention without softmax
+  N, d = Q.shape[0], Q.shape[1]
+
+  i = UOp.range(N, 0)  # output row
+  d_out = UOp.range(d, 1)  # output column
+  j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
+
+  k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
+  qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
+  qk_acc = qk_acc.after(i, j)[0].set(0.0)
+  qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
+  qk_score = qk_acc[0] / (d ** 0.5)
+
+  out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
+  out_acc = out_acc.after(i, d_out)[0].set(0.0)
+  out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
+
+  store = O[i, d_out].store(out_acc[0])
+  return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
+
+# **** backward callbacks ****
+
+def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
+  out, a, b = kernel.src
+  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
+  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
+  return (None, grad_a, grad_b)
+
+def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
+  out, a, b = kernel.src
+  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
+  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
+  return (None, grad_a, grad_b)
+
+# **** tests ****
+
+class TestCustomKernel(unittest.TestCase):
+  def test_simple(self):
+    a = Tensor.ones(16, 16).contiguous()
+    b = Tensor.ones(16, 16).contiguous()
+    c = Tensor.empty(16, 16)
+
+    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
+
+    out = c.flatten().tolist()
+    assert all(x == 2 for x in out), "all 2"
+
+  def test_multioutput(self):
+    a = Tensor.full((16, 16), 3.).contiguous()
+    b = Tensor.full((16, 16), 3.).contiguous()
+    c = Tensor.empty(16, 16)
+    d = Tensor.empty(16, 16)
+
+    c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
+    Tensor.realize(c,d)
+
+    assert all(x == 6 for x in c.flatten().tolist()), "all 6"
+    assert all(x == 9 for x in d.flatten().tolist()), "all 9"
+
+  def test_arange(self):
+    ref = Tensor.arange(100)
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
+  def test_eye(self):
+    ref = Tensor.eye(1024).contiguous().realize()
+    tst = Tensor.empty_like(ref)
+    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
+    self.assertTrue((ref == tst).all().item())
+
+  def test_flip_contract(self):
+    a = Tensor.randn(10,4)
+    b = Tensor.empty_like(a)
+    b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
+    self.assertTrue((a.flip(1) == b).all().item())
+
+  def test_noncontig(self):
+    a = Tensor.ones(16, 16).contiguous()
+    tst = Tensor.empty_like(a)
+    b = a+1
+    b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
+    self.assertTrue((b_p1 == 3).all().item())
+
+  def test_sum(self):
+    # TODO: this only works for float, and silently fails with int
+    a = Tensor([1.0, 2, 3, 4, 5])
+    tst = Tensor.empty(1)
+    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
+    self.assertEqual(b.item(), 15)
+
+  def test_slice_sum(self):
+    A = Tensor.randn(16, 16).contiguous()
+    B = Tensor.empty(16)
+    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
+    self.assertTrue(B.allclose(A.sum(1)))
+
+  def test_gemm(self):
+    N = 16
+    a = Tensor.randn(N, N)
+    b = Tensor.randn(N, N)
+    c = Tensor.empty(N, N)
+
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
+    err = (tst - (a@b)).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_gemm_backward_custom(self): self.test_gemm_backward(True)
+  # NOTE: grad_fxn doesn't work with pyrender
+  @Context(SPEC=1)
+  def test_gemm_backward(self, custom_backward_gemm=False):
+    N = 4
+    a_rand = Tensor.randn(N, 8)
+    b_rand = Tensor.randn(8, N)
+    Tensor.realize(a_rand, b_rand)
+
+    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
+    c = Tensor.empty(N, N)
+    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
+    tst.sum().backward()
+    grad_a, grad_b = a.grad, b.grad
+    Tensor.realize(tst, grad_a, grad_b)
+
+    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
+    ref = (a@b)
+    ref.sum().backward()
+    real_grad_a, real_grad_b = a.grad, b.grad
+    Tensor.realize(ref, real_grad_a, real_grad_b)
+
+    err = (tst - ref).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+    err = (grad_a - real_grad_a).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+    err = (grad_b - real_grad_b).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+  def test_simple_qkv(self):
+    N, d = 8, 4
+    Q = Tensor.randn(N, d)
+    K = Tensor.randn(N, d)
+    V = Tensor.randn(N, d)
+    O = Tensor.empty(N, d)
+
+    O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
+    O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
+
+    Tensor.realize(O_custom, O_ref)
+    err = (O_custom - O_ref).square().max()
+    self.assertLess(err.item(), 1e-6)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -75,7 +75,10 @@ def universal_test_unary(a, dtype, op):
  out: Tensor = op[0](ta)
  tensor_value = out.numpy()
  numpy_value = op[1](ta.numpy())
-  if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value)
+  if dtype in dtypes.fp8s:
+    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
+    if math.isinf(numpy_value): return
+    numpy_value = truncate[dtype](numpy_value)
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
@@ -191,6 +194,7 @@ class TestDTypeALU(unittest.TestCase):
         strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
  @unittest.skipIf(Device.DEFAULT == "PYTHON", "TODO: fix cast inf to int32 in PYTHON")
+  @unittest.skip("broken on Mac")
  def test_float_midcast_int32(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.float32, dtypes.int32)

  @unittest.skip("broken. TODO: fix it")
@@ -51,7 +51,7 @@ class TestFusionOp(unittest.TestCase):
    a = Tensor(val)
    for _ in range(24): a = Tensor.stack(a, a)[0]
    sched = a.schedule()
-    self.assertEqual(len(sched), 1)
+    self.assertEqual(len(sched), 0)
    self.assertLess(time.perf_counter()-st, 2.0)

  def test_recursive_reshape(self):
@@ -52,7 +52,6 @@ class TestImageDType(unittest.TestCase):
    assert isinstance(it.uop.base.realized.dtype, ImageDType)
    np.testing.assert_equal(tst, it.numpy())

-  @unittest.expectedFailure # this isn't supported anymore, CAST to ImageDType stays ImageDType
  def test_image_cast_and_back_collapses(self):
    data = Tensor.randn(9*27*4).realize()
    tst = data.numpy()
@@ -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,10 +38,26 @@ 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])
-      ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.ENDRANGE and u.src[0] in range_in_acc)]
+      ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.END and u.src[0] in range_in_acc)]
      for i,u in enumerate(ranges):
        if skip and i in skip: continue
        assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
@@ -155,6 +171,7 @@ class TestLinearizer(unittest.TestCase):
    assert stores[1].src[1].dtype == dtypes.float
    assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())

+  @unittest.skipIf(Device.DEFAULT=="CPU", "CPU splits the cat so cant upcast")
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
@@ -205,7 +222,7 @@ class TestLinearizer(unittest.TestCase):
    # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
    uops = get_program(ast, opts=opt).uops
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
-    end_range = [i for i, x in enumerate(uops) if x.op is Ops.ENDRANGE][0]
+    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
    for u in uops:
      if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace is AddrSpace.REG:
@@ -214,8 +231,8 @@ class TestLinearizer(unittest.TestCase):
        else:
          assert u.src[1].op in GroupOp.ALU
          assert begin_range < uops.index(u) < end_range
-      # children of STORE are placed after ENDRANGE
-      if any(x.op is Ops.STORE and x.src[1].op in GroupOp.ALU for x in u.src):
+      # children of END are placed after ENDRANGE
+      if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):
        assert end_range < uops.index(u)

  def test_grouped_dims(self):
@@ -261,6 +278,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)

@@ -285,6 +304,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,)
@@ -393,14 +433,15 @@ class TestLinearizer(unittest.TestCase):
    uops = get_program(ast, opts=opt).uops
    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
-    barrier = [u for u in uops if u.op is Ops.BARRIER][0]
+    barrier = [u for u in uops if u.op is Ops.BARRIER]
+    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
-    assert len([u for u in uops if u.op is Ops.IF and u.src[-1] == barrier]) == 1
+    assert len([u for u in uops if u.op is Ops.IF])

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -430,6 +471,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):
@@ -16,14 +16,14 @@ class TestLinearizerFailure(unittest.TestCase):
    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
-    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True))).load()
+    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
    c9 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(47040000), arg=2, src=())
-    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True))).load()
+    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11, c1, c2, c3)
+    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = get_program(ast, Device["METAL"].renderer)

@@ -12,11 +12,11 @@ class TestLinearizerFailures(unittest.TestCase):
    c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
    c4 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(163840), arg=1, src=())
    c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920)))).load()
+    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=2, src=())
-    c8 = c7.index(c3).load()
+    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
-    c10 = c0.index(c3).store(c9, c1, c2)
+    c10 = c0.index(c3).store(c9).end(c1, c2)
    ast = c10.sink()
    get_program(ast)

@@ -596,6 +596,12 @@ class TestMultiTensor(unittest.TestCase):
    # ast are the same on devices
    self.assertEqual(len(set(asts)), 1)

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

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

  def test_std_mean_loaded_nan(self):
-    helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
-                                  lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
+    with warnings.catch_warnings():
+      warnings.filterwarnings("ignore", message="std_mean\\(\\): degrees of freedom is <= 0")
+      helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
+                                    lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
  def test_softmax(self):
    helper_test_op([(45,65)], torch.nn.Softmax(dim=1), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
    helper_test_op([(45)], torch.nn.Softmax(dim=0), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
@@ -2602,17 +2604,13 @@ class TestOps(unittest.TestCase):
      lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(111,28)),
      lambda x: Tensor.avg_pool2d(x, kernel_size=(111,28)), rtol=1e-5)

-  def test_avg_pool3d_failure(self):
-    with Context(NOOPT=0):
-      helper_test_op([(1,1,16,16,16)],
-        lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
-        lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), rtol=1e-5, forward_only=True)
-
-  def test_avg_pool3d_noopt(self):
-    with Context(NOOPT=1):
-      helper_test_op([(1,1,16,16,16)],
-        lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
-        lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), rtol=1e-5, forward_only=True)
+  def test_avg_pool3d(self):
+    # TODO: AMD_LLVM has larger atol
+    # TODO: PYTHON=1 backward hangs?
+    atol = 1e-2 if AMD_LLVM else 1e-6
+    helper_test_op([(1,1,16,16,16)],
+      lambda x: torch.nn.functional.avg_pool3d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False),
+      lambda x: Tensor.avg_pool2d(x, kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False), atol=atol, rtol=1e-5, forward_only=True)

  def test_interpolate_linear(self):
    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
@@ -2824,13 +2822,13 @@ class TestOps(unittest.TestCase):
  @slow_test
  def test_slice_fancy_indexing_list_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[0],b,c,d,:], lambda x: x[[0],j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(0,),b,c,d,:], lambda x: x[(0,),j,k,o,:])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[[0]]],b,c,d,[[1]]], lambda x: x[[[[0]]],j,k,o,[[1]]])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[[1,0,-1],b,c,d,:], lambda x: x[[1,0,-1],j,k,o,:])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[1,2,3],...], lambda x: x[i,j,k,[1,2,3],...])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(1,0,-1),b,c,d,:], lambda x: x[(1,0,-1),j,k,o,:])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,(1,2,3),...], lambda x: x[i,j,k,(1,2,3),...])
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,b,c,[[1],[2],[3]],...], lambda x: x[i,j,k,[[1],[2],[3]],...])
-    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,[2,1,0],c,[-2,1,0],e], lambda x: x[i,[2,1,0],k,[-2,1,0],p])
+    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])

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

  def test_slice_fancy_indexing_errors(self):
    a = Tensor.ones(10,11,12)
@@ -5,7 +5,6 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
-from extra.torch_muon import SingleDeviceMuon as TorchMuon

 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -58,12 +57,11 @@ class TestOptim(unittest.TestCase):
  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
-  #TODO: use torch.muon when it comes out
-  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, TorchMuon, steps, opts, atol, rtol)
+  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, torch.optim.Muon, steps, opts, atol, rtol)

  def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
  def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
-  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
+  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 1e-2, 5e-4)

  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
@@ -87,27 +85,34 @@ class TestOptim(unittest.TestCase):
  def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
    self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)

-  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
-  def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
-  def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
+  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)

  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
-  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
+  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 3e-3, 0)
  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
-  def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
-  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-5, 0)
-  def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 0.5e-1, 1e-1)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
+  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-3, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 5e-2, 1e-1)

-  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-6, 0)
-  def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
-  def test_muon_ns_params(self): self._test_muon(1, {'lr': 0.001,'ns_params': (2.0,-1.5,0.5)}, 1e-6, 0)
-  def test_muon_high_lr_ns_params(self): self._test_muon(1, {'lr': 10,'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
+  def test_muon_ns_coefficients(self): self._test_muon(1, {'lr': 0.001,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  # TODO: disabled due to big atol
+  # def test_muon_high_lr_ns_coefficients(self): self._test_muon(1, {'lr': 10,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)

-  def test_muon_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 0)
-  def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_params(self):
-    self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+  def test_muon_momentum_wd_ns_steps_ns_coefficients(self):
+    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-4, 0)
+  # TODO: disabled due to big atol
+  # def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_coefficients(self):
+  #   self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)

  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
@@ -92,7 +92,9 @@ class TestProfiler(unittest.TestCase):
    #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"

  def test_profile_multidev(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")
+
    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
    buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()

@@ -109,7 +111,8 @@ class TestProfiler(unittest.TestCase):
      assert evs[0].is_copy, "kernel should be copy"

  def test_profile_multidev_transfer(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")

    buf1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:0").realize()
    with helper_collect_profile(TestProfiler.d0, d1) as profile:
@@ -122,7 +125,8 @@ class TestProfiler(unittest.TestCase):

  @unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
  def test_profile_graph(self):
-    d1 = Device[f"{Device.DEFAULT}:1"]
+    try: d1 = Device[f"{Device.DEFAULT}:1"]
+    except Exception as e: self.skipTest(f"second device not available {e}")

    def f(a):
      x = (a + 1).realize()
@@ -145,7 +149,9 @@ class TestProfiler(unittest.TestCase):
  @unittest.skipIf(CI or not issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "skip CI")
  def test_dev_jitter_matrix(self):
    dev_cnt = 6
-    devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
+    try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
+    except Exception as e: self.skipTest(f"multiple devices not available {e}")
+
    for dev in devs: dev.synchronize()
    for dev in devs: dev._at_profile_finalize()

@@ -193,6 +199,7 @@ class TestProfiler(unittest.TestCase):
    #self.assertLess(e1.st, e2.st)
    #self.assertGreater(e1.en-e1.st, e2.en-e2.st)

+  @unittest.skipIf(not CI, "this test is flaky locally")
  @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
  def test_graph(self):
    from test.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs
@@ -219,7 +226,7 @@ class TestProfiler(unittest.TestCase):
    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
    self.assertEqual(len(exec_points), len(range_events), 2)
-    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
+    self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)

 if __name__ == "__main__":
@@ -1,7 +1,55 @@
 import unittest
-from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters, CI
+from tinygrad import Tensor, nn, Device
+from tinygrad.helpers import Context, GlobalCounters, CI, getenv, PCONTIG, DEBUG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
+from tinygrad.codegen.opt import OptOps, Opt
+from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
+class TestDoubleMatmul(unittest.TestCase):
+  def setUp(self):
+    with Context(DEBUG=0):
+      self.a, self.b, self.c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
+      self.ref = (self.a @ self.b @ self.c).realize()
+
+  def _test(self, opts):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+      out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
+
+    with Context(DEBUG=0):
+      err = (out-self.ref).square()
+      self.assertLess(err.max().item(), 1e-4)
+      self.assertLess(err.mean().item(), 1e-6)
+
+  def test_baseline(self): self._test(())
+  def test_upcast_0(self): self._test((Opt(OptOps.UPCAST, 0, 4),))
+  def test_upcast_1(self): self._test((Opt(OptOps.UPCAST, 1, 4),))
+  def test_upcast_2(self): self._test((Opt(OptOps.UPCAST, 2, 4),))
+  def test_upcast_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_01_mismatch(self): self._test((Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 4)))
+  def test_upcast_02(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 2, 4)))
+  def test_upcast_12(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4)))
+
+  def test_unroll_0(self): self._test((Opt(OptOps.UNROLL, 0, 4),))
+  def test_unroll_1(self): self._test((Opt(OptOps.UNROLL, 1, 4),))
+  def test_unroll_01(self): self._test((Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_0_unroll_0(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_1_unroll_0(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)))
+  def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4)))
+
+  def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4)))
+
+  def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2)))
+  def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2)))
+
+  def test_upcast_01_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
+  def test_upcast_12_unroll_01(self):
+    self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))

 class TestRangeifyAssign(unittest.TestCase):
  def test_assign_permuted(self):
@@ -28,6 +76,83 @@ class TestRangeifyEdgeCase(unittest.TestCase):
    res = Tensor.cat(a, c, dim=0)
    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)

+if getenv("BIG") > 2:
+  # llama 8B (8192)
+  BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
+elif getenv("BIG") > 1:
+  # llama 8B
+  BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+elif getenv("BIG") > 0:
+  # bigger
+  BS, HEADS, SEQLEN, EMB = 4, 32, 128, 128
+else:
+  BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+def fa():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+  GlobalCounters.reset()
+  return q.scaled_dot_product_attention(k, v)
+
+def fa_bw():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0):
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
+    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+    attn_output.weight.requires_grad_().realize()
+    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+  GlobalCounters.reset()
+  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+  out = attn_output(attn)
+  loss = (out - target).square().mean()
+  loss.backward()
+  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
+  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
+  ret = [out, q.grad, k.grad, v.grad]
+  Tensor.realize(*ret)
+  return ret
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
+class TestPcontig(unittest.TestCase):
+  def test_flash_attention_bw(self):
+    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
+      grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(PCONTIG=0, DEBUG=2):
+      cmp_grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=0):
+      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
+    mse = sum(mses)
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention(self, opts=None):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=2):
+      cmp = fa().realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=0):
+      mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention_opt(self):
+    opts = ()
+    # columns in top matrix
+    opts += (Opt(OptOps.UPCAST, 0, 4),)
+    # columns in bottom matrix
+    opts += (Opt(OptOps.UPCAST, 3, 4),)
+    # rows in all the matrix
+    opts += (Opt(OptOps.UPCAST, 4, 4),)
+    self.test_flash_attention(opts)
+
 # *** non CI rangeify tests below this line ***

 N = 256
@@ -215,33 +340,6 @@ class TestRangeify(unittest.TestCase):
    out = blk._feed_forward(x)
    out.realize()

-  @unittest.skip("RANGEIFY=0 does nothing")
-  def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
-
-    # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
-
-    # llama 8B
-    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
-
-    def fa():
-      Tensor.manual_seed(1337)
-      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
-
-    with Context(DEBUG=4):
-      GlobalCounters.reset()
-      ret = fa()
-    with Context(RANGEIFY=0):
-      with Context(DEBUG=2):
-        GlobalCounters.reset()
-        cmp = fa()
-      with Context(DEBUG=0):
-        mse = ((cmp-ret)**2).sum().item()
-    print(f"mse: {mse}")
-    self.assertLessEqual(mse, 1e-6)
-
 # contiguous + reduce can support ranges?

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -1,5 +1,4 @@
 import unittest
-from typing import List, cast
 import numpy as np
 from tinygrad.device import Buffer, Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
@@ -15,15 +14,15 @@ from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.codegen import full_rewrite
 from tinygrad.engine.realize import lower_schedule_item

-def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None):
+def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
  for x in inputs: x.realize()
  # NOTE: we only toposort the stores
-  uops: List[UOp] = []
-  def _recursive_add(uop:UOp) -> List[UOp]: return flatten([_recursive_add(x) for x in uop.src])+[uop]
+  uops: list[UOp] = []
+  def _recursive_add(uop:UOp) -> list[UOp]: return flatten([_recursive_add(x) for x in uop.src])+[uop]
  uops = dedup(flatten(_recursive_add(st) for st in stores))
  outbufs = [Buffer(Device.DEFAULT, sz:=(1 if local_size is None else prod(local_size)), (dtype:=u.src[1].dtype), \
      initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
-  inbufs = [cast(UOp,x.uop).base.buffer for x in inputs]
+  inbufs = [x.uop.base.buffer for x in inputs]
  src = Device[Device.DEFAULT].renderer.render(uops)
  ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
                                  src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
@@ -35,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  a = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0)
  b = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 1)
  idx = UOp.const(dtypes.int, 0)
-  ld = UOp(Ops.LOAD, dtype, (b.index(idx),))
+  ld = b.index(idx)
  alu = ld.alu(alu_op, *alu_src_uops)
  store = UOp.store(a.index(idx), alu)
  sink = UOp(Ops.SINK, dtypes.void, (store,))
@@ -47,7 +46,7 @@ class TestRendererFailures(unittest.TestCase):
  def test_gated_store_with_alu(self):
    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, gate_alu), UOp.const(dtypes.int, 1)))
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
    ret = _test_uop_result([], uops, local_size=[4, 1, 1])[0]
@@ -58,7 +57,7 @@ class TestRendererFailures(unittest.TestCase):
    a = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0+lidx1*4, gate_alu_0&gate_alu_1), UOp.const(dtypes.int, 1)))
+    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
    uops = full_rewrite(sink, Device[Device.DEFAULT].renderer)
    ret = _test_uop_result([], uops, local_size=[4, 2, 1])[0]
@@ -333,7 +333,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 4)
+    schedule = check_schedule([out0, out1], 3)
    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
    self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?

@@ -370,6 +370,7 @@ class TestSchedule(unittest.TestCase):

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
  # should contiguous dedup?
+  @unittest.skip("we do the exact opposite now")
  def test_dedup_contiguous(self):
    a = Tensor.ones(4).contiguous()
    b = Tensor.ones(4).contiguous()
@@ -446,7 +447,7 @@ class TestSchedule(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
  def test_fold_conv_batchnorm_optim(self):
    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 13)]:
+    for optim, cnt in [(nn.optim.Adam, 27), (nn.optim.SGD, 7)]:
      with self.subTest(optim=optim.__name__):
        with Tensor.train():
          img = Tensor.ones(1,3,4,4)
@@ -710,7 +711,7 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(b.buffer.numpy(), [12])

  # unlike schedule, kernelize can be called multiple times on a Tensor
-  def test_double_kerenlize(self):
+  def test_double_kernelize(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10)
    c = (a+b)
@@ -759,7 +760,7 @@ class TestSchedule(unittest.TestCase):

  def test_pow_neg_05_is_rsqrt(self):
    t = Tensor([1.0, 2.0, 3.0]) ** -0.5
-    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIP, Ops.SQRT])
+    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIPROCAL, Ops.SQRT])

  def test_pow_2_has_1_mul(self):
    t = Tensor([1.0, 2.0, 3.0]) ** Tensor(2.0)
@@ -1041,13 +1042,12 @@ class TestSchedule(unittest.TestCase):
      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
      np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)

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

  def test_ugly_reduceop_pairing(self):
    Tensor.manual_seed(0)
@@ -1220,7 +1220,7 @@ class TestSchedule(unittest.TestCase):
      _realize_weights(layer)
      opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4)
      layer(x).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 16)
+      check_schedule(opt.schedule_step(), 19)

  def test_adam_conv_fuse(self):
    with Tensor.train():
@@ -1230,7 +1230,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4)
      opt.zero_grad()
      c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 16)
+      check_schedule(opt.schedule_step(), 19)

  def test_adam_2convs_fuse(self):
    with Tensor.train():
@@ -1241,7 +1241,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 18)
+      check_schedule(opt.schedule_step(), 21)

  def test_sgd_conv_fuse(self):
    with Tensor.train():
@@ -1274,7 +1274,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 15)
+      check_schedule(opt.schedule_step(), 13)

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

+  def test_fuse_arange_avg_pool2d_ceil_mode(self):
+    x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True)
+    sched = check_schedule(x, 1)
+    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1)
+
+  def test_fuse_arange_pad_circular_mode_bw(self):
+    x = Tensor.empty(1,1,5,5,5)
+    out = x.pad((1,2,3,5,1,2), mode="circular")
+    g = out.sum().gradient(x)[0]
+    sched = check_schedule(g, 1)
+    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0)
+
  # TODO like openpilot with imagef
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_base_change_expand_expand(self):
@@ -1863,7 +1875,7 @@ class TestSchedule(unittest.TestCase):
    yt = Tensor.randn(BS, 10).realize()
    with Context(SPLIT_REDUCEOP=0):
      loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-      run_schedule(check_schedule(loss, 5))
+      run_schedule(check_schedule(loss, 4))
      loss_fused = loss.numpy()
    loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
    np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -2076,6 +2088,11 @@ class TestCopyFolding(unittest.TestCase):
    check_schedule(b, 0, filter_sink=False)
    assert b.item() == 1

+  def test_one_hot_with_copy(self):
+    y = Tensor([1, 2, 3]).to("CPU")
+    x = y.one_hot(10)
+    check_schedule(x, 3, filter_sink=False)
+
  def test_const_copy_multi(self):
    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"])
    check_schedule(x, 0, filter_sink=False)
@@ -2085,7 +2102,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(3).realize()
    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU")
-    run_schedule(check_schedule(b, 2, filter_sink=False)) # TODO: 0?
+    run_schedule(check_schedule(b, 0, filter_sink=False))
    self.assertListEqual(b.tolist(), [0, 0, 0])
    self.assertEqual(b.device, "CPU")

@@ -2157,8 +2174,8 @@ class TestCopyFolding(unittest.TestCase):
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
-    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
+    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
+    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
    b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
@@ -2261,7 +2278,7 @@ class TestContiguous(unittest.TestCase):
  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 2) # TODO: should be 1?
+    check_schedule(b, 1)

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

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

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

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

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

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
@@ -287,7 +287,6 @@ class TestSymbolicOps(unittest.TestCase):
        symbolic = symbolic_result[:].numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)

-  @unittest.expectedFailure
  def test_conv2d_ceildiv_edge_case(self):
    v = Variable('v', 11, 50_000)
    val = 39601
@@ -295,9 +294,10 @@ class TestSymbolicOps(unittest.TestCase):
    weight = Tensor.randn(256, 22, 12)

    result = x.conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
-    var_val = {v: val}
+    var_val = {v.expr: val}
    shape = tuple(sym_infer(s, var_val) for s in result.shape)
-    self.assertEqual(shape, (1, 256, 6600))  # TODO: fails if ceildiv is incorrect
+    with self.assertRaises(AssertionError):
+      self.assertEqual(shape, (1, 256, 6600))  # TODO: fails if ceildiv is incorrect
    # TODO: test output is correct

 if __name__ == '__main__':
@@ -810,6 +810,7 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(len(si.metadata), 1)
    self.assertEqual(si.metadata[0].name, "relu")

+  @unittest.skip("this no longer works")
  def test_assign(self):
    x = Tensor.empty(10, 10).realize()
    x.assign(Tensor.ones(10, 10).contiguous())
@@ -839,12 +840,11 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
    self.assertTrue(y.grad.uop.metadata[0].backward)
    si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    self.assertEqual(len(si.metadata), 4, f"failed with {si.metadata}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"__mul__", "sigmoid", "relu"})
-    bw = [m for m in si.metadata if m.backward]
-    self.assertEqual(len(bw), 2)
-    self.assertEqual(bw[0].name, "__mul__")
-    self.assertEqual(bw[1].name, "sigmoid")
+    #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
+    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
+    #bw = [m for m in si.metadata if m.backward]
+    #self.assertEqual(len(bw), 1)
+    #self.assertEqual(bw[0].name, "sigmoid")

 class TestIdxUpcast(unittest.TestCase):
  def _find_op(self, ast: UOp, op: Ops):
@@ -920,5 +920,38 @@ class TestIdxUpcast(unittest.TestCase):
    a = Tensor.empty(2**11, 2**11, 1, dtype=dtypes.int8).permute((2, 0, 1)).expand((2**9+10, -1, -1)).contiguous()
    a.realize()

+class TestTensorUnique(unittest.TestCase):
+  def test_empty_bufs_unique(self):
+    a = Tensor.empty(10, 10).contiguous()
+    b = Tensor.empty(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique_sep(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a)
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_zeros_bufs_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = Tensor.zeros(10, 10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_eye_bufs_unique(self):
+    a = Tensor.eye(10).contiguous()
+    b = Tensor.eye(10).contiguous()
+    Tensor.realize(a,b)
+    self.assertIsNot(a.uop.buffer, b.uop.buffer)
+
+  def test_times_2_not_unique(self):
+    a = Tensor.zeros(10, 10).contiguous()
+    b = a * 2
+    c = a * 2
+    Tensor.realize(b,c)
+    self.assertIs(b.uop.buffer, c.uop.buffer)
+
 if __name__ == '__main__':
  unittest.main()
@@ -3,7 +3,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.engine.realize import run_schedule
-from tinygrad.uop.ops import Ops, UOp, UPat
+from tinygrad.uop.ops import UOp
 from tinygrad.helpers import SPLIT_REDUCEOP

 class TestTensorUOp(unittest.TestCase):
@@ -93,7 +93,6 @@ class TestTensorUOp(unittest.TestCase):
    out.realize()
    self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist())

-reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, allow_any_len=True, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE))))))
@unittest.skipUnless(SPLIT_REDUCEOP, "only for SPLIT_REDUCEOP")
 class TestReduceOp(unittest.TestCase):
  def test_no_split_reduce_kernel(self):
@@ -101,23 +100,18 @@ class TestReduceOp(unittest.TestCase):
    a = a.sum()
    sched = a.schedule()
    assert len(sched) == 1
-    assert reduce_kernel.match(sched[0].ast, {})

  def test_split_reduce_kernel_dim0(self):
    a = Tensor.rand(256, 255).realize()
    a = a.sum()
    sched = a.schedule()
    assert len(sched) == 2
-    for s in sched:
-      assert reduce_kernel.match(s.ast, {})

  def test_split_reduce_kernel_dim1(self):
    a = Tensor.rand(255, 256).realize()
    a = a.sum()
    sched = a.schedule()
    assert len(sched) == 2
-    for s in sched:
-      assert reduce_kernel.match(s.ast, {})

 if __name__ == "__main__":
  unittest.main()
@@ -1,12 +1,11 @@
-from typing import List
 import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, KernelInfo
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
 from tinygrad.uop.symbolic import sym
-from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
 from tinygrad.codegen.late.expander import expander
+from test.test_uops import to_uops_list

 simple_pm = PatternMatcher([
  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
@@ -15,12 +14,6 @@ simple_pm = PatternMatcher([
  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
 ])

-def to_uops_list(u:List[UOp]) -> List[UOp]:
-  # we strip the SINK here for legacy reasons
-  ret = full_rewrite(UOp.sink(*u, arg=KernelInfo(opts_to_apply=())))
-  assert ret[-1].op is Ops.SINK
-  return ret[:-1]
-
 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
@@ -270,6 +263,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out])
    self.assertEqual(len([x for x in uops if x.op is Ops.VECTORIZE]), 0)

+  @unittest.skip("this test isn't valid uops")
  def test_gep_vec_fold(self):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
@@ -313,9 +307,10 @@ class TestUOpGraph(unittest.TestCase):
    for vec_size in [2, 4, 8]:
      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
      vec = UOp(Ops.VECTORIZE, dtypes.float.vec(vec_size), tuple(consts))
-      uops = to_uops_list([UOp(Ops.GEP, dtypes.float, (vec,), (i,)) for i in range(vec_size)])
-      for uop, const in zip(uops, consts):
-        self.assertEqual(uop, const)
+      with Context(SPEC=0):
+        uops = to_uops_list([UOp(Ops.GEP, dtypes.float, (vec,), (i,)) for i in range(vec_size)])
+        for uop, const in zip(uops, consts):
+          self.assertEqual(uop, const)

  @unittest.skip("no longer testable standalone")
  def test_wmma_vectorize_fold(self):
@@ -381,7 +376,7 @@ class TestUOpGraph(unittest.TestCase):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(), arg=0)
    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
    idx = UOp.const(dtypes.int, 0)
-    ld = UOp(Ops.LOAD, dtypes.int, (d1.index(idx),))
+    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
    uops = to_uops_list([out])
@@ -391,7 +386,7 @@ class TestUOpGraph(unittest.TestCase):
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), arg=0)
    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=1)
    idx = UOp.const(dtypes.int, 0)
-    ld = UOp(Ops.LOAD, dtypes.int, (d1.index(idx),))
+    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
    out = UOp(Ops.STORE, dtypes.void, (d0.index(idx), alu))
    uops = to_uops_list([out])
@@ -413,7 +408,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
      d0 = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), arg=0)
-      v = UOp(Ops.LOAD, dt, (d0.index(UOp.const(dtypes.int, 0)),))
+      v = d0.index(UOp.const(dtypes.int, 0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")

@@ -425,7 +420,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
-    ld = d0.index(ridx0.valid(ridx0<50)).load()
+    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
    uops = to_uops_list([w])
    for u in uops:
@@ -435,7 +430,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
-    ld = d0.index(ridx0.valid((ridx0<50).logical_not())).load()
+    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
    for u in uops:
@@ -446,7 +441,7 @@ class TestUOpGraph(unittest.TestCase):
    ridx0 = UOp.range(100, 0)
    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    gate_idx = ridx0.valid((ridx0<50))
-    ld = d0.index(gate_idx).load().cast(dtypes.float)
+    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
    uops = to_uops_list([w])
    for u in uops:
@@ -459,65 +454,65 @@ class TestUOpGraph(unittest.TestCase):
    idx = d0.index(ridx0)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
-    st = idx.store(val, ridx0)
+    st = idx.store(val).end(ridx0)
    uops = to_uops_list([st])
    for u in uops:
      assert u.op is not Ops.WHERE
      if u.op is Ops.STORE: assert u.src[1].arg==5

  def test_load_idx_becomes_int(self):
-    # These loads wont overflow int since we know from the gate that the value is bounded
-    r0 = UOp.range(10, 0)
-    d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0)
-    d1 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 1)
-    l0 = UOp(Ops.LOAD, dtypes.long, (d0.index(UOp.const(dtypes.int, 0)),)).cast(dtypes.index)
-    idx = l0 * 600
-    valid = (l0<-1).ne(True)&(l0<3000)
-    l1 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
-    uops = to_uops_list([l1])
+    # mnist indexing with split reduceop
+    # Make sure we are not doign math on the loaded index, which would promote it to long
+    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(128000), arg=0, src=())
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c3 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(512), arg=1, src=())
+    c4 = c3.index(c1)
+    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c7 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(60000), arg=2, src=())
+    c8 = c7.index(c6)
+    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
+    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    uops = to_uops_list([c10])
    for u in uops:
-      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
-
-    valid = (10*r0<5-l0).ne(True)&(l0<3000)
-    l2 = UOp(Ops.LOAD, dtypes.long, (d1.index(idx.valid(valid)),))
-    uops = to_uops_list([l2])
-    for u in uops:
-      if u.op is Ops.INDEX: self.assertEqual(u.src[1].dtype, dtypes.int)
+      self.assertNotEqual(u.dtype, dtypes.long)

  def test_in_out_of_bounds_access(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 0)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 0), ptr=True),))
      to_uops_list([ld0])
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 15)),))
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 15), ptr=True),))
      to_uops_list([ld1])
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 7)),))
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 7), ptr=True),))
      to_uops_list([ld1])

-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 42)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 42), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

  def test_in_out_of_bounds_access_symbolic(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 1, 10)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 1, 10), ptr=True),))
      to_uops_list([ld0])
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 15)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 15), ptr=True),))
      to_uops_list([ld0])

-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

  def test_in_out_of_bounds_access_gated_store(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
      v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v, v<16), UOp.const(dtypes.int, 0)))
+      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v.valid(v<16)), UOp.const(dtypes.int, 0)))
      to_uops_list([st0])

-      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
+      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v.valid(v<20)), v))
      with self.assertRaises(RuntimeError): to_uops_list([st1])

+  @unittest.skip("if not allowed in graph")
  def test_in_bounds_access_gated_local(self):
    with Context(IGNORE_OOB=0):
      # Define buffers
@@ -536,7 +531,7 @@ class TestUOpGraph(unittest.TestCase):
      if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))

      # Load from local memory (after the IF/barrier)
-      local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx), if_barrier))
+      local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx, ptr=True), if_barrier))

      # Store to global memory
      global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
@@ -547,18 +542,18 @@ class TestUOpGraph(unittest.TestCase):
      ridx = UOp.range(20, 0)
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid((0<=i)&(i<16)), ptr=True),))
      to_uops_list([ld0])
      glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
      ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
      i = (ldfloat+3.14).cast(dtypes.int)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16))),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16)), ptr=True),))

  def test_load_cast_to_bool(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx, ridx.cast(dtypes.bool).logical_not()),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx.valid(ridx.cast(dtypes.bool).logical_not()), ptr=True),))
      to_uops_list([ld0])

  @unittest.skip("Bool load is not supported yet")
@@ -567,48 +562,48 @@ class TestUOpGraph(unittest.TestCase):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask),)))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
      to_uops_list([ld0])

  def test_out_of_bounds_off_by_one_access(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 16)),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 16), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

  def test_in_out_bounds_access_with_mask(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, (5<gidx0)&(gidx0<16)),))
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<16),))
+      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid((5<gidx0)&(gidx0<16)), ptr=True),))
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<16), ptr=True),))
      to_uops_list([ld0, ld1])

-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<17),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<17), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

  def test_in_out_of_bounds_access_symbolic_mask(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      i = Variable("i", 1, 80)
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, i<10),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<10), ptr=True),))
      to_uops_list([ld0])
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, i<15),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<15), ptr=True),))
      to_uops_list([ld0])

-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, i<20),))
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<20), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld0])

  def test_in_out_of_bounds_access_index_load(self):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
-      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 42),), "gidx0")
-      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0, gidx0<8),)).cast(dtypes.index)
-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<32)),))
+      gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
+      ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<8), ptr=True),)).cast(dtypes.index)
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<32)), ptr=True),))
      to_uops_list([ld1])

-      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(ld0*2, (ld0>=0)&(ld0<64)),))
+      ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<64)), ptr=True),))
      with self.assertRaises(RuntimeError): to_uops_list([ld1])

  def test_bounds_with_loaded_bool(self):
@@ -616,8 +611,8 @@ class TestUOpGraph(unittest.TestCase):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
      glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 0)
      gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
-      ld0 = glbl0.index(gidx0).load()
-      ld1 = glbl1.index(gidx0.valid(ld0)).load()
+      ld0 = glbl0.index(gidx0, ptr=True).load()
+      ld1 = glbl1.index(gidx0.valid(ld0), ptr=True).load()
      with self.assertRaises(RuntimeError): to_uops_list([ld1])

  def test_fold_gated_load(self):
@@ -625,38 +620,38 @@ class TestUOpGraph(unittest.TestCase):
    glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
    glbl2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 2)
    idx = UOp.const(dtypes.int, 0)
-    ld0 = UOp(Ops.LOAD, dtypes.int, (glbl1.index(UOp.invalid()),))
-    ld1 = UOp(Ops.LOAD, dtypes.int, (glbl2.index(idx, UOp.const(dtypes.bool, True)),))
+    ld0 = glbl1.index(UOp.invalid())
+    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(idx), ld1+ld0))])
    ld0 = uops[-1].src[-1]
    # the gate and invalid value are deleted from ld1
-    self.assertEqual(ld0, UOp.load(glbl2.index(idx), dtype=dtypes.int))
+    self.assertEqual(ld0, UOp.load(glbl2.index(idx, ptr=True), dtype=dtypes.int))

  def test_fold_gated_load_local(self):
    glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    smem = UOp(Ops.DEFINE_LOCAL, dtypes.int.ptr(size=18, addrspace=AddrSpace.LOCAL), (), "temp")
    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
-    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx), UOp.load(glbl0.index(lidx), dtype=dtypes.int)))
+    st = UOp(Ops.STORE, dtypes.void, (smem.index(lidx, ptr=True), glbl0.index(lidx, ptr=True).load()))
    barrier = UOp(Ops.BARRIER, dtypes.void, (st, ))
-    ld0 = UOp(Ops.LOAD, dtypes.int, (smem.index(UOp.invalid()), barrier))
-    ld1 = UOp(Ops.LOAD, dtypes.int, (smem.index(lidx+2, UOp.const(dtypes.bool, True)), barrier))
+    ld0 = smem.after(barrier).index(UOp.invalid())
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0.index(lidx), ld1+ld0))])

    ld0 = uops[-1].src[-1]
    # the gate and invalid value are deleted from ld1
-    self.assertEqual(ld0.src[0], smem.index(lidx+2))
+    self.assertEqual(ld0.src[0], smem.after(barrier).index(lidx+2, ptr=True))

  def test_fold_gated_store(self):
    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
    idx0 = UOp.const(dtypes.int, 0)
    idx1 = UOp.const(dtypes.int, 0)
    val = UOp.const(dtypes.int, 42)
-    st0 = glbl.index(UOp.invalid()).store(val)
-    st1 = glbl.index(idx0, UOp.const(dtypes.bool, True)).store(val)
+    st0 = glbl.index(UOp.invalid(), ptr=True).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True)), ptr=True).store(val)
    uops = to_uops_list([st0, st1])
    # only the second store happens
    self.assertEqual(len(uops), 5)
-    self.assertEqual(uops[-1], glbl.index(idx1).store(val))
+    self.assertEqual(uops[-1], glbl.index(idx1, ptr=True).store(val))

  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
@@ -665,19 +660,6 @@ class TestUOpGraph(unittest.TestCase):
    bad_gate = UOp.const(dtypes.int, 1)
    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])

-  def test_switched_range_order(self):
-    glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
-    cf = UOp.const(dtypes.float, 0.0)
-    r1 = UOp.range(2, 0)
-    r2 = UOp.range(2, 1)
-    alu = UOp(Ops.MUL, dtypes.int, (r2, r1))
-    store = UOp(Ops.STORE, dtypes.void, (glbl.index(alu), cf))
-    uops = to_uops_list([store])
-    ranges = [x for x in uops if x.op is Ops.RANGE]
-    endranges = [x for x in uops if x.op is Ops.ENDRANGE]
-    # ranges are closed in the right order
-    self.assertEqual(endranges[-1].src[0], ranges[0])
-
@track_rewrites()
 def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)

@@ -739,7 +721,7 @@ class TestExpander(unittest.TestCase):
    self.assertTupleEqual(sink.src[0].arg, (0,2,1,3,4,6,5,7))

  def test_contract_no_expand(self):
-    e1 = UOp(Ops.DEFINE_VAR, dtypes.int)
+    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
    sink = expander_rewrite(con)
    assert sink.op is Ops.VECTORIZE and len(sink.src) == 2
@@ -828,60 +810,6 @@ class TestExpander(unittest.TestCase):
    sink = expander_rewrite(sink)
    print(sink)

-class TestIFUOps(unittest.TestCase):
-  def test_create_ifs(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    idx = UOp.const(dtypes.int, 0)
-    st = UOp(Ops.STORE, dtypes.void, (sbuf.index(idx), UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbuf = UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, 0)), barrier))
-    store = UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, 0), gate), lbuf))
-    sink = UOp(Ops.SINK, dtypes.void, (store,))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-    for st in sink.src:
-      self.assertEqual(len(st.src), 2)
-
-  def test_expand_ifs_one_gate(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "gidx0")<1
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    st = UOp(Ops.STORE, dtypes.void, (sbuf, lidx, UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbufs = [UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, i)), barrier)) for i in range(4)]
-    stores = [UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, i), gate), lbufs[i])) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-    for st in sink.src:
-      self.assertEqual(len(st.src), 2)
-
-  # this will be fixed with the merge gated stores bounty
-  @unittest.expectedFailure
-  def test_expand_ifs_dumb(self):
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    stores = [UOp(Ops.STORE, dtypes.void, (buf, UOp.const(dtypes.int, i), UOp.const(dtypes.float, i), gate)) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-    for st in sink.src:
-      self.assertEqual(len(st.src), 2)
-
 class TestUOpTags(unittest.TestCase):
  def test_inc_by_one(self):
    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
@@ -2,23 +2,29 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import CI, DEBUG, getenv, Timing
+from tinygrad.helpers import CI, DEBUG, getenv, Timing, Context
 from tinygrad.dtype import dtypes, DType, AddrSpace
 from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401
-from tinygrad.uop.spec import spec
+from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
+from tinygrad.uop.spec import shared_spec
 from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.engine.realize import CompiledRunner, get_program, get_runner, ExecItem
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer

-def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
+def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
+  sink = UOp.group(*u)
+  for r in sink.ranges: sink = sink.end(r)
+  # we strip the SINK here for legacy reasons
+  ret = full_rewrite(sink.sink(arg=KernelInfo(opts_to_apply=())), ren)
+  assert ret[-1].op is Ops.SINK
+  return ret[:-1]

 def _uops_to_prg(uops_list):
-  uops = full_rewrite(ast:=UOp.sink(*uops_list), opts=Device[Device.DEFAULT].renderer)
+  uops = full_rewrite(ast:=UOp.sink(*uops_list), ren=Device[Device.DEFAULT].renderer)
  src = Device[Device.DEFAULT].renderer.render(uops)
  has_local = Device[Device.DEFAULT].renderer.has_local
  return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
@@ -33,9 +39,9 @@ def _test_single_value(vals, op, dts):
  output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
  buf_loads = [uop(uops, Ops.DEFINE_GLOBAL, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
-  loads = (uop(uops, Ops.LOAD, dtype, [buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0))]) for i, dtype in enumerate(dts))
+  loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
  alu = uop(uops, op, output_dtype, loads)
-  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
+  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
  prg = _uops_to_prg([out])
@@ -50,7 +56,7 @@ def _test_single_value_const(vals, op, dts):
  buf_store = uop(uops, Ops.DEFINE_GLOBAL, output_dtype.ptr(), (), 0)
  loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
  alu = uop(uops, op, output_dtype, loads)
-  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  prg = _uops_to_prg([out])
  prg.exec([buf])
@@ -109,7 +115,7 @@ class TestFloatUOps(TestUOps):
  def test_log2(self): self._test_uop_fxn(Ops.LOG2, lambda a: math.log2(a) if a > 0 else float('-inf' if a==0 else 'nan'))
  @unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
  def test_sin(self): self._test_uop_fxn(Ops.SIN, lambda a: math.sin(a))
-  def test_recip(self): self._test_uop_fxn(Ops.RECIP, lambda a: 1/a if a != 0 else float('inf'))
+  def test_recip(self): self._test_uop_fxn(Ops.RECIPROCAL, lambda a: 1/a if a != 0 else float('inf'))
  def test_sqrt(self): self._test_uop_fxn(Ops.SQRT, lambda a: math.sqrt(a) if a >= 0 else float('nan'))

  def test_add(self): self._test_bop_fxn(Ops.ADD, lambda a,b: a+b)
@@ -212,18 +218,18 @@ class TestExecALU(TestUOps):
    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (7, -3)), -2)
    self.assertEqual(exec_alu(Ops.IDIV, dtypes.int8, (-50, 6)), -8)

-    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIP, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
-    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIP, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (3.0,)))), 2+(1.0/3.0))
+    np.testing.assert_allclose(exec_alu(Ops.MUL, dtypes.float32, (7.0, exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3.0,)))), -2-(1.0/3.0))

  def test_recip(self):
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (8,)), 1/8)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (7,)), 1/7)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (-3,)), 1/-3)
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (-50,)), 1/-50)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (8,)), 1/8)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (7,)), 1/7)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-3,)), 1/-3)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (-50,)), 1/-50)

-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, ((34**2),)), 1/(34**2))
-    np.testing.assert_allclose(exec_alu(Ops.RECIP, dtypes.float32, (10,)), 1/10)
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((32+521+3),)), 1/(32+521+3))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, ((34**2),)), 1/(34**2))
+    np.testing.assert_allclose(exec_alu(Ops.RECIPROCAL, dtypes.float32, (10,)), 1/10)

  def test_bool_cmplt(self):
    self.assertEqual(exec_alu(Ops.CMPLT, dtypes.bool, (False, False)), False)
@@ -271,7 +277,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
    gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    gate = gidx0<UOp.const(dtypes.int, 1)
-    idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, gidx0 * UOp.const(dtypes.int, 2), gate))
+    idx = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem, (gidx0 * UOp.const(dtypes.int, 2)).valid(gate)))
    val = UOp.const(dtypes.float, 42.0)
    store = UOp(Ops.STORE, dtypes.void, (idx, val))
    uops = to_uops_list([store])
@@ -288,7 +294,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
    gidx0 = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'gidx0')
    idx = gidx0 * UOp.const(dtypes.int, 2)
-    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx, gidx0<UOp.const(dtypes.int, 1)))
+    idx0 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem0, idx.valid(gidx0<UOp.const(dtypes.int, 1))))
    idx1 = UOp(Ops.INDEX, dtypes.float.ptr(), (gmem1, idx))
    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
@@ -302,6 +308,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertIs(gated_uops[-1].op, Ops.STORE)

  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
+  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
    gmem0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
    gmem1 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 1)
@@ -331,7 +338,7 @@ class TestLocalAccess(unittest.TestCase):
    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
-    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
+    sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True),))
    self.assertEqual(_test_uops_result(dtypes.float32, uops, sres), 42)

  # NOTE: webgpu specific, since only webgpu performs bitpacking
@@ -341,7 +348,7 @@ class TestLocalAccess(unittest.TestCase):
    smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
    st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
    barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
-    sres = uop(uops, Ops.LOAD, dtypes.uint8, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), barr))
+    sres = smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0))
    self.assertEqual(_test_uops_result(dtypes.uint8, uops, sres), 42)

  # NOTE: webgpu specific, since only webgpu performs bitpacking
@@ -351,7 +358,7 @@ class TestLocalAccess(unittest.TestCase):
    size = 16
    for dtype in _dtypes:
      temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, addrspace=AddrSpace.LOCAL), (), 'smem')
-      uops = to_uops_list([temp], opts=Device[Device.DEFAULT].renderer)
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
      out = Device[Device.DEFAULT].renderer.render(uops)
      # half is supported in wgsl, so it doesn't have to be packed
      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
@@ -375,10 +382,10 @@ class TestAssembly(unittest.TestCase):
    g1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
    c1 = UOp(Ops.CONST, dtypes.int, (), 2)
    c2 = UOp(Ops.CONST, dtypes.int, (), 3)
-    l1 = UOp(Ops.LOAD, dtypes.int, (g1.index(c1),))
+    l1 = g1.index(c1)
    a1 = UOp(Ops.MUL, dtypes.int, (l1, c1))
    a2 = UOp(Ops.MUL, dtypes.int, (l1, c2))
-    uops = to_uops_list([a1,a2], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a1,a2], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHL, ops)
@@ -388,9 +395,9 @@ class TestAssembly(unittest.TestCase):
    for dt in (dtypes.int32, dtypes.uint32):
      g = UOp(Ops.DEFINE_GLOBAL, dt.ptr(), (), 0)
      c = UOp(Ops.CONST, dt, (), 2)
-      l = UOp(Ops.LOAD, dt, (g.index(c),))
+      l = g.index(c)
      a = UOp(Ops.IDIV, dt, (l, c))
-      uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      Device[Device.DEFAULT].renderer.render(uops)
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
@@ -399,16 +406,16 @@ class TestAssembly(unittest.TestCase):
  def test_fast_idiv_and_mod(self):
    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
    c = UOp(Ops.CONST, dtypes.uint, (), 3)
-    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
+    l = g.index(c)
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
-    uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
    self.assertNotIn(Ops.IDIV, ops)

    b = UOp(Ops.MOD, dtypes.uint, (l, c))
-    uops = to_uops_list([b], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([b], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
@@ -421,7 +428,7 @@ class TestAssembly(unittest.TestCase):
    c = UOp(Ops.CONST, dtypes.uint, (), 7)
    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
    a = UOp(Ops.IDIV, dtypes.uint, (l, c))
-    uops = to_uops_list([a], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.SHR, ops)
@@ -429,7 +436,7 @@ class TestAssembly(unittest.TestCase):

  def test_fast_idiv_remove_powers_of_two(self):
    ridx = UOp.range(2**20, 0)
-    uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
+    uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
    ops = [x.op for x in uops]
    # this requires shifting out the powers of two before doing fast_idiv
    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
@@ -451,9 +458,8 @@ class TestAssembly(unittest.TestCase):
  def test_use_cmpeq(self):
    g = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(), (), 0)
    c = UOp(Ops.CONST, dtypes.uint, (), 7)
-    l = UOp(Ops.LOAD, dtypes.uint, (g.index(c),))
-    comp = l.ne(c).ne(True)
-    uops = to_uops_list([comp], opts=Device[Device.DEFAULT].renderer)
+    comp = g.index(c).ne(c).ne(True)
+    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
    self.assertIn(Ops.CMPEQ, ops)
@@ -511,8 +517,8 @@ class TestUOpStr(unittest.TestCase):

 class TestUPatHelpers(unittest.TestCase):
  def test_location(self):
-    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
-    self.assertEqual(spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "math.py")
+    self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
    test_upat = UPat(Ops.CONST, dtypes.bool)
    self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])
    test_upat_named = test_upat.named("test_name")
@@ -540,11 +546,92 @@ class TestUopsObject(unittest.TestCase):

 class TestUOpRender(unittest.TestCase):
  def test_render_vectorize_same(self):
-    u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
-    self.assertEqual(u.render(), "{0, ...}")
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(simplify=False), "{0, ...}")
  def test_render_vectorize_different(self):
-    u = UOp(Ops.VECTORIZE, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
-    self.assertEqual(u.render(), "{0,1,2}")
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(simplify=False), "{0,1,2}")
+  def test_render_vectorize_same_simplified(self):
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
+    self.assertEqual(u.render(), "0")
+  def test_render_vectorize_different_simplified(self):
+    u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
+    self.assertEqual(u.render(), "(0, 1, 2)")
+
+class TestZeroRange(unittest.TestCase):
+  def test_reduce_variable(self):
+    for i in range(3,-1,-1):
+      v = UOp.variable("i", 0, 5).bind(i)
+      out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
+      self.assertEqual(out.item(), i)
+
+class TestUOpPrograms(unittest.TestCase):
+  def _run(self, prog:UOp, *tensors:Tensor):
+    ExecItem(get_runner(Device.DEFAULT, prog), [t.uop.buffer for t in tensors]).run(wait=True)
+
+  def test_simple(self):
+    out = Tensor.empty(10,10,dtype=dtypes.int)
+
+    ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
+    i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
+    prog = ptr[i,j].set(42).end(i,j)
+    self._run(prog.sink(), out)
+
+    with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
+
+  def test_matmul(self):
+    a = Tensor.randn(10,10)
+    b = Tensor.randn(10,10)
+    c = Tensor.empty(10,10)
+    ref = (a@b)
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    # C[i,j] = sum_k A[i,k] * B[k,j]
+    # Shapes: A[M,K], B[K,N], C[M,N]
+    M = N = K = 10
+    DT = dtypes.float32
+
+    # Placeholders (bind slots explicitly)
+    A = UOp.placeholder((M, K), DT, slot=0)
+    B = UOp.placeholder((K, N), DT, slot=1)
+    C = UOp.placeholder((M, N), DT, slot=2)
+
+    # Axes: i,j are spatial; k is a reduction axis over the shared dim K
+    i = UOp.range(M, axis_id=0)                             # rows of A/C
+    j = UOp.range(N, axis_id=1)                             # cols of B/C
+    k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE)  # reduction over K
+
+    # Zero-init: write a scalar 0 to each (i,j).
+    C = C[i, j].set(0.0)
+
+    # Accumulate: C_after(k) enforces the dependency along the reduction axis
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j])
+
+    # Finalize the loop nest / schedule in (i, j, k) order
+    prog = C.end(i, j, k)
+
+    # run program
+    # TODO: make this work with opts_to_apply
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
+
+  def test_matmul_relu(self):
+    a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
+    ref = (a@b).relu()
+    with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
+
+    A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
+    i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
+
+    C = C[i, j].set(0.0)
+    C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
+    C = C[i, j].set(C[i, j].maximum(0.0))
+
+    prog = C.end(i, j)
+
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)

 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -141,8 +141,8 @@ class TestUOpsStats(unittest.TestCase):
    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
-    u1 = UOp(Ops.LOAD, dtypes.int, (globl.index(o1),))
-    u2 = UOp(Ops.LOAD, dtypes.int, (globl.index(o2),))
+    u1 = globl.index(o1)
+    u2 = globl.index(o2)
    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
    u4 = UOp(Ops.MUL, dtypes.int, (u1,u2))
    u5 = UOp(Ops.ADD, dtypes.int, (u4,u3))
@@ -151,8 +151,8 @@ class TestUOpsStats(unittest.TestCase):
    globl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), tuple())
    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
-    u1 = UOp(Ops.LOAD, dtypes.int, (globl.index(o1),))
-    u2 = UOp(Ops.LOAD, dtypes.int, (globl.index(o2),))
+    u1 = globl.index(o1)
+    u2 = globl.index(o2)
    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
    u4 = UOp(Ops.MULACC, dtypes.int, (u1,u2,u3))
    uops_fma = full_rewrite(u4.sink())
@@ -1,76 +0,0 @@
-import unittest, random
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import print_uops, UOp, Ops
-from tinygrad.codegen.late.linearize import block_reorder
-from tinygrad.renderer.cstyle import OpenCLRenderer
-
-def is_toposorted(lst:list[UOp]):
-  seen = set()
-  for u in lst:
-    if any(p not in seen for p in u.src): return False
-    seen.add(u)
-  return True
-
-class TestBlockReorder(unittest.TestCase):
-  def _test_randomize(self, golden:list[UOp]):
-    # test random order is always same
-    for _ in range(50):
-      # shuffle and form a valid toposort
-      lst = golden[:]
-      random.shuffle(lst)
-      topolst = []
-      for u in lst:
-        for p in u.toposort():
-          if p not in topolst: topolst.append(p)
-      assert is_toposorted(topolst)
-
-      for x,y in zip(golden, this_order:=block_reorder(topolst)):
-        if x is not y:
-          print_uops(golden)
-          print_uops(this_order)
-        self.assertIs(x, y)
-
-  def _test_render(self, golden:list[UOp]):
-    return OpenCLRenderer().render(golden)
-
-  def test_loads(self):
-    a = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=0)
-    b = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=1)
-    c = UOp(Ops.DEFINE_GLOBAL, dtype=dtypes.float.ptr(), arg=2)
-    v1 = UOp(Ops.SPECIAL, dtype=dtypes.int, src=(UOp.const(dtypes.int, 4),), arg="gidx0")
-    v2 = UOp(Ops.SPECIAL, dtype=dtypes.int, src=(UOp.const(dtypes.int, 4),), arg="gidx1")
-    v1 = v1*27
-    v2 = v2*4
-    loads = [
-      a.index(v1).load(dtype=dtypes.float),
-      a.index(v1+1).load(dtype=dtypes.float),
-      a.index(v1+2).load(dtype=dtypes.float),
-      a.index(v1+3).load(dtype=dtypes.float),
-      b.index(v2).load(dtype=dtypes.float),
-      b.index(v2+1).load(dtype=dtypes.float),
-      b.index(v2+2).load(dtype=dtypes.float),
-      b.index(v2+3).load(dtype=dtypes.float)]
-    #random.shuffle(loads)
-    sink = c.store(sum(loads)).sink()
-
-    # determine golden order
-    golden = block_reorder(list(sink.toposort()))
-
-    # render for test
-    print(self._test_render(golden))
-    #print_uops(golden)
-
-    # assert the loads are in this order
-    self.assertListEqual([g.src[0].src[1].render() for g in golden if g.op is Ops.LOAD],
-                         ['(gidx1*4)', '((gidx1*4)+1)', '((gidx1*4)+2)', '((gidx1*4)+3)',
-                          '(gidx0*27)', '((gidx0*27)+1)', '((gidx0*27)+2)', '((gidx0*27)+3)'])
-
-    # assert math is after loads
-    first_math = [i for i,g in enumerate(golden) if g.op is Ops.ADD and g.dtype == dtypes.float][0]
-    assert not any(x.op is Ops.LOAD for x in golden[first_math:])
-
-    # confirm the sort is stable
-    self._test_randomize(golden)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -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
@@ -433,17 +433,15 @@ class TestDiskTensorMovement(unittest.TestCase):
    t = Tensor(self.fn)
    self.assertListEqual(t[16:18].tolist(), [16,17])

-  # TODO: fix this! at least assert on it
-  @unittest.expectedFailure
  def test_slice_read_cat(self):
    t = Tensor(self.fn)
-    self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
+    with self.assertRaises(AssertionError):
+      self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])

-  # TODO: fix this! at least assert on it
-  @unittest.expectedFailure
  def test_slice_sum(self):
    t = Tensor(self.fn)
-    self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
+    with self.assertRaises(AssertionError):
+      self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])

 if __name__ == "__main__":
  unittest.main()
@@ -1,11 +1,12 @@
 import unittest, math
 from tinygrad import dtypes
-from tinygrad.helpers import all_same
+from tinygrad.helpers import all_same, Context
 from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
 from tinygrad.codegen import full_rewrite_to_sink
 from hypothesis import given, strategies as strat

 # Helper function to apply the graph rewrite
+@Context(SPEC=0)
 def apply_rewrite(expr):
  return full_rewrite_to_sink(expr.sink()).src[0]

@@ -305,19 +306,19 @@ class TestRecurse(unittest.TestCase):
    graph_rewrite(a, pm, bottom_up=True)

  def test_inf_loop(self):
-    a = UOp.variable('a', 0, 10)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm)

  def test_inf_loop_bottom_up(self):
-    a = UOp.variable('a', 0, 10)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
@@ -3,6 +3,8 @@ import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import CI
+from tinygrad.uop.ops import UOp
+from tinygrad.engine.jit import TinyJit

@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -72,5 +74,52 @@ class TestKeccak(unittest.TestCase):
    data = b"\x00" * 1000
    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))

+  def test_variable_bs(self):
+    data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
+
+    bs = UOp.variable("bs", 1, 4096).bind(1)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(1, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(2)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(2, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(3)
+    data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
+    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
+
+  def test_variable_bs_jit(self):
+    def f(data):
+      return data.keccak()
+    jit_f = TinyJit(f)
+
+    data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
+
+    # initialize jit
+    for _ in range(3):
+      bs = UOp.variable("bs", 1, 4096).bind(4096)
+      _ = jit_f(data.shrink_to(bs, data.shape[-1]))
+
+    bs = UOp.variable("bs", 1, 4096).bind(1)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(1, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(2)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(2, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+
+    bs = UOp.variable("bs", 1, 4096).bind(3)
+    data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
+    out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(3, 32)
+    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
+    self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
+
 if __name__ == "__main__":
  unittest.main()
@@ -99,6 +99,11 @@ class TestStripParens(unittest.TestCase):
  def test_simple(self): self.assertEqual("1+2", strip_parens("(1+2)"))
  def test_nested(self): self.assertEqual("1+(2+3)", strip_parens("(1+(2+3))"))
  def test_casted_no_strip(self): self.assertEqual("(int)(1+2)", strip_parens("(int)(1+2)"))
+  def test_unmatched_parens(self): self.assertEqual("((c35+c39>>23&255)+-127).cast(dtypes.float)",
+    strip_parens("((c35+c39>>23&255)+-127).cast(dtypes.float)"))
+  def test_single_paren_left(self): self.assertEqual("(abc", strip_parens("(abc"))
+  def test_single_paren_right(self): self.assertEqual("abc)", strip_parens("abc)"))
+  def test_parens_at_different_depths(self): self.assertEqual("(a+(b))*(c)", strip_parens("(a+(b))*(c)"))

 class TestProd(unittest.TestCase):
  def test_empty(self): self.assertEqual(1, prod(tuple()))
--- a/Show More
+++ b/Show More