add a gate to rewrite if there's no rules [pr]

2025-11-30 17:28:58 -08:00
145 changed files with 10106 additions and 13768 deletions
@@ -298,7 +298,7 @@ runs:
    - name: Install mesa (linux)
      if: inputs.mesa == 'true' && runner.os == 'Linux'
      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
    - name: Install mesa (macOS)
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
@@ -116,9 +116,11 @@ jobs:
    - name: Verify Qualcomm autogen
      run: |
        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
+        mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
-        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
+        python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
+        diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.py
        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
    - name: Verify libusb autogen
      run: |
@@ -54,9 +54,9 @@ jobs:
    - name: Print macOS version
      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
@@ -64,7 +64,7 @@ jobs:
    - name: Run SDXL
      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
-      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
@@ -132,10 +132,6 @@ jobs:
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -326,27 +322,29 @@ jobs:
      run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+    # TODO: too slow
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 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=110 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=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    # - name: Run 10 CIFAR training steps w HALF
+    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    # - name: Run 10 CIFAR training steps w BF16
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    # - name: Run full CIFAR training steps w 6 GPUS
+    #   run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -437,8 +435,9 @@ jobs:
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -528,21 +527,22 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+    # TODO: too slow
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 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=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    # - name: Run 10 CIFAR training steps w HALF
+    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    # - name: Run 10 CIFAR training steps w BF16
    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS
+    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -592,13 +592,13 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Run MLPerf resnet eval
      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -627,14 +627,18 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
+    # - name: openpilot compile3 0.9.9 driving_vision
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    # - name: openpilot compile3 0.9.9 driving_policy
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    # - name: openpilot compile3 0.9.9 dmonitoring
+    #   run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.10.0 driving_policy
      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.0 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
@@ -704,8 +708,10 @@ jobs:
      run: |
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    # TODO: too slow
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    # TODO: enable
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
@@ -766,10 +772,11 @@ jobs:
        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+    # TODO: too slow
+    # - name: Run full CIFAR training w 1 GPU
+    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '13'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -71,7 +71,9 @@ jobs:
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
-      run: python docs/abstractions3.py
+      run: |
+        python docs/abstractions2.py
+        python docs/abstractions3.py
    - name: Test README
      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Test Quickstart
@@ -261,9 +263,7 @@ jobs:
    - name: Check Device.DEFAULT
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
-      run: |
-        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
+      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
    # TODO: too slow
@@ -289,8 +289,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 20000 lines
-      run: MAX_LINE_COUNT=20000 python sz.py
+    - name: Repo line count < 19000 lines
+      run: MAX_LINE_COUNT=19000 python sz.py

  spec:
    strategy:
@@ -447,7 +447,7 @@ jobs:
        with:
          key: onnxoptl
          deps: testing
-          pydeps: "tensorflow==2.19"
+          pydeps: "tensorflow==2.15.1 tensorflow_addons"
          python-version: '3.11'
          opencl: 'true'
      - name: Test ONNX (CL)
@@ -465,7 +465,7 @@ jobs:
      - name: Test Bert training
        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -972,33 +972,3 @@ jobs:
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
-
-# ****** Compile-only Tests ******
-
-  compiletests:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-${{ matrix.backend }}
-          deps: testing_minimal
-          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
-      - name: Set env
-        shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/test_ops.py --durations=20
@@ -1,213 +0,0 @@
-# Claude Code Guide for tinygrad
-
-## Architecture Overview
-
-tinygrad compiles tensor operations into optimized kernels. The pipeline:
-
-1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
-2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
-3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
-4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
-5. **Runtime** (`runtime/`) - Device-specific execution
-
-## Key Concepts
-
-### UOp (Universal Operation)
-Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
- `op`: The operation type (Ops enum)
- `dtype`: Data type
- `src`: Tuple of source UOps
- `arg`: Operation-specific argument
- `tag`: Optional tag for graph transformations
-
-UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
-
-### PatternMatcher
-Used extensively for graph transformations:
-```python
-pm = PatternMatcher([
-  (UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
-])
-result = graph_rewrite(uop, pm)
-```
-
-### Schedule Cache
-Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
-
-## Directory Structure
-
-```
-tinygrad/
-├── tensor.py          # Tensor class, user API
-├── device.py          # Buffer, device management
-├── dtype.py           # Data types
-├── helpers.py         # Utilities, environment vars
-├── uop/
-│   ├── ops.py         # UOp class, Ops enum, PatternMatcher
-│   ├── spec.py        # UOp type verification
-│   └── symbolic.py    # Symbolic math simplification
-├── engine/
-│   ├── schedule.py    # Schedule creation, caching
-│   ├── realize.py     # Tensor realization
-│   ├── jit.py         # JIT compilation
-│   └── memory.py      # Memory planning
-├── schedule/
-│   ├── rangeify.py    # Convert movements to ranges
-│   └── indexing.py    # Index calculations
-├── codegen/
-│   ├── kernel.py      # Kernel optimization
-│   └── uopgraph.py    # UOp graph transformations
-├── renderer/          # Code generation (CUDA, Metal, etc.)
-└── runtime/           # Device backends
-```
-
-## Testing
-
-```bash
-# Run specific test
-python -m pytest test/unit/test_schedule_cache.py -xvs
-
-# Run with timeout
-python -m pytest test/test_symbolic_ops.py -x --timeout=60
-
-# Debug with print
-DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
-
-# Visualize UOp graphs
-VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
-```
-
-## Common Environment Variables
-
- `DEBUG=1-4` - Increasing verbosity
- `VIZ=1` - Enable graph visualization
- `SPEC=1` - Enable UOp spec verification
- `NOOPT=1` - Disable optimizations
- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
-
-## Debugging Tips
-
-1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
-2. **Check schedule**: `tensor.schedule()` returns list of ScheduleItems
-3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
-4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
-
-## Workflow Rules
-
- **NEVER commit without explicit user approval** - always show the diff and wait for approval
- Run tests before proposing commits
- Test with `SPEC=2` when modifying UOp-related code
-
-## Style Notes
-
- 2-space indentation, 150 char line limit
- PatternMatchers should be defined at module level (slow to construct)
- Prefer `graph_rewrite` over manual graph traversal
- UOp methods like `.replace()` preserve tags unless explicitly changed
- Use `.rtag(value)` to add tags to UOps
-
-## Lessons Learned
-
-### UOp ucache Behavior
-UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
- Two UOps with same structure are identical (`is` comparison works)
-
-### Spec Validation
-When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
-```bash
-SPEC=2 python3 test/unit/test_something.py
-```
-Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
-
-### Schedule Cache Key Normalization
-The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
- `pm_post_sched_cache`: restores original BIND from context
- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
-
-### Avoiding Extra Work
- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
- PatternMatchers are slow to construct - define at module level, not in functions
-
-### Testing LLM Changes
-```bash
-# Quick smoke test
-echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
-
-# Check cache hits (should see "cache hit" after warmup)
-echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
-
-# Test with beam search
-echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
-```
-
-## Common Patterns
-
-### Graph Transformation
-```python
-def my_transform(ctx, x):
-  # Return new UOp or None to skip
-  return x.replace(arg=new_arg)
-
-pm = PatternMatcher([
-  (UPat(Ops.SOMETHING, name="x"), my_transform),
-])
-result = graph_rewrite(input_uop, pm, ctx={})
-```
-
-### Finding Variables
-```python
-# Get all variables in a UOp graph
-variables = uop.variables()
-
-# Get bound variable values
-var, val = bind_uop.unbind()
-```
-
-### Shape Handling
-```python
-# Shapes can be symbolic (contain UOps)
-shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
-```
-
-## Performance Optimization
-
-When optimizing tinygrad internals:
-
-1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
-
-2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
-
-3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
-
-4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
-
-5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
-
-6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
-
-7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
-
-8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
-
-9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
-
-## Pattern Matching Profiling
-
-Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
-
-```bash
-TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
-```
-
-Output format: `matches / attempts -- match_time / total_time ms -- location`
-
-Key patterns to watch (from ResNet50 benchmark):
- `split_load_store`: ~146ms, 31% match rate - does real work
- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
-
-Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
@@ -0,0 +1,135 @@
+# tinygrad is a tensor library, and as a tensor library it has multiple parts
+# 1. a "runtime". this allows buffer management, compilation, and running programs
+# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
+# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
+# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
+
+
+print("******** first, the runtime ***********")
+
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
+
+cpu = CPUDevice()
+
+# allocate some buffers
+out = cpu.allocator.alloc(4)
+a = cpu.allocator.alloc(4)
+b = cpu.allocator.alloc(4)
+
+# load in some values (little endian)
+cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+
+# compile a program to a binary
+lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+
+# create a runtime for the program
+fxn = cpu.runtime("add", lib)
+
+# run the program
+fxn(out, a, b)
+
+# check the data out
+print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
+assert val == 5
+
+
+print("******** second, the Device ***********")
+
+DEVICE = "CPU"   # NOTE: you can change this!
+
+import struct
+from tinygrad.dtype import dtypes
+from tinygrad.device import Buffer, Device
+from tinygrad.uop.ops import UOp, Ops
+
+# allocate some buffers + load in values
+out = Buffer(DEVICE, 1, dtypes.int32).allocate()
+a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+# NOTE: a._buf is the same as the return from cpu.allocator.alloc
+
+# describe the computation
+idx = UOp.const(dtypes.index, 0)
+buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
+buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
+alu = buf_1.index(idx) + buf_2.index(idx)
+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,))
+
+# convert the computation to a "linearized" format (print the format)
+from tinygrad.engine.realize import get_program, CompiledRunner
+program = get_program(s, Device[DEVICE].renderer)
+
+# compile a program (and print the source)
+fxn = CompiledRunner(program)
+print(fxn.p.src)
+# NOTE: fxn.clprg is the CPUProgram
+
+# run the program
+fxn.exec([out, a, b])
+
+# check the data out
+assert out.as_buffer().cast('I')[0] == 5
+
+
+print("******** third, the UOp ***********")
+
+from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.schedule.rangeify import get_rangeify_map
+
+# allocate some values + load in values
+a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+
+# describe the computation
+out = a + b
+s = UOp(Ops.SINK, dtypes.void, (out,))
+
+# group the computation into kernels
+becomes_map = get_rangeify_map(s)
+
+# the compute maps to an assign
+assign = becomes_map[a+b].base
+
+# the first source is the output buffer (data)
+assert assign.src[0].op is Ops.BUFFER
+# the second source is the kernel (compute)
+assert assign.src[1].op is Ops.KERNEL
+
+# schedule the kernel graph in a linear list
+s = UOp(Ops.SINK, dtypes.void, (assign,))
+sched, _ = create_schedule_with_vars(s)
+assert len(sched) == 1
+
+# DEBUGGING: print the compute ast
+print(sched[-1].ast)
+# NOTE: sched[-1].ast is the same as st_0 above
+
+# the output will be stored in a new buffer
+out = assign.buf_uop
+assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
+print(out)
+
+# run that schedule
+run_schedule(sched)
+
+# check the data out
+assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
+
+
+print("******** fourth, the Tensor ***********")
+
+from tinygrad import Tensor
+
+a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
+b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
+out = a + b
+
+# check the data out
+print(val:=out.item())
+assert val == 5
@@ -21,7 +21,7 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
-  opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
+  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))

  @TinyJit
  @Tensor.train()
@@ -1,108 +0,0 @@
-import itertools
-from typing import Callable
-from tinygrad import nn, Tensor, dtypes, Device, TinyJit
-from tinygrad.helpers import getenv, trange, partition
-
-class Model:
-  def __init__(self):
-    self.layers: list[Callable[[Tensor], Tensor]] = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu,
-      nn.BatchNorm(32), Tensor.max_pool2d,
-      nn.Conv2d(32, 64, 3), Tensor.relu,
-      nn.Conv2d(64, 64, 3), Tensor.relu,
-      nn.BatchNorm(64), Tensor.max_pool2d,
-      lambda x: x.flatten(1), nn.Linear(576, 10)]
-
-  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
-
-# TODO: refactor this into optim/onnx
-def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
-  b1_t *= b1
-  b2_t *= b2
-  m.assign(b1 * m + (1.0 - b1) * g)
-  v.assign(b2 * v + (1.0 - b2) * (g * g))
-  m_hat = m / (1.0 - b1_t)
-  v_hat = v / (1.0 - b2_t)
-  return lr * (m_hat / (v_hat.sqrt() + eps))
-
-if __name__ == "__main__":
-  BS = getenv("BS", 512)
-  ACC_STEPS = getenv("ACC_STEPS", 8)
-
-  X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
-  model = Model()
-
-  params = nn.state.get_parameters(model)
-
-  # init params, set requires grad on the ones we need gradients of
-  for x in params:
-    if x.requires_grad is None: x.requires_grad_()
-    x.replace(x.contiguous())
-  Tensor.realize(*params)
-
-  # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.requires_grad)
-  print(f"params: {len(params)} buffers: {len(buffers)}")
-
-  # optim params
-  pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
-  adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
-
-  # create loss and grads. init all state so the JIT works on microbatch
-  for x in params: x.assign(x.detach())
-  loss = Tensor.zeros(tuple()).contiguous()
-  grads = Tensor.zeros(pos_params[-1]).contiguous()
-  Tensor.realize(*params, *buffers, *adam_params, loss, grads)
-
-  @TinyJit
-  @Tensor.train()
-  def microbatch():
-    samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
-    for t in params: t.grad = None
-    # divide by ACC_STEPS at the loss
-    uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
-    ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
-    for t in params: t.grad = None
-    # concat the grads and assign them
-    loss.assign(loss + uloss)
-    grads.assign(grads + ugrads)
-    Tensor.realize(*params, *buffers, loss, grads)
-
-  @TinyJit
-  def optimizer():
-    # run optimizer (on CPU, where adam params live)
-    delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
-
-    # update the params, copying back the delta one at a time to avoid OOM
-    # NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
-    for j,tt in enumerate(params):
-      tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
-
-    # realize everything, zero out loss and grads
-    loss.assign(Tensor.zeros_like(loss))
-    grads.assign(Tensor.zeros_like(grads))
-    Tensor.realize(*params, *adam_params, loss, grads)
-
-  @TinyJit
-  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
-
-  test_acc = float('nan')
-  for i in (t:=trange(getenv("STEPS", 70))):
-    # microbatch sets the gradients
-    for _ in range(ACC_STEPS): microbatch()
-
-    # get the loss before the optimizer clears it
-    # this is already realized so this isn't a schedule
-    loss_item = loss.item()
-
-    # run the optimizer
-    optimizer()
-
-    # eval
-    if i%10 == 9: test_acc = get_test_acc().item()
-    t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")
@@ -223,13 +223,13 @@ def get_mlperf_bert_model():

 def get_fake_data_bert(BS:int):
  return {
-    "input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
-    "next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
+    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
+    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
  }

 def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
@@ -59,7 +59,9 @@ class EmbeddingBert(nn.Embedding):
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
+    # TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
+    # Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
+    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()

 class LayerNormBert:
  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
@@ -918,6 +918,40 @@ def train_rnnt():
  # TODO: RNN-T
  pass

+@TinyJit
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
+  optimizer.zero_grad()
+
+  for i in range(grad_acc):
+    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
+    # NOTE: these two have different names
+    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
+    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
+
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+    # TODO: OOM without this realize with large grad_acc
+    Tensor.realize(*[p.grad for p in optimizer.params])
+
+  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
+  for p in optimizer.params:
+    p.grad = p.grad / loss_scaler
+    global_norm += p.grad.float().square().sum()
+  global_norm = global_norm.sqrt().contiguous()
+  for p in optimizer.params:
+    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
+
+  optimizer.step()
+  scheduler.step()
+  # TODO: no to("CPU") here because it blocks and messes the python time
+  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
+  return loss, global_norm, optimizer.optimizers[0].lr
+
@TinyJit
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
@@ -980,8 +1014,7 @@ def train_bert():
  # ** hyperparameters **
  BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  # TODO: implement grad accumulation + mlperf logging
-  assert grad_acc == 1
+  # TODO: mlperf logging
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
@@ -1040,8 +1073,8 @@ def train_bert():

  # ** Optimizer **
  parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
-  parameters_wd = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
+  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
  optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)

@@ -1098,38 +1131,12 @@ def train_bert():
  # ** train loop **
  wc_start = time.perf_counter()

-  i, train_data = start_step, next(train_it)
+  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]

  if RUNMLPERF:
    if MLLOGGER:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})

-  @TinyJit
-  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
-                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
-    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-      else: t.to_(GPUS[0])
-    optimizer_group.zero_grad()
-
-    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-    (loss * loss_scaler).backward()
-
-    global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer_group[0].device)
-    for p in optimizer_group.params:
-      p.grad = p.grad / loss_scaler
-      global_norm += p.grad.float().square().sum()
-    global_norm = global_norm.sqrt().contiguous()
-    for p in optimizer_group.params:
-      p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
-
-    optimizer_group.step()
-    scheduler_group.step()
-    # TODO: no to("CPU") here because it blocks and messes the python time
-    Tensor.realize(loss, global_norm, optimizer_group.optimizers[0].lr)
-    return loss, global_norm, optimizer_group.optimizers[0].lr
-
  while train_data is not None and i < train_steps and not achieved:
    if getenv("TRAIN", 1):
      Tensor.training = True
@@ -1137,12 +1144,16 @@ def train_bert():
      st = time.perf_counter()
      GlobalCounters.reset()
      with WallTimeEvent(BenchEvent.STEP):
-        loss, global_norm, lr = train_step_bert(
-          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
+        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
+        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)

        pt = time.perf_counter()
-        next_data = next(train_it)
+
+        try:
+          next_data = [next(train_it) for _ in range(grad_acc)]
+        except StopIteration:
+          next_data = None
+
        dt = time.perf_counter()

        device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
@@ -1177,8 +1188,8 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
-        # TODO: this hangs on tiny green after 90 minutes of training
+      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
+        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
@@ -1213,7 +1224,7 @@ def train_bert():
          return

      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE") and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1289,7 +1300,6 @@ def train_llama3():
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  assert grad_acc == 1, f"{grad_acc=} is not supported"
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1364,17 +1374,20 @@ def train_llama3():

  @TinyJit
  @Tensor.train()
-  def train_step(model, tokens:Tensor):
+  def train_step(model, tokens:Tensor, grad_acc:int):
    optim.zero_grad()
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    loss.backward()
+    # grad acc
+    for batch in tokens.split(tokens.shape[0]//grad_acc):
+      if (DP := getenv("DP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+        batch = batch.shard(device, 0)
+      if (MP := getenv("MP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+        batch = batch.shard(device)
+      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
+      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
+      loss.backward()
+      Tensor.realize(*[p.grad for p in optim.params])
    # L2 norm grad clip
    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1413,14 +1426,14 @@ def train_llama3():

  def get_train_iter():
    if getenv("FAKEDATA", 0):
-      return fake_data(BS, SAMPLES)
+      return fake_data(GBS, SAMPLES)
    else:
      if SMALL:
        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
      else:
        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))

  def get_eval_iter():
    if getenv("FAKEDATA", 0):
@@ -1438,7 +1451,7 @@ def train_llama3():
  for tokens in tqdm(iter, total=SAMPLES//GBS):
    t = time.perf_counter()
    GlobalCounters.reset()
-    loss, lr = train_step(model, tokens)
+    loss, lr = train_step(model, tokens, grad_acc)
    loss = loss.float().item()

    i += 1
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -115,7 +115,7 @@ if __name__ == "__main__":

  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
    if not args.fakeweights:
-      default_weights_url = 'https://huggingface.co/sd2-community/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+      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
@@ -1,14 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
-from tinygrad.runtime.support.am.amdev import AMDev
-
-if __name__ == "__main__":
-  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1])])
-  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
-  amdevs = []
-  with Context(DEBUG=2):
-    for pcidev in pcidevs:
-      amdevs.append(AMDev(pcidev, reset_mode=True))
-    for amdev in amdevs: amdev.smu.mode1_reset()
@@ -1,59 +1,48 @@
 import re, ctypes, sys, importlib
-from tinygrad.helpers import getenv

 from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
-    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
+  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
+    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
    self._run_discovery()
    self._build_regs()

 amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
+
 from tinygrad.runtime.ops_amd import PCIIface

-def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
-  register_map = register_names or {}
+def parse_amdgpu_logs(log_content, register_names=None):
+  register_map = register_names

+  final = ""
  def replace_register(match):
-    reg = match.group(1)
-    return f"Reading register {register_map.get(int(reg, 16), reg)}"
+    register = match.group(1)
+    return f"Reading register {register_map.get(int(register, base=16), register)}"

-  processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
+  pattern = r'Reading register (0x[0-9a-fA-F]+)'
+
+  processed_log = re.sub(pattern, replace_register, log_content)

  def replace_register_2(match):
-    reg = match.group(1)
-    return f"Writing register {register_map.get(int(reg, 16), reg)}"
-
-  processed_log = re.sub(r'Writing register (0x[0-9a-fA-F]+)', replace_register_2, processed_log)
-
-  # remove timing prefix
-  processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
-
-  # keep only xcc=0 lines (but keep lines with no xcc at all)
-  if only_xcc0:
-    kept = []
-    for line in processed_log.splitlines(True):
-      if "xcc=" not in line or re.search(r'\bxcc=0\b', line): kept.append(line)
-    processed_log = "".join(kept)
+    register = match.group(1)
+    return f"Writing register {register_map.get(int(register, base=16), register)}"

+  pattern = r'Writing register (0x[0-9a-fA-F]+)'
+  processed_log = re.sub(pattern, replace_register_2, processed_log)
  return processed_log

 def main():
-  only_xcc0 = bool(getenv("ONLY_XCC0", 0))
-
  reg_names = {}
  dev = PCIIface(None, 0)
  for x, y in dev.dev_impl.__dict__.items():
    if isinstance(y, AMRegister):
-      for xcc, addr in y.addr.items():
-        reg_names[addr] = f"{x}, xcc={xcc}"
+      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
-    log_content = f.read()
+    log_content = log_content_them = f.read()

-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names)

  with open(sys.argv[2], 'w') as f:
    f.write(processed_log)
@@ -62,4 +51,5 @@ if __name__ == '__main__':
  if len(sys.argv) != 3:
    print("Usage: <input_file_path> <output_file_path>")
    sys.exit(1)
-  main()
+
+  main()
@@ -1,136 +0,0 @@
-import os, sys, struct
-sys.path.append(os.getcwd())
-# PROFILE=1 to use
-#os.environ["PROFILE"] = "1"
-os.environ["SQTT"] = "1"
-os.environ["SQTT_ITRACE_SE_MASK"] = "1"
-os.environ["SQTT_LIMIT_SE"] = "1"
-import xml.etree.ElementTree as ET
-
-from tinygrad import nn, Tensor, Device
-from tinygrad.helpers import get_single_element
-from tinygrad.engine.realize import lower_schedule
-from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-def disassemble(text, root:ET.Element):
-  i = 0
-  while i < len(text):
-    ins = struct.unpack("I", text[i:i+4])[0]
-
-    # 1. Get the encoding
-    did_match = False
-    for enc_el in root.findall("./ISA/Encodings/Encoding"):
-      mask = enc_el.findtext("EncodingIdentifierMask")
-      assert len(mask)%32 == 0
-      bit_mask = int(mask, 2)
-      iden = [int(x.text, 2) for x in enc_el.find("EncodingIdentifiers").findall("EncodingIdentifier")]
-      for ide in iden:
-        if ins&bit_mask == ide:
-          did_match = True
-          break
-      if did_match: break
-    if not did_match: raise RuntimeError(f"unknown instruction {ins:08X}")
-    if len(mask) >= 64: ins = (struct.unpack("I", text[i+4:i+8])[0]<<32) | ins
-    if len(mask) >= 96: ins = (struct.unpack("I", text[i+8:i+12])[0]<<64) | ins
-    encoding_name = enc_el.findtext("EncodingName")
-
-    #print(ET.tostring(enc_el).decode())
-
-    # 2. Parse the Fields for this Encoding
-    field_data = {}
-    for field in enc_el.findall("MicrocodeFormat/BitMap/Field"):
-      # Fields can be split into multiple ranges (RangeCount > 1)
-      ranges = sorted(field.findall("BitLayout/Range"), key=lambda x: int(x.attrib.get('Order')))
-      val = 0
-      current_shift = 0
-      for rng in ranges:
-        width = int(rng.find("BitCount").text)
-        chunk = (ins >> int(rng.find("BitOffset").text)) & ((1 << width) - 1)
-        val |= (chunk << current_shift)
-        current_shift += width
-      field_data[field.find("FieldName").text] = val
-    # this is already used
-    del field_data["ENCODING"]
-
-    # 3. Extract the instruction
-    did_match = False
-    for ins_el in root.findall("./ISA/Instructions/Instruction"):
-      ins_name = ins_el.findtext("InstructionName")
-      for ins_enc in ins_el.findall("InstructionEncodings/InstructionEncoding"):
-        if ins_enc.findtext("EncodingName") == encoding_name:
-          opcode = int(ins_enc.findtext("Opcode"))
-          if "OP" in field_data and opcode == field_data["OP"]:
-            did_match = True
-            del field_data["OP"]
-            break
-        if did_match: break
-      if did_match: break
-
-    #print(ET.tostring(ins_enc).decode())
-    #print()
-    #print(field_data)
-    if not did_match:
-      print(f"{i:4X} : {ins:16x} -- {encoding_name}")
-    elif did_match:
-      params = []
-      #print(ET.tostring(ins_el).decode())
-
-      # 4. Extract the opcodes
-      for op_ins in ins_enc.findall("Operands/Operand"):
-        op_type = op_ins.findtext("OperandType")
-        op_size = op_ins.findtext("OperandSize")
-        op_fmt = op_ins.findtext("DataFormatName")
-        op_field_name = op_ins.findtext("FieldName")
-        if op_field_name is None: continue
-        assert op_field_name in field_data
-        # loop through operands for compare
-        for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-          test_op_type = op_el.findtext("OperandTypeName")
-          val_dict = {}
-          for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-            val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-          if op_type == test_op_type:
-            if field_data[op_field_name] in val_dict:
-              print(op_type, op_size, op_fmt)
-              params.append(val_dict[field_data[op_field_name]])
-            else:
-              params.append(f"{op_type}({field_data[op_field_name]})")
-            del field_data[op_field_name]
-            #print(op_type, op_size, op_fmt, op_el, op_field_name,
-            #      field_data[op_field_name],
-            #      val_dict.get(field_data[op_field_name], "<UNK>"))
-            #print(ET.tostring(op_el).decode())
-
-      print(f"{i:4X} : {ins:16x} -- {ins_name.lower()} {', '.join(params)}", field_data)
-
-    # advance
-    i += len(mask) // 8
-
-  #print(ET.tostring(root).decode())
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  with open("/tmp/rdna35.xml", "wb") as f: f.write(bytes(xml_str))
-  root = ET.fromstring(xml_str)
-
-  a = Tensor.empty(16)+1
-  for si, ei in lower_schedule(a.schedule()):
-    # get text
-    _, hdr, _ = elf_loader(ei.prg.lib)
-    text = get_single_element([x for x in hdr if x.name==".text"]).content
-
-    # llvm disassembler
-    Device["AMD"].compiler.disassemble(ei.prg.lib)
-
-    # run program
-    ei.run()
-
-  sqtt_events = [e for e in Device["AMD"].profile_events if isinstance(e, ProfileSQTTEvent)]
-  for e in sqtt_events[0:1]: # only the first SE
-    parse_sqtt_print_packets(e.blob)
-
-  disassemble(text[:0x40], root)
@@ -1,15 +0,0 @@
-from tinygrad import Tensor, nn
-import xml.etree.ElementTree as ET
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  root = ET.fromstring(xml_str)
-
-  for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-    op_name = op_el.findtext("OperandTypeName")
-    val_dict = {}
-    for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-      val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-    print(op_name, val_dict)
@@ -1,7 +1,7 @@
 import argparse, os, hashlib
 from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
 from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
-from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
+from tinygrad import Tensor, dtypes, Device, Variable

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
@@ -9,6 +9,8 @@ if __name__ == "__main__":
  parser.add_argument("--output_dir", type=str, default="extra/hevc/out")
  args = parser.parse_args()

+  os.makedirs(args.output_dir, exist_ok=True)
+
  if args.input_file == "":
    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
    hevc_tensor = Tensor.from_url(url, device="CPU")
@@ -25,38 +27,26 @@ if __name__ == "__main__":
  frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]

  # move all needed data to gpu
-  #all_slices = []
-  with Timing("copy to gpu: "):
+  all_slices = []
+  with Timing("prep slices to gpu: "):
    opaque_nv = opaque.to("NV").contiguous().realize()
-    hevc_tensor = hevc_tensor.to("NV")
+
+    for i, (offset, sz, frame_pos, history_sz, _) in enumerate(frame_info):
+      all_slices.append(hevc_tensor[offset:offset+sz].to("NV").contiguous().realize())
+
+    Device.default.synchronize()

  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
  max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)
+  pos = Variable("pos", 0, max_hist + 1)

-  # define variables
-  v_pos = Variable("pos", 0, max_hist + 1)
-  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
-  v_sz = Variable("sz", 0, hevc_tensor.numel())
-  v_i = Variable("i", 0, len(frame_info)-1)
-
-  @TinyJit
-  def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
-    return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
-
-  # warm up
-  history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
-  for i in range(3):
-    hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
-    decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
-
+  history = []
  out_images = []
  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
    for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
-      history = history[-max_hist:] if max_hist > 0 else []
-      # TODO: this shrink should work as a slice
-      hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
+      history = history[-history_sz:] if history_sz > 0 else []

-      outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
+      outimg = all_slices[i].decode_hevc_frame(pos.bind(frame_pos), out_image_size, opaque_nv[i], history).realize()
      out_images.append(outimg)
      if is_hist: history.append(outimg)

@@ -77,7 +67,5 @@ if __name__ == "__main__":
        assert img.data() == decoded_frames[i], f"Frame {i} does not match reference decoder!"
        print(f"Frame {i} matches reference decoder!")
    else:
-      if len(args.output_dir):
-        os.makedirs(args.output_dir, exist_ok=True)
-        img = to_bgr(img, h, w, luma_w, chroma_off).realize()
-        cv2.imwrite(f"{args.output_dir}/out_frame_{i:04d}.png", img.numpy())
+      img = to_bgr(img, h, w, luma_w, chroma_off).realize()
+      cv2.imwrite(f"{args.output_dir}/out_frame_{i:04d}.png", img.numpy())
@@ -322,7 +322,6 @@ def parse_hevc_file_headers(dat:bytes, device="NV"):
    if hdr.slice_type == avcodec.HEVC_SLICE_B: device_ctx.initreflistidxl1 = (ctypes.c_uint8 * 16)(*[idx for _,idx in after_list + before_list])

    locl_ctx_bytes = bytes(device_ctx)
-    locl_ctx_bytes += b'\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x00\x00' # blackwell extension
    locl_ctx_bytes += bytes(0x200 - len(locl_ctx_bytes)) # pad to 512 bytes

    pic_width_in_ctbs = ceildiv(sps.pic_width_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size))
@@ -66,7 +66,7 @@ def ioctl(fd, request, argp):
    print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
    if name == "AMDKFD_IOC_SVM":
      out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
-      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
+      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.kfd_ioctl_svm_attr_type__enumvalues[out[i].type]:40s}: {out[i].value:#x}")
  else:
    print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
          f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")
@@ -29,9 +29,8 @@ rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw

 # create QCOM tensor with the externally managed buffer
 x = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
-y = (x + 1).reshape(-1).tolist()
-print(y[:10])
-assert y == [i + 1 for i in range(64)]
+y = (x + 1).numpy()
+print(y)

 # all calculations are done, save to free the object
 cl.clReleaseMemObject(cl_buf)
@@ -50,7 +49,7 @@ for i in range(4):
  cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
  rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20]

-  y = calc(x = Tensor.from_blob(rawbuf_ptr, (2, 2), dtype=dtypes.int, device='QCOM')).tolist()
+  y = calc(x = Tensor.from_blob(rawbuf_ptr, (2, 2), dtype=dtypes.int, device='QCOM')).numpy()
  print(f'jit {i}\n', y)

  # all calculations are done, save to free the object
@@ -81,19 +80,8 @@ rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw
 # dtypes.imageh = cl.cl_image_format(cl.CL_RGBA, cl.CL_HALF_FLOAT)
 # dtypes.imagef = cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT)
 x = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
-y = (x + 1).tolist()
-print(y[:10])
+y = (x + 1).numpy()
+print(y)

 # all calculations are done, save to free the object
 cl.clReleaseMemObject(cl_img)
-
-# from numpy
-import numpy as np
-
-YUV_SIZE = 50
-a_np = (32*np.random.randn(YUV_SIZE).astype(np.float32) + 128).clip(0,255).astype(np.uint8)
-a = Tensor.from_blob(a_np.ctypes.data, (YUV_SIZE,), dtype=dtypes.uint8, device='QCOM').realize()
-
-print(a.numpy()[:10], a_np[:10])
-assert np.all(a.numpy() == a_np)
-assert np.all((a - 1).numpy() == a_np - 1)
@@ -1,5 +1,5 @@
 use half::f16;
-use num_traits::{float::FloatCore, PrimInt, Unsigned, clamp};
+use num_traits::{float::FloatCore, PrimInt, Unsigned};

 pub fn bits<T>(word: T, hi: usize, lo: usize) -> T where T: PrimInt + Unsigned {
    assert!(hi >= lo);
@@ -48,7 +48,6 @@ impl IEEEClass<u64> for f64 {
 pub trait VOPModifier<T> {
    fn negate(&self, pos: usize, modifier: usize) -> T;
    fn absolute(&self, pos: usize, modifier: usize) -> T;
-    fn clmp(&self, cm: bool) -> T;
 }
 impl<T> VOPModifier<T> for T
 where
@@ -66,11 +65,6 @@ where
            _ => *self,
        }
    }
-    fn clmp(&self, cm:bool) -> T {
-        if !cm { return *self }
-        let r = clamp(*self, T::zero(), T::one());
-        if r == T::zero() { T::zero() } else { r }
-    }
 }

 pub fn extract_mantissa(x: f64) -> f64 {
@@ -1024,7 +1024,7 @@ impl<'a> Thread<'a> {
                    let vdst = (instr & 0xff) as usize;
                    let abs = ((instr >> 8) & 0x7) as usize;
                    let opsel = ((instr >> 11) & 0xf) as usize;
-                    let cm = ((instr >> 15) & 0x1) != 0;
+                    let cm = (instr >> 15) & 0x1;

                    let s = |n: usize| ((instr >> n) & 0x1ff) as usize;
                    let src = (s(32), s(41), s(50));
@@ -1032,9 +1032,7 @@ impl<'a> Thread<'a> {
                    let omod = (instr >> 59) & 0x3;
                    let neg = ((instr >> 61) & 0x7) as usize;
                    assert_eq!(omod, 0);
-                    if op != 272 && cm {
-                        return todo_instr!(op); // TODO: add VOP3 clamp for all ops
-                    }
+                    assert_eq!(cm, 0);
                    assert_eq!(opsel, 0);

                    match op {
@@ -1268,7 +1266,7 @@ impl<'a> Thread<'a> {
                            }

                            let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 543 | 551 | 567 | 606 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
                                    let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                    let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                    let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1277,26 +1275,12 @@ impl<'a> Thread<'a> {
                                        260 => s0 - s1,
                                        261 => s1 - s0,
                                        264 => s0 * s1,
-                                        272 => f32::max(s0, s1).clmp(cm),
+                                        272 => f32::max(s0, s1),
                                        299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                        426 => s0.recip(),
                                        430 => 1.0 / f32::sqrt(s0),
                                        531 => f32::mul_add(s0, s1, s2),
                                        537 => f32::min(f32::min(s0, s1), s2),
-                                        543 => {
-                                            if s0.is_nan() || s1.is_nan() || s2.is_nan() {
-                                                f32::min(f32::min(s0, s1), s2)
-                                            } else {
-                                                let max = f32::max(f32::max(s0, s1), s2);
-                                                if max == s0 {
-                                                    f32::max(s1, s2)
-                                                } else if max == s1 {
-                                                    f32::max(s0, s2)
-                                                } else {
-                                                    f32::max(s0, s1)
-                                                }
-                                            }
-                                        },
                                        540 => f32::max(f32::max(s0, s1), s2),
                                        551 => s2 / s1,
                                        567 => {
@@ -1306,7 +1290,6 @@ impl<'a> Thread<'a> {
                                                false => ret,
                                            }
                                        }
-                                        606 => f32::min(f32::max(s0, s1), s2),
                                        796 => s0 * 2f32.powi(s1.to_bits() as i32),
                                        // cnd_mask isn't a float only ALU but supports neg
                                        257 => {
@@ -119,6 +119,7 @@ OPNAME = {
  0xb: "VALU",
  0xd: "VALU",
  0xe: "VALU",
+  0x10: "__END",
  0x21: "VMEM_LOAD",
  0x22: "VMEM_LOAD",
  0x24: "VMEM_STORE",
@@ -479,7 +480,6 @@ def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -
        nib = (byte >> (offset & 4)) & 0xF
        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
        offset += 4
-      if offset != target: break  # don't parse past the end

    # 2) Decode token from low 8 bits
    opcode = STATE_TO_OPCODE[reg & 0xFF]
@@ -67,7 +67,10 @@ class OccEvent(WaveSlot):
  time:int
  start:int

-RunKey = tuple[str, int]
+@dataclasses.dataclass(frozen=True)
+class RunKey:
+  prg:str
+  tag:int

 class _ROCParseCtx:
  def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
@@ -83,7 +86,7 @@ class _ROCParseCtx:

  def next_sqtt(self):
    x = next(self.sqtt_evs, None)
-    self.active_run = (x.kern, x.exec_tag) if x is not None else None
+    self.active_run = RunKey(x.kern, x.exec_tag) if x is not None else None
    self.active_se = x.se if x is not None else None
    self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
    return self.active_blob
@@ -138,7 +141,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:

  @rocprof.rocprof_trace_decoder_isa_callback_t
  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run)[0]][pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run).prg][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
@@ -158,14 +161,6 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
  t.join()
  return ROCParseCtx

-def print_pmc(events:list[ProfilePMCEvent]) -> None:
-  from tinygrad.viz.serve import unpack_pmc
-  from tabulate import tabulate
-  for e in events:
-    print("**", e.kern)
-    data = unpack_pmc(e)
-    print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
-
 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)))
@@ -175,4 +170,13 @@ if __name__ == "__main__":
  rctx = decode(profile)
  print('SQTT:', rctx.inst_execs.keys())

-  print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])
+  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
@@ -1,87 +0,0 @@
-import os
-os.environ["PROFILE"] = "1"
-os.environ["PMC"] = "1"
-
-import unittest
-import functools, contextlib
-import numpy as np
-from tinygrad import Tensor, Context, Device
-from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.runtime.ops_amd import ProfilePMCEvent
-from tinygrad.engine.realize import get_runner
-from tinygrad.viz.serve import unpack_pmc
-from extra.sqtt.roc import print_pmc
-
-def copy_kernel(B, A, stride=1):
-  n_threads = 32
-  assert A.size >= n_threads, f"{A.size} is too small, min size {n_threads}"
-  g = UOp.range(A.size//n_threads, 0, AxisType.GLOBAL)
-  l = UOp.range(n_threads, 1, AxisType.LOCAL)
-  i = g * n_threads + l
-  index = (i * stride) % A.size
-  return B[index].store(A[index]).sink(arg=KernelInfo(name=f"copy_{A.size}_stride_{stride}", opts_to_apply=()))
-
-def lds_kernel(offset:UOp, size:int, inst:str) -> UOp:
-  tid = UOp.range(offset.size, 0, AxisType.LOCAL)
-  dst = UOp.placeholder((size,), dtypes.float32, 1, AddrSpace.REG)
-  #lds = UOp.placeholder((1024,), dtypes.float32, 2, AddrSpace.LOCAL)
-  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_waitcnt(0);')
-  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_barrier();', src=(u,))
-  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_sched_barrier(0);', src=(u,))
-  u = UOp(Ops.CUSTOM, arg=f'asm volatile("{inst} '+'%0, %1" : "=v"({0}) : "v"({1}));', src=(dst, offset[tid], u))
-  return UOp.sink(u, arg=KernelInfo(name="test_lds", opts_to_apply=()))
-
-dev = Device[Device.DEFAULT]
-
-@contextlib.contextmanager
-def save_pmc():
-  # clear the old traces
-  dev.profile_events.clear()
-  pmc:list[ProfilePMCEvent] = []
-  yield pmc
-  for e in dev.profile_events:
-    if isinstance(e, ProfilePMCEvent): pmc.append(e)
-
-@unittest.skipIf(dev.device != "AMD", "tests PMC counters on AMD")
-class TestPMC(unittest.TestCase):
-  @Context(IGNORE_OOB=0)
-  def test_copy(self, stride:int=1):
-    N = 1 << 25 # ~134MB
-    a = Tensor(np.arange(N, dtype=np.uint32)+1).realize()
-    b = Tensor(np.zeros(N, dtype=np.uint32)).realize()
-    b = Tensor.custom_kernel(b, a, fxn=functools.partial(copy_kernel, stride=stride))[0]
-    with save_pmc() as pmc:
-      b.realize()
-    print_pmc(pmc)
-    np.testing.assert_equal(a.numpy(), b.numpy())
-
-  def test_copy_uncoalesced(self): return self.test_copy(stride=17)
-
-  # test with two threads issuing ds_reads at different offsets
-  def test_ds_read(self, size=1, inst='ds_read_b32'):
-    test_banks = 256
-    offsets = [Tensor([0, b*4]) for b in range(1, test_banks)]
-    with Context(DEBUG=0): Tensor.realize(*offsets)
-    k = Tensor.custom_kernel(offsets[0], fxn=functools.partial(lds_kernel, size=size, inst=inst))[0]
-    # sample all kernels
-    with save_pmc() as pmc_events:
-      runner = get_runner(Device.DEFAULT, k.schedule()[0].ast)
-      # TODO: llvm eliminates lds definition from the ELF, is there another way to pin lds size?
-      runner._prg.group_segment_size = 1024
-      for offset in offsets: runner([offset.uop.buffer])
-    # find read offsets that created bank conflicts from the pmc counters
-    found:list[Tensor] = []
-    for i,e in enumerate(pmc_events):
-      pmc = unpack_pmc(e)["rows"]
-      # SQ on gfx9, renamed to SQC after gfx10
-      val = next(total for name,total,_all_instances in pmc if name in {"SQ_LDS_BANK_CONFLICT", "SQC_LDS_BANK_CONFLICT"})
-      if val > 0: found.append(offsets[i])
-    print("Found bank conflicts at offsets:", [s.numpy() for s in found])
-
-  def test_ds_read_b64(self): self.test_ds_read(2, 'ds_read_b64')
-
-  def test_ds_read_b128(self): self.test_ds_read(4, 'ds_read_b128')
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,166 +0,0 @@
-import math
-
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-
-from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.kernel import Kernel
-from extra.thunder.tiny.tk.tiles import GL, TileLayout
-
-NUM_WORKERS = 1
-Q_BLOCK_SIZE = 16
-KV_BLOCK_SIZE = 16
-
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
-  if len(xq.shape) == 3: xq, xk, xv = xq.unsqueeze(0), xk.unsqueeze(0), xv.unsqueeze(0)
-
-  odtype = xq.dtype
-  xq, xk, xv = xq.transpose(1, 2).cast(dtypes.bfloat16), xk.transpose(1, 2).cast(dtypes.bfloat16), xv.transpose(1, 2).cast(dtypes.bfloat16)
-
-  _, N_, _, D_ = xq.shape
-  block_size = max(Q_BLOCK_SIZE, KV_BLOCK_SIZE)
-  assert D_ % block_size == 0, f"embedding dimension must be multiple of block size, got {D_=} {block_size=}"
-
-  # pad to multiple of block size
-  xq = xq.pad(((0, 0), (0, (block_size - (xq.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
-  xk = xk.pad(((0, 0), (0, (block_size - (xk.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
-  xv = xv.pad(((0, 0), (0, (block_size - (xv.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
-
-  B, N, H, D = xq.shape
-  H_KV = xk.shape[2]
-  GROUP_SIZE = H // H_KV
-  print(f"Flash Attention {B=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
-
-  def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, mu:UOp) -> UOp:
-    with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
-      warp = ker.warp
-
-      o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(mu, ker), GL(l_vecu, ker)
-
-      head = ker.blockIdx_x
-      head_kv = head // GROUP_SIZE
-      batch = ker.blockIdx_z
-      q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
-
-      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-
-      q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
-      q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
-      q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
-      k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
-      k_reg_transposed = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
-      v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
-      o_reg = ker.rt((D, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
-      o_reg_transposed = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
-      att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
-      att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
-      mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
-      mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
-
-      max_vec_last = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      max_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      norm_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-      scale_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
-
-      max_vec = warp.neg_inf(max_vec)
-      norm_vec = warp.zero(norm_vec)
-      o_reg = warp.zero(o_reg)
-      scale_vec = warp.ones(scale_vec)
-
-      # load q tile
-      q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
-      q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
-      q_reg = warp.copy(q_reg, q_reg_fl)
-      q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
-
-      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
-        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
-        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
-
-        k_reg = warp.load(k_reg, k_smem)
-        v_reg = warp.load(v_reg, v_smem)
-
-        # mma qk^t
-        att_block = warp.zero(att_block.after(kv_idx))
-        k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
-        att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
-
-        # apply attention mask
-        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
-        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
-        att_block += mask_reg_transposed
-
-        # softmax
-        max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
-        max_vec = warp.row_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)
-
-        scale_vec = warp.map(scale_vec.after(max_vec_last, max_vec), lambda _, idx: max_vec_last[*idx] - max_vec[*idx])
-        scale_vec = scale_vec.exp2()
-
-        o_reg *= scale_vec
-        norm_vec *= scale_vec
-
-        att_block -= max_vec
-        att_block = att_block.exp2()
-
-        norm_vec = warp.row_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)
-
-        # mma av
-        att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
-        o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
-      o_reg = ker.endrange()
-      norm_vec = norm_vec.after(o_reg)
-      max_vec = max_vec.after(o_reg)
-
-      o_reg /= norm_vec
-
-      o_reg_transposed = warp.transpose(o_reg_transposed, o_reg)
-      o = warp.store(o, o_reg_transposed, (batch, q_seq, head, 0), (), axis=1)
-
-      norm_vec = norm_vec.after(o)
-      max_vec = max_vec.after(o)
-
-      max_vec *= math.log(2)
-      norm_vec = norm_vec.log2() * math.log(2)
-      norm_vec += max_vec
-      l_vec = warp.store(l_vec, norm_vec, (batch, head, 0, q_seq), (), axis=2)
-      o = o.after(l_vec)
-
-      return ker.finish()
-
-  def custom_backward_q(out_qu:UOp, gradu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
-    return UOp.sink(arg=KernelInfo(name="fa_custom_backward_q"))
-
-  def custom_backward_kv(out_ku:UOp, out_vu:UOp, gradu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
-    return UOp.sink(arg=KernelInfo(name="fa_custom_backward_kv"))
-
-  if is_causal:
-    if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-    attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.bool).tril()
-  if attn_mask is not None:
-    if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
-  else:
-    attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.float32)
-
-  attn = Tensor.empty_like(xq)
-  l_vec = Tensor.empty(B, H, 1, N, requires_grad=False, device=xq.device, dtype=dtypes.float32).detach()
-
-  def grad(grad:UOp, kernel:UOp) -> tuple[None, None, UOp, UOp, UOp, None]:
-    grad_q = Tensor.empty_like(q := Tensor(kernel.src[2]))
-    grad_k = Tensor.empty_like(k := Tensor(kernel.src[3]))
-    grad_v = Tensor.empty_like(v := Tensor(kernel.src[4]))
-    mask = Tensor(kernel.src[5])
-
-    delta_vec = (Tensor(grad) * attn).sum(-1).unsqueeze(-2).detach()
-
-    print(l_vec.numpy())
-
-    grad_q = Tensor.custom_kernel(grad_q, Tensor(grad), q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
-    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, Tensor(grad), q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
-    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
-
-  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
-  attn = attn[:, :N_, :, :D_]
-
-  return attn.transpose(1, 2).cast(odtype)
@@ -7,7 +7,7 @@ from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.helpers import getenv, prod

 from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout, VecLayout
+from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout

 class Group:
  def __init__(self, warps:int, ker):
@@ -66,10 +66,7 @@ class Group:
    for height in self.ker.range(src.shape[-3], track=False):
      for width in self.ker.range(src.shape[-2], track=False):
        for inner in self.ker.range(src.shape[-1], track=False):
-          src_load = src[height, width, inner]
-          if src.dtype.base != dst.dtype.base:
-            src_load = src_load.cast(dst.dtype.base)
-          dst_store = dst[width, height, inner].store(src_load).end(height, width, inner)
+          dst_store = dst[width, height, inner].store(src[height, width, inner]).end(height, width, inner)

    self.ker.push_store(dst_store, dst)
    return dst.after(dst_store).reshape(dst.shape)
@@ -80,9 +77,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
    else:  raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -110,9 +107,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
    else:  raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -140,9 +137,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
    else:  raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -170,9 +167,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
    else:  raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -203,9 +200,9 @@ class Group:
    Group.map_rid += len(a.shape)

    if op.__code__.co_argcount == 1:
-      to_store = op(a[*rngs_for_shape]) # type: ignore
+      to_store = op(a[*rngs_for_shape])
    else:
-      to_store = op(a[*rngs_for_shape], rngs_for_shape) # type: ignore
+      to_store = op(a[*rngs_for_shape], rngs_for_shape)

    a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)

@@ -287,7 +284,7 @@ class Group:
  def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0):
    dst, src = cast(UOp, dst), cast(UOp, src)
    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = dst.dtype, src.dtype
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
    if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
      laneid = self.ker.laneid
      rt, st = cast(RT, dst), cast(ST, src)
@@ -341,7 +338,7 @@ class Group:
                src_load = src_load.cast(dst.dtype.base)
              dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
              dst_store = dst_store.end(height, width, outer, inner).barrier()
-    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
+    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace ==AddrSpace.GLOBAL:
      srcf = src.flatten()
      row_stride = prod(src.shape[axis+1:])

@@ -374,28 +371,8 @@ class Group:
            if src.dtype.base != dst.dtype.base:
              src_load = src_load.cast(dst.dtype.base)
            dst_store = dst[*dst_idxs, height, width, inner].store(src_load).end(height, width, inner)
-    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV):
-      srcf = src.flatten()
-      row_stride = prod(src.shape[axis+1:])
-
-      laneid = self.ker.laneid
-      rv = cast(RV, dst)
-      reductions = rv.base_shape.rows
-
-      assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
-
-      idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
-      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
-
-      for outer in self.ker.range(dst.shape[-2]):
-        src_i += outer * reductions + (laneid % reductions)
-
-        src_load = srcf[src_i]
-        if src.dtype.base != dst.dtype.base:
-          src_load = src_load.cast(dst.dtype.base)
-        dst_store = dst[outer, 0].store(src_load).end(outer)
    else:
-      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}")
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")

    self.ker.push_store(dst_store, dst)
    return dst.after(dst_store).reshape(dst.shape)
@@ -403,30 +380,8 @@ class Group:
  def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0):
    dst, src = cast(UOp, dst), cast(UOp, src)
    assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = dst.dtype, src.dtype
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
-      laneid = self.ker.laneid
-      st, rt = cast(ST, dst), cast(RT, src)
-      elements_per_thread = rt.base_shape.elements_per_thread
-
-      for height in self.ker.range(src.shape[-3], track=False):
-        for width in self.ker.range(src.shape[-2], track=False):
-          for inner in self.ker.range(elements_per_thread, track=False):
-            if rt.layout != st.layout:
-              row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
-              col = laneid % rt.base_shape.cols
-            else:
-              row = laneid % rt.base_shape.rows
-              col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
-
-            srow, scol = cast(ST, dst).swizzle(row, col)
-
-            src_load = src[*src_idxs, height, width, inner]
-            if src.dtype.base != dst.dtype.base:
-              src_load = src_load.cast(dst.dtype.base)
-            dst_store = dst[*idxs[:-2], height, width, srow, scol].store(src_load)
-            dst_store = dst_store.end(height, width, inner)
-    elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
+    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL:
      dstf = dst.flatten()
      row_stride = prod(dst.shape[axis+1:])

@@ -459,28 +414,8 @@ class Group:
            if src.dtype.base != dst.dtype.base:
              src_load = src_load.cast(dst.dtype.base)
            dst_store = dstf[dst_i].store(src_load).end(height, width, inner)
-    elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RV):
-      dstf = dst.flatten()
-      row_stride = prod(dst.shape[axis+1:])
-
-      laneid = self.ker.laneid
-      rv = cast(RV, src)
-      reductions = rv.base_shape.rows
-
-      assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
-
-      idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
-      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
-
-      for outer in self.ker.range(src.shape[-2]):
-        dst_i += outer * reductions + (laneid % reductions)
-
-        src_load = src[outer, 0]
-        if src.dtype.base != dst.dtype.base:
-          src_load = src_load.cast(dst.dtype.base)
-        dst_store = dstf[dst_i].store(src_load).end(outer)
    else:
-      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}")
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")

    self.ker.push_store(dst_store, dst)
    return dst.after(dst_store).reshape(dst.shape)
@@ -5,34 +5,33 @@ from extra.thunder.tiny.tk.group import Group
 from extra.thunder.tiny.tk.tiles import GL, ST_16X16, ST_16X16_SWIZZLED, ST, RT_16X16, RT, RV, TileLayout, VecLayout

 class _tk_range:
-  def __init__(self, start:int, end:int, step:int, axis_type:AxisType, rid:int):
+  user_rid = 0
+  def __init__(self, start:int, end:int, step:int, axis_type:AxisType):
    self.start, self.end, self.step = start, end, step
-    self.axis_type, self.rid, self.done = axis_type, rid, False
+    self.axis_type, self.done = axis_type, False
  def __iter__(self): return self
  def __next__(self):
    if not self.done:
      self.done = True
-      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
+      _tk_range.user_rid += 1
+      self._rng = UOp.range(self.end // self.step, _tk_range.user_rid-1, axis_type=self.axis_type) * self.step + self.start
      return self._rng
    raise StopIteration

 class Kernel(AbstractContextManager):
-  def __init__(self, name:str, grid_size:tuple[int, int, int], block_size:int):
-    self.name = name
-
+  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
    self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
    self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
    self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
    self.threadIdx_x = UOp.special(block_size, "lidx0")

-    self.range_stack: list[_tk_range] = []
-    self.store_stack: list[tuple[UOp, UOp]] = []
+    self.range_stack = []
+    self.store_stack = []

    self.global_slot = 0
    self.shared_slot = 0
    self.register_slot = 0
-    self.range_id = 0
-    self.allocs: dict[tuple[str, tuple], UOp] = {}
+    self.allocs = {}

  @property
  def warpid(self): return self.threadIdx_x // WARP_THREADS
@@ -50,8 +49,7 @@ class Kernel(AbstractContextManager):

  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
    if end == 0: start, end = 0, start
-    rng = _tk_range(start, end, step, axis_type, self.range_id)
-    self.range_id += 1
+    rng = _tk_range(start, end, step, axis_type)
    if track: self.range_stack.append(rng)
    return rng

@@ -91,7 +89,7 @@ class Kernel(AbstractContextManager):
    if hasattr(last_store, '_uop'): uop = last_store._uop
    else: uop = last_store

-    return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
+    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()

  def endrange(self):
    last_store = self.store_stack.pop()
@@ -36,7 +36,7 @@ def autowrap(source_cls, blacklist=None):
        def proxy(*args, **kwargs):
          return wrap(val(*unwrap(args), **unwrap(kwargs)), self)
        return proxy
-      if name in UOp.__slots__: return val # type: ignore
+      if name in UOp.__slots__: return val
      return wrap(val, self)
    cls.__getattr__ = __getattr__

@@ -250,12 +250,11 @@ class RT(TileMathMixin):

@autowrap(UOp)
 class RV(TileMathMixin):
-  def __init__(self, uop:UOp, length:int, layout:VecLayout, base_shape:RTBaseShape, ker):
-    self._uop, self.ker = uop, ker
-    self.length, self.layout, self.base_shape = length, layout, base_shape
+  def __init__(self, uop:UOp, layout:VecLayout, ker):
+    self._uop, self.layout, self.ker = uop, layout, ker

  def ruop(self, uop:UOp):
-    return RV(uop, self.length, self.layout, self.base_shape, self.ker)
+    return RV(uop, self.layout, self.ker)

  @classmethod
  def create(cls, length, dtype:DType, layout:VecLayout, base_shape:RTBaseShape, ker):
@@ -267,6 +266,6 @@ class RV(TileMathMixin):
        outer_dim = tiles

    uop = ker.alloc((outer_dim, inner_dim), dtype, AddrSpace.REG)
-    return RV(uop, length, layout, base_shape, ker)
+    return RV(uop, layout, ker)

 ALL_TILES = UOp | GL | ST | RT | RV
@@ -3,13 +3,6 @@ import sys, os, zlib, struct, hashlib
 from tinygrad.helpers import DEBUG, getenv, fetch
 from tinygrad.runtime.support.usb import USB3

-SUPPORTED_CONTROLLERS = [
-  (0x174C, 0x2464),
-  (0x174C, 0x2463),
-  (0xADD1, 0x0001),
-]
-if getenv("USBDEV", ""): SUPPORTED_CONTROLLERS.insert(0, (int(x, 16) for x in getenv("USBDEV", "").split(":")))
-
 def patch(input_filepath, file_hash, patches):
  with open(input_filepath, 'rb') as infile: data = bytearray(infile.read())

@@ -47,14 +40,10 @@ if not os.path.exists(file_path):
 patches = [(0x2a0d + 1 + 4, b'\x0a', b'\x05')]
 patched_fw = patch(file_path, file_hash, patches)

-dev = None
-for vendor, device in SUPPORTED_CONTROLLERS:
-  try:
-    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
-    break
-  except RuntimeError: pass
-if dev is None:
-  raise RuntimeError('Could not open controller. You can set USBDEV environment variable to your device\'s vendor and device ID (e.g., USBDEV="174C:2464")')
+vendor, device = [int(x, base=16) for x in getenv("USBDEV", "174C:2464").split(":")]
+try: dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
+except RuntimeError as e:
+  raise RuntimeError(f'{e}. You can set USBDEV environment variable to your device\'s vendor and device ID (e.g., USBDEV="174C:2464")') from e

 config1 = bytes([
  0xFF, 0xFF, 0xFF, 0xFF, 0x41, 0x41, 0x41, 0x41, 0x42, 0x42, 0x42, 0x42, 0x30, 0x30, 0x36, 0x30,
@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt

-NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot","Christopher Milan"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]
@@ -69,14 +69,17 @@ testing_minimal = [
  "hypothesis",
  "z3-solver",
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate"]
 testing = [
-  "tinygrad[testing_unit]",
+  "tinygrad[testing_minimal]",
  "pillow",
-  "onnx==1.19.0",
+  "onnx==1.18.0",
  "onnx2torch",
  "onnxruntime",
  "opencv-python",
+  "tabulate",
+  "tqdm",
+  "safetensors",
  "transformers",
  "sentencepiece",
  "tiktoken",
@@ -32,8 +32,6 @@ class FakeAM:
    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
  def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
  def paddr2mc(self, paddr:int) -> int: return paddr
-  def paddr2xgmi(self, paddr:int) -> int: return paddr
-  def xgmi2paddr(self, xgmi_paddr:int) -> int: return xgmi_paddr

 #  * PTE format:
 #  * 63:59 reserved
@@ -170,19 +170,6 @@ backend_test.exclude('test_scan_*')
 backend_test.exclude('test_split_to_sequence_*')
 backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121

-# TODO: not yet implemented
-backend_test.exclude('test_tensorscatter_*')
-backend_test.exclude('test_l1normalization_*')
-backend_test.exclude('test_l2normalization_*')
-backend_test.exclude('test_lpnormalization_*')
-backend_test.exclude('test_einsum_scalar_cpu')
-backend_test.exclude('test_mod_mixed_sign_float16_cpu')
-backend_test.exclude('test_qlinearmatmul_2D_uint8_float16_cpu')
-backend_test.exclude('test_qlinearmatmul_3D_uint8_float16_cpu')
-backend_test.exclude('test_attention_3d_*')
-backend_test.exclude('test_attention_4d_*')
-
-
 # rest of the failing tests
 backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # tf_crop_and_resize not implemented
 backend_test.exclude('test_resize_tf_crop_and_resize_axes_2_3_cpu') # tf_crop_and_resize not implemented
@@ -197,6 +184,9 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
 backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
 backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported

+backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
+backend_test.exclude('test_scatternd_max_cpu') # max not yet supported
+
 # regression from removing StrEnum in Domain
 backend_test.exclude('test_adam_cpu')
 backend_test.exclude('test_gradient_of_add_and_mul_cpu')
@@ -2,7 +2,7 @@
 import unittest, math
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.optimizers import Lamb
+import tensorflow_addons as tfa
 from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup

@@ -88,8 +88,6 @@ def create_tiny_lars(params, lr, skip_list=False):
  if skip_list: return OptimizerGroup(LARS([params[0]], lr), SGD([params[1]], lr, classic=True, weight_decay=0., momentum=.9))
  return LARS(params, lr)
 def create_tf_lars(lr, skip_list=False): return LARSOptimizer(lr, skip_list=["W"] if skip_list else None)
-def create_tf_lamb(lr=0.001, b1=0.9, b2=0.999, eps=1e-7, weight_decay=0.0):
-  return Lamb(learning_rate=float(lr), beta_1=b1, beta_2=b2, epsilon=eps, weight_decay=weight_decay)

 def create_tiny_polylr(optim, initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
  assert power == 2
@@ -114,7 +112,7 @@ class ExternalTestOptim(unittest.TestCase):
                   step_tf(tensorflow_optim, steps=steps, kwargs=opts, scheduler=tf_sched, schedopts=schedopts, do_optim=do_optim)):
      np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)

-  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, create_tf_lamb, steps, opts, atol, rtol)
+  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, tfa.optimizers.LAMB, steps, opts, atol, rtol)
  def _test_lars(self, steps, opts, atol, rtol): self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol)
  def _test_lars_polylr(self, steps, opts, schedopts, atol, rtol, do_optim=True):
    self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol,
@@ -1,57 +0,0 @@
-import os, sys, time, multiprocessing
-
-N = int(os.environ.get("NPROC", str(os.cpu_count())))
-DEVICE = os.environ.get("DEV", "AMD")
-
-# this tests the total number of processes that can be running tinygrad at a time
-def proc(i, device, stop_evt):
-  from tinygrad import Tensor
-
-  try:
-    a = Tensor.ones(2, device=device).contiguous()
-    b = Tensor.ones(2, device=device).contiguous()
-    c = (a + b).realize()
-    assert c.tolist() == [2, 2]
-  except Exception as e:
-    # fail if it fails
-    print(f"[child {i:2d}] tinygrad op failed: {e}", file=sys.stderr)
-    # non-zero exit code propagated back to parent
-    sys.exit(1)
-
-  # TODO: wait here for global exit if success. fail if it fails
-  # -> We wait on a global Event shared from the parent.
-  print(f"[child {i:2d}] success")
-  stop_evt.wait()
-  # Normal successful exit
-  sys.exit(0)
-
-if __name__ == "__main__":
-  print(f"testing {N} concurrent tinygrad processes")
-
-  # global exit event, shared by all children
-  stop_evt = multiprocessing.Event()
-  procs = []
-
-  # launch n proc of proc 1 per 200 ms
-  for i in range(N):
-    p = multiprocessing.Process(target=proc, args=(i, DEVICE, stop_evt), name=f"tinygrad-proc-{i}")
-    p.start()
-    procs.append(p)
-    time.sleep(0.1)  # 100 ms between launches
-
-  # signal global exit
-  time.sleep(0.5)
-  stop_evt.set()
-
-  # join all children
-  for p in procs: p.join()
-
-  # check for failures
-  failed = [p for p in procs if p.exitcode != 0]
-  if failed:
-    print(f"{len(failed)} / {len(procs)} processes failed "
-          f"with exit codes: {[p.exitcode for p in failed]}", file=sys.stderr)
-    sys.exit(1)
-
-  print(f"All {len(procs)} tinygrad processes ran successfully")
-  sys.exit(0)
@@ -1,6 +1,5 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.engine.schedule import schedule_cache
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general
@@ -69,7 +68,6 @@ if __name__ == "__main__":
    t()

    # these caches will keep uops alive
-    schedule_cache.clear()
    method_cache.clear()
    apply_movement_op.cache_clear()
    _apply_reshape.cache_clear()
@@ -29,6 +29,7 @@ from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import linalg_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.training import training_ops
 from tensorflow.python.ops import state_ops


@@ -146,7 +147,20 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
    return scaled_lr, grad

  def _apply_dense(self, grad, var, apply_state=None):
-    return self._resource_apply_dense(grad, var, apply_state)
+    var_device, var_dtype = var.device, var.dtype.base_dtype
+    coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                    or self._fallback_apply_state(var_device, var_dtype))
+
+    scaled_lr, grad = self.compute_lr(grad, var, coefficients)
+    mom = self.get_slot(var, "momentum")
+    return training_ops.apply_momentum(
+        var,
+        mom,
+        math_ops.cast(1.0, var.dtype.base_dtype),
+        grad * scaled_lr,
+        self.momentum,
+        use_locking=False,
+        use_nesterov=self.use_nesterov)

  def _resource_apply_dense(self, grad, var, apply_state=None):
    var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -180,13 +194,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                    or self._fallback_apply_state(var_device, var_dtype))

    mom = self.get_slot(var, "momentum")
-    return tf.raw_ops.SparseApplyMomentum(
-        var=var,
-        accum=mom,
-        lr=coefficients["learning_rate"],
-        grad=grad.values,
-        indices=grad.indices,
-        momentum=self.momentum,
+    return training_ops.sparse_apply_momentum(
+        var,
+        mom,
+        coefficients["learning_rate"],
+        grad.values,
+        grad.indices,
+        self.momentum,
        use_locking=False,
        use_nesterov=self.use_nesterov)

@@ -196,13 +210,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                    or self._fallback_apply_state(var_device, var_dtype))

    mom = self.get_slot(var, "momentum")
-    return tf.raw_ops.ResourceSparseApplyKerasMomentum(
-        var=var.handle,
-        accum=mom.handle,
-        lr=coefficients["learning_rate"],
-        grad=grad,
-        indices=indices,
-        momentum=self.momentum,
+    return training_ops.resource_sparse_apply_keras_momentum(
+        var.handle,
+        mom.handle,
+        coefficients["learning_rate"],
+        grad,
+        indices,
+        self.momentum,
        use_locking=False,
        use_nesterov=self.use_nesterov)

@@ -14,6 +14,7 @@ try:
  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, BEAM
+  from tinygrad.device import Device
 except ImportError as e:
  print(repr(e))
  exit(int(ASSERT_DIFF))
@@ -51,10 +52,12 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
    return "\n".join([f"{len(asts)} kernels", *asts])
  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
  # 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
  p2 = get_program(input_ast, renderer=renderer)
  def to_str(ret:ProgramSpec) -> str:
    # PYTHON renderer pickles UOps, first unpickle and decode here
@@ -1,4 +1,4 @@
-import time, struct, functools
+import time, struct
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
@@ -60,14 +60,5 @@ def not_support_multi_device():
  # CL and CUDA don't support multi device if in CI
  return CI and REAL_DEV in ("CL", "CUDA")

-def needs_second_gpu(fn):
-  @functools.wraps(fn)
-  def wrapper(self, *args, **kwargs):
-    # check if there's a second GPU, if not, skip multi tests
-    try: Tensor.zeros(10, device=f"{Device.DEFAULT}:1").contiguous().realize()
-    except Exception as e: self.skipTest(f"second device not available: {e}")
-    return fn(self, *args, **kwargs)
-  return wrapper
-
 # NOTE: This will open REMOTE if it's the default device
 REAL_DEV = (Device.DEFAULT if Device.DEFAULT != "REMOTE" else Device['REMOTE'].properties.real_device)
@@ -113,7 +113,7 @@ class TestEnd2End(unittest.TestCase):

  def test_bn_linear(self):
    BS, K = 2, 1
-    eps = 1e-12  # torch asserts if this is 0
+    eps = 0
    X = Tensor([1,0]).reshape(BS, K, 1, 1)
    Y = Tensor([-1,0]).reshape(BS, K, 1, 1)
    class LinTiny:
@@ -28,16 +28,13 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
      model(*early_gen)
      Device[Device.DEFAULT].synchronize()
      tms.append(time.perf_counter_ns() - st)
-    mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
+    mem_used = GlobalCounters.mem_used - global_mem_used

    # TODO: jit should expose this correctly with graph
    kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
    print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
-    assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
-    assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
-    if kernels_used:
-      assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
-      assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used"
+    assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB - {mem_used/1e9:.2} GB used"
+    assert not kernels_used or kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
    if all_jitted:
      assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501

@@ -64,7 +61,7 @@ class TestRealWorld(unittest.TestCase):
    derandomize_model(model)
    @TinyJit
    def test(t, t2): return model(t, Tensor([801]), t2).realize()
-    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
+    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32),Tensor.randn(1, 77, params["ctx_dim"])), test, 18.0, 515)

  def test_unet_resblock(self):
    model = [ResBlock(16, 24, 16) for _ in range(4)]
@@ -73,7 +70,7 @@ class TestRealWorld(unittest.TestCase):
    def test(t, t2):
      for l in model: t = l(t, t2)
      return t.realize()
-    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.0002, 37)
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.01, 37)

  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
  def test_llama(self):
@@ -85,7 +82,7 @@ class TestRealWorld(unittest.TestCase):
    @TinyJit
    def test(t): return model(t, 0).realize()
    # TODO: test first token vs rest properly
-    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
+    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.27, 168, all_jitted=True)

  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
  def test_gpt2(self):
@@ -115,7 +112,7 @@ class TestRealWorld(unittest.TestCase):
        loss.backward()
        optimizer.step()

-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
+      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):
@@ -125,7 +122,7 @@ class TestRealWorld(unittest.TestCase):
      model = SpeedyResNet(Tensor.ones((12,3,2,2)))
      @TinyJit
      def run(X): return model(X)
-      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
+      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, (1.0/48)*BS, 126)

  @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
  def test_train_cifar(self):
@@ -142,7 +139,7 @@ class TestRealWorld(unittest.TestCase):
        loss.backward()
        optimizer.step()

-      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
+      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, (1.0/48)*BS, 126)

  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
  def test_train_cifar_hyp(self):
@@ -179,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, 400)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 427)

 if __name__ == '__main__':
  unittest.main()
@@ -24,7 +24,7 @@ class TestFloat4(unittest.TestCase):
    s = c.schedule()[0]
    realized_ast = s.ast
    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)

    assert TestFloat4.count_float4(program.uops) == (2, 1)

@@ -35,8 +35,7 @@ class TestFloat4(unittest.TestCase):
    c = a + b

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
    assert TestFloat4.count_float4(uops) == (4, 2)

  @unittest.skipUnless(Device.DEFAULT in {"CPU"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -47,8 +46,7 @@ class TestFloat4(unittest.TestCase):
      c = a + b

      s = c.schedule()[0]
-      return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                         opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops

    sizes = [12, 8, 16]
    shifts = [3, 2, 4]
@@ -66,7 +64,7 @@ class TestFloat4(unittest.TestCase):
    s = c.schedule()[0]
    realized_ast = s.ast
    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)

    assert TestFloat4.count_float4(program.uops) == (0, 1)

@@ -77,8 +75,7 @@ class TestFloat4(unittest.TestCase):
    c = a + b

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops

    assert TestFloat4.count_float4(uops) == (0, 2)

@@ -90,8 +87,7 @@ class TestFloat4(unittest.TestCase):
      c = a + b

      s = c.schedule()[0]
-      return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                         opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
+      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops

    sizes = [13, 9, 17]
    shifts = [3, 2, 4]
@@ -109,7 +105,7 @@ class TestFloat4(unittest.TestCase):
    # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops

    assert TestFloat4.count_float4(uops) == (0, 0)

@@ -123,8 +119,7 @@ class TestFloat4(unittest.TestCase):
    # UPDATE: now we do this fusion

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops

    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}

@@ -137,7 +132,7 @@ class TestFloat4(unittest.TestCase):
    # since the top axis is not contiguous.

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops

    assert TestFloat4.count_float4(uops) == (0, 1)

@@ -149,7 +144,7 @@ class TestFloat4(unittest.TestCase):
    # should float4 b but not a

    s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops

    assert TestFloat4.count_float4(uops) == (1, 1)

@@ -13,7 +13,7 @@ class TestArange(unittest.TestCase):
    GlobalCounters.reset()
    sched = tensor.schedule()
    self.assertEqual(len(sched), 1)
-    p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
+    p = get_program(sched[-1].ast)
    ExecItem(CompiledRunner(p), [tensor.uop.buffer]).run()
    np.testing.assert_equal(tensor.numpy(), desired)
    return p.estimates.ops
@@ -36,7 +36,7 @@ class TestArange(unittest.TestCase):
    with Context(NOOPT=1):
      t = Tensor.ones(256, 256).contiguous().realize()
      sched = t.triu().schedule()
-      p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
+      p = get_program(sched[-1].ast)
      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)

 DSET, DDIM = 2048, 32
@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import temp, CI, CPU_LVP
+from tinygrad.helpers import temp

 N = 200  # has to be bigger than the cache to fail

@@ -180,7 +180,6 @@ class TestAssign(unittest.TestCase):
    new = a + times_a
    np.testing.assert_allclose(new.numpy(), 8)

-  @unittest.skipIf(CI and CPU_LVP, "flaky in CI")
  def test_double_assign(self):
    a = Tensor.ones(4).contiguous().realize()
    a += 1
@@ -155,17 +155,12 @@ class TestCustomKernel(unittest.TestCase):
    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_sum_int(self):
-    a = Tensor([1, 2, 3, 4, 5])
-    tst = Tensor.empty(1, dtype=a.dtype)
-    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)
@@ -17,6 +17,8 @@ pytestmark = pytest.mark.filterwarnings("ignore")
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")

+if Device.DEFAULT == "CPU": core_dtypes.remove(dtypes.bfloat16)  # NOTE: this is for teenygrad, don't remove
+
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
  if not is_dtype_supported(dtype): return []
  # dont cast internal dtypes
@@ -433,8 +435,6 @@ class TestOpsBFloat16(unittest.TestCase):
    data = [60000.0, 70000.0, 80000.0]
    np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())

-  # some CPUs there is no native bfloat16 sqrt
-  @unittest.skipIf(Device.DEFAULT == "CPU", "no approximation")
  def test_no_approximation(self):
    data = [326.0, 339.0, 10603200512.0]
    expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
@@ -442,4 +442,3 @@ class TestOpsBFloat16(unittest.TestCase):

 if __name__ == '__main__':
  unittest.main()
-
@@ -8,8 +8,6 @@ from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import ExecItem, BufferXfer, get_runner, CompiledRunner

-from test.helpers import needs_second_gpu
-
 np.random.seed(1337)
 Tensor.manual_seed(1337)
 BUF_SIZE = 4096 if CI else 4096 * 128
@@ -156,7 +154,6 @@ class TestGraph(unittest.TestCase):

    helper_test_graphs(Device[d0].graph, graphs)

-  @needs_second_gpu
  def test_copies_2_devs(self):
    self.skip_if_not_multigraph()

@@ -170,7 +167,6 @@ class TestGraph(unittest.TestCase):

    helper_test_graphs(Device[d0].graph, graphs)

-  @needs_second_gpu
  def test_copies_after_graph_global(self):
    self.skip_if_not_multigraph()

@@ -219,7 +215,6 @@ class TestGraph(unittest.TestCase):

    helper_test_graphs(Device[d0].graph, graphs)

-  @needs_second_gpu
  def test_graph_after_copies_devs(self):
    self.skip_if_not_multigraph()

@@ -3,7 +3,7 @@ import unittest, functools
 import numpy as np

 from hypothesis import given, settings, strategies as strat
-from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV, needs_second_gpu
+from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV
 from tinygrad.tensor import Tensor
 from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
@@ -439,7 +439,6 @@ class TestJit(unittest.TestCase):
      ja = jf(a)
      np.testing.assert_allclose(a.numpy(), ja.numpy(), atol=1e-4, rtol=1e-5)

-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_jitted_transfers(self):
    d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@@ -473,7 +472,6 @@ class TestJit(unittest.TestCase):
      np.testing.assert_allclose((a.numpy()+b.numpy()), zc.numpy(), atol=1e-4, rtol=1e-5)
      np.testing.assert_allclose((a.numpy()*b.numpy()), wc.numpy(), atol=1e-4, rtol=1e-5)

-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_jitted_view(self):
    d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@@ -501,18 +499,6 @@ class TestJit(unittest.TestCase):
    b = f(Tensor([2.0]))
    assert abs((a - b).item()) > 0.5

-  def test_jit_init_with_empty_different_size(self):
-    @TinyJit
-    def f(x:Tensor) -> Tensor: return (x + 1).realize()
-
-    f(Tensor.empty(1))
-    f(Tensor.empty(1))
-    # TODO: this should fail since input has a different size
-    f(Tensor(2.0)).item()
-    # TODO: this should not fail, and should return 3
-    with self.assertRaises(AssertionError):
-      f(Tensor([2.0])).item()
-
@unittest.skip("Pending multioutput implementation #3607")
 class TestMultioutputJit(unittest.TestCase):
  def _test(self, f):
@@ -45,7 +45,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  @unittest.expectedFailure
@@ -53,7 +53,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -63,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.empty(16)
    out = img.conv2d(w, b)
    ast = helper_linearizer_opt(out)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    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
@@ -84,7 +84,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -92,15 +92,16 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

+  @unittest.expectedFailure # TODO: investigate
  def test_two_nested_range_alt_indexing(self):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
    ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -113,7 +114,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> STORE
    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -123,7 +124,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -134,7 +135,7 @@ class TestLinearizer(unittest.TestCase):
    # these are of size 3 to avoid float4 coalesce
    r = a[:-1] + a[1:]

-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -146,7 +147,7 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = a.expand([2]) + b.expand([2])

-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

@@ -155,8 +156,7 @@ class TestLinearizer(unittest.TestCase):
    x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
    r = Tensor.conv2d(x,w,padding=1).relu()

-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
    accs = [u for u in uops if u.op is Ops.DEFINE_REG]
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
@@ -179,7 +179,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(r.schedule()[-1].ast, opts=opts_to_apply)

    stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

@@ -194,7 +194,7 @@ class TestLinearizer(unittest.TestCase):
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

@@ -204,14 +204,14 @@ class TestLinearizer(unittest.TestCase):
      if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
        realized_ast = a.schedule()[-1].ast
-        program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
+        program = get_program(realized_ast, opts=[])
        local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
      realized_ast = c.schedule()[-1].ast
-      program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
+      program = get_program(realized_ast, opts=[])
      local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
      self.assertEqual(local[0].dtype.base, expected_dtype)

@@ -239,7 +239,7 @@ class TestLinearizer(unittest.TestCase):
    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
    ast = helper_linearizer_opt(r, [opt])
    # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
+    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.END][0]
    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -353,7 +353,7 @@ class TestLinearizer(unittest.TestCase):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
    ast = helper_linearizer_opt(t+1)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(ast, opts=[]).uops
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -386,13 +386,13 @@ class TestLinearizer(unittest.TestCase):
    sched_copy = sched[:]
    run_schedule(sched)
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
-    program = get_program(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=())
+    program = get_program(sched_copy[-1].ast, opts=())
    assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"

  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      ast = helper_linearizer_opt(t)
-      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+      uops = get_program(ast).uops
      # ignore kernel optimized IF statements for now
      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
        uops = uops[:uops.index(if_op)]
@@ -425,7 +425,7 @@ class TestLinearizer(unittest.TestCase):
    out = x.matmul(y)
    with Context(TC=0):
      ast = helper_linearizer_opt(out)
-      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+      uops = get_program(ast).uops
    # check that the float4 cast collapses
    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
@@ -436,7 +436,7 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in get_program(ast).uops if u.op is Ops.STORE][0]
    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -449,7 +449,7 @@ class TestLinearizer(unittest.TestCase):
            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
    ast = helper_linearizer_opt(out, opts=[opt])
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
+    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]
@@ -470,7 +470,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    ast = helper_linearizer_opt(r)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast).uops
    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
@@ -517,7 +517,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
  device = real_bufs[0].device
  wanna_output = [np.array(x).flatten() for x in wanna_output]

-  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, opts=opts), device=device))

  def check_opt(opts):
    prg = get_prg(opts=opts)
@@ -3,7 +3,6 @@ import unittest
 from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
-from tinygrad.device import Device

 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
@@ -19,7 +18,7 @@ class TestLinearizerFailures(unittest.TestCase):
    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).end(c1, c2)
    ast = c10.sink()
-    get_program(ast, renderer=Device[Device.DEFAULT].renderer)
+    get_program(ast)

 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()
@@ -7,7 +7,7 @@ from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import REAL_DEV, not_support_multi_device, needs_second_gpu
+from test.helpers import REAL_DEV, not_support_multi_device

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -35,9 +35,6 @@ def _test_allreduce(t:Tensor):

@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestMultiTensor(unittest.TestCase):
-  @needs_second_gpu
-  def setUp(self): pass
-
  def test_to(self):
    X = Tensor.ones(256).contiguous().realize()
    X.to_(devices_2)
@@ -57,12 +54,6 @@ class TestMultiTensor(unittest.TestCase):
      assert lb.shape == (128,)
    (X + X).realize()

-  def test_shard_empty(self):
-    GlobalCounters.reset()
-    X = Tensor.empty(256).shard(devices_2, 0).realize()
-    assert GlobalCounters.kernel_count == 0
-    (X + X).realize()
-
  def _test_shard_op(self, op, out, n=4):
    t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
    r = op(t).realize()
@@ -836,7 +827,6 @@ class TestMultiTensor(unittest.TestCase):

@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestHandleData(unittest.TestCase):
-  @needs_second_gpu
  def test_copied_to_device(self):
    device = (d0, d1, d2, d3)
    t = Tensor([1, 2, 3, 4]).shard(device).realize()
@@ -861,9 +851,6 @@ class TestHandleData(unittest.TestCase):

@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
-  @needs_second_gpu
-  def setUp(self): pass
-
  # shrink a multitensor on sharded axis
  def test_shrink_bad_args(self):
    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
@@ -985,9 +972,6 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestBatchNorm(unittest.TestCase):
-  @needs_second_gpu
-  def setUp(self): pass
-
  def test_unsynced_backprop_conv_bn(self):
    with Tensor.train():
      from extra.lr_scheduler import OneCycleLR
@@ -1142,11 +1126,9 @@ def helper_test_shard_op(shps, fxn, atol=1e-6, rtol=1e-3):

@unittest.skipIf(not_support_multi_device(), "no multi")
 class TestTensorOps(unittest.TestCase):
-  @needs_second_gpu
  def test_interpolate(self):
    helper_test_shard_op([(4,16,16),(4,24,24)], lambda x: Tensor.interpolate(x, (19,19)))

-  @needs_second_gpu
  def test_bitcast(self):
    helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))

@@ -1189,7 +1171,6 @@ class TestMultiRamUsage(unittest.TestCase):

@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
-  @needs_second_gpu
  def test_from_npy(self):
    t = Tensor(np.arange(100, dtype=np.uint32))
    ll = t.shard((d0, d1), axis=0) + 1
@@ -1199,9 +1180,6 @@ class TestMultiFromUnrenderable(unittest.TestCase):
 class TestMultiAssign(unittest.TestCase):
  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))

-  @needs_second_gpu
-  def setUp(self): pass
-
  def test_multi_assign_realized(self):
    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
@@ -1264,7 +1242,6 @@ class TestMultiAssign(unittest.TestCase):

@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
-  @needs_second_gpu
  def test_transformer(self):
    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))

@@ -9,7 +9,7 @@ from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
 from tinygrad.engine.realize import run_schedule
-from test.helpers import not_support_multi_device, needs_second_gpu
+from test.helpers import not_support_multi_device

@unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
 class TestNN(unittest.TestCase):
@@ -481,21 +481,6 @@ class TestNN(unittest.TestCase):
    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())

-  #https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/torch/nn/modules/module.py#L2425
-  def test_load_conv_num_batches_tracked(self):
-    layer = BatchNorm(sz=1, track_running_stats=False)
-    state_dict = {
-      'weight': Tensor.ones(1),
-      'bias': Tensor.ones(1),
-      'num_batches_tracked': Tensor.ones(1),
-    }
-    load_state_dict(layer, state_dict)
-    state_dict['num_batches_tracked'] = Tensor.empty()
-    load_state_dict(layer, state_dict)
-    layer.num_batches_tracked = Tensor.ones(1)
-    load_state_dict(layer, state_dict)
-
-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_load_state_dict_sharded_model(self):
    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
@@ -534,7 +519,6 @@ class TestNN(unittest.TestCase):
    np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
    np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())

-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_load_state_dict_sharded_model_dict_same_axis(self):
    devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM, EMULATE
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -16,7 +16,6 @@ if CI:

 FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
 PRINT_TENSORS = getenv("PRINT_TENSORS", 0)
-COMPILE_ONLY = Device.DEFAULT == "NULL" and not EMULATE

 def slow_test(test_func):
  return unittest.skipIf(getenv("SKIP_SLOW_TEST"), "Skipping slow test")(test_func)
@@ -39,7 +38,6 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
  tinygrad_fp = time.monotonic() - st

  def compare(s, tinygrad_output, torch_output, atol, rtol):
-    if COMPILE_ONLY: return
    if PRINT_TENSORS: print(s, tinygrad_output, torch_output)
    try:
      assert tinygrad_output.shape == torch_output.shape, f"shape mismatch: tinygrad={tinygrad_output.shape} | torch={torch_output.shape}"
@@ -423,9 +421,8 @@ class TestOps(unittest.TestCase):
  def test_isinf(self):
    val = [float('-inf'), 0., float('inf'), float('nan'), 1.1]
    helper_test_op(None, torch.isinf, Tensor.isinf, vals=[val], forward_only=True)
-    if not COMPILE_ONLY:
-      np.testing.assert_equal(Tensor(val).isinf(detect_positive=True, detect_negative=False).numpy(), [False, False, True, False, False])
-      np.testing.assert_equal(Tensor(val).isinf(detect_positive=False, detect_negative=True).numpy(), [True, False, False, False, False])
+    np.testing.assert_equal(Tensor(val).isinf(detect_positive=True, detect_negative=False).numpy(), [False, False, True, False, False])
+    np.testing.assert_equal(Tensor(val).isinf(detect_positive=False, detect_negative=True).numpy(), [True, False, False, False, False])

  def test_isnan(self):
    helper_test_op(None, torch.isnan, Tensor.isnan, vals=[[float('-inf'), 0., float('inf'), float('nan'), 1.1]], forward_only=True)
@@ -597,7 +594,7 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x//2, forward_only=True, vals=[[3, 4, 5]])
    helper_test_op(None, functools.partial(torch.div, rounding_mode="trunc"), Tensor.idiv, forward_only=True,
                   vals=[[-4, 7, 5, 4, -7, 8], [2, -3, 8, -2, 3, 5]])
-    if is_dtype_supported(dtypes.uint64) and not COMPILE_ONLY:
+    if is_dtype_supported(dtypes.uint64):
      x = Tensor(2**64 - 1, dtype=dtypes.uint64).idiv(1)
      np.testing.assert_equal(x.numpy(), 2**64 - 1)

@@ -682,7 +679,6 @@ class TestOps(unittest.TestCase):
    # float to power of int
    helper_test_op(None, lambda x: 0.7**x, vals=[[-2,-1,0,1,2,3]], forward_only=True)

-  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
  def test_pow_const_direct(self):
    # x ** c
    def get_tiny_gradient(x, c):
@@ -1074,7 +1070,8 @@ class TestOps(unittest.TestCase):
  @slow_test
  def test_cummax(self):
    helper_test_op([()], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
-    self.helper_test_exception([()], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
+    # TODO: torch allows this?
+    # self.helper_test_exception([()], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
    helper_test_op([(20,)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
    self.helper_test_exception([(20,)], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
    self.helper_test_exception([(20,)], lambda x: torch.cummax(x, dim=-2).values, lambda x: Tensor.cummax(x, axis=-2), expected=IndexError)
@@ -1091,9 +1088,8 @@ class TestOps(unittest.TestCase):
    # check if it returns the first index for multiple occurences
    helper_test_op(None, lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[2, 2]])
    helper_test_op(None, lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[1, 2, 2]])
-    if not COMPILE_ONLY:
-      np.testing.assert_equal(Tensor([2,2]).argmax().numpy(), 0)
-      np.testing.assert_equal(Tensor([1,2,2]).argmax().numpy(), 1)
+    np.testing.assert_equal(Tensor([2,2]).argmax().numpy(), 0)
+    np.testing.assert_equal(Tensor([1,2,2]).argmax().numpy(), 1)
    helper_test_op([(10,20)], lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True)
    helper_test_op([(10,20)], lambda x: x.argmax(0, False).type(torch.int32), lambda x: x.argmax(0, False), forward_only=True)
    helper_test_op([(10,20)], lambda x: x.argmax(1, False).type(torch.int32), lambda x: x.argmax(1, False), forward_only=True)
@@ -1111,9 +1107,8 @@ class TestOps(unittest.TestCase):
    # check if it returns the first index for multiple occurences
    helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
    helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[3, 2, 2]])
-    if not COMPILE_ONLY:
-      np.testing.assert_equal(Tensor([2,2]).argmin().numpy(), 0)
-      np.testing.assert_equal(Tensor([3,2,2]).argmin().numpy(), 1)
+    np.testing.assert_equal(Tensor([2,2]).argmin().numpy(), 0)
+    np.testing.assert_equal(Tensor([3,2,2]).argmin().numpy(), 1)
    helper_test_op([(10,20)], lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True)
    helper_test_op([(10,20)], lambda x: x.argmin(0, False).type(torch.int32), lambda x: x.argmin(0, False), forward_only=True)
    helper_test_op([(10,20)], lambda x: x.argmin(1, False).type(torch.int32), lambda x: x.argmin(1, False), forward_only=True)
@@ -1161,13 +1156,12 @@ class TestOps(unittest.TestCase):
                          lambda x: x.topk(4, dim, largest, sorted_).indices.type(torch.int32),
                          lambda x: x.topk(4, dim, largest, sorted_)[1], forward_only=True)
    # repeated values
-    if not COMPILE_ONLY:
-      value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
-      np.testing.assert_equal(value.numpy(), [1, 1, 1])
-      np.testing.assert_equal(indices.numpy(), [0, 1, 3])
-      value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3, largest=False)
-      np.testing.assert_equal(value.numpy(), [0, 0, 0])
-      np.testing.assert_equal(indices.numpy(), [2, 4, 6])
+    value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
+    np.testing.assert_equal(value.numpy(), [1, 1, 1])
+    np.testing.assert_equal(indices.numpy(), [0, 1, 3])
+    value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3, largest=False)
+    np.testing.assert_equal(value.numpy(), [0, 0, 0])
+    np.testing.assert_equal(indices.numpy(), [2, 4, 6])
    self.helper_test_exception([(4)], lambda x: x.topk(5), expected=(RuntimeError, ValueError))

  @slow_test
@@ -1319,7 +1313,6 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
  @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and CPU_LLVM) or IMAGE
  or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "not precise enough")
  def test_gemm_fp16(self):
    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
  def test_gemm(self):
@@ -1730,7 +1723,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(7,5,10)], lambda x: x[1:5:2, 3, ::4])
    helper_test_op([(7,5,10)], lambda x: x[1:5:2, None, None, 3, None, ::4])

-  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
  def test_slice_negative_strides(self):
    # Torch doesn't support slicing with negative steps
    a = np.random.randn(10, 10, 10).astype(np.float32)
@@ -2737,9 +2729,6 @@ class TestOps(unittest.TestCase):

  def test_clip(self):
    helper_test_op([(45,65)], lambda x: x.clip(-2.3, 1.2))
-    # NOTE: torch set backward to 1 at the boundaries
-    # https://github.com/pytorch/pytorch/blob/7a41b66367c38d0af3e8a90f7be48d6b281e7bca/tools/autograd/derivatives.yaml#L421
-    helper_test_op(None, lambda x: x.clip(-2.5, 1.5), vals=[[-3.0, -2.5, 0, 1.5, 2]])
    helper_test_op([(45,65)], lambda x: x.clip(0, 0))
    helper_test_op([(45,65)], lambda x: x.clip(10, 100))
    helper_test_op([(45,65)], lambda x: x.clip(0, 0.1))
@@ -2760,7 +2749,6 @@ class TestOps(unittest.TestCase):
    n = Tensor([1, float("nan")]).max().numpy()
    assert math.isnan(n.item()), f"{n.item()} is not nan"

-  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
  def test_inf_where(self):
    x = Tensor.full((3, 3), float("inf"))
    n = (x < 0).where(x, 1).numpy()
@@ -3177,7 +3165,6 @@ class TestOps(unittest.TestCase):

  @unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
-  @unittest.skipIf(Device.DEFAULT == "QCOM", "QCOM fails with: Resource deadlock avoided")
  def test_masked_select(self):
    helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
    helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)
@@ -1,6 +1,6 @@
 import numpy as np
 import unittest
-from tinygrad import Tensor, Device
+from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
@@ -17,7 +17,7 @@ class TestOptGemm(unittest.TestCase):
    t = self.a.T @ self.b.T
    # TODO: this should be a generic test helper
    si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts))
+    run = CompiledRunner(get_program(si.ast, opts=opts))
    ExecItem(run, si.bufs).run()
    test = si.bufs[0].numpy().reshape(self.res.shape)
    np.testing.assert_allclose(self.res, test, atol=1e-4)
@@ -13,7 +13,7 @@ class TestOpts(unittest.TestCase):
    s = out.schedule()
    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
-      prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
+      prg = get_program(s[-1].ast)
      self.assertIn('float4', prg.src)

 if __name__ == '__main__':
@@ -38,7 +38,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=

 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
  si = out.schedule()[-1]
-  prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
+  prg = get_program(si.ast, opts=opts)
  if replace_src is not None:
    old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
    prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
@@ -7,7 +7,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import lower_schedule, CompiledRunner
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from test.helpers import not_support_multi_device, needs_second_gpu
+from test.helpers import not_support_multi_device

 import numpy as np
 import torch
@@ -141,7 +141,6 @@ class TestRandomness(unittest.TestCase):
    r = Tensor.rand(10).numpy()
    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)

-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_threefry_tensors_cnt(self):
    Tensor.manual_seed(1337)
@@ -161,7 +160,6 @@ class TestRandomness(unittest.TestCase):
    assert len(Tensor._device_rng_counters) == 0
    assert len(Tensor._device_seeds) == 0

-  @needs_second_gpu
  @unittest.skipIf(not_support_multi_device(), "no multi")
  def test_threefry_same_kernels(self):
    Tensor.manual_seed(0)
@@ -12,7 +12,8 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.schedule.rangeify import Kernel
+from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
+from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule

 class KernelCountException(Exception): pass
@@ -23,11 +24,13 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
  else:
    assert isinstance(t, UOp), f"can't schedule {t}"
-    sched = Tensor(t).schedule()
+    sink = UOp.sink(t) if t.op is not Ops.SINK else t
+    becomes_map = get_rangeify_map(sink)
+    sched, _ = create_schedule_with_vars(sink.substitute(becomes_map))
  # test lowering all the ScheduleItems to ExecItems
  kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
  if kernel_cnt != allowed:
-    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
    if DEBUG >= 3:
      for i,s in enumerate(sched):
        print("kernel", i+1)
@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
 from tinygrad.helpers import Context
-from test.helpers import REAL_DEV, needs_second_gpu
+from test.helpers import REAL_DEV

@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
 class TestSubBuffer(unittest.TestCase):
@@ -41,7 +41,6 @@ class TestSubBuffer(unittest.TestCase):
    out = (vt + 100).tolist()
    assert out == [102, 103]

-  @needs_second_gpu
  @unittest.skipIf(REAL_DEV not in {"CUDA", "NV", "AMD"}, "only NV, AMD, CUDA")
  def test_subbuffer_transfer(self):
    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
@@ -70,15 +70,15 @@ class TestTinygrad(unittest.TestCase):
    out = out.log_softmax()
    out = out.mul(m).add(m).sum()
    out.backward()
-    xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
+    xgrad,wgrad = x.grad, W.grad
    out.backward()
-    xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
+    xgrad2,wgrad2 = x.grad, W.grad
    out.backward() # no need to retain again since we will not re-run backward
-    xgrad3, wgrad3 = x.grad.numpy(), W.grad.numpy()
-    np.testing.assert_allclose(xgrad3, xgrad * 3., atol=1e-6)
-    np.testing.assert_allclose(wgrad3, wgrad * 3., atol=1e-6)
-    np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6)
-    np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6)
+    xgrad3,wgrad3 = x.grad, W.grad
+    np.testing.assert_allclose(xgrad3.numpy(), xgrad.numpy() * 3., atol=1e-6)
+    np.testing.assert_allclose(wgrad3.numpy(), wgrad.numpy() * 3., atol=1e-6)
+    np.testing.assert_allclose(xgrad2.numpy(), xgrad.numpy() * 2., atol=1e-6)
+    np.testing.assert_allclose(wgrad2.numpy(), wgrad.numpy() * 2., atol=1e-6)

  def test_second_order_backward_pass(self):
    def test_pytorch():
@@ -829,7 +829,6 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertEqual(len(si.metadata), 3)
    self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})

-  @unittest.skip("metadata is no longer promised to be exact with schedulecache")
  def test_complex_backward(self):
    x = Tensor.rand(3, requires_grad=True).realize()
    y = Tensor.rand(3, requires_grad=True).realize()
@@ -842,13 +841,11 @@ class TestTensorMetadata(unittest.TestCase):
    self.assertTrue(y.grad.uop.metadata[0].backward)
    si = Tensor.schedule(out, x.grad, y.grad)[-1]
    #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
-    # skip numpy, this is schedule cache
-    self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"sigmoid", "relu"})
+    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")

-  @unittest.skip("metadata is no longer promised to be exact with schedulecache")
  def test_tracemeta_0(self):
    with Context(TRACEMETA=0):
      x = Tensor.rand(3, requires_grad=True)
@@ -660,18 +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_after_end(self):
-    r = UOp.range(10, 0)
-
-    c = r + 1
-    self.assertIn(r, c.ranges)
-
-    e = UOp.const(dtypes.void, None).end(r)
-    self.assertNotIn(r, e.ranges)
-
-    a = c.after(e)
-    self.assertNotIn(r, a.ranges)
-
@track_rewrites()
 def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)

@@ -175,13 +175,13 @@ class TestStatsOptimized(unittest.TestCase):
    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats

  def test_gemm(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
+    p = get_program(self.ast_gemm, opts=[])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)

  def test_gemm_tc_unroll(self):
    try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
    except KernelOptError:
      raise unittest.SkipTest("no tensor cores")
    print(p.src)
@@ -190,19 +190,18 @@ class TestStatsOptimized(unittest.TestCase):
  # this is a good lesson about why UPCASTing is a good idea

  def test_gemm_one_upcasted(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4)])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)

  def test_gemm_upcasted(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
-                    opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
    self.check_gemm(p)
    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)

  def test_gemm_upcasted_locals(self):
    try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
                                           Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
@@ -211,7 +210,7 @@ class TestStatsOptimized(unittest.TestCase):

  def test_gemm_group(self):
    try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 4)])
+      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
    SZ = N*N*4
@@ -220,14 +219,14 @@ class TestStatsOptimized(unittest.TestCase):
    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)

  def test_reduce(self):
-    p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
+    p = get_program(self.ast_reduce, opts=[])
    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
    self.assertEqual(p.estimates.ops, N*N)
    self.assertEqual(p.estimates.mem, N*N*4 + 4)

  def test_reduce_group(self):
    try:
-      p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 50)])
+      p = get_program(self.ast_reduce, opts=[Opt(OptOps.GROUP, 0, 50)])
    except KernelOptError:
      raise unittest.SkipTest("no locals")
    # NOTE: these are wrong, they don't respect the if statement
@@ -1,14 +1,14 @@
 import unittest

-from tinygrad import Device
-from tinygrad.helpers import fetch
+from tinygrad import Tensor, Device
 from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu

 class TestHevc(unittest.TestCase):
  def test_hevc_parser(self):
    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
-    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+    hevc_tensor = Tensor.from_url(url, device="CPU")

+    dat = bytes(hevc_tensor.data())
    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat, device=Device.DEFAULT)

    def _test_common(frame, bts):
@@ -1,9 +1,8 @@
-import unittest, math, time
+import unittest, math

 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import ExecItem, get_runner
-from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import CI
 import numpy as np

@@ -13,16 +12,11 @@ from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout

@unittest.skipIf(CI or Device.DEFAULT not in ["AMD"], "only amd")
 class TestTK(unittest.TestCase):
-  def setUp(self):
-    arch = Device["AMD"].arch
-    if not arch.startswith("gfx9"):
-      self.skipTest(f"arch {arch} not supported")
-
  @unittest.skipIf(CI, "no wmma in ci")
  def test_simple_matmul(self):
    N = 8192
    BLOCK_SIZE = 64
-    with Kernel("simple_matmul", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
      warp = ker.warp

      c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -31,16 +25,14 @@ class TestTK(unittest.TestCase):

      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)

      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
-      c_reg_col = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)

      col, row = ker.blockIdx_x, ker.blockIdx_y

-      c_reg_col = warp.zero(c_reg_col)
+      c_reg = warp.zero(c_reg)
      for tile in ker.range(N // BLOCK_SIZE):
        a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
        b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
@@ -48,11 +40,8 @@ class TestTK(unittest.TestCase):
        a_reg = warp.load(a_reg, a_smem)
        b_reg = warp.load(b_reg, b_smem)

-        c_reg_col = warp.mma_AB(c_reg_col, a_reg, b_reg)
-      c_reg_col = ker.endrange()
-
-      c_smem = warp.store(c_smem, c_reg_col)
-      c_reg = warp.load(c_reg, c_smem)
+        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
+      c_reg = ker.endrange()

      c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)

@@ -76,7 +65,7 @@ class TestTK(unittest.TestCase):
  def test_simple_matmul_transposed(self):
    N = 8192
    BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128
-    with Kernel("simple_matmul_transposed", (N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
+    with Kernel((N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
      warp = ker.warp

      c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -124,7 +113,7 @@ class TestTK(unittest.TestCase):
  def test_load_store(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("load_store", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -157,50 +146,11 @@ class TestTK(unittest.TestCase):

    np.testing.assert_allclose(b.numpy(), ref.numpy())

-  def test_load_store_local_hop(self):
-    N = 64
-    BLOCK_SIZE = 32
-    with Kernel("load_store_local_hop", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
-      warp = ker.warp
-
-      b = ker.gl((1, 1, N, N), dtypes.float32)
-      a = ker.gl((1, 1, N, N), dtypes.float32)
-
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-
-      col, row = ker.blockIdx_x, ker.blockIdx_y
-
-      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
-      a_reg = warp.load(a_reg, a_smem)
-      b_reg = warp.copy(b_reg, a_reg)
-      b_smem = warp.store(b_smem, b_reg)
-      b_reg = warp.load(b_reg, b_smem)
-      b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
-
-      sink = ker.finish()
-
-    with Context(DEBUG=0):
-      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
-      b = Tensor.empty(1, 1, N, N, dtype="float32")
-      Tensor.realize(a, b)
-
-    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
-    for _ in range(5): ei.run(wait=True)
-    b = b.float()
-
-    ref = a.float()
-
-    np.testing.assert_allclose(b.numpy(), ref.numpy())
-
  @unittest.skip("TODO")
  def test_load_store_group(self):
    N = 256
    BLOCK_SIZE = 64
-    with Kernel("load_store_group", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
+    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
      warp = ker.warp
      group = ker.group(2)

@@ -237,7 +187,7 @@ class TestTK(unittest.TestCase):
  def test_add(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("add", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -274,7 +224,7 @@ class TestTK(unittest.TestCase):
  def test_max(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("max", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -319,7 +269,7 @@ class TestTK(unittest.TestCase):
  def test_max_nonsquare(self):
    N, M = 32, 128
    BLOCK_N, BLOCK_M = 16, 64
-    with Kernel("max_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -364,7 +314,7 @@ class TestTK(unittest.TestCase):
  def test_sum(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("sum", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -409,7 +359,7 @@ class TestTK(unittest.TestCase):
  def test_sum_nonsquare(self):
    N, M = 32, 128
    BLOCK_N, BLOCK_M = 16, 64
-    with Kernel("sum_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -454,7 +404,7 @@ class TestTK(unittest.TestCase):
  def test_softmax(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("softmax", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
@@ -472,29 +422,28 @@ class TestTK(unittest.TestCase):
      norm_vec = warp.zero(norm_vec)

      for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg_ = warp.load(a_reg, a_smem_)
+        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg = warp.load(a_reg, a_smem)

-        a_reg_ *= 1.0 / math.log(2)
+        a_reg *= 1.0 / math.log(2)

        max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
-        max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
-        a_reg_ = (a_reg_ - max_vec).exp2()
+        max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
+        a_reg = (a_reg - max_vec).exp2()
        max_vec_last = (max_vec_last - max_vec).exp2()
        norm_vec *= max_vec_last
-        norm_vec = warp.row_reduce(norm_vec, a_reg_, lambda a, b: a + b)
+        norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
      norm_vec = ker.endrange()
-      max_vec = max_vec.after(norm_vec)

      for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg_ = warp.load(a_reg, a_smem_)
+        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg = warp.load(a_reg.after(norm_vec), a_smem)

-        a_reg_ *= 1.0 / math.log(2)
-        a_reg_ = (a_reg_ - max_vec).exp2()
-        a_reg_ /= norm_vec
+        a_reg *= 1.0 / math.log(2)
+        a_reg = (a_reg - max_vec).exp2()
+        a_reg /= norm_vec

-        b = warp.store(b, a_reg_, (0, 0, 0, tile_col), (), axis=2)
+        b = warp.store(b, a_reg, (0, 0, 0, tile_col), (), axis=2)

      sink = ker.finish()

@@ -514,7 +463,7 @@ class TestTK(unittest.TestCase):
  def test_softmax_col(self):
    N = 64
    BLOCK_SIZE = 32
-    with Kernel("softmax_col", (1, 1, 1), WARP_THREADS) as ker:
+    with Kernel((1, 1, 1), WARP_THREADS) as ker:
      warp = ker.warp

      b = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
@@ -532,29 +481,28 @@ class TestTK(unittest.TestCase):
      norm_vec = warp.zero(norm_vec)

      for tile_row in ker.range(N // BLOCK_SIZE):
-        a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
-        a_reg_ = warp.load(a_reg, a_smem_)
+        a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
+        a_reg = warp.load(a_reg, a_smem)

-        a_reg_ *= 1.0 / math.log(2)
+        a_reg *= 1.0 / math.log(2)

        max_vec_last = warp.copy(max_vec_last.after(tile_row), max_vec)
-        max_vec = warp.col_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
-        a_reg_ = (a_reg_ - max_vec).exp2()
+        max_vec = warp.col_reduce(max_vec.after(max_vec_last), a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
+        a_reg = (a_reg - max_vec).exp2()
        max_vec_last = (max_vec_last - max_vec).exp2()
        norm_vec *= max_vec_last
-        norm_vec = warp.col_reduce(norm_vec, a_reg_, lambda a, b: a + b)
+        norm_vec = warp.col_reduce(norm_vec, a_reg, lambda a, b: a + b)
      norm_vec = ker.endrange()
-      max_vec = max_vec.after(norm_vec)

      for tile_row in ker.range(N // BLOCK_SIZE):
-        a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
-        a_reg_ = warp.load(a_reg.after(norm_vec), a_smem_)
+        a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
+        a_reg = warp.load(a_reg.after(norm_vec), a_smem)

-        a_reg_ *= 1.0 / math.log(2)
-        a_reg_ = (a_reg_ - max_vec).exp2()
-        a_reg_ /= norm_vec
+        a_reg *= 1.0 / math.log(2)
+        a_reg = (a_reg - max_vec).exp2()
+        a_reg /= norm_vec

-        b = warp.store(b, a_reg_, (0, 0, tile_row, 0), (), axis=2)
+        b = warp.store(b, a_reg, (0, 0, tile_row, 0), (), axis=2)

      sink = ker.finish()

@@ -573,11 +521,11 @@ class TestTK(unittest.TestCase):

  def test_fa(self):
    NUM_WORKERS = 1
-    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
+    B, N, H, H_KV, D = 1, 8192, 32, 8, 128
    Q_BLOCK_SIZE = 16
    KV_BLOCK_SIZE = 16
    GROUP_SIZE = H // H_KV
-    with Kernel("fa", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
+    with Kernel((H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
      warp = ker.warp

      # kernel
@@ -657,7 +605,6 @@ class TestTK(unittest.TestCase):
        att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
        o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
      o_reg = ker.endrange()
-      norm_vec = norm_vec.after(o_reg)

      o_reg /= norm_vec

@@ -674,12 +621,7 @@ class TestTK(unittest.TestCase):
      Tensor.realize(q, k, v, out)

    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (out, q, k, v)])
-    for _ in range(5):
-      et = ei.run(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): ei.run(wait=True)
    out = out.float()

    q_permuted = q.permute(0, 2, 1, 3)
@@ -688,36 +630,7 @@ class TestTK(unittest.TestCase):
    ref = q_permuted.scaled_dot_product_attention(k_permuted, v_permuted, is_causal=True, enable_gqa=True).float()
    ref = ref.permute(0, 2, 1, 3)

-    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
-
-  def test_fast_fa(self):
-    from extra.thunder.tiny.fa import flash_attention
-
-    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
-
-    with Context(DEBUG=0):
-      q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
-      k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
-      v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
-      Tensor.realize(q, k, v)
-
-    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
-
-    fa_jitted = TinyJit(flash_attention)
-
-    for _ in range(10):
-      st = time.perf_counter()
-      out = fa_jitted(q, k, v, is_causal=True)
-      et = time.perf_counter() - st
-      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")
-    out = out.float().transpose(1, 2)
-
-    ref = q.scaled_dot_product_attention(k, v, is_causal=True, enable_gqa=True).float().transpose(1, 2)
-
-    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-2, rtol=1e-5)

 if __name__ == "__main__":
  unittest.main()
@@ -1,14 +1,8 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
+from tinygrad.apps.llm import apply_rope
 #from tinygrad.engine.realize import run_schedule

-def apply_rope(x:Tensor, start_pos:int):
-  B, H, T, Hd = x.shape
-  precompute_freqs_cis.cache_clear()
-  freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
-  return apply_rope_new(x, freqs_cis)
-
 # TODO: test_scheduler, but just in uint
 class TestAttention(unittest.TestCase):
  def test_half_qkv_buffers(self):
@@ -45,7 +39,7 @@ class TestAttention(unittest.TestCase):
    prune_size = len(rope_prune.captured.jit_cache)

    self.assertGreater(noprune_size, prune_size)
-    self.assertGreaterEqual(noprune_size, 2)
+    self.assertGreaterEqual(noprune_size, 3)
    self.assertEqual(prune_size, 1)

 if __name__ == '__main__':
@@ -44,34 +44,6 @@ class TestAutogen(unittest.TestCase):
      test.argtypes = [Baz]
      self.assertEqual(test(b), b.a + b.b + b.c + b.d)

-  # https://github.com/python/cpython/issues/90914
-  @unittest.skipIf(WIN, "doesn't compile on windows")
-  def test_bitfield_interop(self):
-    class Baz(Struct): pass
-    Baz._fields_ = [(chr(ord('a') + i), ctypes.c_bool, 1) for i in range(8)]
-    src = '''#include <stdbool.h>
-      struct baz {
-        bool a:1;
-        bool b:1;
-        bool c:1;
-        bool d:1;
-        bool e:1;
-        bool f:1;
-        bool g:1;
-        bool h:1;
-      };
-
-      int test(struct baz x) {
-        return x.c;
-      }
-    '''
-    args = ('-x', 'c', '-fPIC', '-shared')
-    with tempfile.NamedTemporaryFile(suffix=".so") as f:
-      subprocess.check_output(('clang',) + args + ('-', '-o', f.name), input=src.encode('utf-8'))
-      test = ctypes.CDLL(f.name).test
-      test.argtypes = [Baz]
-      for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
-
  @unittest.skipIf(WIN, "doesn't compile on windows")
  def test_packed_structs(self):
    NvU32 = ctypes.c_uint32
@@ -110,53 +82,4 @@ class TestAutogen(unittest.TestCase):
    assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
    assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC

-  def test_packed_fields(self):
-    uint8_t = ctypes.c_ubyte
-    uint16_t = ctypes.c_ushort
-    uint32_t = ctypes.c_uint32
-
-    class struct_die_info(Struct): pass
-    struct_die_info._packed_ = True
-    struct_die_info._fields_ = [
-      ('die_id', uint16_t),
-      ('die_offset', uint16_t),
-    ]
-    die_info = struct_die_info
-    class struct_ip_discovery_header(Struct): pass
-    class struct_ip_discovery_header_0(ctypes.Union): pass
-    class struct_ip_discovery_header_0_0(Struct): pass
-    uint8_t = ctypes.c_ubyte
-    struct_ip_discovery_header_0_0._fields_ = [
-      ('base_addr_64_bit', uint8_t,1),
-      ('reserved', uint8_t,7),
-      ('reserved2', uint8_t),
-    ]
-    struct_ip_discovery_header_0._anonymous_ = ['_0']
-    struct_ip_discovery_header_0._packed_ = True
-    struct_ip_discovery_header_0._fields_ = [
-      ('padding', (uint16_t * 1)),
-      ('_0', struct_ip_discovery_header_0_0),
-    ]
-    struct_ip_discovery_header._anonymous_ = ['_0']
-    struct_ip_discovery_header._packed_ = True
-    struct_ip_discovery_header._fields_ = [
-      ('signature', uint32_t),
-      ('version', uint16_t),
-      ('size', uint16_t),
-      ('id', uint32_t),
-      ('num_dies', uint16_t),
-      ('die_info', (die_info * 16)),
-      ('_0', struct_ip_discovery_header_0),
-    ]
-    ip_discovery_header = struct_ip_discovery_header
-
-    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
-    ihdr = ip_discovery_header.from_buffer_copy(hdr)
-
-    assert ctypes.sizeof(ihdr) == 80
-    assert ihdr.signature == 0x53445049
-    assert ihdr.version == 0x0004
-    assert ihdr.num_dies == 1
-    assert ihdr.base_addr_64_bit == 1
-
 if __name__ == "__main__": unittest.main()
@@ -42,10 +42,12 @@ class TestDevice(unittest.TestCase):
                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_LLVM": "1"})
      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_LLVM": "0"})
+      subprocess.run([f'python3 -c "{imports}; {expect_failure}"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "0", "CPU_LLVM": "0"})
      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CC": "LLVM"})
-      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CC": "CLANGJIT"})
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "0"})
+      subprocess.run([f'python3 -c "{imports}; {expect_failure}"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "1", "CPU_LLVM": "1"})
    elif Device.DEFAULT == "AMD":
      from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
      try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
@@ -62,20 +64,14 @@ class TestDevice(unittest.TestCase):
                        shell=True, check=True, env={**os.environ, "DEV": "AMD", "AMD_HIP": "1", "AMD_LLVM": "1"})
    else: self.skipTest("only run on CPU/AMD")

-  @unittest.skipIf((WIN and CI) or (not Device.DEFAULT == "CPU"), "skipping windows test")
-  def test_env_online(self):
-    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
-    try: _, _ = CPULLVMCompiler(), ClangJITCompiler()
-    except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
-
-    with Context(CPU_LLVM=1):
-      inst = Device["CPU"].compiler
-      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
-    with Context(CPU_LLVM=0):
-      self.assertIsInstance(Device["CPU"].compiler, ClangJITCompiler)
-    with Context(CPU_LLVM=1):
-      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
-      assert inst is Device["CPU"].compiler  # cached
+  def test_compiler_envvar(self):
+    d = Device[Device.DEFAULT]
+    dname = Device.DEFAULT.split(':')[0].upper()
+    assert d._get_compiler_envvar(type("Compiler", (), {})) == f"{dname}_COMPILER"
+    assert d._get_compiler_envvar(type("LLVMCompiler", (), {})) == f"{dname}_LLVM"
+    assert d._get_compiler_envvar(type("RandomCompiler", (), {})) == f"{dname}_RANDOM"
+    assert d._get_compiler_envvar(type(f"{dname}Compiler", (), {})) == f"{dname}_{dname}COMPILER" # do not repeat device name alone
+    assert d._get_compiler_envvar(type(f"{dname}LLVMCompiler", (), {})) == f"{dname}_LLVM" # do not repeat device name

 class MockCompiler(Compiler):
  def __init__(self, key): super().__init__(key)
@@ -110,18 +110,6 @@ class TestTensorGradient(unittest.TestCase):
    with self.assertRaises(RuntimeError): x.sum().gradient(x)
    with self.assertRaises(RuntimeError): x.float().sum().gradient(x)

-  def test_multiple_backward(self):
-    x = Tensor([3.], requires_grad=True)
-    (x*2)[0].backward()
-    np.testing.assert_allclose(x.grad.numpy(), [2.0])
-    old_grad = x.grad
-    (x*3)[0].backward()
-    np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0])
-    self.assertIs(x.grad, old_grad)
-    (x*x)[0].backward()
-    np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
-    self.assertIs(x.grad, old_grad)
-
 class TestRealizeMeansRealize(unittest.TestCase):
  def test_randn_realizes(self):
    x = Tensor.randn(2, 3, 64, 64, requires_grad=True).realize()
@@ -163,14 +163,6 @@ class TestFetch(unittest.TestCase):
    fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
          allow_caching=False)

-  def test_fetch_half_and_full_file(self):
-    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
-          headers={"Range": "bytes=0-10"}).read_bytes()
-    assert len(x) == 11, f"{len(x) != 11}"
-    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
-          headers={"Range": "bytes=0-100"}).read_bytes()
-    assert len(x) == 101, f"{len(x) != 101}"
-
 class TestFullyFlatten(unittest.TestCase):
  def test_fully_flatten(self):
    self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])
@@ -176,8 +176,9 @@ class TestIndexing(unittest.TestCase):
    self.assertRaises(IndexError, lambda: reference[0.0, ..., 0.0:2.0])
    self.assertRaises(IndexError, lambda: reference[0.0, :, 0.0])

-    def delitem(): del reference[0]
-    self.assertRaises(TypeError, delitem)
+    # TODO: delitem
+    # def delitem(): del reference[0]
+    # self.assertRaises(TypeError, delitem)

  # TODO setitem backward
  '''
@@ -1,136 +0,0 @@
-import unittest, threading, time
-from unittest.mock import Mock
-
-class TestLLMServer(unittest.TestCase):
-  """Integration tests using the real OpenAI client."""
-
-  @classmethod
-  def setUpClass(cls):
-    cls.mock_tok = Mock()
-    cls.mock_tok.role = Mock(return_value=[100, 101])
-    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
-    cls.mock_tok.decode = Mock(return_value="Hello")
-
-    cls.mock_model = Mock()
-    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
-
-    cls.bos_id = 1
-    cls.eos_id = 999
-
-    import tinygrad.apps.llm as llm_module
-    llm_module.model = cls.mock_model
-    llm_module.tok = cls.mock_tok
-    llm_module.bos_id = cls.bos_id
-    llm_module.eos_id = cls.eos_id
-
-    from tinygrad.apps.llm import Handler
-    from tinygrad.helpers import TCPServerWithReuse
-
-    cls.port = 11435
-    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
-    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
-    cls.server_thread.start()
-    time.sleep(0.1)
-
-    from openai import OpenAI
-    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
-
-  @classmethod
-  def tearDownClass(cls):
-    cls.server.shutdown()
-    cls.server.server_close()
-
-  def test_chat_completion_stream(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=True
-    )
-
-    chunks = list(stream)
-    self.assertGreater(len(chunks), 0)
-    self.assertEqual(chunks[0].choices[0].delta.role, "assistant")
-    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
-
-  def test_openai_response_structure(self):
-    stream = self.client.chat.completions.create(
-      model="test-model",
-      messages=[{"role": "user", "content": "Test"}],
-      stream=True
-    )
-
-    for chunk in stream:
-      self.assertTrue(chunk.id.startswith("chatcmpl-"))
-      self.assertEqual(chunk.object, "chat.completion.chunk")
-      self.assertIsNotNone(chunk.choices)
-      self.assertIsNotNone(chunk.created)
-      self.assertIsInstance(chunk.created, int)
-      self.assertEqual(chunk.model, "test-model")
-
-  def test_stream_with_usage(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=True,
-      stream_options={"include_usage": True}
-    )
-
-    chunks = list(stream)
-    last_chunk = chunks[-1]
-
-    self.assertIsNotNone(last_chunk.usage)
-    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
-    self.assertIsNotNone(last_chunk.usage.completion_tokens)
-    self.assertIsNotNone(last_chunk.usage.total_tokens)
-
-  def test_multi_turn_conversation(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[
-        {"role": "system", "content": "You are helpful."},
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Hi!"},
-        {"role": "user", "content": "How are you?"}
-      ],
-      stream=True
-    )
-
-    chunks = list(stream)
-    self.assertGreater(len(chunks), 0)
-    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
-
-  def test_content_is_streamed(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=True
-    )
-
-    contents = []
-    for chunk in stream:
-      if chunk.choices and chunk.choices[0].delta.content:
-        contents.append(chunk.choices[0].delta.content)
-
-    self.assertGreater(len(contents), 0)
-
-  def test_non_streaming(self):
-    resp = self.client.chat.completions.create(
-      model="test-model",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=False
-    )
-
-    self.assertTrue(resp.id.startswith("chatcmpl-"))
-    self.assertEqual(resp.object, "chat.completion")
-    self.assertEqual(resp.model, "test-model")
-    self.assertIsNotNone(resp.created)
-    self.assertEqual(len(resp.choices), 1)
-    self.assertEqual(resp.choices[0].message.role, "assistant")
-    self.assertIsNotNone(resp.choices[0].message.content)
-    self.assertEqual(resp.choices[0].finish_reason, "stop")
-    self.assertIsNotNone(resp.usage)
-    self.assertIsNotNone(resp.usage.prompt_tokens)
-    self.assertIsNotNone(resp.usage.completion_tokens)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,31 +0,0 @@
-#!/usr/bin/env python
-import io, unittest
-import numpy as np
-from tinygrad import Tensor, fetch
-from tinygrad.nn.state import png_load
-try:
-  from PIL import Image
-except ImportError:
-  raise unittest.SkipTest("PIL not installed")
-
-class TestPNGLoad(unittest.TestCase):
-  def test_real_png(self):
-    # test against a real PNG file (uses only filters 0, 1)
-    fp = fetch('https://upload.wikimedia.org/wikipedia/en/d/d4/Norwegian_Forest_Cat_in_Norway.png')
-    with open(fp, 'rb') as f: png_bytes = f.read()
-    expected = np.array(Image.open(io.BytesIO(png_bytes)))[:, :, :3]
-    result = png_load(Tensor(np.frombuffer(png_bytes, dtype=np.uint8))).numpy()
-    np.testing.assert_array_equal(result, expected)
-
-  def test_roundtrip_png(self):
-    # horizontal stripes pattern uses only filters 0, 1
-    img_array = np.zeros((32, 32, 3), dtype=np.uint8)
-    img_array[::2] = 255  # white stripes on black
-    buf = io.BytesIO()
-    Image.fromarray(img_array).save(buf, format='PNG')
-    png_bytes = buf.getvalue()
-    result = png_load(Tensor(np.frombuffer(png_bytes, dtype=np.uint8))).numpy()
-    np.testing.assert_array_equal(result, img_array)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,47 +0,0 @@
-import unittest
-from tinygrad import Tensor, Variable
-from tinygrad.engine.schedule import schedule_cache
-
-class TestScheduleCache(unittest.TestCase):
-  def test_bound_variable_reuses_cache(self):
-    schedule_cache.clear()
-    v = Variable('v', 1, 100)
-    x = Tensor.ones(10).contiguous().realize()
-
-    # first run with v=5
-    t1 = (x + Tensor(v.bind(5))).sum()
-    self.assertEqual(t1.item(), 60.0)
-    cache_size_after_first = len(schedule_cache)
-
-    # second run with v=10 should reuse cache
-    t2 = (x + Tensor(v.bind(10))).sum()
-    self.assertEqual(t2.item(), 110.0)
-    self.assertEqual(len(schedule_cache), cache_size_after_first)
-
-  def test_bound_variable_var_vals(self):
-    v = Variable('pos', 1, 100)
-    x = Tensor.ones(10).contiguous().realize()
-
-    t = x + Tensor(v.bind(42))
-    _, var_vals = t.schedule_with_vars()
-    self.assertEqual(var_vals, {'pos': 42})
-
-  def test_simple(self):
-    a = Tensor.ones(10).contiguous()
-    b = Tensor.ones(10).contiguous()
-    Tensor.realize(a, b)
-
-    # warm up
-    for _ in range(2):
-      num = (a.sum().contiguous()+b.sum().contiguous()).item()
-      print(num)
-
-    # confirm schedule cache doesn't grow
-    start_len_schedule_cache = len(schedule_cache)
-    for _ in range(3):
-      num = (a.sum().contiguous()+b.sum().contiguous()).item()
-      print(num)
-    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -22,8 +22,7 @@ class TestRawShmBuffer(unittest.TestCase):

  @unittest.skipIf(CI, "CI doesn't like big shared memory")
  def test_e2e_big(self):
-    # bigger than this doesn't work on Linux, maybe this is a limit somewhere?
-    t = Tensor.randn(2048, 128, 8).realize()
+    t = Tensor.randn(2048, 2048, 8).realize()

    # copy to shm
    shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name
@@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
+from tinygrad.uop.symbolic import sym, commutative
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -112,9 +112,8 @@ class TestSymbolic(unittest.TestCase):
    b = Variable("b", 1, 8)
    self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
-    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2))
    self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
-    self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1)
+    self.assertEqual((a*a*3).divide_exact(a*a.const_like(-3)).simplify(), a*-1)
    self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
    self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
    self.assertEqual((a*b*3+a*b*b).divide_exact(a*b).simplify(), b+3)
@@ -323,12 +322,12 @@ class TestSymbolic(unittest.TestCase):
  def test_mod_mod_wrong_sign(self):
    v1=Variable("v1", 0, 128)
    v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -3, 4, "(v1%2*2+(v3+-1)%5+-2)")
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")

  def test_mod_mod_wrong_sign2(self):
    v2=Variable("v2", 0, 8)
    v3=Variable("v3", 0, 4)
-    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -2, 6, "(((v2+((v3+3)%7))+-2)%7)")
+    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")

  def test_mul_mul(self):
    self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
@@ -378,9 +377,9 @@ class TestSymbolic(unittest.TestCase):
  def test_big_mod(self):
    self.helper_test_variable(Variable("a", -20, 20)%10, -9, 9, "(a%10)")
    self.helper_test_variable(Variable("a", -20, 0)%10, -9, 0, "(((a*-1)%10)*-1)")
-    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 1, "(a%10)")
+    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 9, "(a%10)")  # TODO: tighter max
    self.helper_test_variable(Variable("a", 0, 20)%10, 0, 9, "(a%10)")
-    self.helper_test_variable(Variable("a", -1, 20)%10, -1, 9, "(a%10)")
+    self.helper_test_variable(Variable("a", -1, 20)%10, -9, 9, "(a%10)")  # TODO: tighter min

  def test_ge_remove(self):
    self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "False")
@@ -581,12 +580,12 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")
    self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")

+  @unittest.expectedFailure  # TODO: improve nest_div_by_smallest_factor
  def test_sum_div_complex4(self):
    gidx0 = Variable("gidx0", 0, 2)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 12)
-    # TODO: improve nest_div_by_smallest_factor to get ((lidx2+(lidx3*2))//3)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((gidx0+(lidx2*19+lidx3*38)//3)//19)")
+    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")

  def test_sum_mul_distribute(self):
    gidx0 = Variable("gidx0", 0, 7)
@@ -1030,25 +1029,6 @@ class TestSymbolicRealWorld(unittest.TestCase):
    self.assertIn(idx.render(),
      ("(lidx3+((lidx5+1)//16*802816+(lidx5+1)%16*49+gidx0*3211264+gidx1*784+gidx2*8+lidx4*100352)+2207744)",))

-class TestGatedUopGivenValid(unittest.TestCase):
-  def test_invalid_gate_simplifies_index(self):
-    r0 = Variable("r0", 0, 2)
-
-    idx:UOp = (r0 < 3).where((r0 + uconst(-1)) // uconst(3), UOp.invalid())
-    idx = graph_rewrite(idx, pm_simplify_valid)
-    self.assertEqual(idx, (r0 < 3).where(uconst(0), UOp.invalid()))
-
-  def test_invalid_gate_simplifies_vectorize(self):
-    r0 = Variable("r0", 0, 2)
-
-    idx0 = (r0 + uconst(-1)) // uconst(3)
-    idx1 = r0 % uconst(3)
-    idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.index.vec(2), (idx0, idx1)), UOp.invalid())
-    idx = graph_rewrite(idx, pm_simplify_valid)
-    # NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
-    expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
-    self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
-
 class TestBounds(unittest.TestCase):
  def test_unrolled_arange(self):
    # #include <metal_stdlib>
@@ -1,7 +1,6 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv
-from tinygrad.helpers import partition, TCPServerWithReuse, HTTPRequestHandler, DEBUG, Timing, GlobalCounters, stderr_log, colored
+import sys, argparse, typing, re, unicodedata
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers

 class SimpleTokenizer:
  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
@@ -10,7 +9,6 @@ class SimpleTokenizer:
    self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}

    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # TODO: ucat_range is slow
    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
    self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
@@ -26,7 +24,7 @@ class SimpleTokenizer:
    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
    if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'")
    vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
-    normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
+    normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens))

  def _encode_word(self, word:bytes) -> list[int]:
@@ -52,18 +50,15 @@ class SimpleTokenizer:
  def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode()
  def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")

-@functools.cache
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).contiguous()
-
-def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
+def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
  B, H, T, Hd = x.shape
  assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
-  x_pairs     = x.reshape(B, H, T, Hd//2, 2)
-  cos = freqs_cis.reshape(1, 1, T, Hd//2, 2)[..., 0]
-  sin = freqs_cis.reshape(1, 1, T, Hd//2, 2)[..., 1]
+  half = Hd // 2
+  t_start_pos = start_pos if isinstance(start_pos, int) else Tensor(start_pos)
+  angles = (Tensor.arange(T, dtype="float32") + t_start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
+  # contiguous here allows RoPE to be pruned in the JIT
+  cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous()
+  x_pairs = x.reshape(B, H, T, half, 2)
  return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin,
                      x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)

@@ -99,10 +94,8 @@ class TransformerBlock:
    k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
    v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)

-    # TODO: make UOp have SupportsIndex
-    freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context)[start_pos:start_pos+T]  # type: ignore
-    q = apply_rope(q, freqs_cis)
-    k = apply_rope(k, freqs_cis)
+    q = apply_rope(q, start_pos)
+    k = apply_rope(k, start_pos)

    # TODO: remove these kv cache realizes
    if not hasattr(self, "cache_kv"):
@@ -120,8 +113,7 @@ class TransformerBlock:

  def _feed_forward(self, h: Tensor) -> Tensor:
    h_norm = self.ffn_norm(h)
-    # TODO: remove the need for this contiguous
-    gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
+    gated  = self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)
    return h + self.ffn_down(gated)

  def __call__(self, x: Tensor, start_pos: int|UOp):
@@ -147,7 +139,7 @@ class Transformer:
    return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)

  @staticmethod
-  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=True) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None) -> tuple[Transformer, dict]:
    # TODO: remove the need for copy to default device
    kv, state_dict = nn.state.gguf_load(gguf.to(None))

@@ -164,8 +156,7 @@ class Transformer:
                        norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), max_context=max_context)
    nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
    # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
-    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
-    if realize: Tensor.realize(*params)
+    for s in nn.state.get_parameters(model): s.replace(s.contiguous())
    return model, kv

  def generate(self, tokens:list[int], start_pos=0):
@@ -181,92 +172,26 @@ class Transformer:
      yield next_id

 models = {
-  "llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
-  "llama3.2:3b": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
-  "llama3.2:3b-f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
-  "llama3.1:8b": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
+  "1B": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
+  "3B": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
+  "3B_f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
+  "8B": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
 }

-# *** simple OpenAI compatible server on 11434 to match ollama ***
-# OPENAI_BASE_URL=http://localhost:11434/v1 OPENAI_API_KEY=ollama uvx --from gpt-command-line gpt
-
-class Handler(HTTPRequestHandler):
-  def log_request(self, code='-', size='-'): pass
-  def run_model(self, ids:list[int], model_name:str, include_usage=False):
-    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  in:{len(ids):5d}  {colored('--', 'BLACK')}  ")
-    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
-    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
-    out: list[int] = []
-    st = time.perf_counter()
-    for next_id in model.generate(ids):
-      if len(out) == 0: stderr_log(f"prefill:{len(ids)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
-      if next_id == eos_id: break
-      out.append(next_id)
-      yield {"choices": [{"index":0, "delta":{"content":tok.decode([next_id])}, "finish_reason":None}], **tmpl}
-    yield {"choices": [{"index":0, "delta":{},"finish_reason":"stop"}], **tmpl}
-    if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
-    stderr_log(f"out:{len(out):5d}  {colored('--', 'BLACK')}  gen: {len(out)/(time.perf_counter()-pt):4.0f} tok/s\n")
-
-  def do_POST(self):
-    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
-    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
-    if DEBUG >= 1: print(json.dumps(body, indent=2))
-    if self.path == "/v1/chat/completions":
-      # extract tokens
-      ids = [bos_id]
-      for msg in body["messages"]:
-        ids += tok.role(msg["role"])
-        # content can be a str or a list
-        content = msg["content"]
-        if isinstance(content, str): ids += tok.encode(content)
-        elif isinstance(content, list):
-          for c in content:
-            if c["type"] == "text": ids += tok.encode(c["text"])
-            else: raise RuntimeError(f"unhandled type: {c['type']}")
-        else: raise RuntimeError(f"unknown content type: {type(content)}")
-        ids += tok.role("assistant")
-
-      # reply
-      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False))
-      if body.get("stream"): self.stream_json(chunks)
-      else:
-        out = []
-        for c in chunks: out.append(c["choices"][0]["delta"].get("content", "") if c["choices"] else "")
-        self.send_data(json.dumps({**c, "object":"chat.completion",
-          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":"stop"}]}).encode())
-    else:
-      raise RuntimeError(f"unhandled path {self.path}")
-
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
+  parser.add_argument("--size", choices=list(models.keys()), default=list(models.keys())[0], help="Model size")
  parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
-  parser.add_argument("--serve", action="store_true", help="Run OpenAI compatible API")
-  parser.add_argument("--benchmark", action="store_true", help="Benchmark tok/s")
  args = parser.parse_args()

  # load the model
-  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
-  if DEBUG >= 1: print(f"using model {args.model}")
-
-  # do benchmark
-  if args.benchmark:
-    param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
-    gen = model.generate([0], 0)
-    for _ in range(20):
-      GlobalCounters.reset()
-      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s, param {param_bytes/x:7.2f} GB/s"): next(gen)
-    exit(0)
+  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.size]), args.max_context)

  # extract some metadata
  tok = SimpleTokenizer.from_gguf_kv(kv)
  bos_id: int = kv['tokenizer.ggml.bos_token_id']
  eos_id: int = kv['tokenizer.ggml.eos_token_id']

-  # start server
-  if args.serve: TCPServerWithReuse(('', 11434), Handler).serve_forever()
-
  ids: list[int] = [bos_id]
  while 1:
    start_pos = len(ids) - 1
@@ -1,50 +0,0 @@
-# classification in 50 lines
-import sys
-from tinygrad import nn, Tensor
-
-class Bottleneck:
-  expansion = 4
-  def __init__(self, in_c, mid_c, stride=1):
-    out_c = mid_c * self.expansion
-    self.conv1, self.bn1 = nn.Conv2d(in_c, mid_c, 1, bias=False), nn.BatchNorm2d(mid_c)
-    self.conv2, self.bn2 = nn.Conv2d(mid_c, mid_c, 3, stride, 1, bias=False), nn.BatchNorm2d(mid_c)
-    self.conv3, self.bn3 = nn.Conv2d(mid_c, out_c, 1, bias=False), nn.BatchNorm2d(out_c)
-    self.downsample = (stride != 1 or in_c != out_c) and [nn.Conv2d(in_c, out_c, 1, stride, bias=False), nn.BatchNorm2d(out_c)] or []
-
-  def __call__(self, x:Tensor) -> Tensor:
-    identity = x.sequential(self.downsample)
-    x = self.bn1(self.conv1(x)).relu()
-    x = self.bn2(self.conv2(x)).relu()
-    x = self.bn3(self.conv3(x))
-    return (x + identity).relu()
-
-class ResNet50:
-  def __init__(self, num_classes=1000):
-    self.conv1, self.bn1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), nn.BatchNorm2d(64)
-    self.layer1 = self._make_layer(64,  64,  3, 1)
-    self.layer2 = self._make_layer(256, 128, 4, 2)
-    self.layer3 = self._make_layer(512, 256, 6, 2)
-    self.layer4 = self._make_layer(1024,512, 3, 2)
-    self.fc = nn.Linear(2048, num_classes)
-
-  def _make_layer(self, in_c, mid_c, blocks, stride):
-    layers = [Bottleneck(in_c, mid_c, stride)]
-    for _ in range(1, blocks): layers.append(Bottleneck(mid_c * Bottleneck.expansion, mid_c))
-    return layers
-
-  def __call__(self, x:Tensor) -> Tensor:
-    x = self.bn1(self.conv1(x)).relu()
-    # TODO: max_pool2d return type is Tensor | tuple[Tensor, Tensor], this should be type specialised
-    x = x.max_pool2d()  # type: ignore
-    x = x.sequential([*self.layer1, *self.layer2, *self.layer3, *self.layer4])
-    x = x.mean((2, 3))
-    return self.fc(x)
-
-if __name__ == "__main__":
-  test_url = "https://upload.wikimedia.org/wikipedia/en/d/d4/Norwegian_Forest_Cat_in_Norway.png"
-  img = nn.state.png_load(Tensor.from_url(sys.argv[1] if len(sys.argv) > 1 else test_url))
-  model = ResNet50()
-  state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/timm/resnet50.a1_in1k/resolve/main/model.safetensors"))
-  nn.state.load_state_dict(model, state_dict)
-  value = model(img.rearrange("h w c -> 1 c h w").float()/255).argmax().item()
-  print(value, nn.datasets.imagenet_labels()[value])
@@ -7,7 +7,6 @@ from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
-from tinygrad.helpers import ALLOW_TF32
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
@@ -229,7 +228,6 @@ class Scheduler:
      except IndexError:
        raise KernelOptError(f"invalid tensor core choice {tc_select}")
      for tc in tensor_cores:
-        if self.ren.device in ("CUDA", "NV") and tc.dtype_in == dtypes.float and not ALLOW_TF32: continue
        if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
          # tensor cores have three ranges. X, Y, and REDUCE
          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
@@ -92,9 +92,9 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
 # *** external API ***

 # get dictionary of all possible actions
-def get_kernel_actions(s:Scheduler, include_0=True) -> dict[int, Scheduler]:
+def get_kernel_actions(s:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
  acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
-  kernel_actions = actions.copy()
+  kernel_actions = (actions if candidates is None else candidates).copy()

  for i,a in enumerate(kernel_actions):
    if a.axis is not None and a.op is not OptOps.TC:
@@ -1,6 +1,7 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
+from tinygrad.helpers import getenv

@dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -91,7 +92,8 @@ cuda_8168_tf32 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2
  swizzle=((('r0', 'r1', 'l2', 'l3', 'l4'), ('u1', 'r2'), ('l0', 'l1', 'u0')),
           (('r0', 'r1', 'u0', 'l0', 'l1'), ('u1', 'r2'), ('l2', 'l3', 'l4'))))]
 cuda_sm75: list[TensorCore] = cuda_8168_f16
-cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16 + cuda_8168_tf32
+cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16
+if getenv("ALLOW_TF32", 0): cuda_sm80 += cuda_8168_tf32
 cuda_sm89: list[TensorCore] = cuda_sm80 + cuda_81632_f8

 # ***** AMD *****
@@ -88,8 +88,7 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
  # lift x+y out of reduce on lt
  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
  # lift x*y out of reduce
-  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
+  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
  # fold the range
  # bound from below
  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),
@@ -1,11 +1,11 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator
+from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
-from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
-from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
+from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
+from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer

@@ -277,49 +277,29 @@ class Compiler:
    return lib
  def disassemble(self, lib:bytes): pass

-@dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial; ctrl_var:ContextVar|None = None # noqa: E702
-
-@dataclass(frozen=True)
-class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
-
+CompilerPairT = tuple[functools.partial|type[Renderer], functools.partial|type[Compiler]]
 class Compiled:
  profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.

-  def __init__(self, device:str, allocator:Allocator, compilers:CompilerSet|None, runtime, graph=None, group_id=None):
+  def __init__(self, device:str, allocator:Allocator, compilers:Sequence[CompilerPairT]|None, runtime, graph=None, group_id=None):
    self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
+    self.compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])

-    self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler]] = {}
-    for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+    envnames = [self._get_compiler_envvar(c) for r,c in self.compilers]
+    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 1)
+    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 0)

-  @property
-  def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
+    if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
+    for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)

-  @property
-  def compiler(self) -> Compiler: return self._select_compiler_pair()[1]
+    self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers,
+                                                       f"No compiler for {self.device} is available")

-  def _compiler_name(self, c:type[Compiler]|functools.partial) -> str:
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
+    if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")

-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler]:
-    # select forced compiler from global env var.
-    forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
-
-    # add forced compilers from individual env vars.
-    forced_comps |= set(rc for en, rc in self.comp_sets.values() if en is not None and en.value == 1)
-    if len(forced_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {forced_comps}")
-
-    # select remaining compilers (all or forced only)
-    comps = list(rc for en, rc in self.comp_sets.values())
-
-    # remove disabled compilers
-    for en, rc in self.comp_sets.values():
-      if en is not None and en.value == 0 and rc in comps: comps.remove(rc)
-
-    return select_first_inited(list(forced_comps) if len(forced_comps)>0 else comps, f"No compiler for {self.device} is available", self.cached_pair)
+  def _get_compiler_envvar(self, c):
+    compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
+    return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"

  def synchronize(self):
    """
@@ -345,13 +325,11 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
  if device is None: device = Device.DEFAULT
  if dtype == dtypes.bfloat16:
    if device == "METAL": return not CI
-    if device == "CUDA": return not CI and not CUDA_PTX
-    if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
+    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP")
    return device in {"AMD", "PYTHON", "NULL"}
  if dtype in dtypes.fp8s:
-    if device == "CUDA": return not CI and not CUDA_PTX
-    if device == "NV": return not CI and not NV_PTX and not NV_NAK
+    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
    if device == "AMD": return not CI and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
    return device in {"PYTHON", "NULL"}
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
@@ -366,7 +344,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device in ["CUDA", "NV"]: return not CI
    if device == "CPU" and CPU_LLVM: return OSX
    if device == "PYTHON": return sys.version_info >= (3, 12)
-  if dtype == dtypes.float64: return device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
+  if dtype == dtypes.float64: return device != "METAL" and not (OSX and device == "CL")
  return True

 if PROFILE:
@@ -389,23 +367,20 @@ def enumerate_devices_str() -> Generator[str, None, None]:
    compilers_results, any_works = [], False
    try:
      d = Device[device]
-      default_comp_pairs, default_compiler, cc_ctrl_var = d.comp_sets, d.compiler, d.comps_ctrl_var
+      default_renderer, default_compiler = d.renderer, d.compiler
      try:
-        for k,(en,(r,c)) in default_comp_pairs.items():
-          d.comp_sets = {k:(None,(r,c))} # env var set to None, so it doesn't interfere
-          d.comps_ctrl_var = None
+        for r,c in d.compilers:
          try:
-            # d.renderer, d.compiler = r(), c()
+            d.renderer, d.compiler = r(), c()
            with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
            if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(c)} to make default)' if cc_ctrl_var is not None else ''
-            default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
+            default_text = '(default)' if type(default_compiler) is type(d.compiler) else f'({d._get_compiler_envvar(c)}=1 to make default)'
            compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
            any_works = True
          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
      finally:
        # put the defaults back!
-        d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var
+        d.renderer, d.compiler = default_renderer, default_compiler
      result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
    except Exception as e:
      result = f"{colored('FAIL', 'red')} {e}"
@@ -5,7 +5,7 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv,
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
-from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
+from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
 from tinygrad.schedule.rangeify import mop_cleanup
@@ -143,7 +143,7 @@ class MultiGraphRunner(GraphRunner):

 def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
  if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]
-  if isinstance(ei.prg, (BufferCopy, BufferXfer, EncDec)): return [cast(Buffer, ei.bufs[0])]
+  if isinstance(ei.prg, (BufferCopy, BufferXfer)): return [cast(Buffer, ei.bufs[0])]
  return []

 def update_depends(depends:set[Buffer|None], jit_cache:list[ExecItem]):
@@ -318,7 +318,7 @@ class TinyJit(Generic[ReturnType]):

      # memory planning (optional)
      # Exclude buffers involved in transfer ops to preserve parallelism.
-      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec)) for b in ji.bufs}
+      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy)) for b in ji.bufs}
      assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], noopt_buffers, debug_prefix="JIT ")
      jit_cache = [ExecItem(item.prg, [assigned.get(b,b).ensure_allocated() for b in item.bufs if b is not None],
                            item.metadata, item.fixedvars) for item in jit_cache]
@@ -14,7 +14,7 @@ from tinygrad.codegen.opt import Opt
 # **************** Program Creation ****************

@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
  """
  Transform an AST into a ProgramSpec. May trigger BEAM search.

@@ -30,6 +30,7 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
  if DEBUG >= 5: print(pyrender(ast))

  # linearize
+  if renderer is None: renderer = Device.default.renderer
  if opts is not None:
    assert ast.arg is None, "can't apply opts if sink has an arg"
    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
@@ -2,11 +2,10 @@ import time
 from typing import cast
 from dataclasses import dataclass, field, replace
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites
-from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
-from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
+from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten

 # **** ScheduleItem return type

@@ -20,11 +19,12 @@ class ScheduleItem:

 # **** schedule linearizer

-def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
+def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[str, int]]:
  with cpu_profile(TracingKey("toposort sched_sink")):
    # construct the KERNEL children graph based on assigns
    children: dict[UOp, list[UOp]] = {}
    in_degree: dict[UOp, int] = {}
+    var_vals: dict[str, int] = {}
    for u in sched_sink.toposort():
      if u.op is Ops.RANGE:
        in_degree.setdefault(u, 0)
@@ -43,8 +43,14 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
              assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
              children.setdefault(ss.src[1], []).append(k)
              in_degree[k] += 1
-        elif s.op in {Ops.BUFFER, Ops.BIND}:
-          pass  # a BUFFER is already realized, BINDs are handled in complete_create_schedule_with_vars
+        elif s.op is Ops.BUFFER:
+          pass  # a BUFFER is already realized, nothing to do here
+        elif s.op is Ops.BIND:
+          # for RANGE this is in fixedvars
+          if s.src[1].op is not Ops.RANGE:
+            var, val = s.unbind()
+            assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+            var_vals[var.expr] = val
        else:
          raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")

@@ -66,7 +72,7 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
          assert isinstance(base, Buffer), "base can't be MultiBuffer"
          buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
        ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
-        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
+        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and s.src[1].op is Ops.RANGE)
        if any(isinstance(x, MultiBuffer) for x in ubufs):
          assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
          dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
@@ -101,106 +107,37 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
      else:
        real_schedule.append(replace(si, fixedvars=si.fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in si.bound_ranges}, bound_ranges=()))
      sched_ptr += 1
-  return real_schedule
+  return real_schedule, var_vals

 from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map

-def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
-  if (ret:=ctx.get(b, None)) is None:
-    if b.op is Ops.BUFFER:
-      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
-    else:
-      # TODO: flip args in CONST
-      assert b.op is Ops.CONST
-      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
-  return ret
-
-pm_pre_sched_cache = PatternMatcher([
-  # replace input buffers
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
-  # remove unique consts
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
-  # strip value from BIND for cache key normalization, so different values hit same cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
-])
-
-def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
-  if (ret:=ctx.get(b, None)) is None:
-    assert b.op is Ops.BUFFER
-    # if it's not in the cache, create a new buffer
-    ctx[b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
-  return ret
-
-pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer_back),
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.LUNIQUE)), name="b"), replace_input_buffer_back),
-  # restore BIND value stripped in pm_pre_sched_cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)),
-])
-
-schedule_cache: dict[bytes, tuple[UOp, UOp]] = {}
-@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
 def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ScheduleItem], dict[str, int]]:
  # big_sink srcs are all the Tensors
  st = time.perf_counter()

-  # replace all UNIQUE buffers with LUNIQUE, strip BIND values for cache key
-  input_buffers: dict[UOp, UOp] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
-  sched_cache_key = big_sink_cache.key
+  # verify Tensors match the spec
+  if SPEC: type_verify(big_sink, tensor_spec)

-  if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
-    # verify Tensors match the spec (on big_sink, we only need to do this if cache misses)
-    if SPEC: type_verify(big_sink, tensor_spec)
+  # tensor map is what we return
+  tensor_map: dict[UOp, UOp] = {}

-    # hack to preserve metadata
-    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
+  if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
+    tensor_map |= get_multi_map(big_sink)
+    big_sink = big_sink.substitute(tensor_map, name="Apply Multi Map")
+    big_sink = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink.src]))

-    # tensor map is what we return
-    tensor_map: dict[UOp, UOp] = {}
-
-    if any(isinstance(x._device, tuple) for x in big_sink_cache.toposort()):
-      tensor_map |= get_multi_map(big_sink_cache)
-      big_sink_cache = big_sink_cache.substitute(tensor_map, name="Apply Multi Map")
-      big_sink_cache = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink_cache.src]))
-
-    tensor_map |= get_rangeify_map(big_sink_cache)
-    big_sink = big_sink_cache.substitute(tensor_map, name="Apply Kernelize Map")
-
-    # save in schedule cache
-    tensor_map_sink = UOp.sink(*flatten([(k,v) for k,v in tensor_map.items()]))
-    schedule_cache[sched_cache_key] = (big_sink, tensor_map_sink)
-  else:
-    # schedule cache hit
-    del big_sink_cache
-    big_sink, tensor_map_sink = sc_ret
-
-  # replace all the LUNIQUEs with UNIQUEs
-  input_buffers_reverse = {v:k for k,v in input_buffers.items()}
-  big_sink = graph_rewrite(big_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for sched cache")
-  tm_src = graph_rewrite(tensor_map_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for tensor map").src
-  tensor_map = {tm_src[i]:tm_src[i+1] for i in range(0, len(tm_src), 2)}
+  tensor_map |= get_rangeify_map(big_sink)
+  big_sink = big_sink.substitute(tensor_map, name="Apply Kernelize Map")

  # create the schedule
-  schedule = create_schedule(big_sink)
+  schedule, var_vals = create_schedule_with_vars(big_sink)
  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)

-  # extract var_vals from BINDs that were stripped (only if there are kernels)
-  var_vals: dict[str, int] = {}
-  if schedule:
-    for u in input_buffers:
-      if u.op is Ops.BIND:
-        var, val = u.unbind()
-        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
-        var_vals[var.expr] = val
-
  # remove all AFTERs, after scheduling, the tensors are just buffers
  tensor_map |= {u:u.buf_uop for u in big_sink.toposort() if u.op is Ops.AFTER}

  if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
-    print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
-          f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
-          f" | {len(UOpMetaClass.ucache)} uops in cache")
+    print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms ({len(UOpMetaClass.ucache)} uops in cache)")
  return tensor_map, schedule, var_vals
@@ -1,9 +1,8 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
-import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, socketserver, json
+import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
-from http.server import BaseHTTPRequestHandler

 T = TypeVar("T")
 U = TypeVar("U")
@@ -115,14 +114,10 @@ def suppress_finalizing(func):
      if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
  return wrapper

-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg:str, cache:dict|None=None) -> tuple[T,...]|T:
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T:
  excs = []
  for typ in candidates:
-    if cache is not None and typ in cache: return cache[typ]
-    try:
-      x = tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
-      if cache is not None: cache[typ] = x
-      return x
+    try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
    except Exception as e: excs.append(e)
  raise ExceptionGroup(err_msg, excs)

@@ -149,10 +144,6 @@ def getenv(key:str, default:Any=0): return type(default)(os.getenv(key, default)
 def temp(x:str, append_user:bool=False) -> str:
  return (pathlib.Path(tempfile.gettempdir()) / (f"{x}.{getpass.getuser()}" if append_user else x)).as_posix()

-def stderr_log(msg):
-  sys.stderr.write(msg)
-  sys.stderr.flush()
-
 class Context(contextlib.ContextDecorator):
  def __init__(self, **kwargs): self.kwargs = kwargs
  def __enter__(self):
@@ -182,19 +173,14 @@ WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1),
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
-LRU = ContextVar("LRU", 1)
+PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
-# Compilers
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-NV_PTX, CUDA_PTX, NV_NAK, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0)
-NULL_IR3, NULL_NAK = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0)
-AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
-QCOM_CC = ContextVar("QCOM_CC", "")
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", VIZ.value)
@@ -207,8 +193,6 @@ DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
 TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
 # set to 0 to disable the compiler cache
 CCACHE = ContextVar("CCACHE", 1)
-# allow tf32 to be used on NVIDIA GPUs
-ALLOW_TF32 = ContextVar("ALLOW_TF32", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -388,16 +372,14 @@ def _ensure_downloads_dir() -> pathlib.Path:
  return pathlib.Path(cache_dir) / "downloads"

 def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False,
-          allow_caching=not getenv("DISABLE_HTTP_CACHE"), headers:dict[str, str]={}) -> pathlib.Path:
+          allow_caching=not getenv("DISABLE_HTTP_CACHE")) -> pathlib.Path:
  if url.startswith(("/", ".")): return pathlib.Path(url)
  if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name)
-  else:
-    hh = "_"+hashlib.md5(("\n".join(f"{k.strip()}:{v.strip()}" for k,v in sorted(headers.items()))).encode("utf-8")).hexdigest() if headers else ""
-    fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else ""))
+  else: fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + (".gunzip" if gunzip else ""))
  if not fp.is_file() or not allow_caching:
    (_dir := fp.parent).mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.11.0", **headers}), timeout=10) as r:
-      assert r.status in {200, 206}, r.status
+    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.11.0"}), timeout=10) as r:
+      assert r.status == 200, r.status
      length = int(r.headers.get('content-length', 0)) if not gunzip else None
      readfile = gzip.GzipFile(fileobj=r) if gunzip else r
      progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)
@@ -409,33 +391,6 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
      if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
  return fp

-# NOTE: using HTTPServer forces a potentially slow socket.getfqdn
-class TCPServerWithReuse(socketserver.TCPServer):
-  allow_reuse_address = True
-  def __init__(self, server_address, RequestHandlerClass):
-    print(f"*** started server on http://127.0.0.1:{server_address[1]}")
-    super().__init__(server_address, RequestHandlerClass)
-
-class HTTPRequestHandler(BaseHTTPRequestHandler):
-  def send_data(self, data:bytes, content_type:str="application/json", status_code:int=200):
-    self.send_response(status_code)
-    self.send_header("Content-Type", content_type)
-    self.send_header("Content-Length", str(len(data)))
-    self.end_headers()
-    return self.wfile.write(data)
-  def stream_json(self, source:Generator):
-    try:
-      self.send_response(200)
-      self.send_header("Content-Type", "text/event-stream")
-      self.send_header("Cache-Control", "no-cache")
-      self.end_headers()
-      for r in source:
-        self.wfile.write(f"data: {json.dumps(r)}\n\n".encode("utf-8"))
-        self.wfile.flush()
-      self.wfile.write("data: [DONE]\n\n".encode("utf-8"))
-    # pass if client closed connection
-    except (BrokenPipeError, ConnectionResetError): return
-
 # *** Exec helpers

 def system(cmd:str, **kwargs) -> str:
@@ -1,4 +1,3 @@
-import ast
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import tar_extract

@@ -13,8 +12,3 @@ def cifar(device=None):
  train = Tensor.cat(*[tt[f"cifar-10-batches-bin/data_batch_{i}.bin"].reshape(-1, 3073).to(device) for i in range(1,6)])
  test = tt["cifar-10-batches-bin/test_batch.bin"].reshape(-1, 3073).to(device)
  return train[:, 1:].reshape(-1,3,32,32), train[:, 0], test[:, 1:].reshape(-1,3,32,32), test[:, 0]
-
-def imagenet_labels():
-  return ast.literal_eval(Tensor.from_url(
-    "https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt"
-  ).tobytes().decode())
@@ -498,15 +498,13 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)

  # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
-  def _onnx_pads_to_tiny_pads(pads):
-    n = len(pads) // 2
-    return tuple(x for i in range(n-1, -1, -1) for x in (pads[i], pads[i+n]))
+  def _onnx_pads_to_tiny_pads(pads): return tuple(flatten(reversed(list(zip(pads, pads[len(pads)//2:])))))

  AUTO_PAD_OPTIONS = Literal["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"]
  # (padding_height, padding_width) -> (padding_top, padding_left, padding_bottom, padding_right)
  def _auto_pad(pads, auto_pad: AUTO_PAD_OPTIONS):
-    first = [p//2 for p in pads] if auto_pad == "SAME_UPPER" else [p - p//2 for p in pads]
-    return first + [p - f for p, f in zip(pads, first)]
+    if auto_pad == "SAME_UPPER": return [pads[i]//2 for i in range(len(pads))] + [pads[i]-pads[i]//2 for i in range(len(pads))]
+    return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))]

  def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS):
    if auto_pad == "VALID": return [0]*(len(k_)*2)
@@ -628,7 +626,6 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def ThresholdedRelu(X:Tensor, alpha:float=1.0): return (X > alpha).where(X, 0)
  def LogSoftmax(x: Tensor, axis:int=-1): return x.log_softmax(axis)
  def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
-  def Swish(x:Tensor, alpha:float=1.0): return x * (x * alpha).sigmoid()

  # ***** Unary Ops (broadcasted) *****
  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
@@ -649,7 +646,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def Mod(x:Tensor,y:Tensor,fmod=0): return x - x.div(y, rounding_mode="trunc") * y if fmod else x % y

  # ***** Casting Ops *****
-  # TODO: saturate parameter is ignored in Cast and CastLike
+  # TODO: saturate
  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_fallback(OnnxDataType(to).to_dtype(), "Cast op"))
  def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)

@@ -701,8 +698,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
  def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
    axes = axes or list(range(data.ndim))
-    steps = steps or [1] * data.ndim
-    slices = [slice(None)] * data.ndim
+    steps = steps or [1]*data.ndim
+    slices = [slice(0,x,1) for x in data.shape]
    for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
    return data[tuple(slices)]

@@ -713,7 +710,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
          mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
-    value = _resolve_const(constant_value or value)
+    value = constant_value or value
    axes = axes or list(range(x.ndim))
    real_pads = [0] * (x.ndim*2)
    for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
@@ -812,7 +809,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

    input_shape = cast(tuple[int, ...], X.shape[2:])
    if scales is not None: assert all(sc==1 for sc in scales[:-len(input_shape)]), "resizing batch_size dim or channel dim not supported"
-    if sizes is not None: assert tuple(sizes[:-2]) == tuple(X.shape[X.ndim-len(sizes):-2]), "resizing batch_size dim or channel dim not supported"
+    if sizes is not None: assert tuple(sizes[:-2]) == tuple(X.shape[X.ndim-len(sizes):-2]),  "resizing batch_size dim or channel dim not supported"

    scales, sizes = (None if scales is None else scales[-len(input_shape):]), (None if sizes is None else sizes[-len(input_shape):])
    if sizes is not None:
@@ -936,8 +933,11 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    # https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
    assert (segment_ids is None) is (segment_embedding is None)
    assert mask is None and not mask_index_type, "functionality not supported yet"  # TODO
-    input_shape, seq_length = input_ids.shape, input_ids.shape[1]
+    input_shape = input_ids.shape
+    seq_length = input_shape[1]
+    compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
    vocab_size, max_position_embeddings = word_embedding.shape[0], position_embedding.shape[0]
+    type_vocab_size  = (segment_embedding.shape[0] if compute_seg_emb else None)

    def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
      return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
@@ -946,9 +946,10 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape)
    wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding)
    pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding)
+    seg_embedding_res = embedding(segment_ids, type_vocab_size, segment_embedding) if compute_seg_emb else None

    embedding_sum = wrd_embedding_res + pos_embedding_res
-    if segment_embedding is not None: embedding_sum = embedding_sum + embedding(segment_ids, segment_embedding.shape[0], segment_embedding)
+    if seg_embedding_res is not None: embedding_sum = embedding_sum + seg_embedding_res
    out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
    return out, None, embedding_sum
  def MeanVarianceNormalization(x:Tensor, axis:list[int]|None=None):
@@ -1002,7 +1003,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)

  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
-                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None, do_rotary:int=0, mask_filter_value:float=-10000.0,
+                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
    assert not do_rotary and not attention_bias, "TODO"
@@ -1157,7 +1158,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    b_idx = Tensor.arange(b, device=x.device).reshape(b, *(1,)*(indices.ndim - 2)).expand(*indices.shape[:-1])
    ret = x[(b_idx,) + tuple(i.squeeze(-1) for i in indices.split(1, -1))]
    return ret.reshape(*x_shape[:batch_dims], *i_shape[batch_dims:-1], *ret.shape[indices.ndim-1:])
-  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul", "max", "min"]='none'):
+  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul"]='none'):
    assert updates.shape == indices.shape[:-1] + x.shape[cast(int, indices.shape[-1]):]
    x = x.contiguous()
    for index, u in zip(indices.split(1, 0), updates.split(1, 0)):
@@ -1166,8 +1167,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
      if reduction == "none": x[i] = u
      elif reduction == "add": x[i] += u
      elif reduction == "mul": x[i] *= u
-      elif reduction == "max": x[i] = x[i].maximum(u)
-      elif reduction == "min": x[i] = x[i].minimum(u)
+      else: raise NotImplementedError("reduction doesn't support max or min")
    return x

  def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Literal["none", "add", "mul", "min", "max"]="none"):
@@ -1286,7 +1286,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    # Tensor ops
    **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
    "Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
-    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Round", "Erf")},
+    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh",  "Elu", "Celu", "Selu", "Round", "Erf")},
    # Implemented ops
    **{name:obj for name,obj in locals().items() if isinstance(obj, types.FunctionType) and not name.startswith("_") and name[0].isupper()},
    # Version ops
@@ -9,12 +9,13 @@ class Optimizer:
  Base class for all optimizers.
  """
  def __init__(self, params: list[Tensor], lr: float, fused=FUSE_OPTIM):
-    # if requires_grad is None, but being put into an optimizer, set it to True
+    # if it's None, but being put into an optimizer, set it to True
    for x in params:
      if x.requires_grad is None: x.requires_grad = True

    self.params: list[Tensor] = dedup([x for x in params if x.requires_grad])
    assert len(self.params) != 0, "optimizer must have at least one param"
+    self.device = self.params[0].device
    self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized
    self.fused = fused
    # store lr in at least float32 precision
@@ -22,9 +23,6 @@ class Optimizer:
                     dtype=least_upper_dtype(dtypes.default_float, dtypes.float32))
    if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))

-  @property
-  def device(self): return self.params[0].device
-
  def _new_optim_param(self) -> list[Tensor]:
    param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False).contiguous()]
@@ -51,7 +49,6 @@ class Optimizer:
                - help: Consider setting Tensor.training=True before calling Optimizer.step().""")
    if self.fused:
      # optimizer fusion just concatenates all the buffers, runs the _step, then splits them back up
-      # NOTE: contiguous is for speed
      out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
                              [Tensor.cat(*[unwrap(t.grad).contiguous().flatten() for t in self.params], dim=0)])
      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
@@ -1,4 +1,4 @@
-import json, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
+import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -151,8 +151,7 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
        if DEBUG >= 1: print(f"WARNING: not loading {k}")
        continue
      if v.shape != state_dict[k].shape:
-        if {(), (1,)} == {state_dict[k].shape, v.shape}: state_dict[k] = state_dict[k].reshape(v.shape)
-        else: raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
+        raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
      if isinstance(v.device, tuple):
        if isinstance(state_dict[k].device, tuple): v.replace(state_dict[k])
        else: v.replace(state_dict[k].shard(v.device, v.uop.axis))
@@ -162,27 +161,6 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
      ret.append(v)
  return ret

-@accept_filename
-def zip_extract(t: Tensor) -> dict[str, Tensor]:
-  files: dict[str, Tensor] = {}
-  file_offsets: dict[str, tuple[Tensor, int, int]] = {}
-  with zipfile.ZipFile(TensorIO(t), "r") as myzip:
-    for zi in myzip.filelist:
-      file_offset = zi.header_offset+30+t[zi.header_offset+26:zi.header_offset+30].bitcast(dtypes.uint16).to("CPU").sum()
-      file_offsets[zi.filename] = (file_offset, zi.compress_size, zi.compress_type)
-  # sadly, the extra length needs to be read from the local header of each file. this is a limitation of the zip file format
-  Tensor.realize(*[x[0] for x in file_offsets.values()])
-  for filename, (file_offset, compress_size, compress_type) in file_offsets.items():
-    # possible to remove this realize/item? it's slow
-    file_offset_int = int(file_offset.item())
-    files[filename] = t[file_offset_int:file_offset_int+compress_size]
-    match compress_type:
-      case zipfile.ZIP_STORED: pass
-      # TODO: we need a zlib UOp so this can be lazy
-      case zipfile.ZIP_DEFLATED: files[filename] = Tensor(zlib.decompress(files[filename].data(), -15))
-      case _: raise NotImplementedError(f"compression {compress_type} not supported")
-  return files
-
@accept_filename
 def tar_extract(t: Tensor) -> dict[str, Tensor]:
  """
@@ -201,7 +179,6 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:

 # torch support!

-# TODO: this should use tar_extract and zip_extract
@accept_filename
 def torch_load(t:Tensor) -> dict[str, Tensor]:
  """
@@ -383,24 +360,3 @@ def gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
  for name, dims, typ, off in t_infos: state_dict[name] = ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims))

  return kv_data, state_dict
-
-@accept_filename
-def png_load(t:Tensor) -> Tensor:
-  f = io.BufferedReader(TensorIO(t))
-  assert f.read(8) == b'\x89PNG\r\n\x1a\n', "not a PNG"
-  idats = []
-  while (slen:=f.read(4)):
-    typ, dat = f.read(4), f.read(struct.unpack(">I", slen)[0])
-    if DEBUG >= 3: print(len(dat), typ)
-    if typ == b'IHDR':
-      width, height, depth, color_type = struct.unpack(">IIBB", dat[:10])
-      assert depth == 8 and color_type in [2, 6], f"only 8-bit RGB/RGBA PNG supported {depth=} {color_type=}"
-      bpp = 3 if color_type == 2 else 4
-    if typ == b'IDAT': idats.append(dat)
-    f.seek(4, 1)
-  data = Tensor(zlib.decompress(b''.join(idats))).reshape(height, width * bpp + 1)
-  filters, pixels = data[:, 0], data[:, 1:].reshape(height, width, bpp)
-  assert filters.max().item() <= 1, f"only PNG filters 0/1 supported, got {set(filters.tolist())}"  # type: ignore[arg-type]
-  # Sub filter (type 1): each pixel adds the pixel to its left, which is cumsum along width
-  pixels = (filters == 1).reshape(height, 1, 1).where(pixels.cast(dtypes.int16).cumsum(axis=1).bitwise_and(0xff).cast(dtypes.uint8), pixels)
-  return pixels[:, :, :3]
@@ -224,12 +224,8 @@ class ClangRenderer(CStyleLanguage):
                 Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
                 Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
                 Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
-
-  # LLVM legalizes double => half/bf16 cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
-  # there is also no native bfl16 <-> fp16 conversion on those CPUs
+  # LLVM legalizes double => half cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
  extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
-                                 (UPat.var("x", dtypes.float64).cast(dtypes.bfloat16), lambda x: x.cast(dtypes.float32).cast(dtypes.bfloat16)),
-                                 (UPat.var("x", dtypes.bfloat16).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
    (UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu)]) + CStyleLanguage.extra_matcher

  if sys.platform == 'win32':
@@ -377,7 +373,6 @@ class CUDARenderer(CStyleLanguage):
  code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+int(x))}", "l": lambda x: f"threadIdx.{chr(120+int(x))}",
                       "i": lambda x: f"(blockIdx.{chr(120+int(x))}*blockDim.{chr(120+int(x))}+threadIdx.{chr(120+int(x))})"}
  code_for_op = { **CStyleLanguage.code_for_op,
-    Ops.TRUNC: lambda x,dtype: f"htrunc({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"trunc({x})",
    Ops.SIN: lambda x,dtype: f"hsin({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sin({x})",
    Ops.LOG2: lambda x,dtype: f"hlog2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"log2({x})",
    Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})",
@@ -1,18 +1,21 @@
 from typing import Callable, cast, Any
-from tinygrad.dtype import AddrSpace, DType, PtrDType, ImageDType, dtypes
+from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
 from tinygrad.helpers import DEBUG, OSX, unwrap, charptr
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa
-import base64, ctypes, ctypes.util, struct, functools, inspect, contextlib, itertools
+import base64, contextlib, ctypes, ctypes.util, struct, functools, inspect

 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))

-def glsl_type(t:DType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents if isinstance(t, PtrDType) else {
-  **{getattr(dtypes,k):g(f"glsl_type_builtin_{v}") for k,v in [('double','double'),('float','float'),('float16','float16_t'),('bool','uint8_t')]},
-  **{d:g(f"glsl_type_builtin_{'u' * (d in dtypes.uints)}int{str(d.itemsize*8)+'_t' if d.itemsize != 4 else ''}") for d in dtypes.ints}}[t]
+# this is a ridiculous hack, but I can't find a better way to grab the glsl_type objects
+glsl_base = {**{d:g(f"GLSL_TYPE_{'U' if d in dtypes.uints else ''}INT{d.itemsize*8 if d.itemsize != 4 else ''}") for d in dtypes.ints},
+             **{getattr(dtypes,d):g(f"GLSL_TYPE_{d.upper()}") for d in ['double', 'float', 'float16']}, dtypes.bool: mesa.GLSL_TYPE_UINT8}
+def glsl_type(t:DType) -> mesa.struct_glsl_type:
+  if isinstance(t, PtrDType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents
+  return mesa.glsl_get_base_glsl_type(mesa.glsl_type(base_type=glsl_base[t])).contents

 # alu ops, aop[<dtype>][<op>]
 u_aop = { Ops.ADD: "iadd", Ops.MUL: "imul", Ops.IDIV: "udiv", Ops.MOD: "umod", Ops.CMPLT: "ult", Ops.CMPNE: "ine", Ops.CMPEQ: "ieq", Ops.OR: "ior",
@@ -49,7 +52,7 @@ def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=N
      if has_def: mesa.nir_def_init(instr.contents.instr, getattr(instr.contents, "def"), go(nc), go(bs))
      for k, v in go(intrins or {}).items():
        idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic.value].index_map[g(f"NIR_INTRINSIC_{k}")]
-        assert idx > 0, "invalid intrinsic. mesa version mismatch?"
+        assert idx > 0
        instr.contents.const_index[idx - 1] = go(v)
      for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src)
      for k,v in {k:vcomp for k,v in contents.items() if (vcomp:=go(v)) is not None}.items(): setattr(instr.contents, k, go(v))
@@ -67,16 +70,11 @@ def nchannel(b:mesa.nir_builder, src:mesa.nir_def, c:int):
  ctypes.cast(mov.contents.src, ctypes.POINTER(mesa.nir_alu_src))[0] = alu_src
  return mov

-def nimm_set(imm:mesa.nir_def, x, dtype:DType):
-  instr = ctypes.cast(imm.parent_instr, ctypes.POINTER(mesa.nir_load_const_instr))
-  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, x)
-
@nir_instr(nc=1, bs=lambda dtype: 1 if dtype == dtypes.bool else dtype.itemsize * 8)
 def nimm(b:mesa.nir_builder, x, dtype:DType) -> mesa.nir_def:
-  nimm_set(getattr((instr:=mesa.nir_load_const_instr_create(b.shader, 1, 1 if dtype==dtypes.bool else dtype.itemsize * 8)).contents, "def"), x, dtype)
+  instr = mesa.nir_load_const_instr_create(b.shader, 1, 1 if dtype == dtypes.bool else dtype.itemsize * 8)
+  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, x)
  return instr
-@nir_instr(nc=1, bs=lambda dtype: 1 if dtype == dtypes.bool else dtype.itemsize * 8)
-def nundef(b, dtype): return mesa.nir_undef_instr_create(b.shader, 1, 1 if dtype == dtypes.bool else dtype.itemsize * 8)

 deref_var = nir_instr(nc=1, bs=32, modes=lambda var:var.data.mode, type=lambda var:var.type, var=lambda var:ctypes.pointer(var))( # pylint: disable=W0108
  lambda b, var: mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_var))
@@ -92,8 +90,6 @@ nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.itemsize*8/

 ngid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_id))
 nlid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_local_invocation_id))
-ngsz = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_size))
-def nid(b): return nalu(b, "iadd", nalu(b, "imul", ngid(b), ngsz(b)), nlid(b))

 nbarrier = nir_instr(has_def=False, intrins={"EXECUTION_SCOPE":mesa.SCOPE_WORKGROUP})(
  lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_barrier))
@@ -130,16 +126,16 @@ class NIRRenderer(Renderer):
    (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
     lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
    # load/store use pointer arithmetic, and the cast does nothing
-    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.VECTORIZE) else None),
+    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"),
+     lambda x,buf,off: x.replace(src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op != Ops.CAST else None),
    (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
  ])

  def_rewrite = PatternMatcher([
    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
-    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
-    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
-    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
+    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 8)),
+    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 4)),
+    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, ngid(ctx.b) if x.arg[0] == 'g' else nlid(ctx.b), int(x.arg[-1]))),
    (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True, name="x"),
     lambda ctx,x,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
    (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"),
@@ -161,15 +157,13 @@ class NIRRenderer(Renderer):
  def __init__(self): mesa.glsl_type_singleton_init_or_ref()

  def __del__(self):
-    with contextlib.suppress(AttributeError): mesa.glsl_type_singleton_decref()
+    with contextlib.suppress(AttributeError):mesa.glsl_type_singleton_decref()

  @property
  def nir_options(self): raise NotImplementedError("needs nir_options")
-  def param(self, b:mesa.nir_builder, x, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
+  def param(self, b:mesa.nir_builder, dtype:DType, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
  def prerender(self, uops:list[UOp]):
    self.b = mesa.nir_builder_init_simple_shader(mesa.MESA_SHADER_COMPUTE, mesa.nir_shader_compiler_options.from_buffer_copy(self.nir_options), None)
-    self.b.shader.contents.info.workgroup_size_variable = any([u.op == Ops.SPECIAL and u.arg[0] == 'i' for u in uops])
-  def postrender(self, uops:list[UOp]): pass

  def render(self, uops:list[UOp]):
    self.prerender(uops)
@@ -202,7 +196,6 @@ class NIRRenderer(Renderer):
      else:
        if (d:=self.def_rewrite.rewrite(u, ctx=self)) is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
        self.r[u] = cast(mesa.nir_def, d)
-    self.postrender(uops)

    mesa.nir_validate_shader(self.b.shader, b"after render")
    if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(ctypes.CDLL(ctypes.util.find_library('c')),
@@ -216,23 +209,22 @@ class NIRRenderer(Renderer):

    return ret

-class NIRRendererWithOpts(NIRRenderer):
+class NAKRenderer(NIRRenderer):
+  device = "NV"
  def __init__(self, dev=None, nir_options=None):
    self.dev, self._nir_options = dev, nir_options
    super().__init__()

-  def __reduce__(self): return self.__class__, (None, self.nir_options)
+  def __reduce__(self): return NAKRenderer, (None, self.nir_options,)

  @property
  def nir_options(self):
    if self._nir_options is None: self._nir_options = self.dev.compiler.nir_options
    return self._nir_options

-class NAKRenderer(NIRRendererWithOpts):
-  device = "NV"
  param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
    intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
-       lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
+       lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))

 class LVPRenderer(NIRRenderer):
  device = "CPU"
@@ -243,55 +235,9 @@ class LVPRenderer(NIRRenderer):

  param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
    srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
-    setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
+    setattr(self, "param_idx", self.param_idx+sz))(lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))

  def prerender(self, uops:list[UOp]):
    super().prerender(uops)
    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])

-# FIXME: this should be a rewrite rule
-def tovec(b, coord): return nalu(b, "vec4", nchannel(b, coord, 0), nchannel(b, coord, 1), nundef(b, dtypes.int), nundef(b, dtypes.int))
-def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
-nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=lambda val:val.num_components,
-  intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'SRC_TYPE':nfloat(dtype)},
-  srcs=lambda b,img,coord,val:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
-    lambda b,img,coord,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))
-
-_nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'DEST_TYPE':nfloat(dtype)},
-  nc=4, bs=32, num_components=4, srcs=lambda b,img,coord:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
-    lambda b,img,coord,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
-
-class IR3Renderer(NIRRendererWithOpts):
-  device = "QCOM"
-
-  def nload_img(ctx,img,coord):
-    ctx.texs.add(img)
-    return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype)
-
-  def_rewrite = PatternMatcher([
-    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")),
-          allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))),
-     lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img),
-  ]) + NIRRenderer.def_rewrite
-
-  _param = LVPRenderer.param
-  def _param_img(self, x):
-    self.img_idx += 1
-    return nimm(self.b, self.img_idx - 1, dtypes.int)
-
-  def param(self, b, x, sz): return self._param_img(x) if isinstance(x.dtype, ImageDType) else self._param(b, x, sz)
-
-  def prerender(self, uops:list[UOp]):
-    super().prerender(uops)
-    self.texs:set[UOp] = set()
-    self.uops, self.ibo_idx, self.img_idx = uops, 0, 0
-    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
-
-  def postrender(self, uops:list[UOp]):
-    bufs, texs, imgs = [u for u in uops if u.op == Ops.DEFINE_GLOBAL], itertools.count().__next__, itertools.count().__next__
-    for b in filter(lambda b: isinstance(b.dtype, ImageDType), bufs): nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int)
-
-    self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)])
-    self.b.shader.contents.info.num_images = texs() + imgs()
@@ -5,7 +5,6 @@ root = (here:=pathlib.Path(__file__).parent).parents[2]
 nv_src = {"nv_570": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/81fe4fb417c8ac3b9bdcc1d56827d116743892a5.tar.gz",
          "nv_580": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/2af9f1f0f7de4988432d4ae875b5858ffdb09cc2.tar.gz"}
 ffmpeg_src = "https://ffmpeg.org/releases/ffmpeg-8.0.1.tar.gz"
-rocr_src = "https://github.com/ROCm/rocm-systems/archive/refs/tags/rocm-7.1.1.tar.gz"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"

 def load(name, dll, files, **kwargs):
@@ -40,7 +39,7 @@ def __getattr__(nm):
        *[root/"extra/nv_gpu_driver"/s for s in ["clc9b0.h", "clc6c0qmd.h","clcec0qmd.h", "nvdec_drv.h"]], "{}/kernel-open/common/inc/nvmisc.h",
        *[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0070", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
                                                                    "83de", "c6c0", "cdc0"]],
-        *[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "clcfb0", "uvm_ioctl", "uvm_linux_ioctl", "hwref/ampere/ga100/dev_fault"]],
+        *[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "uvm_ioctl", "uvm_linux_ioctl", "hwref/ampere/ga100/dev_fault"]],
        *[f"{{}}/src/nvidia/arch/nvalloc/unix/include/nv{s}.h" for s in ["_escape", "-ioctl", "-ioctl-numbers",
                                                                         "-ioctl-numa", "-unix-nvos-params-wrappers"]],
        *[f"{{}}/src/common/sdk/nvidia/inc/{s}.h" for s in ["alloc/alloc_channel", "nvos", "ctrl/ctrlc36f", "ctrl/ctrlcb33",
@@ -90,15 +89,13 @@ def __getattr__(nm):
        "os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so'", "'/usr/local/lib/libamd_comgr.dylib'", "'/opt/homebrew/lib/libamd_comgr.dylib'"
      ], ["/opt/rocm/include/amd_comgr/amd_comgr.h"], args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"])
    case "hsa": return load("hsa", ["os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so'", "find_library('hsa-runtime64')"], [
-      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
-      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue",
-                                                                              "amd_hsa_kernel_code", "hsa_ext_finalize",
-                                                                              "hsa_ext_image", "hsa_ven_amd_aqlprofile"]]],
-      tarball=rocr_src, args=["-DLITTLEENDIAN_CPU"])
+      f"/opt/rocm/include/hsa/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code", "hsa_ext_finalize",
+                                               "hsa_ext_image", "hsa_ven_amd_aqlprofile"] ], args=["-I/opt/rocm/include"])
    case "amd_gpu": return load("amd_gpu", [], [root/f"extra/hip_gpu_driver/{s}.h" for s in ["sdma_registers", "nvd", "gc_11_0_0_offset",
                                                                                             "sienna_cichlid_ip_offset"]],
                                args=["-I/opt/rocm/include", "-x", "c++"])
    case "kgsl": return load("kgsl", [], [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
+    case "adreno": return load("adreno", [], [root/"extra/qcom_gpu_driver/a6xx.xml.h"])
    case "qcom_dsp":
      return load("qcom_dsp", [], [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
    case "sqtt": return load("sqtt", [], [root/"extra/sqtt/sqtt.h"])
@@ -113,20 +110,18 @@ def __getattr__(nm):
        *[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
        *[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
                                                                     "_flow", "_const"]],
-        *[f"{{}}/src/freedreno/{s}.h" for s in ["common/freedreno_dev_info", "ir3/ir3_compiler", "ir3/ir3_shader", "ir3/ir3_nir"]],
-        "{}/src/compiler/glsl_types.h", "{}/src/util/blob.h", "{}/src/util/ralloc.h", "{}/gen/ir3-isa.h", "{}/gen/builtin_types.h",
-        "{}/gen/a6xx.xml.h", "{}/gen/adreno_pm4.xml.h", "{}/gen/a6xx_enums.xml.h", "{}/gen/a6xx_descriptors.xml.h"], args=lambda:[
+        "{}/src/compiler/glsl_types.h", "{}/src/util/blob.h", "{}/src/util/ralloc.h"], args=lambda:[
          "-DHAVE_ENDIAN_H", "-DHAVE_STRUCT_TIMESPEC", "-DHAVE_PTHREAD", "-DHAVE_FUNC_ATTRIBUTE_PACKED", "-I{}/src", "-I{}/include", "-I{}/gen",
-          "-I{}/src/compiler/nir", "-I{}/src/gallium/auxiliary", "-I{}/src/gallium/include", "-I{}/src/freedreno/common",
-          f"-I{system('llvm-config-20 --includedir')}"],
-        preprocess=lambda path: subprocess.run("\n".join(["mkdir -p gen/util/format", "python3 src/compiler/builtin_types_h.py gen/builtin_types.h",
-          "python3 src/compiler/isaspec/decode.py --xml src/freedreno/isa/ir3.xml --out-c /dev/null --out-h gen/ir3-isa.h",
-          "python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h",
-          *["python3 src/freedreno/registers/gen_header.py --rnn src/freedreno/registers/ --xml " +
-            f"src/freedreno/registers/adreno/{s}.xml c-defines > gen/{s}.xml.h" for s in ["a6xx", "adreno_pm4", "a6xx_enums", "a6xx_descriptors"]],
-          *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
-          *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
-  tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
+          "-I{}/src/compiler/nir", "-I{}/src/gallium/auxiliary", "-I{}/src/gallium/include", f"-I{system('llvm-config-20 --includedir')}"],
+        preprocess=lambda path: subprocess.run("""mkdir -p gen/util/format
+python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
+python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
+python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
+python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
+python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+python3 src/compiler/builtin_types_h.py gen/builtin_types.h""", cwd=path, shell=True, check=True),
+  tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.4/mesa-25.2.4.tar.gz",
  prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
    case "libclang":
      return load("libclang", ["os.getenv('LIBCLANG_PATH', find_library('clang-20'))"],
@@ -136,4 +131,4 @@ def __getattr__(nm):
      return load("metal", ["find_library('Metal')"],[f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
                  ["ComputeCommandEncoder", "ComputePipeline", "CommandQueue", "Device", "IndirectCommandBuffer", "Resource", "CommandEncoder"]],
                  args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
-    case _: raise AttributeError(f"no such autogen: {nm}")
+    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -6,21 +6,18 @@ inc = ["-include", "stdint.h"]

 def __getattr__(nm):
  match nm:
-    case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
-      "discovery", "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]] + \
-      [f"{AMD}/include/{s}.h" for s in ["v9_structs", "soc15_ih_clientid"]], args=inc, tarball=am_src)
+    case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm", "discovery",
+  "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]]+[f"{AMD}/include/soc15_ih_clientid.h"], args=inc, tarball=am_src)
    case "pm4_soc15": return load("am/pm4_soc15", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], tarball=am_src)
    case "pm4_nv": return load("am/pm4_nv", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], tarball=am_src)
    case "sdma_4_0_0": return load("am/sdma_4_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
    case "sdma_5_0_0": return load("am/sdma_5_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
-    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
+    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}//amdgpu/sdma_v6_0_0_pkt_open.h"],
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
    case "smu_v13_0_0": return load("am/smu_v13_0_0",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
-                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
-    case "smu_v13_0_6": return load("am/smu_v13_0_6",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu13_driver_if_v13_0_6"]]
-                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
+                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], tarball=am_src),
    case "smu_v14_0_2": return load("am/smu_v14_0_2", [], [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
                                    "smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -3878,745 +3878,6 @@ AMDGPU_DOORBELL_LAYOUT1_LAST_NON_CP = enum_AMDGPU_DOORBELL_ASSIGNMENT_LAYOUT1.de
 AMDGPU_DOORBELL_LAYOUT1_MAX_ASSIGNMENT = enum_AMDGPU_DOORBELL_ASSIGNMENT_LAYOUT1.define('AMDGPU_DOORBELL_LAYOUT1_MAX_ASSIGNMENT', 488)
 AMDGPU_DOORBELL_LAYOUT1_INVALID = enum_AMDGPU_DOORBELL_ASSIGNMENT_LAYOUT1.define('AMDGPU_DOORBELL_LAYOUT1_INVALID', 65535)

-class struct_v9_sdma_mqd(Struct): pass
-struct_v9_sdma_mqd._fields_ = [
-  ('sdmax_rlcx_rb_cntl', uint32_t),
-  ('sdmax_rlcx_rb_base', uint32_t),
-  ('sdmax_rlcx_rb_base_hi', uint32_t),
-  ('sdmax_rlcx_rb_rptr', uint32_t),
-  ('sdmax_rlcx_rb_rptr_hi', uint32_t),
-  ('sdmax_rlcx_rb_wptr', uint32_t),
-  ('sdmax_rlcx_rb_wptr_hi', uint32_t),
-  ('sdmax_rlcx_rb_wptr_poll_cntl', uint32_t),
-  ('sdmax_rlcx_rb_rptr_addr_hi', uint32_t),
-  ('sdmax_rlcx_rb_rptr_addr_lo', uint32_t),
-  ('sdmax_rlcx_ib_cntl', uint32_t),
-  ('sdmax_rlcx_ib_rptr', uint32_t),
-  ('sdmax_rlcx_ib_offset', uint32_t),
-  ('sdmax_rlcx_ib_base_lo', uint32_t),
-  ('sdmax_rlcx_ib_base_hi', uint32_t),
-  ('sdmax_rlcx_ib_size', uint32_t),
-  ('sdmax_rlcx_skip_cntl', uint32_t),
-  ('sdmax_rlcx_context_status', uint32_t),
-  ('sdmax_rlcx_doorbell', uint32_t),
-  ('sdmax_rlcx_status', uint32_t),
-  ('sdmax_rlcx_doorbell_log', uint32_t),
-  ('sdmax_rlcx_watermark', uint32_t),
-  ('sdmax_rlcx_doorbell_offset', uint32_t),
-  ('sdmax_rlcx_csa_addr_lo', uint32_t),
-  ('sdmax_rlcx_csa_addr_hi', uint32_t),
-  ('sdmax_rlcx_ib_sub_remain', uint32_t),
-  ('sdmax_rlcx_preempt', uint32_t),
-  ('sdmax_rlcx_dummy_reg', uint32_t),
-  ('sdmax_rlcx_rb_wptr_poll_addr_hi', uint32_t),
-  ('sdmax_rlcx_rb_wptr_poll_addr_lo', uint32_t),
-  ('sdmax_rlcx_rb_aql_cntl', uint32_t),
-  ('sdmax_rlcx_minor_ptr_update', uint32_t),
-  ('sdmax_rlcx_midcmd_data0', uint32_t),
-  ('sdmax_rlcx_midcmd_data1', uint32_t),
-  ('sdmax_rlcx_midcmd_data2', uint32_t),
-  ('sdmax_rlcx_midcmd_data3', uint32_t),
-  ('sdmax_rlcx_midcmd_data4', uint32_t),
-  ('sdmax_rlcx_midcmd_data5', uint32_t),
-  ('sdmax_rlcx_midcmd_data6', uint32_t),
-  ('sdmax_rlcx_midcmd_data7', uint32_t),
-  ('sdmax_rlcx_midcmd_data8', uint32_t),
-  ('sdmax_rlcx_midcmd_cntl', uint32_t),
-  ('reserved_42', uint32_t),
-  ('reserved_43', uint32_t),
-  ('reserved_44', uint32_t),
-  ('reserved_45', uint32_t),
-  ('reserved_46', uint32_t),
-  ('reserved_47', uint32_t),
-  ('reserved_48', uint32_t),
-  ('reserved_49', uint32_t),
-  ('reserved_50', uint32_t),
-  ('reserved_51', uint32_t),
-  ('reserved_52', uint32_t),
-  ('reserved_53', uint32_t),
-  ('reserved_54', uint32_t),
-  ('reserved_55', uint32_t),
-  ('reserved_56', uint32_t),
-  ('reserved_57', uint32_t),
-  ('reserved_58', uint32_t),
-  ('reserved_59', uint32_t),
-  ('reserved_60', uint32_t),
-  ('reserved_61', uint32_t),
-  ('reserved_62', uint32_t),
-  ('reserved_63', uint32_t),
-  ('reserved_64', uint32_t),
-  ('reserved_65', uint32_t),
-  ('reserved_66', uint32_t),
-  ('reserved_67', uint32_t),
-  ('reserved_68', uint32_t),
-  ('reserved_69', uint32_t),
-  ('reserved_70', uint32_t),
-  ('reserved_71', uint32_t),
-  ('reserved_72', uint32_t),
-  ('reserved_73', uint32_t),
-  ('reserved_74', uint32_t),
-  ('reserved_75', uint32_t),
-  ('reserved_76', uint32_t),
-  ('reserved_77', uint32_t),
-  ('reserved_78', uint32_t),
-  ('reserved_79', uint32_t),
-  ('reserved_80', uint32_t),
-  ('reserved_81', uint32_t),
-  ('reserved_82', uint32_t),
-  ('reserved_83', uint32_t),
-  ('reserved_84', uint32_t),
-  ('reserved_85', uint32_t),
-  ('reserved_86', uint32_t),
-  ('reserved_87', uint32_t),
-  ('reserved_88', uint32_t),
-  ('reserved_89', uint32_t),
-  ('reserved_90', uint32_t),
-  ('reserved_91', uint32_t),
-  ('reserved_92', uint32_t),
-  ('reserved_93', uint32_t),
-  ('reserved_94', uint32_t),
-  ('reserved_95', uint32_t),
-  ('reserved_96', uint32_t),
-  ('reserved_97', uint32_t),
-  ('reserved_98', uint32_t),
-  ('reserved_99', uint32_t),
-  ('reserved_100', uint32_t),
-  ('reserved_101', uint32_t),
-  ('reserved_102', uint32_t),
-  ('reserved_103', uint32_t),
-  ('reserved_104', uint32_t),
-  ('reserved_105', uint32_t),
-  ('reserved_106', uint32_t),
-  ('reserved_107', uint32_t),
-  ('reserved_108', uint32_t),
-  ('reserved_109', uint32_t),
-  ('reserved_110', uint32_t),
-  ('reserved_111', uint32_t),
-  ('reserved_112', uint32_t),
-  ('reserved_113', uint32_t),
-  ('reserved_114', uint32_t),
-  ('reserved_115', uint32_t),
-  ('reserved_116', uint32_t),
-  ('reserved_117', uint32_t),
-  ('reserved_118', uint32_t),
-  ('reserved_119', uint32_t),
-  ('reserved_120', uint32_t),
-  ('reserved_121', uint32_t),
-  ('reserved_122', uint32_t),
-  ('reserved_123', uint32_t),
-  ('reserved_124', uint32_t),
-  ('reserved_125', uint32_t),
-  ('sdma_engine_id', uint32_t),
-  ('sdma_queue_id', uint32_t),
-]
-class struct_v9_mqd(Struct): pass
-class struct_v9_mqd_0(ctypes.Union): pass
-class struct_v9_mqd_0_0(Struct): pass
-struct_v9_mqd_0_0._fields_ = [
-  ('compute_static_thread_mgmt_se4', uint32_t),
-  ('compute_static_thread_mgmt_se5', uint32_t),
-  ('compute_static_thread_mgmt_se6', uint32_t),
-  ('compute_static_thread_mgmt_se7', uint32_t),
-]
-class struct_v9_mqd_0_1(Struct): pass
-struct_v9_mqd_0_1._fields_ = [
-  ('compute_current_logic_xcc_id', uint32_t),
-  ('compute_restart_cg_tg_id', uint32_t),
-  ('compute_tg_chunk_size', uint32_t),
-  ('compute_restore_tg_chunk_size', uint32_t),
-]
-struct_v9_mqd_0._anonymous_ = ['_0', '_1']
-struct_v9_mqd_0._fields_ = [
-  ('_0', struct_v9_mqd_0_0),
-  ('_1', struct_v9_mqd_0_1),
-]
-class struct_v9_mqd_1(ctypes.Union): pass
-class struct_v9_mqd_1_0(Struct): pass
-struct_v9_mqd_1_0._fields_ = [
-  ('reserved_225', uint32_t),
-  ('reserved_226', uint32_t),
-]
-class struct_v9_mqd_1_1(Struct): pass
-struct_v9_mqd_1_1._fields_ = [
-  ('pm4_target_xcc_in_xcp', uint32_t),
-  ('cp_mqd_stride_size', uint32_t),
-]
-struct_v9_mqd_1._anonymous_ = ['_0', '_1']
-struct_v9_mqd_1._fields_ = [
-  ('_0', struct_v9_mqd_1_0),
-  ('_1', struct_v9_mqd_1_1),
-]
-struct_v9_mqd._anonymous_ = ['_0', '_1']
-struct_v9_mqd._fields_ = [
-  ('header', uint32_t),
-  ('compute_dispatch_initiator', uint32_t),
-  ('compute_dim_x', uint32_t),
-  ('compute_dim_y', uint32_t),
-  ('compute_dim_z', uint32_t),
-  ('compute_start_x', uint32_t),
-  ('compute_start_y', uint32_t),
-  ('compute_start_z', uint32_t),
-  ('compute_num_thread_x', uint32_t),
-  ('compute_num_thread_y', uint32_t),
-  ('compute_num_thread_z', uint32_t),
-  ('compute_pipelinestat_enable', uint32_t),
-  ('compute_perfcount_enable', uint32_t),
-  ('compute_pgm_lo', uint32_t),
-  ('compute_pgm_hi', uint32_t),
-  ('compute_tba_lo', uint32_t),
-  ('compute_tba_hi', uint32_t),
-  ('compute_tma_lo', uint32_t),
-  ('compute_tma_hi', uint32_t),
-  ('compute_pgm_rsrc1', uint32_t),
-  ('compute_pgm_rsrc2', uint32_t),
-  ('compute_vmid', uint32_t),
-  ('compute_resource_limits', uint32_t),
-  ('compute_static_thread_mgmt_se0', uint32_t),
-  ('compute_static_thread_mgmt_se1', uint32_t),
-  ('compute_tmpring_size', uint32_t),
-  ('compute_static_thread_mgmt_se2', uint32_t),
-  ('compute_static_thread_mgmt_se3', uint32_t),
-  ('compute_restart_x', uint32_t),
-  ('compute_restart_y', uint32_t),
-  ('compute_restart_z', uint32_t),
-  ('compute_thread_trace_enable', uint32_t),
-  ('compute_misc_reserved', uint32_t),
-  ('compute_dispatch_id', uint32_t),
-  ('compute_threadgroup_id', uint32_t),
-  ('compute_relaunch', uint32_t),
-  ('compute_wave_restore_addr_lo', uint32_t),
-  ('compute_wave_restore_addr_hi', uint32_t),
-  ('compute_wave_restore_control', uint32_t),
-  ('_0', struct_v9_mqd_0),
-  ('reserved_43', uint32_t),
-  ('reserved_44', uint32_t),
-  ('reserved_45', uint32_t),
-  ('reserved_46', uint32_t),
-  ('reserved_47', uint32_t),
-  ('reserved_48', uint32_t),
-  ('reserved_49', uint32_t),
-  ('reserved_50', uint32_t),
-  ('reserved_51', uint32_t),
-  ('reserved_52', uint32_t),
-  ('reserved_53', uint32_t),
-  ('reserved_54', uint32_t),
-  ('reserved_55', uint32_t),
-  ('reserved_56', uint32_t),
-  ('reserved_57', uint32_t),
-  ('reserved_58', uint32_t),
-  ('reserved_59', uint32_t),
-  ('reserved_60', uint32_t),
-  ('reserved_61', uint32_t),
-  ('reserved_62', uint32_t),
-  ('reserved_63', uint32_t),
-  ('reserved_64', uint32_t),
-  ('compute_user_data_0', uint32_t),
-  ('compute_user_data_1', uint32_t),
-  ('compute_user_data_2', uint32_t),
-  ('compute_user_data_3', uint32_t),
-  ('compute_user_data_4', uint32_t),
-  ('compute_user_data_5', uint32_t),
-  ('compute_user_data_6', uint32_t),
-  ('compute_user_data_7', uint32_t),
-  ('compute_user_data_8', uint32_t),
-  ('compute_user_data_9', uint32_t),
-  ('compute_user_data_10', uint32_t),
-  ('compute_user_data_11', uint32_t),
-  ('compute_user_data_12', uint32_t),
-  ('compute_user_data_13', uint32_t),
-  ('compute_user_data_14', uint32_t),
-  ('compute_user_data_15', uint32_t),
-  ('cp_compute_csinvoc_count_lo', uint32_t),
-  ('cp_compute_csinvoc_count_hi', uint32_t),
-  ('reserved_83', uint32_t),
-  ('reserved_84', uint32_t),
-  ('reserved_85', uint32_t),
-  ('cp_mqd_query_time_lo', uint32_t),
-  ('cp_mqd_query_time_hi', uint32_t),
-  ('cp_mqd_connect_start_time_lo', uint32_t),
-  ('cp_mqd_connect_start_time_hi', uint32_t),
-  ('cp_mqd_connect_end_time_lo', uint32_t),
-  ('cp_mqd_connect_end_time_hi', uint32_t),
-  ('cp_mqd_connect_end_wf_count', uint32_t),
-  ('cp_mqd_connect_end_pq_rptr', uint32_t),
-  ('cp_mqd_connect_end_pq_wptr', uint32_t),
-  ('cp_mqd_connect_end_ib_rptr', uint32_t),
-  ('cp_mqd_readindex_lo', uint32_t),
-  ('cp_mqd_readindex_hi', uint32_t),
-  ('cp_mqd_save_start_time_lo', uint32_t),
-  ('cp_mqd_save_start_time_hi', uint32_t),
-  ('cp_mqd_save_end_time_lo', uint32_t),
-  ('cp_mqd_save_end_time_hi', uint32_t),
-  ('cp_mqd_restore_start_time_lo', uint32_t),
-  ('cp_mqd_restore_start_time_hi', uint32_t),
-  ('cp_mqd_restore_end_time_lo', uint32_t),
-  ('cp_mqd_restore_end_time_hi', uint32_t),
-  ('disable_queue', uint32_t),
-  ('reserved_107', uint32_t),
-  ('gds_cs_ctxsw_cnt0', uint32_t),
-  ('gds_cs_ctxsw_cnt1', uint32_t),
-  ('gds_cs_ctxsw_cnt2', uint32_t),
-  ('gds_cs_ctxsw_cnt3', uint32_t),
-  ('reserved_112', uint32_t),
-  ('reserved_113', uint32_t),
-  ('cp_pq_exe_status_lo', uint32_t),
-  ('cp_pq_exe_status_hi', uint32_t),
-  ('cp_packet_id_lo', uint32_t),
-  ('cp_packet_id_hi', uint32_t),
-  ('cp_packet_exe_status_lo', uint32_t),
-  ('cp_packet_exe_status_hi', uint32_t),
-  ('gds_save_base_addr_lo', uint32_t),
-  ('gds_save_base_addr_hi', uint32_t),
-  ('gds_save_mask_lo', uint32_t),
-  ('gds_save_mask_hi', uint32_t),
-  ('ctx_save_base_addr_lo', uint32_t),
-  ('ctx_save_base_addr_hi', uint32_t),
-  ('dynamic_cu_mask_addr_lo', uint32_t),
-  ('dynamic_cu_mask_addr_hi', uint32_t),
-  ('cp_mqd_base_addr_lo', uint32_t),
-  ('cp_mqd_base_addr_hi', uint32_t),
-  ('cp_hqd_active', uint32_t),
-  ('cp_hqd_vmid', uint32_t),
-  ('cp_hqd_persistent_state', uint32_t),
-  ('cp_hqd_pipe_priority', uint32_t),
-  ('cp_hqd_queue_priority', uint32_t),
-  ('cp_hqd_quantum', uint32_t),
-  ('cp_hqd_pq_base_lo', uint32_t),
-  ('cp_hqd_pq_base_hi', uint32_t),
-  ('cp_hqd_pq_rptr', uint32_t),
-  ('cp_hqd_pq_rptr_report_addr_lo', uint32_t),
-  ('cp_hqd_pq_rptr_report_addr_hi', uint32_t),
-  ('cp_hqd_pq_wptr_poll_addr_lo', uint32_t),
-  ('cp_hqd_pq_wptr_poll_addr_hi', uint32_t),
-  ('cp_hqd_pq_doorbell_control', uint32_t),
-  ('reserved_144', uint32_t),
-  ('cp_hqd_pq_control', uint32_t),
-  ('cp_hqd_ib_base_addr_lo', uint32_t),
-  ('cp_hqd_ib_base_addr_hi', uint32_t),
-  ('cp_hqd_ib_rptr', uint32_t),
-  ('cp_hqd_ib_control', uint32_t),
-  ('cp_hqd_iq_timer', uint32_t),
-  ('cp_hqd_iq_rptr', uint32_t),
-  ('cp_hqd_dequeue_request', uint32_t),
-  ('cp_hqd_dma_offload', uint32_t),
-  ('cp_hqd_sema_cmd', uint32_t),
-  ('cp_hqd_msg_type', uint32_t),
-  ('cp_hqd_atomic0_preop_lo', uint32_t),
-  ('cp_hqd_atomic0_preop_hi', uint32_t),
-  ('cp_hqd_atomic1_preop_lo', uint32_t),
-  ('cp_hqd_atomic1_preop_hi', uint32_t),
-  ('cp_hqd_hq_status0', uint32_t),
-  ('cp_hqd_hq_control0', uint32_t),
-  ('cp_mqd_control', uint32_t),
-  ('cp_hqd_hq_status1', uint32_t),
-  ('cp_hqd_hq_control1', uint32_t),
-  ('cp_hqd_eop_base_addr_lo', uint32_t),
-  ('cp_hqd_eop_base_addr_hi', uint32_t),
-  ('cp_hqd_eop_control', uint32_t),
-  ('cp_hqd_eop_rptr', uint32_t),
-  ('cp_hqd_eop_wptr', uint32_t),
-  ('cp_hqd_eop_done_events', uint32_t),
-  ('cp_hqd_ctx_save_base_addr_lo', uint32_t),
-  ('cp_hqd_ctx_save_base_addr_hi', uint32_t),
-  ('cp_hqd_ctx_save_control', uint32_t),
-  ('cp_hqd_cntl_stack_offset', uint32_t),
-  ('cp_hqd_cntl_stack_size', uint32_t),
-  ('cp_hqd_wg_state_offset', uint32_t),
-  ('cp_hqd_ctx_save_size', uint32_t),
-  ('cp_hqd_gds_resource_state', uint32_t),
-  ('cp_hqd_error', uint32_t),
-  ('cp_hqd_eop_wptr_mem', uint32_t),
-  ('cp_hqd_aql_control', uint32_t),
-  ('cp_hqd_pq_wptr_lo', uint32_t),
-  ('cp_hqd_pq_wptr_hi', uint32_t),
-  ('reserved_184', uint32_t),
-  ('reserved_185', uint32_t),
-  ('reserved_186', uint32_t),
-  ('reserved_187', uint32_t),
-  ('reserved_188', uint32_t),
-  ('reserved_189', uint32_t),
-  ('reserved_190', uint32_t),
-  ('reserved_191', uint32_t),
-  ('iqtimer_pkt_header', uint32_t),
-  ('iqtimer_pkt_dw0', uint32_t),
-  ('iqtimer_pkt_dw1', uint32_t),
-  ('iqtimer_pkt_dw2', uint32_t),
-  ('iqtimer_pkt_dw3', uint32_t),
-  ('iqtimer_pkt_dw4', uint32_t),
-  ('iqtimer_pkt_dw5', uint32_t),
-  ('iqtimer_pkt_dw6', uint32_t),
-  ('iqtimer_pkt_dw7', uint32_t),
-  ('iqtimer_pkt_dw8', uint32_t),
-  ('iqtimer_pkt_dw9', uint32_t),
-  ('iqtimer_pkt_dw10', uint32_t),
-  ('iqtimer_pkt_dw11', uint32_t),
-  ('iqtimer_pkt_dw12', uint32_t),
-  ('iqtimer_pkt_dw13', uint32_t),
-  ('iqtimer_pkt_dw14', uint32_t),
-  ('iqtimer_pkt_dw15', uint32_t),
-  ('iqtimer_pkt_dw16', uint32_t),
-  ('iqtimer_pkt_dw17', uint32_t),
-  ('iqtimer_pkt_dw18', uint32_t),
-  ('iqtimer_pkt_dw19', uint32_t),
-  ('iqtimer_pkt_dw20', uint32_t),
-  ('iqtimer_pkt_dw21', uint32_t),
-  ('iqtimer_pkt_dw22', uint32_t),
-  ('iqtimer_pkt_dw23', uint32_t),
-  ('iqtimer_pkt_dw24', uint32_t),
-  ('iqtimer_pkt_dw25', uint32_t),
-  ('iqtimer_pkt_dw26', uint32_t),
-  ('iqtimer_pkt_dw27', uint32_t),
-  ('iqtimer_pkt_dw28', uint32_t),
-  ('iqtimer_pkt_dw29', uint32_t),
-  ('iqtimer_pkt_dw30', uint32_t),
-  ('iqtimer_pkt_dw31', uint32_t),
-  ('_1', struct_v9_mqd_1),
-  ('reserved_227', uint32_t),
-  ('set_resources_header', uint32_t),
-  ('set_resources_dw1', uint32_t),
-  ('set_resources_dw2', uint32_t),
-  ('set_resources_dw3', uint32_t),
-  ('set_resources_dw4', uint32_t),
-  ('set_resources_dw5', uint32_t),
-  ('set_resources_dw6', uint32_t),
-  ('set_resources_dw7', uint32_t),
-  ('reserved_236', uint32_t),
-  ('reserved_237', uint32_t),
-  ('reserved_238', uint32_t),
-  ('reserved_239', uint32_t),
-  ('queue_doorbell_id0', uint32_t),
-  ('queue_doorbell_id1', uint32_t),
-  ('queue_doorbell_id2', uint32_t),
-  ('queue_doorbell_id3', uint32_t),
-  ('queue_doorbell_id4', uint32_t),
-  ('queue_doorbell_id5', uint32_t),
-  ('queue_doorbell_id6', uint32_t),
-  ('queue_doorbell_id7', uint32_t),
-  ('queue_doorbell_id8', uint32_t),
-  ('queue_doorbell_id9', uint32_t),
-  ('queue_doorbell_id10', uint32_t),
-  ('queue_doorbell_id11', uint32_t),
-  ('queue_doorbell_id12', uint32_t),
-  ('queue_doorbell_id13', uint32_t),
-  ('queue_doorbell_id14', uint32_t),
-  ('queue_doorbell_id15', uint32_t),
-  ('reserved_256', uint32_t),
-  ('reserved_257', uint32_t),
-  ('reserved_258', uint32_t),
-  ('reserved_259', uint32_t),
-  ('reserved_260', uint32_t),
-  ('reserved_261', uint32_t),
-  ('reserved_262', uint32_t),
-  ('reserved_263', uint32_t),
-  ('reserved_264', uint32_t),
-  ('reserved_265', uint32_t),
-  ('reserved_266', uint32_t),
-  ('reserved_267', uint32_t),
-  ('reserved_268', uint32_t),
-  ('reserved_269', uint32_t),
-  ('reserved_270', uint32_t),
-  ('reserved_271', uint32_t),
-  ('reserved_272', uint32_t),
-  ('reserved_273', uint32_t),
-  ('reserved_274', uint32_t),
-  ('reserved_275', uint32_t),
-  ('reserved_276', uint32_t),
-  ('reserved_277', uint32_t),
-  ('reserved_278', uint32_t),
-  ('reserved_279', uint32_t),
-  ('reserved_280', uint32_t),
-  ('reserved_281', uint32_t),
-  ('reserved_282', uint32_t),
-  ('reserved_283', uint32_t),
-  ('reserved_284', uint32_t),
-  ('reserved_285', uint32_t),
-  ('reserved_286', uint32_t),
-  ('reserved_287', uint32_t),
-  ('reserved_288', uint32_t),
-  ('reserved_289', uint32_t),
-  ('reserved_290', uint32_t),
-  ('reserved_291', uint32_t),
-  ('reserved_292', uint32_t),
-  ('reserved_293', uint32_t),
-  ('reserved_294', uint32_t),
-  ('reserved_295', uint32_t),
-  ('reserved_296', uint32_t),
-  ('reserved_297', uint32_t),
-  ('reserved_298', uint32_t),
-  ('reserved_299', uint32_t),
-  ('reserved_300', uint32_t),
-  ('reserved_301', uint32_t),
-  ('reserved_302', uint32_t),
-  ('reserved_303', uint32_t),
-  ('reserved_304', uint32_t),
-  ('reserved_305', uint32_t),
-  ('reserved_306', uint32_t),
-  ('reserved_307', uint32_t),
-  ('reserved_308', uint32_t),
-  ('reserved_309', uint32_t),
-  ('reserved_310', uint32_t),
-  ('reserved_311', uint32_t),
-  ('reserved_312', uint32_t),
-  ('reserved_313', uint32_t),
-  ('reserved_314', uint32_t),
-  ('reserved_315', uint32_t),
-  ('reserved_316', uint32_t),
-  ('reserved_317', uint32_t),
-  ('reserved_318', uint32_t),
-  ('reserved_319', uint32_t),
-  ('reserved_320', uint32_t),
-  ('reserved_321', uint32_t),
-  ('reserved_322', uint32_t),
-  ('reserved_323', uint32_t),
-  ('reserved_324', uint32_t),
-  ('reserved_325', uint32_t),
-  ('reserved_326', uint32_t),
-  ('reserved_327', uint32_t),
-  ('reserved_328', uint32_t),
-  ('reserved_329', uint32_t),
-  ('reserved_330', uint32_t),
-  ('reserved_331', uint32_t),
-  ('reserved_332', uint32_t),
-  ('reserved_333', uint32_t),
-  ('reserved_334', uint32_t),
-  ('reserved_335', uint32_t),
-  ('reserved_336', uint32_t),
-  ('reserved_337', uint32_t),
-  ('reserved_338', uint32_t),
-  ('reserved_339', uint32_t),
-  ('reserved_340', uint32_t),
-  ('reserved_341', uint32_t),
-  ('reserved_342', uint32_t),
-  ('reserved_343', uint32_t),
-  ('reserved_344', uint32_t),
-  ('reserved_345', uint32_t),
-  ('reserved_346', uint32_t),
-  ('reserved_347', uint32_t),
-  ('reserved_348', uint32_t),
-  ('reserved_349', uint32_t),
-  ('reserved_350', uint32_t),
-  ('reserved_351', uint32_t),
-  ('reserved_352', uint32_t),
-  ('reserved_353', uint32_t),
-  ('reserved_354', uint32_t),
-  ('reserved_355', uint32_t),
-  ('reserved_356', uint32_t),
-  ('reserved_357', uint32_t),
-  ('reserved_358', uint32_t),
-  ('reserved_359', uint32_t),
-  ('reserved_360', uint32_t),
-  ('reserved_361', uint32_t),
-  ('reserved_362', uint32_t),
-  ('reserved_363', uint32_t),
-  ('reserved_364', uint32_t),
-  ('reserved_365', uint32_t),
-  ('reserved_366', uint32_t),
-  ('reserved_367', uint32_t),
-  ('reserved_368', uint32_t),
-  ('reserved_369', uint32_t),
-  ('reserved_370', uint32_t),
-  ('reserved_371', uint32_t),
-  ('reserved_372', uint32_t),
-  ('reserved_373', uint32_t),
-  ('reserved_374', uint32_t),
-  ('reserved_375', uint32_t),
-  ('reserved_376', uint32_t),
-  ('reserved_377', uint32_t),
-  ('reserved_378', uint32_t),
-  ('reserved_379', uint32_t),
-  ('reserved_380', uint32_t),
-  ('reserved_381', uint32_t),
-  ('reserved_382', uint32_t),
-  ('reserved_383', uint32_t),
-  ('reserved_384', uint32_t),
-  ('reserved_385', uint32_t),
-  ('reserved_386', uint32_t),
-  ('reserved_387', uint32_t),
-  ('reserved_388', uint32_t),
-  ('reserved_389', uint32_t),
-  ('reserved_390', uint32_t),
-  ('reserved_391', uint32_t),
-  ('reserved_392', uint32_t),
-  ('reserved_393', uint32_t),
-  ('reserved_394', uint32_t),
-  ('reserved_395', uint32_t),
-  ('reserved_396', uint32_t),
-  ('reserved_397', uint32_t),
-  ('reserved_398', uint32_t),
-  ('reserved_399', uint32_t),
-  ('reserved_400', uint32_t),
-  ('reserved_401', uint32_t),
-  ('reserved_402', uint32_t),
-  ('reserved_403', uint32_t),
-  ('reserved_404', uint32_t),
-  ('reserved_405', uint32_t),
-  ('reserved_406', uint32_t),
-  ('reserved_407', uint32_t),
-  ('reserved_408', uint32_t),
-  ('reserved_409', uint32_t),
-  ('reserved_410', uint32_t),
-  ('reserved_411', uint32_t),
-  ('reserved_412', uint32_t),
-  ('reserved_413', uint32_t),
-  ('reserved_414', uint32_t),
-  ('reserved_415', uint32_t),
-  ('reserved_416', uint32_t),
-  ('reserved_417', uint32_t),
-  ('reserved_418', uint32_t),
-  ('reserved_419', uint32_t),
-  ('reserved_420', uint32_t),
-  ('reserved_421', uint32_t),
-  ('reserved_422', uint32_t),
-  ('reserved_423', uint32_t),
-  ('reserved_424', uint32_t),
-  ('reserved_425', uint32_t),
-  ('reserved_426', uint32_t),
-  ('reserved_427', uint32_t),
-  ('reserved_428', uint32_t),
-  ('reserved_429', uint32_t),
-  ('reserved_430', uint32_t),
-  ('reserved_431', uint32_t),
-  ('reserved_432', uint32_t),
-  ('reserved_433', uint32_t),
-  ('reserved_434', uint32_t),
-  ('reserved_435', uint32_t),
-  ('reserved_436', uint32_t),
-  ('reserved_437', uint32_t),
-  ('reserved_438', uint32_t),
-  ('reserved_439', uint32_t),
-  ('reserved_440', uint32_t),
-  ('reserved_441', uint32_t),
-  ('reserved_442', uint32_t),
-  ('reserved_443', uint32_t),
-  ('reserved_444', uint32_t),
-  ('reserved_445', uint32_t),
-  ('reserved_446', uint32_t),
-  ('reserved_447', uint32_t),
-  ('reserved_448', uint32_t),
-  ('reserved_449', uint32_t),
-  ('reserved_450', uint32_t),
-  ('reserved_451', uint32_t),
-  ('reserved_452', uint32_t),
-  ('reserved_453', uint32_t),
-  ('reserved_454', uint32_t),
-  ('reserved_455', uint32_t),
-  ('reserved_456', uint32_t),
-  ('reserved_457', uint32_t),
-  ('reserved_458', uint32_t),
-  ('reserved_459', uint32_t),
-  ('reserved_460', uint32_t),
-  ('reserved_461', uint32_t),
-  ('reserved_462', uint32_t),
-  ('reserved_463', uint32_t),
-  ('reserved_464', uint32_t),
-  ('reserved_465', uint32_t),
-  ('reserved_466', uint32_t),
-  ('reserved_467', uint32_t),
-  ('reserved_468', uint32_t),
-  ('reserved_469', uint32_t),
-  ('reserved_470', uint32_t),
-  ('reserved_471', uint32_t),
-  ('reserved_472', uint32_t),
-  ('reserved_473', uint32_t),
-  ('reserved_474', uint32_t),
-  ('reserved_475', uint32_t),
-  ('reserved_476', uint32_t),
-  ('reserved_477', uint32_t),
-  ('reserved_478', uint32_t),
-  ('reserved_479', uint32_t),
-  ('reserved_480', uint32_t),
-  ('reserved_481', uint32_t),
-  ('reserved_482', uint32_t),
-  ('reserved_483', uint32_t),
-  ('reserved_484', uint32_t),
-  ('reserved_485', uint32_t),
-  ('reserved_486', uint32_t),
-  ('reserved_487', uint32_t),
-  ('reserved_488', uint32_t),
-  ('reserved_489', uint32_t),
-  ('reserved_490', uint32_t),
-  ('reserved_491', uint32_t),
-  ('reserved_492', uint32_t),
-  ('reserved_493', uint32_t),
-  ('reserved_494', uint32_t),
-  ('reserved_495', uint32_t),
-  ('reserved_496', uint32_t),
-  ('reserved_497', uint32_t),
-  ('reserved_498', uint32_t),
-  ('reserved_499', uint32_t),
-  ('reserved_500', uint32_t),
-  ('reserved_501', uint32_t),
-  ('reserved_502', uint32_t),
-  ('reserved_503', uint32_t),
-  ('reserved_504', uint32_t),
-  ('reserved_505', uint32_t),
-  ('reserved_506', uint32_t),
-  ('reserved_507', uint32_t),
-  ('reserved_508', uint32_t),
-  ('reserved_509', uint32_t),
-  ('reserved_510', uint32_t),
-  ('reserved_511', uint32_t),
-]
-class struct_v9_mqd_allocation(Struct): pass
-struct_v9_mqd_allocation._fields_ = [
-  ('mqd', struct_v9_mqd),
-  ('wptr_poll_mem', uint32_t),
-  ('rptr_report_mem', uint32_t),
-  ('dynamic_cu_mask', uint32_t),
-  ('dynamic_rb_mask', uint32_t),
-]
-class struct_v9_ce_ib_state(Struct): pass
-struct_v9_ce_ib_state._fields_ = [
-  ('ce_ib_completion_status', uint32_t),
-  ('ce_constegnine_count', uint32_t),
-  ('ce_ibOffset_ib1', uint32_t),
-  ('ce_ibOffset_ib2', uint32_t),
-  ('ce_chainib_addrlo_ib1', uint32_t),
-  ('ce_chainib_addrlo_ib2', uint32_t),
-  ('ce_chainib_addrhi_ib1', uint32_t),
-  ('ce_chainib_addrhi_ib2', uint32_t),
-  ('ce_chainib_size_ib1', uint32_t),
-  ('ce_chainib_size_ib2', uint32_t),
-]
-class struct_v9_de_ib_state(Struct): pass
-struct_v9_de_ib_state._fields_ = [
-  ('ib_completion_status', uint32_t),
-  ('de_constEngine_count', uint32_t),
-  ('ib_offset_ib1', uint32_t),
-  ('ib_offset_ib2', uint32_t),
-  ('chain_ib_addrlo_ib1', uint32_t),
-  ('chain_ib_addrlo_ib2', uint32_t),
-  ('chain_ib_addrhi_ib1', uint32_t),
-  ('chain_ib_addrhi_ib2', uint32_t),
-  ('chain_ib_size_ib1', uint32_t),
-  ('chain_ib_size_ib2', uint32_t),
-  ('preamble_begin_ib1', uint32_t),
-  ('preamble_begin_ib2', uint32_t),
-  ('preamble_end_ib1', uint32_t),
-  ('preamble_end_ib2', uint32_t),
-  ('chain_ib_pream_addrlo_ib1', uint32_t),
-  ('chain_ib_pream_addrlo_ib2', uint32_t),
-  ('chain_ib_pream_addrhi_ib1', uint32_t),
-  ('chain_ib_pream_addrhi_ib2', uint32_t),
-  ('draw_indirect_baseLo', uint32_t),
-  ('draw_indirect_baseHi', uint32_t),
-  ('disp_indirect_baseLo', uint32_t),
-  ('disp_indirect_baseHi', uint32_t),
-  ('gds_backup_addrlo', uint32_t),
-  ('gds_backup_addrhi', uint32_t),
-  ('index_base_addrlo', uint32_t),
-  ('index_base_addrhi', uint32_t),
-  ('sample_cntl', uint32_t),
-]
-class struct_v9_gfx_meta_data(Struct): pass
-struct_v9_gfx_meta_data._fields_ = [
-  ('ce_payload', struct_v9_ce_ib_state),
-  ('reserved1', (uint32_t * 54)),
-  ('de_payload', struct_v9_de_ib_state),
-  ('DeIbBaseAddrLo', uint32_t),
-  ('DeIbBaseAddrHi', uint32_t),
-  ('reserved2', (uint32_t * 931)),
-]
 enum_soc15_ih_clientid = CEnum(ctypes.c_uint32)
 SOC15_IH_CLIENTID_IH = enum_soc15_ih_clientid.define('SOC15_IH_CLIENTID_IH', 0)
 SOC15_IH_CLIENTID_ACP = enum_soc15_ih_clientid.define('SOC15_IH_CLIENTID_ACP', 1)
--- a/Show More
+++ b/Show More