cleaner

* nbio7_9

* psp

* gmc

* gfx

* sdma

* ih

* linter

* linter

* minor

* finish

* add missing

* do not allow warm boot for now

.github/actions/setup-tinygrad/action.yml

@@ -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/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+      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
     - name: Install mesa (macOS)
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash

.github/workflows/autogen.yml

@@ -116,11 +116,9 @@ 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, adreno, qcom_dsp"
+        python3 -c "from tinygrad.runtime.autogen import kgsl, 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: |

.github/workflows/benchmark.yml

@@ -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=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      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
     - 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 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 NOCLANG=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,6 +132,10 @@ 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)
@@ -318,31 +322,31 @@ jobs:
     # TODO: too slow
     # - name: Fuzz Padded Tensor Core GEMM (PTX)
     #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    - name: HEVC Decode Benchmark
+      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=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
+      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
     # - 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=66 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=72 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)
@@ -433,9 +437,8 @@ 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
-    # 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 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
@@ -525,22 +528,21 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: Train MNIST
       run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-    # TODO: too slow
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
-    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=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
     # - 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)
@@ -590,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=66 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=72 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)
@@ -625,18 +627,14 @@ 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
@@ -706,10 +704,8 @@ jobs:
       run: |
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    # TODO: enable
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
     # - name: Run 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)
@@ -770,11 +766,10 @@ jobs:
         NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Test LLAMA-3
       run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    #- name: Run 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 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

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
+  CACHE_VERSION: '15'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -71,9 +71,7 @@ jobs:
     - name: Test Docs Build
       run: python -m mkdocs build --strict
     - name: Test Docs
-      run: |
-        python docs/abstractions2.py
-        python docs/abstractions3.py
+      run: 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
@@ -263,7 +261,9 @@ 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 -m pytest -n=auto test/unit/ --durations=20
+      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
     - 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 < 19000 lines
-      run: MAX_LINE_COUNT=19000 python sz.py
+    - name: Repo line count < 20000 lines
+      run: MAX_LINE_COUNT=20000 python sz.py
 
   spec:
     strategy:
@@ -447,7 +447,7 @@ jobs:
         with:
           key: onnxoptl
           deps: testing
-          pydeps: "tensorflow==2.15.1 tensorflow_addons"
+          pydeps: "tensorflow==2.19"
           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=8 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=1 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,3 +972,33 @@ 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

CLAUDE.md

@@ -0,0 +1,213 @@
+# 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.

docs/abstractions2.py

@@ -1,135 +0,0 @@
-# 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

examples/beautiful_mnist.py

@@ -21,7 +21,7 @@ if __name__ == "__main__":
   X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
 
   model = Model()
-  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
+  opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
 
   @TinyJit
   @Tensor.train()

examples/gradaccum_mnist.py

@@ -0,0 +1,108 @@
+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}%")

examples/mlperf/helpers.py

@@ -223,13 +223,13 @@ def get_mlperf_bert_model():
 
 def get_fake_data_bert(BS:int):
   return {
-    "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"),
+    "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(),
   }
 
 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:

examples/mlperf/initializers.py

@@ -59,9 +59,7 @@ 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)
-    # 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()
+    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
 
 class LayerNormBert:
   def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):

examples/mlperf/model_train.py

@@ -918,40 +918,6 @@ 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):
@@ -1014,7 +980,8 @@ 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: mlperf logging
+  # TODO: implement grad accumulation + mlperf logging
+  assert grad_acc == 1
   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))
@@ -1073,8 +1040,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 = [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)
+  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)
   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)
 
@@ -1131,12 +1098,38 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
+  i, train_data = start_step, next(train_it)
 
   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
@@ -1144,16 +1137,12 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        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)
+        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"])
 
         pt = time.perf_counter()
-
-        try:
-          next_data = [next(train_it) for _ in range(grad_acc)]
-        except StopIteration:
-          next_data = None
-
+        next_data = next(train_it)
         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)}"
@@ -1188,8 +1177,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", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+      elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
+        # TODO: this hangs on tiny green after 90 minutes of training
         train_step_bert.captured.free_intermediates()
       eval_lm_losses = []
       eval_clsf_losses = []
@@ -1224,7 +1213,7 @@ def train_bert():
           return
 
       if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE") 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)
@@ -1300,6 +1289,7 @@ 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)
@@ -1374,20 +1364,17 @@ def train_llama3():
 
   @TinyJit
   @Tensor.train()
-  def train_step(model, tokens:Tensor, grad_acc:int):
+  def train_step(model, tokens:Tensor):
     optim.zero_grad()
-    # 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])
+    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()
     # 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
@@ -1426,14 +1413,14 @@ def train_llama3():
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
-      return fake_data(GBS, SAMPLES)
+      return fake_data(BS, SAMPLES)
     else:
       if SMALL:
         from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
       else:
         from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
 
   def get_eval_iter():
     if getenv("FAKEDATA", 0):
@@ -1451,7 +1438,7 @@ def train_llama3():
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
-    loss, lr = train_step(model, tokens, grad_acc)
+    loss, lr = train_step(model, tokens)
     loss = loss.float().item()
 
     i += 1

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh

@@ -2,7 +2,7 @@
 
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh

@@ -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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000

examples/sdv2.py

@@ -115,7 +115,7 @@ if __name__ == "__main__":
 
   with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
     if not args.fakeweights:
-      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+      default_weights_url = 'https://huggingface.co/sd2-community/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

extra/amdpci/hive_reset.py

@@ -0,0 +1,14 @@
+#!/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()

extra/amdpci/proclogs.py

@@ -1,48 +1,59 @@
 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, devfmt, vram, doorbell, mmio, dma_regions=None):
-    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+  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')
     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):
-  register_map = register_names
+def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+  register_map = register_names or {}
 
-  final = ""
   def replace_register(match):
-    register = match.group(1)
-    return f"Reading register {register_map.get(int(register, base=16), register)}"
+    reg = match.group(1)
+    return f"Reading register {register_map.get(int(reg, 16), reg)}"
 
-  pattern = r'Reading register (0x[0-9a-fA-F]+)'
-
-  processed_log = re.sub(pattern, replace_register, log_content)
+  processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
 
   def replace_register_2(match):
-    register = match.group(1)
-    return f"Writing register {register_map.get(int(register, base=16), register)}"
+    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)
 
-  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 inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+      for xcc, addr in y.addr.items():
+        reg_names[addr] = f"{x}, xcc={xcc}"
 
   with open(sys.argv[1], 'r') as f:
-    log_content = log_content_them = f.read()
+    log_content = f.read()
 
-  processed_log = parse_amdgpu_logs(log_content, reg_names)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)
@@ -51,5 +62,4 @@ if __name__ == '__main__':
   if len(sys.argv) != 3:
     print("Usage: <input_file_path> <output_file_path>")
     sys.exit(1)
-
-  main()
+  main()

extra/assembly/rdna3/generate.py

@@ -0,0 +1,136 @@
+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)

extra/assembly/rdna3/parse.py

@@ -0,0 +1,15 @@
+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)

extra/gemm/amd_uop_matmul.py

@@ -1,3 +1,4 @@
+import numpy as np
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
 from tinygrad.engine.realize import ExecItem, get_runner
@@ -140,15 +141,14 @@ def hand_spec_kernel3():
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
 def test_matmul(sink:UOp, N=N):
-  with Context(DEBUG=0):
-    a = Tensor.randn(N, N)
-    b = Tensor.randn(N, N)
-    hc = Tensor.empty(N, N)
-    Tensor.realize(a, b, hc)
+  rng = np.random.default_rng()
+  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
+  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
+  hc = Tensor.empty(N, N)
+  Tensor.realize(a, b, hc)
 
   ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
 
-  GlobalCounters.reset()
   ets = []
   with Context(DEBUG=2):
     for _ in range(run_count):

extra/hevc/.gitignore

@@ -0,0 +1 @@
+out/

extra/hevc/decode.py

@@ -0,0 +1,83 @@
+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
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--input_file", type=str, default="")
+  parser.add_argument("--output_dir", type=str, default="extra/hevc/out")
+  args = parser.parse_args()
+
+  if args.input_file == "":
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    hevc_tensor = Tensor.from_url(url, device="CPU")
+  else:
+    hevc_tensor = Tensor.empty(os.stat(args.input_file).st_size, dtype=dtypes.uint8, device=f"disk:{args.input_file}").to("CPU")
+
+  dat = bytes(hevc_tensor.data())
+  dat_hash = hashlib.md5(dat).hexdigest()
+
+  with Timing("prep infos: "):
+    dat_nv = hevc_tensor.to("NV")
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
+
+  frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
+
+  # move all needed data to gpu
+  #all_slices = []
+  with Timing("copy to gpu: "):
+    opaque_nv = opaque.to("NV").contiguous().realize()
+    hevc_tensor = hevc_tensor.to("NV")
+
+  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)
+
+  # 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)
+
+  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)),))
+
+      outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
+      out_images.append(outimg)
+      if is_hist: history.append(outimg)
+
+    Device.default.synchronize()
+
+  if getenv("VALIDATE", 0):
+    import pickle
+    if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
+      url = "https://github.com/nimlgen/hevc_validate_set/raw/refs/heads/main/decoded_frames_b813bfdbec194fd17fdf0e3ceb8cea1c.pkl"
+      decoded_frames = pickle.load(fetch(url).open("rb"))
+    else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
+  else: import cv2
+
+  for i, img in tqdm(enumerate(out_images)):
+    if getenv("VALIDATE", 0):
+      if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
+        img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
+        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())

extra/hevc/hevc.py

@@ -0,0 +1,450 @@
+import dataclasses, enum, argparse, os, itertools, time, ctypes
+from typing import Any
+from tinygrad import Tensor, dtypes, Device, TinyJit
+from tinygrad.helpers import DEBUG, round_up, ceildiv, Timing, prod
+from tinygrad.runtime.autogen import avcodec, nv_570 as nv_gpu
+
+class BitReader:
+  def __init__(self, data:bytes): self.reader, self.current_bits, self.bits, self.read_bits, self.total = iter(data), 0, 0, 0, len(data) * 8
+  def empty(self): return self.read_bits == self.total and self.current_bits == 0
+  def peak_bits(self, n):
+    while self.current_bits < n:
+      self.bits = (self.bits << 8) | next(self.reader)
+      self.current_bits += 8
+      self.read_bits += 8
+    return (self.bits >> (self.current_bits - n)) & ((1 << n) - 1)
+  def _next_bits(self, n):
+    val = self.peak_bits(n)
+    self.bits &= (1 << (self.current_bits - n)) - 1
+    self.current_bits -= n
+    return val
+
+  def u(self, n): return self._next_bits(n)
+
+  # 9.2 Parsing process for 0-th order Exp-Golomb codes
+  def ue_v(self):
+    leading_zero_bits = -1
+    while True:
+      bit = self.u(1)
+      leading_zero_bits += 1
+      if bit == 1: break
+
+    part = self.u(leading_zero_bits)
+
+    if leading_zero_bits == 0: return 0
+    return (1 << leading_zero_bits) - 1 + part
+
+  # 9.2.2 Mapping process for signed Exp-Golomb codes
+  def se_v(self):
+    k = self.ue_v()
+    return (-1 ** (k + 1)) * (k // 2)
+
+# 7.3.1.1 General NAL unit syntax
+def _hevc_get_rbsp(dat:bytes, off=0) -> bytes:
+  rbsp = bytes()
+  while off < len(dat):
+    if off + 2 < len(dat) and dat[off:off+3] == b'\x00\x00\x03':
+      rbsp += bytes([0, 0])
+      off += 3
+    else:
+      rbsp += bytes([dat[off]])
+      off += 1
+  return rbsp
+
+class HevcSlice:
+  # 7.3.3 Profile, tier and level syntax
+  def profile_tier_level(self, r:BitReader, enable:bool, max_sub_layers:int):
+    assert enable and max_sub_layers == 0, "no sublayers supported"
+    self._notimpl_profile_tier_level = r.u(88)
+    self.general_level_idc = r.u(8)
+
+  # 7.3.7 Short-term reference picture set syntax
+  def st_ref_pic_set(self, r:BitReader, stRpsIdx:int, num_short_term_ref_pic_sets:int=0, sps=None):
+    inter_ref_pic_set_prediction_flag = r.u(1) if stRpsIdx != 0 else 0
+
+    if inter_ref_pic_set_prediction_flag:
+      if stRpsIdx == num_short_term_ref_pic_sets:
+        delta_idx_minus1 = r.ue_v()
+      delta_rps_sign = r.u(1)
+      abs_delta_rps_minus1 = r.ue_v()
+
+      NumDeltaPocs = sps.num_negative_pics + sps.num_positive_pics
+      for i in range(NumDeltaPocs + 1):
+        used_by_curr_pic_flag = r.u(1)
+        if not used_by_curr_pic_flag:
+          use_delta_flag = r.u(1)
+    else:
+      self.num_negative_pics = r.ue_v()
+      self.num_positive_pics = r.ue_v()
+      for i in range(self.num_negative_pics):
+        delta_poc_s0_minus1 = r.ue_v()
+        used_by_curr_pic_s0_flag = r.u(1)
+      for i in range(self.num_positive_pics):
+        delta_poc_s1_minus1 = r.ue_v()
+        used_by_curr_pic_s1_flag = r.u(1)
+
+# 7.3.2.2 Sequence parameter set RBSP syntax
+class SPS(HevcSlice):
+  def __init__(self, r:BitReader):
+    self.sps_video_parameter_set_id = r.u(4)
+    self.sps_max_sub_layers_minus1 = r.u(3)
+    self.sps_temporal_id_nesting_flag = r.u(1)
+
+    self.profile_tier_level(r, True, self.sps_max_sub_layers_minus1)
+
+    self.sps_seq_parameter_set_id = r.ue_v()
+    self.chroma_format_idc = r.ue_v()
+    self.separate_colour_plane_flag = r.u(1) if self.chroma_format_idc == 3 else 0
+    self.pic_width_in_luma_samples = r.ue_v()
+    self.pic_height_in_luma_samples = r.ue_v()
+    self.conformance_window_flag = r.u(1)
+
+    if self.conformance_window_flag:
+      self.conf_win_left_offset = r.ue_v()
+      self.conf_win_right_offset = r.ue_v()
+      self.conf_win_top_offset = r.ue_v()
+      self.conf_win_bottom_offset = r.ue_v()
+    else: self.conf_win_left_offset = self.conf_win_right_offset = self.conf_win_top_offset = self.conf_win_bottom_offset = 0
+
+    self.bit_depth_luma = r.ue_v() + 8
+    self.bit_depth_chroma = r.ue_v() + 8
+    self.log2_max_pic_order_cnt_lsb_minus4 = r.ue_v()
+    self.sps_sub_layer_ordering_info_present_flag = r.u(1)
+    self.sps_max_dec_pic_buffering, self.sps_max_num_reorder_pics, self.sps_max_latency_increase_plus1 = [], [], []
+    for i in range((0 if self.sps_sub_layer_ordering_info_present_flag else self.sps_max_sub_layers_minus1), self.sps_max_sub_layers_minus1 + 1):
+      self.sps_max_dec_pic_buffering.append(r.ue_v() + 1)
+      self.sps_max_num_reorder_pics.append(r.ue_v())
+      self.sps_max_latency_increase_plus1.append(r.ue_v())
+    self.log2_min_luma_coding_block_size = r.ue_v() + 3
+    self.log2_max_luma_coding_block_size = self.log2_min_luma_coding_block_size + r.ue_v()
+    self.log2_min_transform_block_size = r.ue_v() + 2
+    self.log2_max_transform_block_size = self.log2_min_transform_block_size + r.ue_v()
+    self.max_transform_hierarchy_depth_inter = r.ue_v()
+    self.max_transform_hierarchy_depth_intra = r.ue_v()
+    if scaling_list_enabled_flag := r.u(1):
+      if sps_scaling_list_data_present_flag := r.u(1): assert False, "scaling_list_data parsing not implemented"
+    self.amp_enabled_flag = r.u(1)
+    self.sample_adaptive_offset_enabled_flag = r.u(1)
+    self.pcm_enabled_flag = r.u(1)
+    assert self.pcm_enabled_flag == 0, "pcm not implemented"
+    self.num_short_term_ref_pic_sets = r.ue_v()
+    for i in range(self.num_short_term_ref_pic_sets):
+      self.st_ref_pic_set(r, i, self.num_short_term_ref_pic_sets)
+    self.long_term_ref_pics_present_flag = r.u(1)
+    if self.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented"
+    self.sps_temporal_mvp_enabled_flag = r.u(1)
+    self.strong_intra_smoothing_enabled_flag = r.u(1)
+
+# 7.3.2.3 Picture parameter set RBSP syntax
+class PPS(HevcSlice):
+  def __init__(self, r:BitReader):
+    self.pps_pic_parameter_set_id = r.ue_v()
+    self.pps_seq_parameter_set_id = r.ue_v()
+    self.dependent_slice_segments_enabled_flag = r.u(1)
+    self.output_flag_present_flag = r.u(1)
+    self.num_extra_slice_header_bits = r.u(3)
+    self.sign_data_hiding_enabled_flag = r.u(1)
+    self.cabac_init_present_flag = r.u(1)
+    self.num_ref_idx_l0_default_active = r.ue_v() + 1
+    self.num_ref_idx_l1_default_active = r.ue_v() + 1
+    self.init_qp = r.se_v() + 26
+    self.constrained_intra_pred_flag = r.u(1)
+    self.transform_skip_enabled_flag = r.u(1)
+    self.cu_qp_delta_enabled_flag = r.u(1)
+    if self.cu_qp_delta_enabled_flag: self.diff_cu_qp_delta_depth = r.ue_v()
+
+    self.pps_cb_qp_offset = r.se_v()
+    self.pps_cr_qp_offset = r.se_v()
+    self.pps_slice_chroma_qp_offsets_present_flag = r.u(1)
+    self.weighted_pred_flag = r.u(1)
+    self.weighted_bipred_flag = r.u(1)
+    self.transquant_bypass_enabled_flag = r.u(1)
+    self.tiles_enabled_flag = r.u(1)
+    self.entropy_coding_sync_enabled_flag = r.u(1)
+    if self.tiles_enabled_flag:
+      self.num_tile_columns_minus1 = r.ue_v()
+      self.num_tile_rows_minus1 = r.ue_v()
+      self.uniform_spacing_flag = r.u(1)
+      self.column_width_minus1, self.row_height_minus1 = [], []
+      if not self.uniform_spacing_flag:
+        for i in range(self.num_tile_columns_minus1): self.column_width_minus1.append(r.ue_v())
+        for i in range(self.num_tile_rows_minus1): self.row_height_minus1.append(r.ue_v())
+      self.loop_filter_across_tiles_enabled_flag = r.u(1)
+    self.loop_filter_across_slices_enabled_flag = r.u(1)
+    self.deblocking_filter_control_present_flag = r.u(1)
+    if self.deblocking_filter_control_present_flag: assert False, "deblocking_filter parsing not implemented"
+    self.scaling_list_data_present_flag = r.u(1)
+    if self.scaling_list_data_present_flag: assert False, "scaling_list_data parsing not implemented"
+    self.lists_modification_present_flag = r.u(1)
+    self.log2_parallel_merge_level = r.ue_v() + 2
+
+# 7.3.6 Slice segment header syntax
+class SliceSegment(HevcSlice):
+  def __init__(self, r:BitReader, nal_unit_type:int, sps:SPS, pps:PPS):
+    self.first_slice_segment_in_pic_flag = r.u(1)
+    if nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23:
+      self.no_output_of_prior_pics_flag = r.u(1)
+    self.slice_pic_parameter_set_id = r.ue_v()
+    if not self.first_slice_segment_in_pic_flag:
+      if pps.dependent_slice_segments_enabled_flag:
+        self.dependent_slice_segment_flag = r.u(1)
+      self.slice_segment_address = r.ue_v()
+    self.dependent_slice_segment_flag = 0
+    if not self.dependent_slice_segment_flag:
+      r.u(pps.num_extra_slice_header_bits) # extra bits ignored
+      self.slice_type = r.ue_v()
+
+      self.sw_skip_start = r.read_bits - r.current_bits
+      self.pic_output_flag = r.u(1) if pps.output_flag_present_flag else 0
+      self.colour_plane_id = r.u(2) if sps.separate_colour_plane_flag else 0
+
+      if nal_unit_type != avcodec.HEVC_NAL_IDR_W_RADL and nal_unit_type != avcodec.HEVC_NAL_IDR_N_LP:
+        self.slice_pic_order_cnt_lsb = r.u(sps.log2_max_pic_order_cnt_lsb_minus4 + 4)
+
+        self.short_term_ref_pic_set_sps_flag = r.u(1)
+        if not self.short_term_ref_pic_set_sps_flag:
+          self.short_term_ref_pics_in_slice_start = r.read_bits - r.current_bits
+          self.st_ref_pic_set(r, sps.num_short_term_ref_pic_sets, sps=sps)
+          self.short_term_ref_pics_in_slice_end = r.read_bits - r.current_bits
+        elif sps.num_short_term_ref_pic_sets > 1: assert False, "short_term_ref_pic_set parsing not implemented"
+
+        if sps.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented"
+
+        self.sw_skip_end = r.read_bits - r.current_bits
+        self.slice_temporal_mvp_enabled_flag = r.u(1) if sps.sps_temporal_mvp_enabled_flag else 0
+      else: self.slice_pic_order_cnt_lsb, self.sw_skip_end = 0, self.sw_skip_start
+
+      if sps.sample_adaptive_offset_enabled_flag:
+        slice_sao_luma_flag = r.u(1)
+        ChromaArrayType = sps.chroma_format_idc if sps.separate_colour_plane_flag == 0 else 0
+        slice_sao_chroma_flag = r.u(1) if ChromaArrayType != 0 else 0
+
+      if self.slice_type in {avcodec.HEVC_SLICE_B, avcodec.HEVC_SLICE_B}:
+        if num_ref_idx_active_override_flag := r.u(1):
+          num_ref_idx_l0_active_minus1 = r.ue_v()
+          num_ref_idx_l1_active_minus1 = r.ue_v() if self.slice_type == avcodec.HEVC_SLICE_B else 0
+
+def fill_sps_into_dev_context(device_ctx, sps:SPS):
+  device_ctx.chroma_format_idc = sps.chroma_format_idc
+  device_ctx.pic_width_in_luma_samples = sps.pic_width_in_luma_samples
+  device_ctx.pic_height_in_luma_samples = sps.pic_height_in_luma_samples
+  device_ctx.bit_depth_luma = sps.bit_depth_luma
+  device_ctx.bit_depth_chroma = sps.bit_depth_chroma
+  device_ctx.log2_max_pic_order_cnt_lsb_minus4 = sps.log2_max_pic_order_cnt_lsb_minus4
+  device_ctx.log2_min_luma_coding_block_size = sps.log2_min_luma_coding_block_size
+  device_ctx.log2_max_luma_coding_block_size = sps.log2_max_luma_coding_block_size
+  device_ctx.log2_min_transform_block_size = sps.log2_min_transform_block_size
+  device_ctx.log2_max_transform_block_size = sps.log2_max_transform_block_size
+  device_ctx.amp_enabled_flag = sps.amp_enabled_flag
+  device_ctx.pcm_enabled_flag = sps.pcm_enabled_flag
+  device_ctx.sample_adaptive_offset_enabled_flag = sps.sample_adaptive_offset_enabled_flag
+  device_ctx.sps_temporal_mvp_enabled_flag = sps.sps_temporal_mvp_enabled_flag
+  device_ctx.strong_intra_smoothing_enabled_flag = sps.strong_intra_smoothing_enabled_flag
+
+def fill_pps_into_dev_context(device_ctx, pps:PPS):
+  device_ctx.sign_data_hiding_enabled_flag = pps.sign_data_hiding_enabled_flag
+  device_ctx.cabac_init_present_flag = pps.cabac_init_present_flag
+  device_ctx.num_ref_idx_l0_default_active = pps.num_ref_idx_l0_default_active
+  device_ctx.num_ref_idx_l1_default_active = pps.num_ref_idx_l1_default_active
+  device_ctx.init_qp = pps.init_qp
+  device_ctx.cu_qp_delta_enabled_flag = pps.cu_qp_delta_enabled_flag
+  device_ctx.diff_cu_qp_delta_depth = getattr(pps, 'diff_cu_qp_delta_depth', 0)
+  device_ctx.pps_cb_qp_offset = pps.pps_cb_qp_offset
+  device_ctx.pps_cr_qp_offset = pps.pps_cr_qp_offset
+  device_ctx.pps_slice_chroma_qp_offsets_present_flag = pps.pps_slice_chroma_qp_offsets_present_flag
+  device_ctx.weighted_pred_flag = pps.weighted_pred_flag
+  device_ctx.weighted_bipred_flag = pps.weighted_bipred_flag
+  device_ctx.transquant_bypass_enabled_flag = pps.transquant_bypass_enabled_flag
+  device_ctx.tiles_enabled_flag = pps.tiles_enabled_flag
+  device_ctx.entropy_coding_sync_enabled_flag = pps.entropy_coding_sync_enabled_flag
+  device_ctx.loop_filter_across_slices_enabled_flag = pps.loop_filter_across_slices_enabled_flag
+  device_ctx.deblocking_filter_control_present_flag = pps.deblocking_filter_control_present_flag
+  device_ctx.scaling_list_data_present_flag = pps.scaling_list_data_present_flag
+  device_ctx.lists_modification_present_flag = pps.lists_modification_present_flag
+  device_ctx.log2_parallel_merge_level = pps.log2_parallel_merge_level
+  device_ctx.loop_filter_across_tiles_enabled_flag = getattr(pps, 'loop_filter_across_tiles_enabled_flag', 0)
+
+def parse_hevc_file_headers(dat:bytes, device="NV"):
+  res = []
+  nal_unit_start = 1
+  history:list[tuple[int, int, int]] = []
+  device_ctx = nv_gpu.nvdec_hevc_pic_s(gptimer_timeout_value=92720000, tileformat=1, sw_start_code_e=1, pattern_id=2)
+  nal_infos = []
+  ctx_bytes = bytes()
+  align_ctx_bytes_size = 0x300
+
+  def _flush_picture():
+    nonlocal res, history, device_ctx, nal_infos, ctx_bytes, align_ctx_bytes_size
+
+    if not len(nal_infos): return
+
+    hdr, nal_unit_type = nal_infos[0][0]
+    assert all(nal_unit_type == x[0][1] for x in nal_infos), "all NAL units in a picture must be of the same type"
+
+    device_ctx.curr_pic_idx = next(i for i in range(16) if all(d[0] != i for d in history))
+
+    if nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP}:
+      history = []
+
+    device_ctx.num_ref_frames = len(history)
+    device_ctx.IDR_picture_flag = int(nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP})
+    device_ctx.RAP_picture_flag = int(nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23)
+    device_ctx.RefDiffPicOrderCnts=(ctypes.c_int16 * 16)()
+    device_ctx.colMvBuffersize = (round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64) // 16) // 256
+    device_ctx.framestride=(ctypes.c_uint32 * 2)(round_up(sps.pic_width_in_luma_samples, 64), round_up(sps.pic_width_in_luma_samples, 64))
+    device_ctx.sw_hdr_skip_length = hdr.sw_skip_end - hdr.sw_skip_start
+    device_ctx.num_bits_short_term_ref_pics_in_slice = max(0, device_ctx.sw_hdr_skip_length - 9)
+    device_ctx.stream_len = sum(x[2] for x in nal_infos)
+
+    if pps.tiles_enabled_flag:
+      device_ctx.num_tile_columns = pps.num_tile_columns_minus1 + 1
+      device_ctx.num_tile_rows = pps.num_tile_rows_minus1 + 1
+
+    device_ctx.num_short_term_ref_pic_sets = sps.num_short_term_ref_pic_sets
+
+    luma_h_rounded = round_up(sps.pic_height_in_luma_samples, 64)
+    device_ctx.HevcSaoBufferOffset = (608 * luma_h_rounded) >> 8
+    device_ctx.HevcBsdCtrlOffset = ((device_ctx.HevcSaoBufferOffset<<8) + 4864 * luma_h_rounded) >> 8
+
+    device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset = ((device_ctx.HevcBsdCtrlOffset<<8) + 152 * luma_h_rounded) >> 8
+    device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset = ((device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset<<8) + 2000 * luma_h_rounded) >> 8
+    device_ctx.v3.HevcSliceEdgeOffset = device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset
+
+    before_list, after_list = [], []
+    for pic_idx, poc, _ in history:
+      device_ctx.RefDiffPicOrderCnts[pic_idx] = hdr.slice_pic_order_cnt_lsb - poc
+      if hdr.slice_pic_order_cnt_lsb < poc: after_list.append((poc - hdr.slice_pic_order_cnt_lsb, pic_idx))
+      else: before_list.append((hdr.slice_pic_order_cnt_lsb - poc, pic_idx))
+    before_list.sort()
+    after_list.sort()
+
+    device_ctx.initreflistidxl0 = (ctypes.c_uint8 * 16)(*[idx for _,idx in before_list + after_list])
+    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))
+    pic_height_in_ctbs = ceildiv(sps.pic_height_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size))
+    # append tile sizes 0x200
+    if pps.tiles_enabled_flag and pps.uniform_spacing_flag:
+      assert device_ctx.num_tile_columns == 1 and device_ctx.num_tile_rows == 1, "not implemented: uniform spacing with multiple tiles"
+      locl_ctx_bytes += pic_width_in_ctbs.to_bytes(2, "little") + pic_height_in_ctbs.to_bytes(2, "little")
+    else:
+      if pps.tiles_enabled_flag and not getattr(pps, 'uniform_spacing_flag', 0):
+        column_width = [cw_minus1 + 1 for cw_minus1 in pps.column_width_minus1[0:pps.num_tile_columns_minus1]]
+        row_height = [rh_minus1 + 1 for rh_minus1 in pps.row_height_minus1[0:pps.num_tile_rows_minus1]]
+      else:
+        column_width = []
+        row_height = []
+
+      column_width.append(pic_width_in_ctbs - sum(column_width))
+      row_height.append(pic_height_in_ctbs - sum(row_height))
+
+      for c in column_width:
+        for r in row_height: locl_ctx_bytes += c.to_bytes(2, "little") + r.to_bytes(2, "little")
+
+    luma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64)
+    chroma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up((sps.pic_height_in_luma_samples + 1) // 2, 64)
+    is_hist = nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL}
+
+    res.append((nal_infos[0][1], device_ctx.stream_len, device_ctx.curr_pic_idx, len(history), is_hist))
+
+    locl_ctx_bytes += (align_ctx_bytes_size - len(locl_ctx_bytes)) * b'\x00'
+    ctx_bytes += locl_ctx_bytes
+
+    if nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL}:
+      history.append((device_ctx.curr_pic_idx, hdr.slice_pic_order_cnt_lsb, None))
+
+    if len(history) >= sps.sps_max_dec_pic_buffering[0]:
+      # remove the oldest poc
+      history.pop(0)
+
+    nal_infos = []
+
+  cnt = 0
+  while nal_unit_start < len(dat):
+    assert dat[nal_unit_start:nal_unit_start+3] == b"\x00\x00\x01", "NAL unit start code not found"
+
+    pos = dat.find(b"\x00\x00\x01", nal_unit_start + 3)
+    nal_unit_len = (pos if pos != -1 else len(dat)) - nal_unit_start
+
+    # 7.3.1.1 General NAL unit syntax
+    nal_unit_type = (dat[nal_unit_start+3] >> 1) & 0x3F
+    slice_dat = dat[nal_unit_start+5:nal_unit_start+nal_unit_len]
+
+    if nal_unit_type == avcodec.HEVC_NAL_SPS:
+      sps = SPS(BitReader(_hevc_get_rbsp(slice_dat)))
+      fill_sps_into_dev_context(device_ctx, sps)
+    elif nal_unit_type == avcodec.HEVC_NAL_PPS:
+      pps = PPS(BitReader(_hevc_get_rbsp(slice_dat)))
+      fill_pps_into_dev_context(device_ctx, pps)
+    elif nal_unit_type in {avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_TRAIL_N}:
+      hdr = SliceSegment(BitReader(slice_dat), nal_unit_type, sps, pps)
+
+      if hdr.first_slice_segment_in_pic_flag == 1: _flush_picture()
+      nal_infos.append(((hdr, nal_unit_type), nal_unit_start, nal_unit_len))
+
+    nal_unit_start += nal_unit_len
+  _flush_picture()
+
+  w = sps.pic_width_in_luma_samples - 2 * (sps.conf_win_left_offset + sps.conf_win_right_offset)
+  h = sps.pic_height_in_luma_samples - 2 * (sps.conf_win_top_offset  + sps.conf_win_bottom_offset)
+  chroma_off = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64)
+  opaque = Tensor(ctx_bytes, device=device).reshape(len(res), align_ctx_bytes_size)
+  return opaque, res, w, h, sps.pic_width_in_luma_samples, sps.pic_height_in_luma_samples, chroma_off
+
+def _addr_table(h, w, w_aligned):
+  GOB_W, GOB_H = 64, 8
+  GOB_SIZE = GOB_W * GOB_H
+  BLOCK_H_GOBS = 2
+
+  xs = Tensor.arange(w, dtype=dtypes.uint32).reshape(1, w)
+  ys = Tensor.arange(h, dtype=dtypes.uint32).reshape(h, 1)
+
+  gob_x = xs // GOB_W
+  gob_y = ys // GOB_H
+  super_block_y = gob_y // BLOCK_H_GOBS
+  gob_y_in_block = gob_y  % BLOCK_H_GOBS
+  stride_gobs = w_aligned // GOB_W
+
+  base = ((super_block_y * stride_gobs + gob_x) * BLOCK_H_GOBS + gob_y_in_block) * GOB_SIZE
+
+  lx, ly = xs % GOB_W, ys % GOB_H
+  swiz = (lx & 0x0F) | ((ly & 0x03) << 4) | ((lx & 0x10) << 2) | ((ly & 0x04) << 5) | ((lx & 0x20) << 3)
+  return (base + swiz).reshape(-1)
+
+def nv12_to_bgr_from_planes(luma: Tensor, chroma: Tensor, h: int, w: int) -> Tensor:
+  Y = luma.reshape(h, w).cast(dtypes.float32)
+
+  uv = chroma.reshape(h // 2, w // 2, 2).cast(dtypes.float32)
+  U_small = uv[..., 0]
+  V_small = uv[..., 1]
+
+  U = U_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w)
+  V = V_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w)
+
+  C = Y - 16.0
+  D = U - 128.0
+  E = V - 128.0
+
+  R = 1.1643835616438356 * C + 1.5960267857142858 * E
+  G = 1.1643835616438356 * C - 0.39176229009491365 * D - 0.8129676472377708 * E
+  B = 1.1643835616438356 * C + 2.017232142857143  * D
+
+  R = R.maximum(0.0).minimum(255.0)
+  G = G.maximum(0.0).minimum(255.0)
+  B = B.maximum(0.0).minimum(255.0)
+
+  return Tensor.stack([B, G, R], dim=2).cast(dtypes.uint8)
+
+def untile_nv12(src:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor:
+  luma = src.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))]
+  chroma = src.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))]
+  return luma.cat(chroma).realize()
+
+def to_bgr(tensor:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor:
+  luma = tensor.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))]
+  chroma = tensor.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))]
+  return nv12_to_bgr_from_planes(luma, chroma, h, w).realize()

extra/hip_gpu_driver/hip_ioctl.py

@@ -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.kfd_ioctl_svm_attr_type__enumvalues[out[i].type]:40s}: {out[i].value:#x}")
+      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}")
   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 "")

extra/mmapeak/mmapeak.py

@@ -48,7 +48,7 @@ if __name__=="__main__":
   COMPILER = HIPCompiler(DEV.arch)
   if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
     if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
-    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 40
+    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 32
     launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
     launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
     launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))

extra/nv_gpu_driver/clc9b0.h

@@ -0,0 +1,603 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef clc9b0_h_
+#define clc9b0_h_
+
+#include "nvtypes.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define NVC9B0_VIDEO_DECODER                                                       (0x0000C9B0)
+
+#define NVC9B0_NOP                                                              (0x00000100)
+#define NVC9B0_NOP_PARAMETER                                                    31:0
+#define NVC9B0_PM_TRIGGER                                                       (0x00000140)
+#define NVC9B0_PM_TRIGGER_V                                                     31:0
+#define NVC9B0_SET_APPLICATION_ID                                               (0x00000200)
+#define NVC9B0_SET_APPLICATION_ID_ID                                            31:0
+#define NVC9B0_SET_APPLICATION_ID_ID_MPEG12                                     (0x00000001)
+#define NVC9B0_SET_APPLICATION_ID_ID_VC1                                        (0x00000002)
+#define NVC9B0_SET_APPLICATION_ID_ID_H264                                       (0x00000003)
+#define NVC9B0_SET_APPLICATION_ID_ID_MPEG4                                      (0x00000004)
+#define NVC9B0_SET_APPLICATION_ID_ID_VP8                                        (0x00000005)
+#define NVC9B0_SET_APPLICATION_ID_ID_CTR64                                      (0x00000006)
+#define NVC9B0_SET_APPLICATION_ID_ID_HEVC                                       (0x00000007)
+#define NVC9B0_SET_APPLICATION_ID_ID_NEW_H264                                   (0x00000008)
+#define NVC9B0_SET_APPLICATION_ID_ID_VP9                                        (0x00000009)
+#define NVC9B0_SET_APPLICATION_ID_ID_PASS1                                      (0x0000000A)
+#define NVC9B0_SET_APPLICATION_ID_ID_HEVC_PARSER                                (0x0000000C)
+#define NVC9B0_SET_APPLICATION_ID_ID_UCODE_TEST                                 (0x0000000D)
+#define NVC9B0_SET_APPLICATION_ID_ID_HWDRM_PR_DECRYPTAUDIO                      (0x0000000E)
+#define NVC9B0_SET_APPLICATION_ID_ID_HWDRM_PR_DECRYPTAUDIOMULTIPLE              (0x0000000F)
+#define NVC9B0_SET_APPLICATION_ID_ID_HWDRM_PR_PREPROCESSENCRYPTEDDATA           (0x00000010)
+#define NVC9B0_SET_APPLICATION_ID_ID_VP9_WITH_PARSER                            (0x00000011)
+#define NVC9B0_SET_APPLICATION_ID_ID_AVD                                        (0x00000012)
+#define NVC9B0_SET_APPLICATION_ID_ID_HW_DRM_PR4_DECRYPTCONTENTMULTIPLE          (0x00000013)
+#define NVC9B0_SET_APPLICATION_ID_ID_DHKE                                       (0x00000020)
+#define NVC9B0_SET_WATCHDOG_TIMER                                               (0x00000204)
+#define NVC9B0_SET_WATCHDOG_TIMER_TIMER                                         31:0
+#define NVC9B0_SEMAPHORE_A                                                      (0x00000240)
+#define NVC9B0_SEMAPHORE_A_UPPER                                                7:0
+#define NVC9B0_SEMAPHORE_B                                                      (0x00000244)
+#define NVC9B0_SEMAPHORE_B_LOWER                                                31:0
+#define NVC9B0_SEMAPHORE_C                                                      (0x00000248)
+#define NVC9B0_SEMAPHORE_C_PAYLOAD                                              31:0
+#define NVC9B0_CTX_SAVE_AREA                                                    (0x0000024C)
+#define NVC9B0_CTX_SAVE_AREA_OFFSET                                             31:0
+#define NVC9B0_CTX_SWITCH                                                       (0x00000250)
+#define NVC9B0_CTX_SWITCH_OP                                                    1:0
+#define NVC9B0_CTX_SWITCH_OP_CTX_UPDATE                                         (0x00000000)
+#define NVC9B0_CTX_SWITCH_OP_CTX_SAVE                                           (0x00000001)
+#define NVC9B0_CTX_SWITCH_OP_CTX_RESTORE                                        (0x00000002)
+#define NVC9B0_CTX_SWITCH_OP_CTX_FORCERESTORE                                   (0x00000003)
+#define NVC9B0_CTX_SWITCH_CTXID_VALID                                           2:2
+#define NVC9B0_CTX_SWITCH_CTXID_VALID_FALSE                                     (0x00000000)
+#define NVC9B0_CTX_SWITCH_CTXID_VALID_TRUE                                      (0x00000001)
+#define NVC9B0_CTX_SWITCH_RESERVED0                                             7:3
+#define NVC9B0_CTX_SWITCH_CTX_ID                                                23:8
+#define NVC9B0_CTX_SWITCH_RESERVED1                                             31:24
+#define NVC9B0_SET_SEMAPHORE_PAYLOAD_LOWER                                      (0x00000254)
+#define NVC9B0_SET_SEMAPHORE_PAYLOAD_LOWER_PAYLOAD_LOWER                        31:0
+#define NVC9B0_SET_SEMAPHORE_PAYLOAD_UPPER                                      (0x00000258)
+#define NVC9B0_SET_SEMAPHORE_PAYLOAD_UPPER_PAYLOAD_UPPER                        31:0
+#define NVC9B0_SET_MONITORED_FENCE_SIGNAL_ADDRESS_BASE_A                        (0x0000025C)
+#define NVC9B0_SET_MONITORED_FENCE_SIGNAL_ADDRESS_BASE_A_LOWER                  31:0
+#define NVC9B0_SET_MONITORED_FENCE_SIGNAL_ADDRESS_BASE_B                        (0x00000260)
+#define NVC9B0_SET_MONITORED_FENCE_SIGNAL_ADDRESS_BASE_B_UPPER                  31:0
+#define NVC9B0_EXECUTE                                                          (0x00000300)
+#define NVC9B0_EXECUTE_NOTIFY                                                   0:0
+#define NVC9B0_EXECUTE_NOTIFY_DISABLE                                           (0x00000000)
+#define NVC9B0_EXECUTE_NOTIFY_ENABLE                                            (0x00000001)
+#define NVC9B0_EXECUTE_NOTIFY_ON                                                1:1
+#define NVC9B0_EXECUTE_NOTIFY_ON_END                                            (0x00000000)
+#define NVC9B0_EXECUTE_NOTIFY_ON_BEGIN                                          (0x00000001)
+#define NVC9B0_EXECUTE_PREDICATION                                              2:2
+#define NVC9B0_EXECUTE_PREDICATION_DISABLE                                      (0x00000000)
+#define NVC9B0_EXECUTE_PREDICATION_ENABLE                                       (0x00000001)
+#define NVC9B0_EXECUTE_PREDICATION_OP                                           3:3
+#define NVC9B0_EXECUTE_PREDICATION_OP_EQUAL_ZERO                                (0x00000000)
+#define NVC9B0_EXECUTE_PREDICATION_OP_NOT_EQUAL_ZERO                            (0x00000001)
+#define NVC9B0_EXECUTE_AWAKEN                                                   8:8
+#define NVC9B0_EXECUTE_AWAKEN_DISABLE                                           (0x00000000)
+#define NVC9B0_EXECUTE_AWAKEN_ENABLE                                            (0x00000001)
+#define NVC9B0_SEMAPHORE_D                                                      (0x00000304)
+#define NVC9B0_SEMAPHORE_D_STRUCTURE_SIZE                                       1:0
+#define NVC9B0_SEMAPHORE_D_STRUCTURE_SIZE_ONE                                   (0x00000000)
+#define NVC9B0_SEMAPHORE_D_STRUCTURE_SIZE_FOUR                                  (0x00000001)
+#define NVC9B0_SEMAPHORE_D_STRUCTURE_SIZE_TWO                                   (0x00000002)
+#define NVC9B0_SEMAPHORE_D_AWAKEN_ENABLE                                        8:8
+#define NVC9B0_SEMAPHORE_D_AWAKEN_ENABLE_FALSE                                  (0x00000000)
+#define NVC9B0_SEMAPHORE_D_AWAKEN_ENABLE_TRUE                                   (0x00000001)
+#define NVC9B0_SEMAPHORE_D_OPERATION                                            17:16
+#define NVC9B0_SEMAPHORE_D_OPERATION_RELEASE                                    (0x00000000)
+#define NVC9B0_SEMAPHORE_D_OPERATION_RESERVED_0                                 (0x00000001)
+#define NVC9B0_SEMAPHORE_D_OPERATION_RESERVED_1                                 (0x00000002)
+#define NVC9B0_SEMAPHORE_D_OPERATION_TRAP                                       (0x00000003)
+#define NVC9B0_SEMAPHORE_D_FLUSH_DISABLE                                        21:21
+#define NVC9B0_SEMAPHORE_D_FLUSH_DISABLE_FALSE                                  (0x00000000)
+#define NVC9B0_SEMAPHORE_D_FLUSH_DISABLE_TRUE                                   (0x00000001)
+#define NVC9B0_SEMAPHORE_D_TRAP_TYPE                                            23:22
+#define NVC9B0_SEMAPHORE_D_TRAP_TYPE_UNCONDITIONAL                              (0x00000000)
+#define NVC9B0_SEMAPHORE_D_TRAP_TYPE_CONDITIONAL                                (0x00000001)
+#define NVC9B0_SEMAPHORE_D_TRAP_TYPE_CONDITIONAL_EXT                            (0x00000002)
+#define NVC9B0_SEMAPHORE_D_PAYLOAD_SIZE                                         24:24
+#define NVC9B0_SEMAPHORE_D_PAYLOAD_SIZE_32BIT                                   (0x00000000)
+#define NVC9B0_SEMAPHORE_D_PAYLOAD_SIZE_64BIT                                   (0x00000001)
+#define NVC9B0_SET_PREDICATION_OFFSET_UPPER                                     (0x00000308)
+#define NVC9B0_SET_PREDICATION_OFFSET_UPPER_OFFSET                              7:0
+#define NVC9B0_SET_PREDICATION_OFFSET_LOWER                                     (0x0000030C)
+#define NVC9B0_SET_PREDICATION_OFFSET_LOWER_OFFSET                              31:0
+#define NVC9B0_SET_AUXILIARY_DATA_BUFFER                                        (0x00000310)
+#define NVC9B0_SET_AUXILIARY_DATA_BUFFER_OFFSET                                 31:0
+#define NVC9B0_SET_CONTROL_PARAMS                                               (0x00000400)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE                                    3:0
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_MPEG1                              (0x00000000)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_MPEG2                              (0x00000001)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_VC1                                (0x00000002)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_H264                               (0x00000003)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_MPEG4                              (0x00000004)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_DIVX3                              (0x00000004)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_VP8                                (0x00000005)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_HEVC                               (0x00000007)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_VP9                                (0x00000009)
+#define NVC9B0_SET_CONTROL_PARAMS_CODEC_TYPE_AV1                                (0x0000000A)
+#define NVC9B0_SET_CONTROL_PARAMS_GPTIMER_ON                                    4:4
+#define NVC9B0_SET_CONTROL_PARAMS_RET_ERROR                                     5:5
+#define NVC9B0_SET_CONTROL_PARAMS_ERR_CONCEAL_ON                                6:6
+#define NVC9B0_SET_CONTROL_PARAMS_ERROR_FRM_IDX                                 12:7
+#define NVC9B0_SET_CONTROL_PARAMS_MBTIMER_ON                                    13:13
+#define NVC9B0_SET_CONTROL_PARAMS_EC_INTRA_FRAME_USING_PSLC                     14:14
+#define NVC9B0_SET_CONTROL_PARAMS_IGNORE_SOME_FIELDS_CRC_CHECK                  15:15
+#define NVC9B0_SET_CONTROL_PARAMS_EVENT_TRACE_LOGGING_ON                        16:16
+#define NVC9B0_SET_CONTROL_PARAMS_ALL_INTRA_FRAME                               17:17
+#define NVC9B0_SET_CONTROL_PARAMS_TESTRUN_ENV                                   19:18
+#define NVC9B0_SET_CONTROL_PARAMS_TESTRUN_ENV_TRACE3D_RUN                       (0x00000000)
+#define NVC9B0_SET_CONTROL_PARAMS_TESTRUN_ENV_PROD_RUN                          (0x00000001)
+#define NVC9B0_SET_CONTROL_PARAMS_HINT_DUMP_EN                                  20:20
+#define NVC9B0_SET_CONTROL_PARAMS_RESERVED                                      25:21
+#define NVC9B0_SET_CONTROL_PARAMS_NVDECSIM_SKIP_SCP                             26:26
+#define NVC9B0_SET_CONTROL_PARAMS_ENABLE_ENCRYPT                                27:27
+#define NVC9B0_SET_CONTROL_PARAMS_ENCRYPTMODE                                   31:28
+#define NVC9B0_SET_DRV_PIC_SETUP_OFFSET                                         (0x00000404)
+#define NVC9B0_SET_DRV_PIC_SETUP_OFFSET_OFFSET                                  31:0
+#define NVC9B0_SET_IN_BUF_BASE_OFFSET                                           (0x00000408)
+#define NVC9B0_SET_IN_BUF_BASE_OFFSET_OFFSET                                    31:0
+#define NVC9B0_SET_PICTURE_INDEX                                                (0x0000040C)
+#define NVC9B0_SET_PICTURE_INDEX_INDEX                                          31:0
+#define NVC9B0_SET_SLICE_OFFSETS_BUF_OFFSET                                     (0x00000410)
+#define NVC9B0_SET_SLICE_OFFSETS_BUF_OFFSET_OFFSET                              31:0
+#define NVC9B0_SET_COLOC_DATA_OFFSET                                            (0x00000414)
+#define NVC9B0_SET_COLOC_DATA_OFFSET_OFFSET                                     31:0
+#define NVC9B0_SET_HISTORY_OFFSET                                               (0x00000418)
+#define NVC9B0_SET_HISTORY_OFFSET_OFFSET                                        31:0
+#define NVC9B0_SET_DISPLAY_BUF_SIZE                                             (0x0000041C)
+#define NVC9B0_SET_DISPLAY_BUF_SIZE_SIZE                                        31:0
+#define NVC9B0_SET_HISTOGRAM_OFFSET                                             (0x00000420)
+#define NVC9B0_SET_HISTOGRAM_OFFSET_OFFSET                                      31:0
+#define NVC9B0_SET_NVDEC_STATUS_OFFSET                                          (0x00000424)
+#define NVC9B0_SET_NVDEC_STATUS_OFFSET_OFFSET                                   31:0
+#define NVC9B0_SET_DISPLAY_BUF_LUMA_OFFSET                                      (0x00000428)
+#define NVC9B0_SET_DISPLAY_BUF_LUMA_OFFSET_OFFSET                               31:0
+#define NVC9B0_SET_DISPLAY_BUF_CHROMA_OFFSET                                    (0x0000042C)
+#define NVC9B0_SET_DISPLAY_BUF_CHROMA_OFFSET_OFFSET                             31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET0                                         (0x00000430)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET0_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET1                                         (0x00000434)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET1_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET2                                         (0x00000438)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET2_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET3                                         (0x0000043C)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET3_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET4                                         (0x00000440)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET4_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET5                                         (0x00000444)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET5_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET6                                         (0x00000448)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET6_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET7                                         (0x0000044C)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET7_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET8                                         (0x00000450)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET8_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET9                                         (0x00000454)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET9_OFFSET                                  31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET10                                        (0x00000458)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET10_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET11                                        (0x0000045C)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET11_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET12                                        (0x00000460)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET12_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET13                                        (0x00000464)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET13_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET14                                        (0x00000468)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET14_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET15                                        (0x0000046C)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET15_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET16                                        (0x00000470)
+#define NVC9B0_SET_PICTURE_LUMA_OFFSET16_OFFSET                                 31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET0                                       (0x00000474)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET0_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET1                                       (0x00000478)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET1_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET2                                       (0x0000047C)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET2_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET3                                       (0x00000480)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET3_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET4                                       (0x00000484)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET4_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET5                                       (0x00000488)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET5_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET6                                       (0x0000048C)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET6_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET7                                       (0x00000490)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET7_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET8                                       (0x00000494)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET8_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET9                                       (0x00000498)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET9_OFFSET                                31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET10                                      (0x0000049C)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET10_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET11                                      (0x000004A0)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET11_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET12                                      (0x000004A4)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET12_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET13                                      (0x000004A8)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET13_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET14                                      (0x000004AC)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET14_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET15                                      (0x000004B0)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET15_OFFSET                               31:0
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET16                                      (0x000004B4)
+#define NVC9B0_SET_PICTURE_CHROMA_OFFSET16_OFFSET                               31:0
+#define NVC9B0_SET_PIC_SCRATCH_BUF_OFFSET                                       (0x000004B8)
+#define NVC9B0_SET_PIC_SCRATCH_BUF_OFFSET_OFFSET                                31:0
+#define NVC9B0_SET_EXTERNAL_MVBUFFER_OFFSET                                     (0x000004BC)
+#define NVC9B0_SET_EXTERNAL_MVBUFFER_OFFSET_OFFSET                              31:0
+#define NVC9B0_SET_SUB_SAMPLE_MAP_OFFSET                                        (0x000004C0)
+#define NVC9B0_SET_SUB_SAMPLE_MAP_OFFSET_OFFSET                                 31:0
+#define NVC9B0_SET_SUB_SAMPLE_MAP_IV_OFFSET                                     (0x000004C4)
+#define NVC9B0_SET_SUB_SAMPLE_MAP_IV_OFFSET_OFFSET                              31:0
+#define NVC9B0_SET_INTRA_TOP_BUF_OFFSET                                         (0x000004C8)
+#define NVC9B0_SET_INTRA_TOP_BUF_OFFSET_OFFSET                                  31:0
+#define NVC9B0_SET_TILE_SIZE_BUF_OFFSET                                         (0x000004CC)
+#define NVC9B0_SET_TILE_SIZE_BUF_OFFSET_OFFSET                                  31:0
+#define NVC9B0_SET_FILTER_BUFFER_OFFSET                                         (0x000004D0)
+#define NVC9B0_SET_FILTER_BUFFER_OFFSET_OFFSET                                  31:0
+#define NVC9B0_SET_CRC_STRUCT_OFFSET                                            (0x000004D4)
+#define NVC9B0_SET_CRC_STRUCT_OFFSET_OFFSET                                     31:0
+#define NVC9B0_SET_PR_SSM_CONTENT_INFO_BUF_OFFSET                               (0x000004D8)
+#define NVC9B0_SET_PR_SSM_CONTENT_INFO_BUF_OFFSET_OFFSET                        31:0
+#define NVC9B0_H264_SET_MBHIST_BUF_OFFSET                                       (0x00000500)
+#define NVC9B0_H264_SET_MBHIST_BUF_OFFSET_OFFSET                                31:0
+#define NVC9B0_VP8_SET_PROB_DATA_OFFSET                                         (0x00000540)
+#define NVC9B0_VP8_SET_PROB_DATA_OFFSET_OFFSET                                  31:0
+#define NVC9B0_VP8_SET_HEADER_PARTITION_BUF_BASE_OFFSET                         (0x00000544)
+#define NVC9B0_VP8_SET_HEADER_PARTITION_BUF_BASE_OFFSET_OFFSET                  31:0
+#define NVC9B0_HEVC_SET_SCALING_LIST_OFFSET                                     (0x00000580)
+#define NVC9B0_HEVC_SET_SCALING_LIST_OFFSET_OFFSET                              31:0
+#define NVC9B0_HEVC_SET_TILE_SIZES_OFFSET                                       (0x00000584)
+#define NVC9B0_HEVC_SET_TILE_SIZES_OFFSET_OFFSET                                31:0
+#define NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET                                    (0x00000588)
+#define NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET_OFFSET                             31:0
+#define NVC9B0_HEVC_SET_SAO_BUFFER_OFFSET                                       (0x0000058C)
+#define NVC9B0_HEVC_SET_SAO_BUFFER_OFFSET_OFFSET                                31:0
+#define NVC9B0_HEVC_SET_SLICE_INFO_BUFFER_OFFSET                                (0x00000590)
+#define NVC9B0_HEVC_SET_SLICE_INFO_BUFFER_OFFSET_OFFSET                         31:0
+#define NVC9B0_HEVC_SET_SLICE_GROUP_INDEX                                       (0x00000594)
+#define NVC9B0_HEVC_SET_SLICE_GROUP_INDEX_OFFSET                                31:0
+#define NVC9B0_VP9_SET_PROB_TAB_BUF_OFFSET                                      (0x000005C0)
+#define NVC9B0_VP9_SET_PROB_TAB_BUF_OFFSET_OFFSET                               31:0
+#define NVC9B0_VP9_SET_CTX_COUNTER_BUF_OFFSET                                   (0x000005C4)
+#define NVC9B0_VP9_SET_CTX_COUNTER_BUF_OFFSET_OFFSET                            31:0
+#define NVC9B0_VP9_SET_SEGMENT_READ_BUF_OFFSET                                  (0x000005C8)
+#define NVC9B0_VP9_SET_SEGMENT_READ_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_VP9_SET_SEGMENT_WRITE_BUF_OFFSET                                 (0x000005CC)
+#define NVC9B0_VP9_SET_SEGMENT_WRITE_BUF_OFFSET_OFFSET                          31:0
+#define NVC9B0_VP9_SET_TILE_SIZE_BUF_OFFSET                                     (0x000005D0)
+#define NVC9B0_VP9_SET_TILE_SIZE_BUF_OFFSET_OFFSET                              31:0
+#define NVC9B0_VP9_SET_COL_MVWRITE_BUF_OFFSET                                   (0x000005D4)
+#define NVC9B0_VP9_SET_COL_MVWRITE_BUF_OFFSET_OFFSET                            31:0
+#define NVC9B0_VP9_SET_COL_MVREAD_BUF_OFFSET                                    (0x000005D8)
+#define NVC9B0_VP9_SET_COL_MVREAD_BUF_OFFSET_OFFSET                             31:0
+#define NVC9B0_VP9_SET_FILTER_BUFFER_OFFSET                                     (0x000005DC)
+#define NVC9B0_VP9_SET_FILTER_BUFFER_OFFSET_OFFSET                              31:0
+#define NVC9B0_VP9_PARSER_SET_PIC_SETUP_OFFSET                                  (0x000005E0)
+#define NVC9B0_VP9_PARSER_SET_PIC_SETUP_OFFSET_OFFSET                           31:0
+#define NVC9B0_VP9_PARSER_SET_PREV_PIC_SETUP_OFFSET                             (0x000005E4)
+#define NVC9B0_VP9_PARSER_SET_PREV_PIC_SETUP_OFFSET_OFFSET                      31:0
+#define NVC9B0_VP9_PARSER_SET_PROB_TAB_BUF_OFFSET                               (0x000005E8)
+#define NVC9B0_VP9_PARSER_SET_PROB_TAB_BUF_OFFSET_OFFSET                        31:0
+#define NVC9B0_VP9_SET_HINT_DUMP_BUF_OFFSET                                     (0x000005EC)
+#define NVC9B0_VP9_SET_HINT_DUMP_BUF_OFFSET_OFFSET                              31:0
+#define NVC9B0_PASS1_SET_CLEAR_HEADER_OFFSET                                    (0x00000600)
+#define NVC9B0_PASS1_SET_CLEAR_HEADER_OFFSET_OFFSET                             31:0
+#define NVC9B0_PASS1_SET_RE_ENCRYPT_OFFSET                                      (0x00000604)
+#define NVC9B0_PASS1_SET_RE_ENCRYPT_OFFSET_OFFSET                               31:0
+#define NVC9B0_PASS1_SET_VP8_TOKEN_OFFSET                                       (0x00000608)
+#define NVC9B0_PASS1_SET_VP8_TOKEN_OFFSET_OFFSET                                31:0
+#define NVC9B0_PASS1_SET_INPUT_DATA_OFFSET                                      (0x0000060C)
+#define NVC9B0_PASS1_SET_INPUT_DATA_OFFSET_OFFSET                               31:0
+#define NVC9B0_PASS1_SET_OUTPUT_DATA_SIZE_OFFSET                                (0x00000610)
+#define NVC9B0_PASS1_SET_OUTPUT_DATA_SIZE_OFFSET_OFFSET                         31:0
+#define NVC9B0_AV1_SET_PROB_TAB_READ_BUF_OFFSET                                 (0x00000640)
+#define NVC9B0_AV1_SET_PROB_TAB_READ_BUF_OFFSET_OFFSET                          31:0
+#define NVC9B0_AV1_SET_PROB_TAB_WRITE_BUF_OFFSET                                (0x00000644)
+#define NVC9B0_AV1_SET_PROB_TAB_WRITE_BUF_OFFSET_OFFSET                         31:0
+#define NVC9B0_AV1_SET_SEGMENT_READ_BUF_OFFSET                                  (0x00000648)
+#define NVC9B0_AV1_SET_SEGMENT_READ_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_AV1_SET_SEGMENT_WRITE_BUF_OFFSET                                 (0x0000064C)
+#define NVC9B0_AV1_SET_SEGMENT_WRITE_BUF_OFFSET_OFFSET                          31:0
+#define NVC9B0_AV1_SET_COL_MV0_READ_BUF_OFFSET                                  (0x00000650)
+#define NVC9B0_AV1_SET_COL_MV0_READ_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_AV1_SET_COL_MV1_READ_BUF_OFFSET                                  (0x00000654)
+#define NVC9B0_AV1_SET_COL_MV1_READ_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_AV1_SET_COL_MV2_READ_BUF_OFFSET                                  (0x00000658)
+#define NVC9B0_AV1_SET_COL_MV2_READ_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_AV1_SET_COL_MVWRITE_BUF_OFFSET                                   (0x0000065C)
+#define NVC9B0_AV1_SET_COL_MVWRITE_BUF_OFFSET_OFFSET                            31:0
+#define NVC9B0_AV1_SET_GLOBAL_MODEL_BUF_OFFSET                                  (0x00000660)
+#define NVC9B0_AV1_SET_GLOBAL_MODEL_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_AV1_SET_FILM_GRAIN_BUF_OFFSET                                    (0x00000664)
+#define NVC9B0_AV1_SET_FILM_GRAIN_BUF_OFFSET_OFFSET                             31:0
+#define NVC9B0_AV1_SET_TILE_STREAM_INFO_BUF_OFFSET                              (0x00000668)
+#define NVC9B0_AV1_SET_TILE_STREAM_INFO_BUF_OFFSET_OFFSET                       31:0
+#define NVC9B0_AV1_SET_SUB_STREAM_ENTRY_BUF_OFFSET                              (0x0000066C)
+#define NVC9B0_AV1_SET_SUB_STREAM_ENTRY_BUF_OFFSET_OFFSET                       31:0
+#define NVC9B0_AV1_SET_HINT_DUMP_BUF_OFFSET                                     (0x00000670)
+#define NVC9B0_AV1_SET_HINT_DUMP_BUF_OFFSET_OFFSET                              31:0
+#define NVC9B0_H264_SET_SCALING_LIST_OFFSET                                     (0x00000680)
+#define NVC9B0_H264_SET_SCALING_LIST_OFFSET_OFFSET                              31:0
+#define NVC9B0_H264_SET_VLDHIST_BUF_OFFSET                                      (0x00000684)
+#define NVC9B0_H264_SET_VLDHIST_BUF_OFFSET_OFFSET                               31:0
+#define NVC9B0_H264_SET_EDOBOFFSET0                                             (0x00000688)
+#define NVC9B0_H264_SET_EDOBOFFSET0_OFFSET                                      31:0
+#define NVC9B0_H264_SET_EDOBOFFSET1                                             (0x0000068C)
+#define NVC9B0_H264_SET_EDOBOFFSET1_OFFSET                                      31:0
+#define NVC9B0_H264_SET_EDOBOFFSET2                                             (0x00000690)
+#define NVC9B0_H264_SET_EDOBOFFSET2_OFFSET                                      31:0
+#define NVC9B0_H264_SET_EDOBOFFSET3                                             (0x00000694)
+#define NVC9B0_H264_SET_EDOBOFFSET3_OFFSET                                      31:0
+#define NVC9B0_SET_CONTENT_INITIAL_VECTOR(b)                                    (0x00000C00 + (b)*0x00000004)
+#define NVC9B0_SET_CONTENT_INITIAL_VECTOR_VALUE                                 31:0
+#define NVC9B0_SET_CTL_COUNT                                                    (0x00000C10)
+#define NVC9B0_SET_CTL_COUNT_VALUE                                              31:0
+#define NVC9B0_SET_UPPER_SRC                                                    (0x00000C14)
+#define NVC9B0_SET_UPPER_SRC_OFFSET                                             7:0
+#define NVC9B0_SET_LOWER_SRC                                                    (0x00000C18)
+#define NVC9B0_SET_LOWER_SRC_OFFSET                                             31:0
+#define NVC9B0_SET_UPPER_DST                                                    (0x00000C1C)
+#define NVC9B0_SET_UPPER_DST_OFFSET                                             7:0
+#define NVC9B0_SET_LOWER_DST                                                    (0x00000C20)
+#define NVC9B0_SET_LOWER_DST_OFFSET                                             31:0
+#define NVC9B0_SET_BLOCK_COUNT                                                  (0x00000C24)
+#define NVC9B0_SET_BLOCK_COUNT_VALUE                                            31:0
+#define NVC9B0_PR_SET_REQUEST_BUF_OFFSET                                        (0x00000D00)
+#define NVC9B0_PR_SET_REQUEST_BUF_OFFSET_OFFSET                                 31:0
+#define NVC9B0_PR_SET_REQUEST_BUF_SIZE                                          (0x00000D04)
+#define NVC9B0_PR_SET_REQUEST_BUF_SIZE_SIZE                                     31:0
+#define NVC9B0_PR_SET_RESPONSE_BUF_OFFSET                                       (0x00000D08)
+#define NVC9B0_PR_SET_RESPONSE_BUF_OFFSET_OFFSET                                31:0
+#define NVC9B0_PR_SET_RESPONSE_BUF_SIZE                                         (0x00000D0C)
+#define NVC9B0_PR_SET_RESPONSE_BUF_SIZE_SIZE                                    31:0
+#define NVC9B0_PR_SET_REQUEST_MESSAGE_BUF_OFFSET                                (0x00000D10)
+#define NVC9B0_PR_SET_REQUEST_MESSAGE_BUF_OFFSET_OFFSET                         31:0
+#define NVC9B0_PR_SET_RESPONSE_MESSAGE_BUF_OFFSET                               (0x00000D14)
+#define NVC9B0_PR_SET_RESPONSE_MESSAGE_BUF_OFFSET_OFFSET                        31:0
+#define NVC9B0_PR_SET_LOCAL_DECRYPT_BUF_OFFSET                                  (0x00000D18)
+#define NVC9B0_PR_SET_LOCAL_DECRYPT_BUF_OFFSET_OFFSET                           31:0
+#define NVC9B0_PR_SET_LOCAL_DECRYPT_BUF_SIZE                                    (0x00000D1C)
+#define NVC9B0_PR_SET_LOCAL_DECRYPT_BUF_SIZE_SIZE                               31:0
+#define NVC9B0_PR_SET_CONTENT_DECRYPT_INFO_BUF_OFFSET                           (0x00000D20)
+#define NVC9B0_PR_SET_CONTENT_DECRYPT_INFO_BUF_OFFSET_OFFSET                    31:0
+#define NVC9B0_PR_SET_REENCRYPTED_BITSTREAM_BUF_OFFSET                          (0x00000D24)
+#define NVC9B0_PR_SET_REENCRYPTED_BITSTREAM_BUF_OFFSET_OFFSET                   31:0
+#define NVC9B0_DH_KE_SET_CHALLENGE_BUF_OFFSET                                   (0x00000E00)
+#define NVC9B0_DH_KE_SET_CHALLENGE_BUF_OFFSET_OFFSET                            31:0
+#define NVC9B0_DH_KE_SET_RESPONSE_BUF_OFFSET                                    (0x00000E04)
+#define NVC9B0_DH_KE_SET_RESPONSE_BUF_OFFSET_OFFSET                             31:0
+#define NVC9B0_SET_SESSION_KEY(b)                                               (0x00000F00 + (b)*0x00000004)
+#define NVC9B0_SET_SESSION_KEY_VALUE                                            31:0
+#define NVC9B0_SET_CONTENT_KEY(b)                                               (0x00000F10 + (b)*0x00000004)
+#define NVC9B0_SET_CONTENT_KEY_VALUE                                            31:0
+#define NVC9B0_PM_TRIGGER_END                                                   (0x00001114)
+#define NVC9B0_PM_TRIGGER_END_V                                                 31:0
+
+#define NVC9B0_ERROR_NONE                                                       (0x00000000)
+#define NVC9B0_OS_ERROR_EXECUTE_INSUFFICIENT_DATA                               (0x00000001)
+#define NVC9B0_OS_ERROR_SEMAPHORE_INSUFFICIENT_DATA                             (0x00000002)
+#define NVC9B0_OS_ERROR_INVALID_METHOD                                          (0x00000003)
+#define NVC9B0_OS_ERROR_INVALID_DMA_PAGE                                        (0x00000004)
+#define NVC9B0_OS_ERROR_UNHANDLED_INTERRUPT                                     (0x00000005)
+#define NVC9B0_OS_ERROR_EXCEPTION                                               (0x00000006)
+#define NVC9B0_OS_ERROR_INVALID_CTXSW_REQUEST                                   (0x00000007)
+#define NVC9B0_OS_ERROR_APPLICATION                                             (0x00000008)
+#define NVC9B0_OS_ERROR_SW_BREAKPT                                              (0x00000009)
+#define NVC9B0_OS_INTERRUPT_EXECUTE_AWAKEN                                      (0x00000100)
+#define NVC9B0_OS_INTERRUPT_BACKEND_SEMAPHORE_AWAKEN                            (0x00000200)
+#define NVC9B0_OS_INTERRUPT_CTX_ERROR_FBIF                                      (0x00000300)
+#define NVC9B0_OS_INTERRUPT_LIMIT_VIOLATION                                     (0x00000400)
+#define NVC9B0_OS_INTERRUPT_LIMIT_AND_FBIF_CTX_ERROR                            (0x00000500)
+#define NVC9B0_OS_INTERRUPT_HALT_ENGINE                                         (0x00000600)
+#define NVC9B0_OS_INTERRUPT_TRAP_NONSTALL                                       (0x00000700)
+#define NVC9B0_H264_VLD_ERR_SEQ_DATA_INCONSISTENT                               (0x00004001)
+#define NVC9B0_H264_VLD_ERR_PIC_DATA_INCONSISTENT                               (0x00004002)
+#define NVC9B0_H264_VLD_ERR_SLC_DATA_BUF_ADDR_OUT_OF_BOUNDS                     (0x00004100)
+#define NVC9B0_H264_VLD_ERR_BITSTREAM_ERROR                                     (0x00004101)
+#define NVC9B0_H264_VLD_ERR_CTX_DMA_ID_CTRL_IN_INVALID                          (0x000041F8)
+#define NVC9B0_H264_VLD_ERR_SLC_HDR_OUT_SIZE_NOT_MULT256                        (0x00004200)
+#define NVC9B0_H264_VLD_ERR_SLC_DATA_OUT_SIZE_NOT_MULT256                       (0x00004201)
+#define NVC9B0_H264_VLD_ERR_CTX_DMA_ID_FLOW_CTRL_INVALID                        (0x00004203)
+#define NVC9B0_H264_VLD_ERR_CTX_DMA_ID_SLC_HDR_OUT_INVALID                      (0x00004204)
+#define NVC9B0_H264_VLD_ERR_SLC_HDR_OUT_BUF_TOO_SMALL                           (0x00004205)
+#define NVC9B0_H264_VLD_ERR_SLC_HDR_OUT_BUF_ALREADY_VALID                       (0x00004206)
+#define NVC9B0_H264_VLD_ERR_SLC_DATA_OUT_BUF_TOO_SMALL                          (0x00004207)
+#define NVC9B0_H264_VLD_ERR_DATA_BUF_CNT_TOO_SMALL                              (0x00004208)
+#define NVC9B0_H264_VLD_ERR_BITSTREAM_EMPTY                                     (0x00004209)
+#define NVC9B0_H264_VLD_ERR_FRAME_WIDTH_TOO_LARGE                               (0x0000420A)
+#define NVC9B0_H264_VLD_ERR_FRAME_HEIGHT_TOO_LARGE                              (0x0000420B)
+#define NVC9B0_H264_VLD_ERR_HIST_BUF_TOO_SMALL                                  (0x00004300)
+#define NVC9B0_VC1_VLD_ERR_PIC_DATA_BUF_ADDR_OUT_OF_BOUND                       (0x00005100)
+#define NVC9B0_VC1_VLD_ERR_BITSTREAM_ERROR                                      (0x00005101)
+#define NVC9B0_VC1_VLD_ERR_PIC_HDR_OUT_SIZE_NOT_MULT256                         (0x00005200)
+#define NVC9B0_VC1_VLD_ERR_PIC_DATA_OUT_SIZE_NOT_MULT256                        (0x00005201)
+#define NVC9B0_VC1_VLD_ERR_CTX_DMA_ID_CTRL_IN_INVALID                           (0x00005202)
+#define NVC9B0_VC1_VLD_ERR_CTX_DMA_ID_FLOW_CTRL_INVALID                         (0x00005203)
+#define NVC9B0_VC1_VLD_ERR_CTX_DMA_ID_PIC_HDR_OUT_INVALID                       (0x00005204)
+#define NVC9B0_VC1_VLD_ERR_SLC_HDR_OUT_BUF_TOO_SMALL                            (0x00005205)
+#define NVC9B0_VC1_VLD_ERR_PIC_HDR_OUT_BUF_ALREADY_VALID                        (0x00005206)
+#define NVC9B0_VC1_VLD_ERR_PIC_DATA_OUT_BUF_TOO_SMALL                           (0x00005207)
+#define NVC9B0_VC1_VLD_ERR_DATA_INFO_IN_BUF_TOO_SMALL                           (0x00005208)
+#define NVC9B0_VC1_VLD_ERR_BITSTREAM_EMPTY                                      (0x00005209)
+#define NVC9B0_VC1_VLD_ERR_FRAME_WIDTH_TOO_LARGE                                (0x0000520A)
+#define NVC9B0_VC1_VLD_ERR_FRAME_HEIGHT_TOO_LARGE                               (0x0000520B)
+#define NVC9B0_VC1_VLD_ERR_PIC_DATA_OUT_BUF_FULL_TIME_OUT                       (0x00005300)
+#define NVC9B0_MPEG12_VLD_ERR_SLC_DATA_BUF_ADDR_OUT_OF_BOUNDS                   (0x00006100)
+#define NVC9B0_MPEG12_VLD_ERR_BITSTREAM_ERROR                                   (0x00006101)
+#define NVC9B0_MPEG12_VLD_ERR_SLC_DATA_OUT_SIZE_NOT_MULT256                     (0x00006200)
+#define NVC9B0_MPEG12_VLD_ERR_CTX_DMA_ID_CTRL_IN_INVALID                        (0x00006201)
+#define NVC9B0_MPEG12_VLD_ERR_CTX_DMA_ID_FLOW_CTRL_INVALID                      (0x00006202)
+#define NVC9B0_MPEG12_VLD_ERR_SLC_DATA_OUT_BUF_TOO_SMALL                        (0x00006203)
+#define NVC9B0_MPEG12_VLD_ERR_DATA_INFO_IN_BUF_TOO_SMALL                        (0x00006204)
+#define NVC9B0_MPEG12_VLD_ERR_BITSTREAM_EMPTY                                   (0x00006205)
+#define NVC9B0_MPEG12_VLD_ERR_INVALID_PIC_STRUCTURE                             (0x00006206)
+#define NVC9B0_MPEG12_VLD_ERR_INVALID_PIC_CODING_TYPE                           (0x00006207)
+#define NVC9B0_MPEG12_VLD_ERR_FRAME_WIDTH_TOO_LARGE                             (0x00006208)
+#define NVC9B0_MPEG12_VLD_ERR_FRAME_HEIGHT_TOO_LARGE                            (0x00006209)
+#define NVC9B0_MPEG12_VLD_ERR_SLC_DATA_OUT_BUF_FULL_TIME_OUT                    (0x00006300)
+#define NVC9B0_CMN_VLD_ERR_PDEC_RETURNED_ERROR                                  (0x00007101)
+#define NVC9B0_CMN_VLD_ERR_EDOB_FLUSH_TIME_OUT                                  (0x00007102)
+#define NVC9B0_CMN_VLD_ERR_EDOB_REWIND_TIME_OUT                                 (0x00007103)
+#define NVC9B0_CMN_VLD_ERR_VLD_WD_TIME_OUT                                      (0x00007104)
+#define NVC9B0_CMN_VLD_ERR_NUM_SLICES_ZERO                                      (0x00007105)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_DATA_BUF_ADDR_OUT_OF_BOUND                     (0x00008100)
+#define NVC9B0_MPEG4_VLD_ERR_BITSTREAM_ERROR                                    (0x00008101)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_HDR_OUT_SIZE_NOT_MULT256                       (0x00008200)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_DATA_OUT_SIZE_NOT_MULT256                      (0x00008201)
+#define NVC9B0_MPEG4_VLD_ERR_CTX_DMA_ID_CTRL_IN_INVALID                         (0x00008202)
+#define NVC9B0_MPEG4_VLD_ERR_CTX_DMA_ID_FLOW_CTRL_INVALID                       (0x00008203)
+#define NVC9B0_MPEG4_VLD_ERR_CTX_DMA_ID_PIC_HDR_OUT_INVALID                     (0x00008204)
+#define NVC9B0_MPEG4_VLD_ERR_SLC_HDR_OUT_BUF_TOO_SMALL                          (0x00008205)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_HDR_OUT_BUF_ALREADY_VALID                      (0x00008206)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_DATA_OUT_BUF_TOO_SMALL                         (0x00008207)
+#define NVC9B0_MPEG4_VLD_ERR_DATA_INFO_IN_BUF_TOO_SMALL                         (0x00008208)
+#define NVC9B0_MPEG4_VLD_ERR_BITSTREAM_EMPTY                                    (0x00008209)
+#define NVC9B0_MPEG4_VLD_ERR_FRAME_WIDTH_TOO_LARGE                              (0x0000820A)
+#define NVC9B0_MPEG4_VLD_ERR_FRAME_HEIGHT_TOO_LARGE                             (0x0000820B)
+#define NVC9B0_MPEG4_VLD_ERR_PIC_DATA_OUT_BUF_FULL_TIME_OUT                     (0x00051E01)
+#define NVC9B0_DEC_ERROR_MPEG12_APPTIMER_EXPIRED                                (0xDEC10001)
+#define NVC9B0_DEC_ERROR_MPEG12_MVTIMER_EXPIRED                                 (0xDEC10002)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_TOKEN                                   (0xDEC10003)
+#define NVC9B0_DEC_ERROR_MPEG12_SLICEDATA_MISSING                               (0xDEC10004)
+#define NVC9B0_DEC_ERROR_MPEG12_HWERR_INTERRUPT                                 (0xDEC10005)
+#define NVC9B0_DEC_ERROR_MPEG12_DETECTED_VLD_FAILURE                            (0xDEC10006)
+#define NVC9B0_DEC_ERROR_MPEG12_PICTURE_INIT                                    (0xDEC10100)
+#define NVC9B0_DEC_ERROR_MPEG12_STATEMACHINE_FAILURE                            (0xDEC10101)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_CTXID_PIC                               (0xDEC10901)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_CTXID_UCODE                             (0xDEC10902)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_CTXID_FC                                (0xDEC10903)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_CTXID_SLH                               (0xDEC10904)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_UCODE_SIZE                              (0xDEC10905)
+#define NVC9B0_DEC_ERROR_MPEG12_INVALID_SLICE_COUNT                             (0xDEC10906)
+#define NVC9B0_DEC_ERROR_VC1_APPTIMER_EXPIRED                                   (0xDEC20001)
+#define NVC9B0_DEC_ERROR_VC1_MVTIMER_EXPIRED                                    (0xDEC20002)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_TOKEN                                      (0xDEC20003)
+#define NVC9B0_DEC_ERROR_VC1_SLICEDATA_MISSING                                  (0xDEC20004)
+#define NVC9B0_DEC_ERROR_VC1_HWERR_INTERRUPT                                    (0xDEC20005)
+#define NVC9B0_DEC_ERROR_VC1_DETECTED_VLD_FAILURE                               (0xDEC20006)
+#define NVC9B0_DEC_ERROR_VC1_TIMEOUT_POLLING_FOR_DATA                           (0xDEC20007)
+#define NVC9B0_DEC_ERROR_VC1_PDEC_PIC_END_UNALIGNED                             (0xDEC20008)
+#define NVC9B0_DEC_ERROR_VC1_WDTIMER_EXPIRED                                    (0xDEC20009)
+#define NVC9B0_DEC_ERROR_VC1_ERRINTSTART                                        (0xDEC20010)
+#define NVC9B0_DEC_ERROR_VC1_IQT_ERRINT                                         (0xDEC20011)
+#define NVC9B0_DEC_ERROR_VC1_MC_ERRINT                                          (0xDEC20012)
+#define NVC9B0_DEC_ERROR_VC1_MC_IQT_ERRINT                                      (0xDEC20013)
+#define NVC9B0_DEC_ERROR_VC1_REC_ERRINT                                         (0xDEC20014)
+#define NVC9B0_DEC_ERROR_VC1_REC_IQT_ERRINT                                     (0xDEC20015)
+#define NVC9B0_DEC_ERROR_VC1_REC_MC_ERRINT                                      (0xDEC20016)
+#define NVC9B0_DEC_ERROR_VC1_REC_MC_IQT_ERRINT                                  (0xDEC20017)
+#define NVC9B0_DEC_ERROR_VC1_DBF_ERRINT                                         (0xDEC20018)
+#define NVC9B0_DEC_ERROR_VC1_DBF_IQT_ERRINT                                     (0xDEC20019)
+#define NVC9B0_DEC_ERROR_VC1_DBF_MC_ERRINT                                      (0xDEC2001A)
+#define NVC9B0_DEC_ERROR_VC1_DBF_MC_IQT_ERRINT                                  (0xDEC2001B)
+#define NVC9B0_DEC_ERROR_VC1_DBF_REC_ERRINT                                     (0xDEC2001C)
+#define NVC9B0_DEC_ERROR_VC1_DBF_REC_IQT_ERRINT                                 (0xDEC2001D)
+#define NVC9B0_DEC_ERROR_VC1_DBF_REC_MC_ERRINT                                  (0xDEC2001E)
+#define NVC9B0_DEC_ERROR_VC1_DBF_REC_MC_IQT_ERRINT                              (0xDEC2001F)
+#define NVC9B0_DEC_ERROR_VC1_PICTURE_INIT                                       (0xDEC20100)
+#define NVC9B0_DEC_ERROR_VC1_STATEMACHINE_FAILURE                               (0xDEC20101)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_CTXID_PIC                                  (0xDEC20901)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_CTXID_UCODE                                (0xDEC20902)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_CTXID_FC                                   (0xDEC20903)
+#define NVC9B0_DEC_ERROR_VC1_INVAILD_CTXID_SLH                                  (0xDEC20904)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_UCODE_SIZE                                 (0xDEC20905)
+#define NVC9B0_DEC_ERROR_VC1_INVALID_SLICE_COUNT                                (0xDEC20906)
+#define NVC9B0_DEC_ERROR_H264_APPTIMER_EXPIRED                                  (0xDEC30001)
+#define NVC9B0_DEC_ERROR_H264_MVTIMER_EXPIRED                                   (0xDEC30002)
+#define NVC9B0_DEC_ERROR_H264_INVALID_TOKEN                                     (0xDEC30003)
+#define NVC9B0_DEC_ERROR_H264_SLICEDATA_MISSING                                 (0xDEC30004)
+#define NVC9B0_DEC_ERROR_H264_HWERR_INTERRUPT                                   (0xDEC30005)
+#define NVC9B0_DEC_ERROR_H264_DETECTED_VLD_FAILURE                              (0xDEC30006)
+#define NVC9B0_DEC_ERROR_H264_ERRINTSTART                                       (0xDEC30010)
+#define NVC9B0_DEC_ERROR_H264_IQT_ERRINT                                        (0xDEC30011)
+#define NVC9B0_DEC_ERROR_H264_MC_ERRINT                                         (0xDEC30012)
+#define NVC9B0_DEC_ERROR_H264_MC_IQT_ERRINT                                     (0xDEC30013)
+#define NVC9B0_DEC_ERROR_H264_REC_ERRINT                                        (0xDEC30014)
+#define NVC9B0_DEC_ERROR_H264_REC_IQT_ERRINT                                    (0xDEC30015)
+#define NVC9B0_DEC_ERROR_H264_REC_MC_ERRINT                                     (0xDEC30016)
+#define NVC9B0_DEC_ERROR_H264_REC_MC_IQT_ERRINT                                 (0xDEC30017)
+#define NVC9B0_DEC_ERROR_H264_DBF_ERRINT                                        (0xDEC30018)
+#define NVC9B0_DEC_ERROR_H264_DBF_IQT_ERRINT                                    (0xDEC30019)
+#define NVC9B0_DEC_ERROR_H264_DBF_MC_ERRINT                                     (0xDEC3001A)
+#define NVC9B0_DEC_ERROR_H264_DBF_MC_IQT_ERRINT                                 (0xDEC3001B)
+#define NVC9B0_DEC_ERROR_H264_DBF_REC_ERRINT                                    (0xDEC3001C)
+#define NVC9B0_DEC_ERROR_H264_DBF_REC_IQT_ERRINT                                (0xDEC3001D)
+#define NVC9B0_DEC_ERROR_H264_DBF_REC_MC_ERRINT                                 (0xDEC3001E)
+#define NVC9B0_DEC_ERROR_H264_DBF_REC_MC_IQT_ERRINT                             (0xDEC3001F)
+#define NVC9B0_DEC_ERROR_H264_PICTURE_INIT                                      (0xDEC30100)
+#define NVC9B0_DEC_ERROR_H264_STATEMACHINE_FAILURE                              (0xDEC30101)
+#define NVC9B0_DEC_ERROR_H264_INVALID_CTXID_PIC                                 (0xDEC30901)
+#define NVC9B0_DEC_ERROR_H264_INVALID_CTXID_UCODE                               (0xDEC30902)
+#define NVC9B0_DEC_ERROR_H264_INVALID_CTXID_FC                                  (0xDEC30903)
+#define NVC9B0_DEC_ERROR_H264_INVALID_CTXID_SLH                                 (0xDEC30904)
+#define NVC9B0_DEC_ERROR_H264_INVALID_UCODE_SIZE                                (0xDEC30905)
+#define NVC9B0_DEC_ERROR_H264_INVALID_SLICE_COUNT                               (0xDEC30906)
+#define NVC9B0_DEC_ERROR_MPEG4_APPTIMER_EXPIRED                                 (0xDEC40001)
+#define NVC9B0_DEC_ERROR_MPEG4_MVTIMER_EXPIRED                                  (0xDEC40002)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_TOKEN                                    (0xDEC40003)
+#define NVC9B0_DEC_ERROR_MPEG4_SLICEDATA_MISSING                                (0xDEC40004)
+#define NVC9B0_DEC_ERROR_MPEG4_HWERR_INTERRUPT                                  (0xDEC40005)
+#define NVC9B0_DEC_ERROR_MPEG4_DETECTED_VLD_FAILURE                             (0xDEC40006)
+#define NVC9B0_DEC_ERROR_MPEG4_TIMEOUT_POLLING_FOR_DATA                         (0xDEC40007)
+#define NVC9B0_DEC_ERROR_MPEG4_PDEC_PIC_END_UNALIGNED                           (0xDEC40008)
+#define NVC9B0_DEC_ERROR_MPEG4_WDTIMER_EXPIRED                                  (0xDEC40009)
+#define NVC9B0_DEC_ERROR_MPEG4_ERRINTSTART                                      (0xDEC40010)
+#define NVC9B0_DEC_ERROR_MPEG4_IQT_ERRINT                                       (0xDEC40011)
+#define NVC9B0_DEC_ERROR_MPEG4_MC_ERRINT                                        (0xDEC40012)
+#define NVC9B0_DEC_ERROR_MPEG4_MC_IQT_ERRINT                                    (0xDEC40013)
+#define NVC9B0_DEC_ERROR_MPEG4_REC_ERRINT                                       (0xDEC40014)
+#define NVC9B0_DEC_ERROR_MPEG4_REC_IQT_ERRINT                                   (0xDEC40015)
+#define NVC9B0_DEC_ERROR_MPEG4_REC_MC_ERRINT                                    (0xDEC40016)
+#define NVC9B0_DEC_ERROR_MPEG4_REC_MC_IQT_ERRINT                                (0xDEC40017)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_ERRINT                                       (0xDEC40018)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_IQT_ERRINT                                   (0xDEC40019)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_MC_ERRINT                                    (0xDEC4001A)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_MC_IQT_ERRINT                                (0xDEC4001B)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_REC_ERRINT                                   (0xDEC4001C)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_REC_IQT_ERRINT                               (0xDEC4001D)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_REC_MC_ERRINT                                (0xDEC4001E)
+#define NVC9B0_DEC_ERROR_MPEG4_DBF_REC_MC_IQT_ERRINT                            (0xDEC4001F)
+#define NVC9B0_DEC_ERROR_MPEG4_PICTURE_INIT                                     (0xDEC40100)
+#define NVC9B0_DEC_ERROR_MPEG4_STATEMACHINE_FAILURE                             (0xDEC40101)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_CTXID_PIC                                (0xDEC40901)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_CTXID_UCODE                              (0xDEC40902)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_CTXID_FC                                 (0xDEC40903)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_CTXID_SLH                                (0xDEC40904)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_UCODE_SIZE                               (0xDEC40905)
+#define NVC9B0_DEC_ERROR_MPEG4_INVALID_SLICE_COUNT                              (0xDEC40906)
+
+#ifdef __cplusplus
+};     /* extern "C" */
+#endif
+#endif // clc9b0_h

extra/qcom_gpu_driver/qcom_opencl_interop.py

@@ -29,8 +29,9 @@ 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).numpy()
-print(y)
+y = (x + 1).reshape(-1).tolist()
+print(y[:10])
+assert y == [i + 1 for i in range(64)]
 
 # all calculations are done, save to free the object
 cl.clReleaseMemObject(cl_buf)
@@ -49,7 +50,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')).numpy()
+  y = calc(x = Tensor.from_blob(rawbuf_ptr, (2, 2), dtype=dtypes.int, device='QCOM')).tolist()
   print(f'jit {i}\n', y)
 
   # all calculations are done, save to free the object
@@ -80,8 +81,19 @@ 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).numpy()
-print(y)
+y = (x + 1).tolist()
+print(y[:10])
 
 # 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)

extra/remu/src/helpers.rs

@@ -1,5 +1,5 @@
 use half::f16;
-use num_traits::{float::FloatCore, PrimInt, Unsigned};
+use num_traits::{float::FloatCore, PrimInt, Unsigned, clamp};
 
 pub fn bits<T>(word: T, hi: usize, lo: usize) -> T where T: PrimInt + Unsigned {
     assert!(hi >= lo);
@@ -48,6 +48,7 @@ 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
@@ -65,6 +66,11 @@ 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 {

extra/remu/src/thread.rs

@@ -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;
+                    let cm = ((instr >> 15) & 0x1) != 0;
 
                     let s = |n: usize| ((instr >> n) & 0x1ff) as usize;
                     let src = (s(32), s(41), s(50));
@@ -1032,7 +1032,9 @@ impl<'a> Thread<'a> {
                     let omod = (instr >> 59) & 0x3;
                     let neg = ((instr >> 61) & 0x7) as usize;
                     assert_eq!(omod, 0);
-                    assert_eq!(cm, 0);
+                    if op != 272 && cm {
+                        return todo_instr!(op); // TODO: add VOP3 clamp for all ops
+                    }
                     assert_eq!(opsel, 0);
 
                     match op {
@@ -1266,7 +1268,7 @@ impl<'a> Thread<'a> {
                             }
 
                             let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 543 | 551 | 567 | 606 | 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);
@@ -1275,12 +1277,26 @@ impl<'a> Thread<'a> {
                                         260 => s0 - s1,
                                         261 => s1 - s0,
                                         264 => s0 * s1,
-                                        272 => f32::max(s0, s1),
+                                        272 => f32::max(s0, s1).clmp(cm),
                                         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 => {
@@ -1290,6 +1306,7 @@ 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 => {

extra/sqtt/attempt_sqtt_parse.py

@@ -119,7 +119,6 @@ OPNAME = {
   0xb: "VALU",
   0xd: "VALU",
   0xe: "VALU",
-  0x10: "__END",
   0x21: "VMEM_LOAD",
   0x22: "VMEM_LOAD",
   0x24: "VMEM_STORE",
@@ -480,6 +479,7 @@ 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]

extra/sqtt/rgptool.py

@@ -166,6 +166,7 @@ class RGP:
           se=ev.se,
           itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
           blob=merged_sqtt_events[ev.se].blob + ev.blob,
+          exec_tag=0,
         )
     sqtt_events = list(merged_sqtt_events.values())
 

extra/sqtt/roc.py

@@ -67,12 +67,14 @@ class OccEvent(WaveSlot):
   time:int
   start:int
 
+RunKey = tuple[str, int]
+
 class _ROCParseCtx:
   def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
     self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
     self.disasms:dict[str, dict[int, tuple[str, int]]] = {}
-    self.inst_execs:dict[str, list[WaveExec]] = {}
-    self.occ_events:dict[str, list[OccEvent]] = {}
+    self.inst_execs:dict[RunKey, list[WaveExec]] = {}
+    self.occ_events:dict[RunKey, list[OccEvent]] = {}
 
     for prog in prog_evs:
       arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
@@ -81,14 +83,14 @@ class _ROCParseCtx:
 
   def next_sqtt(self):
     x = next(self.sqtt_evs, None)
-    self.active_kern = x.kern if x is not None else None
+    self.active_run = (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
 
   def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
     if DEBUG >= 5: print(f"OCC {ev.time=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.wave_id=} {ev.start=}")
-    self.occ_events.setdefault(unwrap(self.active_kern), []).append(OccEvent(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.time, ev.start))
+    self.occ_events.setdefault(unwrap(self.active_run), []).append(OccEvent(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.time, ev.start))
 
   def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
     if DEBUG >= 5: print(f"WAVE {ev.wave_id=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.contexts=} {ev.begin_time=} {ev.end_time=}")
@@ -98,7 +100,7 @@ class _ROCParseCtx:
     insts_blob = bytearray(sz:=ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t))
     ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
 
-    self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
+    self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
                                                                              ev.end_time, insts_blob))
 
 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
@@ -136,7 +138,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_kern)][pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run)[0]][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
@@ -156,6 +158,14 @@ 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)))
@@ -165,13 +175,4 @@ if __name__ == "__main__":
   rctx = decode(profile)
   print('SQTT:', rctx.inst_execs.keys())
 
-  for ev in profile:
-    if not isinstance(ev, ProfilePMCEvent): continue
-    print(f"PMC Event: dev={ev.device} kern={ev.kern}")
-    ptr = 0
-    for s in ev.sched:
-      view = memoryview(ev.blob).cast('Q')
-      print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
-        ptr += 1
+  print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])

extra/sqtt/test_pmc.py

@@ -0,0 +1,87 @@
+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()

extra/sqtt/test_timing.py

@@ -2,7 +2,9 @@ import os
 os.environ["PYTHONPATH"] = "."
 os.environ["SQTT"] = "1"
 if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["VIZ"] = "1"
+os.environ["PROFILE"] = "1"
+# VIZ=1 to launch server
+# os.environ["VIZ"] = "1"
 os.environ["AMD_LLVM"] = "0"
 
 import unittest
@@ -129,5 +131,13 @@ class TestTiming(unittest.TestCase):
     for w in waves:
       print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
 
+  def test_ones(self):
+    N = getenv("N", 4096)
+    CNT = getenv("CNT", 2)
+    with save_sqtt() as sqtt:
+      for _ in range(CNT):
+        Tensor.ones(N, N).contiguous().realize()
+    self.assertEqual(len(sqtt), CNT)
+
 if __name__ == "__main__":
   unittest.main()

extra/thunder/tiny/fa.py

@@ -0,0 +1,166 @@
+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)

extra/thunder/tiny/tk/group.py

@@ -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
+from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout, VecLayout
 
 class Group:
   def __init__(self, warps:int, ker):
@@ -66,7 +66,10 @@ 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):
-          dst_store = dst[width, height, inner].store(src[height, width, inner]).end(height, width, inner)
+          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)
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
@@ -77,9 +80,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))), ())
+      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
     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))), ())
+      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
     else:  raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -107,9 +110,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))), ())
+      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
     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))), ())
+      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
     else:  raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -137,9 +140,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))), ())
+      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
     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))), ())
+      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
     else:  raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -167,9 +170,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))), ())
+      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
     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))), ())
+      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
     else:  raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -200,9 +203,9 @@ class Group:
     Group.map_rid += len(a.shape)
 
     if op.__code__.co_argcount == 1:
-      to_store = op(a[*rngs_for_shape])
+      to_store = op(a[*rngs_for_shape]) # type: ignore
     else:
-      to_store = op(a[*rngs_for_shape], rngs_for_shape)
+      to_store = op(a[*rngs_for_shape], rngs_for_shape) # type: ignore
 
     a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
 
@@ -284,7 +287,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 = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    dst_dtype, src_dtype = dst.dtype, 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)
@@ -338,7 +341,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:
+    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
       srcf = src.flatten()
       row_stride = prod(src.shape[axis+1:])
 
@@ -371,8 +374,28 @@ 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")
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
@@ -380,8 +403,30 @@ 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 = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL:
+    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):
       dstf = dst.flatten()
       row_stride = prod(dst.shape[axis+1:])
 
@@ -414,8 +459,28 @@ 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")
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)

extra/thunder/tiny/tk/kernel.py

@@ -5,33 +5,34 @@ 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:
-  user_rid = 0
-  def __init__(self, start:int, end:int, step:int, axis_type:AxisType):
+  def __init__(self, start:int, end:int, step:int, axis_type:AxisType, rid:int):
     self.start, self.end, self.step = start, end, step
-    self.axis_type, self.done = axis_type, False
+    self.axis_type, self.rid, self.done = axis_type, rid, False
   def __iter__(self): return self
   def __next__(self):
     if not self.done:
       self.done = True
-      _tk_range.user_rid += 1
-      self._rng = UOp.range(self.end // self.step, _tk_range.user_rid-1, axis_type=self.axis_type) * self.step + self.start
+      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
       return self._rng
     raise StopIteration
 
 class Kernel(AbstractContextManager):
-  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+  def __init__(self, name:str, grid_size:tuple[int, int, int], block_size:int):
+    self.name = name
+
     self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
     self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
     self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
     self.threadIdx_x = UOp.special(block_size, "lidx0")
 
-    self.range_stack = []
-    self.store_stack = []
+    self.range_stack: list[_tk_range] = []
+    self.store_stack: list[tuple[UOp, UOp]] = []
 
     self.global_slot = 0
     self.shared_slot = 0
     self.register_slot = 0
-    self.allocs = {}
+    self.range_id = 0
+    self.allocs: dict[tuple[str, tuple], UOp] = {}
 
   @property
   def warpid(self): return self.threadIdx_x // WARP_THREADS
@@ -49,7 +50,8 @@ 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)
+    rng = _tk_range(start, end, step, axis_type, self.range_id)
+    self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
@@ -89,7 +91,7 @@ class Kernel(AbstractContextManager):
     if hasattr(last_store, '_uop'): uop = last_store._uop
     else: uop = last_store
 
-    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+    return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
 
   def endrange(self):
     last_store = self.store_stack.pop()

extra/thunder/tiny/tk/tiles.py

@@ -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
+      if name in UOp.__slots__: return val # type: ignore
       return wrap(val, self)
     cls.__getattr__ = __getattr__
 
@@ -250,11 +250,12 @@ class RT(TileMathMixin):
 
 @autowrap(UOp)
 class RV(TileMathMixin):
-  def __init__(self, uop:UOp, layout:VecLayout, ker):
-    self._uop, self.layout, self.ker = uop, layout, ker
+  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 ruop(self, uop:UOp):
-    return RV(uop, self.layout, self.ker)
+    return RV(uop, self.length, self.layout, self.base_shape, self.ker)
 
   @classmethod
   def create(cls, length, dtype:DType, layout:VecLayout, base_shape:RTBaseShape, ker):
@@ -266,6 +267,6 @@ class RV(TileMathMixin):
         outer_dim = tiles
 
     uop = ker.alloc((outer_dim, inner_dim), dtype, AddrSpace.REG)
-    return RV(uop, layout, ker)
+    return RV(uop, length, layout, base_shape, ker)
 
 ALL_TILES = UOp | GL | ST | RT | RV

extra/usbgpu/patch.py

@@ -3,6 +3,13 @@ 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())
 
@@ -40,10 +47,14 @@ if not os.path.exists(file_path):
 patches = [(0x2a0d + 1 + 4, b'\x0a', b'\x05')]
 patched_fw = patch(file_path, file_hash, patches)
 
-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
+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")')
 
 config1 = bytes([
   0xFF, 0xFF, 0xFF, 0xFF, 0x41, 0x41, 0x41, 0x41, 0x42, 0x42, 0x42, 0x42, 0x30, 0x30, 0x36, 0x30,

extra/weekly_commits_table.py

@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt
 
-NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot","Christopher Milan"]
 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)]

mkdocs.yml

@@ -119,7 +119,7 @@ plugins:
 - mkdocstrings:
     handlers:
       python:
-        import:
+        inventories:
         - https://docs.python.org/3/objects.inv
         paths: [tinygrad]
         options:

pyproject.toml

@@ -69,17 +69,14 @@ testing_minimal = [
   "hypothesis",
   "z3-solver",
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai"]
 testing = [
-  "tinygrad[testing_minimal]",
+  "tinygrad[testing_unit]",
   "pillow",
-  "onnx==1.18.0",
+  "onnx==1.19.0",
   "onnx2torch",
   "onnxruntime",
   "opencv-python",
-  "tabulate",
-  "tqdm",
-  "safetensors",
   "transformers",
   "sentencepiece",
   "tiktoken",

test/device/test_amd_llvm.py

@@ -1,10 +1,8 @@
 import unittest
-import numpy as np
 from tinygrad import Device
 from tinygrad.device import CompileError
-from tinygrad.helpers import flat_mv
-if Device.DEFAULT=="AMD":
-  from tinygrad.runtime.ops_amd import AMDAllocator, AMDDevice, AMDProgram
+if Device.DEFAULT == "AMD":
+  # NOTE: if you don't gate this, LVP fails on Mac
   from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "Runs only on AMD")
@@ -18,16 +16,8 @@ entry:
   ret void
 }
     '''
-    device = AMDDevice()
     compiler = AMDLLVMCompiler("gfx1100")
-    obj = compiler.compile(src)
-    allocator = AMDAllocator(device)
-    a = allocator.alloc(1*8)
-    prog = AMDProgram(device, "test", obj)
-    prog(a, wait=True)
-    na = np.empty(1, np.uint64)
-    allocator._copyout(flat_mv(na.data), a)
-    assert na == [0x1234567800000005]
+    compiler.compile(src)
 
   def test_compiler_diag_error(self):
     src = """

test/device/test_hcq.py

@@ -224,7 +224,8 @@ class TestHCQ(unittest.TestCase):
   def test_copy_64bit(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
 
-    for sz in [(1 << 32) - 1, (1 << 32), (1 << 32) + 1, (5 << 30), (6 << 30) - 0x4642ee1]:
+    # NOTE: these must be a multiple of 8 for .view(fmt='Q') to work
+    for sz in [(1 << 32) - 8, (1 << 32), (1 << 32) + 8, (5 << 30), (6 << 30) - 0x4642ee0]:
       buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
       buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
 

test/external/external_test_am.py

@@ -32,6 +32,8 @@ 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

test/external/external_test_onnx_backend.py

@@ -170,6 +170,19 @@ 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
@@ -184,9 +197,6 @@ 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')

test/external/external_test_optim.py

@@ -2,7 +2,7 @@
 import unittest, math
 import numpy as np
 import tensorflow as tf
-import tensorflow_addons as tfa
+from tensorflow.keras.optimizers import Lamb
 from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup
 
@@ -88,6 +88,8 @@ 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
@@ -112,7 +114,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, tfa.optimizers.LAMB, steps, opts, atol, rtol)
+  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, create_tf_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,

test/external/external_test_process_count.py

@@ -0,0 +1,57 @@
+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)

test/external/external_uop_gc.py

@@ -1,5 +1,6 @@
 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
@@ -68,10 +69,12 @@ 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()
     fold_divmod_general.cache_clear()
+    UOp.const.cache_clear()
     Tensor._device_seeds.clear()
     Tensor._device_rng_counters.clear()
 

test/external/mlperf_resnet/lars_optimizer.py

@@ -29,7 +29,6 @@ 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
 
 
@@ -147,20 +146,7 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
     return scaled_lr, grad
 
   def _apply_dense(self, grad, var, apply_state=None):
-    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)
+    return self._resource_apply_dense(grad, var, apply_state)
 
   def _resource_apply_dense(self, grad, var, apply_state=None):
     var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -194,13 +180,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                     or self._fallback_apply_state(var_device, var_dtype))
 
     mom = self.get_slot(var, "momentum")
-    return training_ops.sparse_apply_momentum(
-        var,
-        mom,
-        coefficients["learning_rate"],
-        grad.values,
-        grad.indices,
-        self.momentum,
+    return tf.raw_ops.SparseApplyMomentum(
+        var=var,
+        accum=mom,
+        lr=coefficients["learning_rate"],
+        grad=grad.values,
+        indices=grad.indices,
+        momentum=self.momentum,
         use_locking=False,
         use_nesterov=self.use_nesterov)
 
@@ -210,13 +196,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                     or self._fallback_apply_state(var_device, var_dtype))
 
     mom = self.get_slot(var, "momentum")
-    return training_ops.resource_sparse_apply_keras_momentum(
-        var.handle,
-        mom.handle,
-        coefficients["learning_rate"],
-        grad,
-        indices,
-        self.momentum,
+    return tf.raw_ops.ResourceSparseApplyKerasMomentum(
+        var=var.handle,
+        accum=mom.handle,
+        lr=coefficients["learning_rate"],
+        grad=grad,
+        indices=indices,
+        momentum=self.momentum,
         use_locking=False,
         use_nesterov=self.use_nesterov)
 

test/external/process_replay/process_replay.py

@@ -14,7 +14,6 @@ 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))
@@ -52,12 +51,10 @@ 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|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, 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

test/helpers.py

@@ -1,4 +1,4 @@
-import time, struct
+import time, struct, functools
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
@@ -60,5 +60,14 @@ 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)

test/mockgpu/nv/nvdriver.py

@@ -100,6 +100,9 @@ class NVDriver(VirtDriver):
       assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
       struct.hObjectNew = self._alloc_handle()
       self.object_by_handle[struct.hObjectNew] = NVSubDevice(self.object_by_handle[struct.hObjectParent])
+    elif struct.hClass == nv_gpu.NV01_MEMORY_VIRTUAL:
+      assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
+      struct.hObjectNew = self._alloc_handle()
     elif struct.hClass == nv_gpu.TURING_USERMODE_A:
       assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVSubDevice)
       struct.hObjectNew = self._alloc_handle()
@@ -215,6 +218,8 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.NV_ESC_RM_FREE:
       st = nv_gpu.NVOS00_PARAMETERS.from_address(argp)
       self.object_by_handle.pop(st.hObjectOld)
+    elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY_DMA:
+      pass # mappings are same as uvm
     elif nr == nv_gpu.NV_ESC_CARD_INFO:
       for i,gpu in enumerate(self.gpus.values()):
         st = nv_gpu.nv_ioctl_card_info_t.from_address(argp + i * ctypes.sizeof(nv_gpu.nv_ioctl_card_info_t))

test/models/test_end2end.py

@@ -113,7 +113,7 @@ class TestEnd2End(unittest.TestCase):
 
   def test_bn_linear(self):
     BS, K = 2, 1
-    eps = 0
+    eps = 1e-12  # torch asserts if this is 0
     X = Tensor([1,0]).reshape(BS, K, 1, 1)
     Y = Tensor([-1,0]).reshape(BS, K, 1, 1)
     class LinTiny:

test/models/test_real_world.py

@@ -28,13 +28,16 @@ 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
+    mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
 
     # 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/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}"
+    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"
     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
 
@@ -61,7 +64,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, 18.0, 515)
+    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
 
   def test_unet_resblock(self):
     model = [ResBlock(16, 24, 16) for _ in range(4)]
@@ -70,7 +73,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.01, 37)
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.0002, 37)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_llama(self):
@@ -82,7 +85,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.27, 168, all_jitted=True)
+    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_gpt2(self):
@@ -112,7 +115,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 103)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_forward_cifar(self):
@@ -122,7 +125,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, (1.0/48)*BS, 126)
+      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_train_cifar(self):
@@ -139,7 +142,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, (1.0/48)*BS, 126)
+      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_train_cifar_hyp(self):
@@ -176,7 +179,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, 427)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 400)
 
 if __name__ == '__main__':
   unittest.main()

test/models/test_whisper.py

@@ -94,12 +94,14 @@ class TestWhisper(unittest.TestCase):
     self.assertEqual(TRANSCRIPTION_2,  transcriptions[0])
     self.assertEqual(TRANSCRIPTION_1,  transcriptions[1])
 
+  @unittest.skip("file 3 url is broken")
   @unittest.skipIf(CI or (Device.DEFAULT == "CPU" and CPU_LLVM), "too long for CI")
   def test_transcribe_long(self):
     waveform = [load_file_waveform(fetch(TEST_FILE_3_URL))]
     transcription = transcribe_waveform(self.model, self.enc, waveform)
     self.assertWER(transcription, TRANSCRIPTION_3, 0.085)
 
+  @unittest.skip("file 3 url is broken")
   @unittest.skipIf(CI or (Device.DEFAULT == "CPU" and CPU_LLVM), "too long for CI")
   def test_transcribe_long_no_batch(self):
     waveforms = [load_file_waveform(fetch(TEST_FILE_3_URL)), load_file_waveform(TEST_FILE_1)]

test/opt/test_gen_float4.py

@@ -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, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
 
     assert TestFloat4.count_float4(program.uops) == (2, 1)
 
@@ -35,7 +35,8 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
+    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
     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")
@@ -46,7 +47,8 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
+      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
 
     sizes = [12, 8, 16]
     shifts = [3, 2, 4]
@@ -64,7 +66,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, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
 
     assert TestFloat4.count_float4(program.uops) == (0, 1)
 
@@ -75,7 +77,8 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
+    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
 
     assert TestFloat4.count_float4(uops) == (0, 2)
 
@@ -87,7 +90,8 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
+      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
 
     sizes = [13, 9, 17]
     shifts = [3, 2, 4]
@@ -105,7 +109,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, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
 
     assert TestFloat4.count_float4(uops) == (0, 0)
 
@@ -119,7 +123,8 @@ class TestFloat4(unittest.TestCase):
     # UPDATE: now we do this fusion
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    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
 
     assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
 
@@ -132,7 +137,7 @@ class TestFloat4(unittest.TestCase):
     # since the top axis is not contiguous.
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
 
     assert TestFloat4.count_float4(uops) == (0, 1)
 
@@ -144,7 +149,7 @@ class TestFloat4(unittest.TestCase):
     # should float4 b but not a
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 

test/test_arange.py

@@ -13,7 +13,7 @@ class TestArange(unittest.TestCase):
     GlobalCounters.reset()
     sched = tensor.schedule()
     self.assertEqual(len(sched), 1)
-    p = get_program(sched[-1].ast)
+    p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
     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)
+      p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
       self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)
 
 DSET, DDIM = 2048, 32

test/test_assign.py

@@ -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
+from tinygrad.helpers import temp, CI, CPU_LVP
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -180,6 +180,7 @@ 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

test/test_custom_kernel.py

@@ -155,12 +155,17 @@ 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)

test/test_dtype.py

@@ -17,8 +17,6 @@ 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
@@ -435,6 +433,8 @@ 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,3 +442,4 @@ class TestOpsBFloat16(unittest.TestCase):
 
 if __name__ == '__main__':
   unittest.main()
+

test/test_gc.py

@@ -7,13 +7,22 @@ from tinygrad.engine.realize import run_schedule
 from tinygrad.uop.ops import UOp
 from tinygrad.tensor import Tensor
 
+def _allocations_of_type(t):
+  ret = 0
+  for x in gc.get_objects():
+    try:
+      if isinstance(x, t): ret += 1
+    except ReferenceError:
+      pass
+  return ret
+
 def tensors_allocated():
   gc.collect()
-  return sum([isinstance(x, Tensor) for x in gc.get_objects()])
+  return _allocations_of_type(Tensor)
 
 def bufs_allocated():
   gc.collect()
-  return sum([isinstance(x, Buffer) for x in gc.get_objects()])
+  return _allocations_of_type(Buffer)
 
 class TestGC(unittest.TestCase):
 

test/test_graph.py

@@ -8,6 +8,8 @@ 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
@@ -154,6 +156,7 @@ 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()
 
@@ -167,6 +170,7 @@ 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()
 
@@ -215,6 +219,7 @@ 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()
 

test/test_jit.py

@@ -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
+from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV, needs_second_gpu
 from tinygrad.tensor import Tensor
 from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
@@ -439,6 +439,7 @@ 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"
@@ -472,6 +473,7 @@ 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"
@@ -499,6 +501,18 @@ 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):

test/test_linearizer.py

@@ -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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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,16 +92,15 @@ 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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]])
@@ -114,7 +113,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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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
@@ -124,7 +123,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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -135,7 +134,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, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, 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?"
@@ -147,7 +146,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, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, 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"
 
@@ -156,7 +155,8 @@ 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, 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, renderer=Device[Device.DEFAULT].renderer,
+                       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, opts=opts_to_apply)
+    program = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[])
+        program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[])
+      program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=opt).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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, opts=())
+    program = get_program(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, 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).uops
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).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).uops
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).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).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).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, opts=opt).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, 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).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).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, opts=opts), device=device))
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/test_linearizer_failures.py

@@ -3,6 +3,7 @@ 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):
@@ -18,7 +19,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)
+    get_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/test_multitensor.py

@@ -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
+from test.helpers import REAL_DEV, not_support_multi_device, needs_second_gpu
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -35,6 +35,9 @@ 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)
@@ -54,6 +57,12 @@ 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()
@@ -827,6 +836,7 @@ 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()
@@ -851,6 +861,9 @@ 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()
@@ -972,6 +985,9 @@ 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
@@ -1126,9 +1142,11 @@ 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))
 
@@ -1171,6 +1189,7 @@ 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
@@ -1180,6 +1199,9 @@ 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()
@@ -1242,6 +1264,7 @@ 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))
 

test/test_nn.py

@@ -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
+from test.helpers import not_support_multi_device, needs_second_gpu
 
 @unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
 class TestNN(unittest.TestCase):
@@ -481,6 +481,21 @@ 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")
@@ -519,6 +534,7 @@ 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")

test/test_ops.py

@@ -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
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM, EMULATE
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -16,6 +16,7 @@ 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)
@@ -38,6 +39,7 @@ 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}"
@@ -421,8 +423,9 @@ 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)
-    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])
+    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])
 
   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)
@@ -594,7 +597,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):
+    if is_dtype_supported(dtypes.uint64) and not COMPILE_ONLY:
       x = Tensor(2**64 - 1, dtype=dtypes.uint64).idiv(1)
       np.testing.assert_equal(x.numpy(), 2**64 - 1)
 
@@ -679,6 +682,7 @@ 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):
@@ -1070,8 +1074,7 @@ 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))
-    # 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)
+    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)
@@ -1088,8 +1091,9 @@ 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]])
-    np.testing.assert_equal(Tensor([2,2]).argmax().numpy(), 0)
-    np.testing.assert_equal(Tensor([1,2,2]).argmax().numpy(), 1)
+    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)
     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)
@@ -1107,8 +1111,9 @@ 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]])
-    np.testing.assert_equal(Tensor([2,2]).argmin().numpy(), 0)
-    np.testing.assert_equal(Tensor([3,2,2]).argmin().numpy(), 1)
+    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)
     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)
@@ -1156,12 +1161,13 @@ 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
-    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])
+    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])
     self.helper_test_exception([(4)], lambda x: x.topk(5), expected=(RuntimeError, ValueError))
 
   @slow_test
@@ -1313,6 +1319,7 @@ 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):
@@ -1723,6 +1730,7 @@ 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)
@@ -2729,6 +2737,9 @@ 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))
@@ -2749,6 +2760,7 @@ 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()
@@ -3165,6 +3177,7 @@ 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)

test/test_opt_gemm.py

@@ -1,6 +1,6 @@
 import numpy as np
 import unittest
-from tinygrad import Tensor
+from tinygrad import Tensor, Device
 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, opts=opts))
+    run = CompiledRunner(get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, 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)

test/test_opts.py

@@ -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)
+      prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
       self.assertIn('float4', prg.src)
 
 if __name__ == '__main__':

test/test_quantize_onnx.py

@@ -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, opts=opts)
+  prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, 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"))

test/test_randomness.py

@@ -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
+from test.helpers import not_support_multi_device, needs_second_gpu
 
 import numpy as np
 import torch
@@ -141,6 +141,7 @@ 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)
@@ -160,6 +161,7 @@ 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)

test/test_schedule.py

@@ -12,8 +12,7 @@ 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 get_rangeify_map, Kernel
-from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.schedule.rangeify import Kernel
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 
 class KernelCountException(Exception): pass
@@ -24,13 +23,11 @@ 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}"
-    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))
+    sched = Tensor(t).schedule()
   # 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 {len(sched)}")
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
     if DEBUG >= 3:
       for i,s in enumerate(sched):
         print("kernel", i+1)

test/test_subbuffer.py

@@ -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
+from test.helpers import REAL_DEV, needs_second_gpu
 
 @unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
 class TestSubBuffer(unittest.TestCase):
@@ -41,6 +41,7 @@ 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()

test/test_tensor.py

@@ -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, W.grad
+    xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
     out.backward()
-    xgrad2,wgrad2 = x.grad, W.grad
+    xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
     out.backward() # no need to retain again since we will not re-run backward
-    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)
+    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)
 
   def test_second_order_backward_pass(self):
     def test_pytorch():
@@ -829,6 +829,7 @@ 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()
@@ -841,11 +842,13 @@ 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}")
-    self.assertSetEqual(set(m.name for m in si.metadata), {"sigmoid", "relu"})
+    # skip numpy, this is schedule cache
+    self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"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)

test/test_uop_graph.py

@@ -660,6 +660,18 @@ 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)
 

test/test_uops.py

@@ -8,6 +8,7 @@ from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
 from tinygrad.uop.spec import shared_spec
 from tinygrad.renderer import ProgramSpec
+from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner, ExecItem
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
@@ -135,9 +136,9 @@ class TestFloatUOps(TestUOps):
 class TestNonFloatUOps(TestUOps):
   def test_add_int32(self): self._test_bop_fxn(Ops.ADD, lambda a,b: int(a)+int(b), (dtypes.int32, dtypes.int32))
   def test_mul_int32(self): self._test_bop_fxn(Ops.MUL, lambda a,b: int(a)*int(b), (dtypes.int32, dtypes.int32))
-  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "only ptx uses bitshifts")
+  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
   def test_shr_int32(self): self._test_bop_fxn(Ops.SHR, lambda a,b: int(a)>>int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
-  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "only ptx uses bitshifts")
+  @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
   def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)<<int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
   def test_div_int32(self):
     self._test_bop_fxn(Ops.IDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)

test/test_uops_stats.py

@@ -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, opts=[])
+    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, 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, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, 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,18 +190,19 @@ 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, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, 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, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    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)])
     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, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, 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")
@@ -210,7 +211,7 @@ class TestStatsOptimized(unittest.TestCase):
 
   def test_gemm_group(self):
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 4)])
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
@@ -219,14 +220,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, opts=[])
+    p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, 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, opts=[Opt(OptOps.GROUP, 0, 50)])
+      p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 50)])
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     # NOTE: these are wrong, they don't respect the if statement

test/testextra/test_hevc.py

@@ -0,0 +1,65 @@
+import unittest
+
+from tinygrad import Device
+from tinygrad.helpers import fetch
+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()
+
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat, device=Device.DEFAULT)
+
+    def _test_common(frame, bts):
+      self.assertEqual(frame0.pic_width_in_luma_samples, 1952)
+      self.assertEqual(frame0.pic_height_in_luma_samples, 1216)
+      self.assertEqual(frame0.chroma_format_idc, 1)
+      self.assertEqual(frame0.bit_depth_luma, 8)
+      self.assertEqual(frame0.bit_depth_chroma, 8)
+      self.assertEqual(frame0.log2_min_luma_coding_block_size, 3)
+      self.assertEqual(frame0.log2_max_luma_coding_block_size, 5)
+      self.assertEqual(frame0.log2_min_transform_block_size, 2)
+      self.assertEqual(frame0.log2_max_transform_block_size, 5)
+      self.assertEqual(frame0.num_tile_columns, 3)
+      self.assertEqual(frame0.num_tile_rows, 1)
+      self.assertEqual(frame0.colMvBuffersize, 589)
+      self.assertEqual(frame0.HevcSaoBufferOffset, 2888)
+      self.assertEqual(frame0.HevcBsdCtrlOffset, 25992)
+      self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcFltAboveOffset, 26714)
+      self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcSaoAboveOffset, 36214)
+
+      # tiles
+      self.assertEqual(bytes(bts[0x200:0x210]), b'\x18\x00&\x00\x18\x00&\x00\r\x00&\x00\x00\x00\x00\x00')
+
+    frame0 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[0].data())
+    _test_common(frame0, opaque[0].data())
+    self.assertEqual(frame0.stream_len, 148063)
+    self.assertEqual(frame0.IDR_picture_flag, 1)
+    self.assertEqual(frame0.RAP_picture_flag, 1)
+    self.assertEqual(frame0.sw_hdr_skip_length, 0)
+    self.assertEqual(frame0.num_ref_frames, 0)
+
+    frame1 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[1].data())
+    _test_common(frame1, opaque[1].data())
+    self.assertEqual(frame1.stream_len, 57110)
+    self.assertEqual(frame1.IDR_picture_flag, 0)
+    self.assertEqual(frame1.RAP_picture_flag, 0)
+    self.assertEqual(frame1.sw_hdr_skip_length, 9)
+    self.assertEqual(frame1.num_ref_frames, 1)
+    self.assertEqual(list(frame1.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame1.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame1.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+
+    frame3 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[3].data())
+    _test_common(frame3, opaque[3].data())
+    self.assertEqual(frame3.stream_len, 47036)
+    self.assertEqual(frame3.IDR_picture_flag, 0)
+    self.assertEqual(frame3.RAP_picture_flag, 0)
+    self.assertEqual(frame3.sw_hdr_skip_length, 9)
+    self.assertEqual(frame3.num_ref_frames, 1)
+    self.assertEqual(list(frame3.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+if __name__ == "__main__":
+  unittest.main()

test/testextra/test_tk.py

@@ -1,8 +1,9 @@
-import unittest, math
+import unittest, math, time
 
 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
 
@@ -10,13 +11,18 @@ from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
 from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout
 
-@unittest.skipIf(CI and Device.DEFAULT not in ["AMD"], "only amd")
+@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((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel("simple_matmul", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -25,14 +31,16 @@ 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 = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, 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)
 
       col, row = ker.blockIdx_x, ker.blockIdx_y
 
-      c_reg = warp.zero(c_reg)
+      c_reg_col = warp.zero(c_reg_col)
       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)
@@ -40,8 +48,11 @@ class TestTK(unittest.TestCase):
         a_reg = warp.load(a_reg, a_smem)
         b_reg = warp.load(b_reg, b_smem)
 
-        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
-      c_reg = ker.endrange()
+        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 = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
 
@@ -65,7 +76,7 @@ class TestTK(unittest.TestCase):
   def test_simple_matmul_transposed(self):
     N = 8192
     BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128
-    with Kernel((N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
+    with Kernel("simple_matmul_transposed", (N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -113,7 +124,7 @@ class TestTK(unittest.TestCase):
   def test_load_store(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel("load_store", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -146,11 +157,50 @@ 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((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
+    with Kernel("load_store_group", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
       warp = ker.warp
       group = ker.group(2)
 
@@ -187,7 +237,7 @@ class TestTK(unittest.TestCase):
   def test_add(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("add", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -224,7 +274,7 @@ class TestTK(unittest.TestCase):
   def test_max(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("max", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -269,7 +319,7 @@ class TestTK(unittest.TestCase):
   def test_max_nonsquare(self):
     N, M = 32, 128
     BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("max_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -314,7 +364,7 @@ class TestTK(unittest.TestCase):
   def test_sum(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("sum", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -359,7 +409,7 @@ class TestTK(unittest.TestCase):
   def test_sum_nonsquare(self):
     N, M = 32, 128
     BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("sum_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -404,7 +454,7 @@ class TestTK(unittest.TestCase):
   def test_softmax(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("softmax", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
@@ -422,28 +472,29 @@ 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.after(norm_vec), 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 = (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()
 
@@ -463,7 +514,7 @@ class TestTK(unittest.TestCase):
   def test_softmax_col(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("softmax_col", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
@@ -481,28 +532,29 @@ 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()
 
@@ -521,11 +573,11 @@ class TestTK(unittest.TestCase):
 
   def test_fa(self):
     NUM_WORKERS = 1
-    B, N, H, H_KV, D = 1, 8192, 32, 8, 128
+    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
     Q_BLOCK_SIZE = 16
     KV_BLOCK_SIZE = 16
     GROUP_SIZE = H // H_KV
-    with Kernel((H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
+    with Kernel("fa", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
       # kernel
@@ -605,6 +657,7 @@ 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
 
@@ -621,7 +674,12 @@ 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): ei.run(wait=True)
+    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")
     out = out.float()
 
     q_permuted = q.permute(0, 2, 1, 3)
@@ -630,7 +688,36 @@ 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=1e-2, rtol=1e-5)
+    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)
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,8 +1,14 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.apps.llm import apply_rope
+from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
 #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):
@@ -39,7 +45,7 @@ class TestAttention(unittest.TestCase):
     prune_size = len(rope_prune.captured.jit_cache)
 
     self.assertGreater(noprune_size, prune_size)
-    self.assertGreaterEqual(noprune_size, 3)
+    self.assertGreaterEqual(noprune_size, 2)
     self.assertEqual(prune_size, 1)
 
 if __name__ == '__main__':

test/unit/test_autogen.py

@@ -44,6 +44,34 @@ 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
@@ -82,4 +110,53 @@ 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()

test/unit/test_device.py

@@ -42,12 +42,10 @@ 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_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"})
+                        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"})
     elif Device.DEFAULT == "AMD":
       from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
       try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
@@ -64,14 +62,20 @@ 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")
 
-  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
+  @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
 
 class MockCompiler(Compiler):
   def __init__(self, key): super().__init__(key)

test/unit/test_gradient.py

@@ -110,6 +110,18 @@ 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()

test/unit/test_helpers.py

@@ -163,6 +163,14 @@ 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])

test/unit/test_indexing.py

@@ -176,9 +176,8 @@ class TestIndexing(unittest.TestCase):
     self.assertRaises(IndexError, lambda: reference[0.0, ..., 0.0:2.0])
     self.assertRaises(IndexError, lambda: reference[0.0, :, 0.0])
 
-    # TODO: delitem
-    # def delitem(): del reference[0]
-    # self.assertRaises(TypeError, delitem)
+    def delitem(): del reference[0]
+    self.assertRaises(TypeError, delitem)
 
   # TODO setitem backward
   '''

test/unit/test_llm_server.py

@@ -0,0 +1,136 @@
+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()

test/unit/test_microbenchmarks.py

@@ -1,51 +1,99 @@
 import unittest, time
+from tinygrad import dtypes, Tensor, UOp, getenv
 from tinygrad.helpers import Profiling
-from tinygrad.uop.ops import UOp
-from tinygrad.dtype import dtypes
 
-# it's about 1 ms per 1k UOps on M3
-N = 10000
+PYPROFILE = getenv("PYPROFILE")
+class TestBench(unittest.TestCase):
+  @staticmethod
+  def setUpClass():
+    # no fixed cost
+    Tensor.empty(10,10)
+    Tensor.randn(10,10)
 
-class TestMicrobenchmarks(unittest.TestCase):
+  def start_time(self): self.st = time.perf_counter()
   def setUp(self):
-    self.st = time.perf_counter()
+    # it's about 1 ms per 1k UOps on M3
+    if PYPROFILE:
+      self.prof = Profiling()
+      self.prof.__enter__()
+    else:
+      self.prof = None
+    self.N = 10000
+    self.start_time()
+
   def tearDown(self):
     et = (time.perf_counter() - self.st)
-    print(f"{self._testMethodName} {et*1e3:.2f} ms")
+    if self.prof is not None: self.prof.__exit__()
+    print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")
 
   def test_uop_instant_creation(self):
-    for i in range(N): UOp.const(dtypes.int, 100+i)
+    for i in range(self.N): UOp.const(dtypes.int, 100+i)
 
   def test_uop_list_creation(self):
-    [UOp.const(dtypes.int, 100+i) for i in range(N)]
+    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]
 
   def test_uop_add_2n(self):
     a = UOp.const(dtypes.int, 2)
-    for _ in range(N): a = a + a
+    for _ in range(self.N): a = a + a
 
   def test_uop_toposort(self):
     a = UOp.const(dtypes.int, 0)
-    for i in range(N): a = a + UOp.const(dtypes.int, 100+i)
-    self.setUp()
-    self.assertEqual(len(a.toposort()), 2*N+1)
+    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
+    self.start_time()
+    self.assertEqual(len(a.toposort()), 2*self.N+1)
 
   def test_uop_toposort_2n(self):
     a = UOp.const(dtypes.int, 0)
-    for i in range(N): a = a + a
-    self.setUp()
-    self.assertEqual(len(a.toposort()), N+1)
+    for _ in range(self.N): a = a + a
+    self.start_time()
+    self.assertEqual(len(a.toposort()), self.N+1)
 
   def test_uop_simplify(self):
     a = UOp.const(dtypes.int, 2)
-    for _ in range(N): (a+a).simplify()
+    for _ in range(self.N): (a+a).simplify()
 
-class TestMicroprofile(unittest.TestCase):
   def test_uop_simplify_complex(self):
+    self.N //= 10 # this test is slow
     x = UOp.variable("x", 0, 10)
     y = UOp.variable("y", 0, 10)
     expr = (x*2)+5+(x*4)+(y*2)+y
-    with Profiling():
-      for _ in range(1000): expr.simplify()
+    for _ in range(self.N): expr.simplify()
+
+  def test_uop_simplify_div(self):
+    self.N //= 10 # this test is slow
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    z = UOp.variable("z", 0, 10)
+    expr = (x*4+y*8)//(z*2)
+    for _ in range(self.N): expr.simplify()
+
+  def test_uop_chain_free(self):
+    a = UOp.const(dtypes.int, 2)
+    for _ in range(self.N): a = a + a
+    self.start_time()
+    del a
+
+  def test_tensor_zeros(self):
+    self.N //= 10 # this test is slow
+    for _ in range(self.N): Tensor.zeros(10, 10)
+
+  def test_tensor_add(self):
+    self.N //= 10 # this test is slow
+    a = Tensor.zeros(10, 10)
+    b = Tensor.zeros(10, 10)
+    for _ in range(self.N): a+b
+
+  def test_tensor_empty(self):
+    self.N //= 10 # this test is slow
+    for _ in range(self.N): Tensor.empty(10, 10)
+
+  def test_tensor_rand(self):
+    self.N //= 100 # this test is very slow
+    for _ in range(self.N): Tensor.rand(10, 10)
+
+  def test_tensor_randn(self):
+    self.N //= 100 # this test is very slow
+    for _ in range(self.N): Tensor.randn(10, 10)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_png.py

@@ -0,0 +1,31 @@
+#!/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()

test/unit/test_schedule_cache.py

@@ -0,0 +1,47 @@
+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()

test/unit/test_shm_tensor.py

@@ -22,7 +22,8 @@ class TestRawShmBuffer(unittest.TestCase):
 
   @unittest.skipIf(CI, "CI doesn't like big shared memory")
   def test_e2e_big(self):
-    t = Tensor.randn(2048, 2048, 8).realize()
+    # bigger than this doesn't work on Linux, maybe this is a limit somewhere?
+    t = Tensor.randn(2048, 128, 8).realize()
 
     # copy to shm
     shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name

test/unit/test_uop_symbolic.py

@@ -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
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -112,8 +112,9 @@ 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*a.const_like(-3)).simplify(), a*-1)
+    self.assertEqual((a*a*3).divide_exact(a*(-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)
@@ -322,12 +323,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), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -3, 4, "(v1%2*2+(v3+-1)%5+-2)")
 
   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), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")
+    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -2, 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)")
@@ -377,9 +378,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, 9, "(a%10)")  # TODO: tighter max
+    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 1, "(a%10)")
     self.helper_test_variable(Variable("a", 0, 20)%10, 0, 9, "(a%10)")
-    self.helper_test_variable(Variable("a", -1, 20)%10, -9, 9, "(a%10)")  # TODO: tighter min
+    self.helper_test_variable(Variable("a", -1, 20)%10, -1, 9, "(a%10)")
 
   def test_ge_remove(self):
     self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "False")
@@ -580,12 +581,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)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")
+    # 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)")
 
   def test_sum_mul_distribute(self):
     gidx0 = Variable("gidx0", 0, 7)
@@ -1029,6 +1030,25 @@ 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>

tinygrad/apps/llm.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers
+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
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
@@ -9,6 +10,7 @@ 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)|" + \
@@ -24,7 +26,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 = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
+    normal_tokens, special_tokens = 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]:
@@ -50,15 +52,18 @@ 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")
 
-def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
+@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:
   B, H, T, Hd = x.shape
   assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
-  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)
+  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]
   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)
 
@@ -94,8 +99,10 @@ 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)
 
-    q = apply_rope(q, start_pos)
-    k = apply_rope(k, start_pos)
+    # 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)
 
     # TODO: remove these kv cache realizes
     if not hasattr(self, "cache_kv"):
@@ -113,7 +120,8 @@ class TransformerBlock:
 
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
-    gated  = self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)
+    # TODO: remove the need for this contiguous
+    gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
@@ -139,7 +147,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) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=True) -> tuple[Transformer, dict]:
     # TODO: remove the need for copy to default device
     kv, state_dict = nn.state.gguf_load(gguf.to(None))
 
@@ -156,7 +164,8 @@ 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 nn.state.get_parameters(model): s.replace(s.contiguous())
+    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
+    if realize: Tensor.realize(*params)
     return model, kv
 
   def generate(self, tokens:list[int], start_pos=0):
@@ -172,26 +181,92 @@ class Transformer:
       yield next_id
 
 models = {
-  "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",
+  "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",
 }
 
+# *** 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("--size", choices=list(models.keys()), default=list(models.keys())[0], help="Model size")
+  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
   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.size]), args.max_context)
+  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)
 
   # 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

tinygrad/apps/resnet.py

@@ -0,0 +1,50 @@
+# 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])

tinygrad/codegen/late/devectorizer.py

@@ -3,8 +3,8 @@ import functools, operator, itertools
 from collections import defaultdict
 from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, symbolic, invalid_gate
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
 from tinygrad.helpers import getenv, flatten, AMX, prod
 from tinygrad.renderer import Renderer
 
@@ -59,11 +59,16 @@ load_store_indexing = PatternMatcher([
 def expand_index(buf:UOp, vec:UOp):
   if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
   # generate the individual indexes
-  midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)]),
-                       symbolic+load_store_indexing, name=f"index_buf_{buf.arg}")
+  return UOp(Ops.VECTORIZE, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))
+
+def fold_expanded_index(midx:UOp):
+  buf = midx.src[0].src[0]
+  if not all(s.src[0] is buf for s in midx.src): return None
+  if not all(isinstance(s.dtype, PtrDType) for s in midx.src): return None
+
   # extract all the relevant offsets
   offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
-  for i in range(vec.dtype.count):
+  for i in range(len(midx.src)):
     idx: Any = midx.src[i].src[1].get_idx()
     if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
     elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
@@ -75,7 +80,7 @@ def expand_index(buf:UOp, vec:UOp):
 
   # then rewrite everything we can into groups
   ret = []
-  idxs: list[int|None] = [None]*vec.dtype.count
+  idxs: list[int|None] = [None]*len(midx.src)
   global_offset = 0
   for offsets in offsets_rootsrc.values():
     grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
@@ -113,6 +118,7 @@ def gep_on_store(gep:UOp, st:UOp, sto:UOp):
 
 load_store_folding = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
+  (UPat(Ops.VECTORIZE, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index),
   # GEP after LOAD
   (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
    lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),

tinygrad/codegen/opt/postrange.py

@@ -7,6 +7,7 @@ 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
@@ -228,6 +229,7 @@ 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)

tinygrad/codegen/opt/search.py

@@ -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, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
+def get_kernel_actions(s:Scheduler, include_0=True) -> 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 if candidates is None else candidates).copy()
+  kernel_actions = actions.copy()
 
   for i,a in enumerate(kernel_actions):
     if a.axis is not None and a.op is not OptOps.TC:

tinygrad/codegen/opt/tc.py

@@ -1,7 +1,6 @@
 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)
@@ -92,8 +91,7 @@ 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
-if getenv("ALLOW_TF32", 0): cuda_sm80 += cuda_8168_tf32
+cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16 + cuda_8168_tf32
 cuda_sm89: list[TensorCore] = cuda_sm80 + cuda_81632_f8
 
 # ***** AMD *****