fixed and test is real

* viz: tabulate pmc

* linter

* enable nesting

* pmc comes before waves

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

@@ -61,7 +61,7 @@ runs:
       uses: actions/cache@v4
       with:
         path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/pyproject.toml') }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
 
     # **** Caching downloads ****
 
@@ -221,7 +221,7 @@ runs:
         sudo mkdir -p /usr/local/lib
         curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
-          sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
+          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
 
     # **** gpuocelot ****
@@ -278,7 +278,7 @@ runs:
       if: inputs.webgpu == 'true' && runner.os == 'Linux'
       shell: bash
       run: |
-        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo curl -fL https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
         sudo ldconfig
     - name: Install WebGPU dawn (macOS)
       if: inputs.webgpu == 'true' && runner.os == 'macOS'
@@ -298,7 +298,7 @@ runs:
     - name: Install mesa (linux)
       if: inputs.mesa == 'true' && runner.os == 'Linux'
       shell: bash
-      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
     - name: Install mesa (macOS)
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash

.github/workflows/autogen.yml

@@ -13,9 +13,11 @@ on:
   pull_request:
     paths:
     - 'tinygrad/runtime/autogen/**/*'
+    - 'tinygrad/runtime/support/autogen.py'
   workflow_dispatch:
     paths:
     - 'tinygrad/runtime/autogen/**/*'
+    - 'tinygrad/runtime/support/autogen.py'
 
 jobs:
   autogen:

.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
@@ -318,28 +318,28 @@ 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
@@ -433,9 +433,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 +524,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,10 +588,10 @@ 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
@@ -625,12 +623,6 @@ 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
@@ -643,14 +635,14 @@ jobs:
       run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
-    # - name: benchmark MobileNetV2 on DSP
-    #   run: |
-    #     # generate quantized weights
-    #     ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-    #     ln -s /data/home/tiny/tinygrad/testsig-*.so .
-    #     PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
-    #     # benchmark on DSP with NOOPT=1, the devectorizer has issues
-    #     PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+    - name: benchmark MobileNetV2 on DSP
+      run: |
+        # generate quantized weights
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        # benchmark on DSP with NOOPT=1, the devectorizer has issues
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -706,10 +698,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 +760,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/szdiff.yml

@@ -56,15 +56,15 @@ jobs:
         uses: actions/checkout@v4
         with:
           path: base
-      - name: Set up Python 3.10
+      - name: Set up Python 3.12
         uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: '3.12'
       - name: Count Line Diff
         run: |
-          pip install tabulate
           BASE="$GITHUB_WORKSPACE/base"
           PR="$GITHUB_WORKSPACE/pr"
+          pip install tabulate $BASE
           cp "$BASE/sz.py" .
           echo "loc_content<<EOF" >> "$GITHUB_ENV"
           python sz.py "$BASE" "$PR" >> "$GITHUB_ENV"

.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: '14'
   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
@@ -86,65 +84,67 @@ jobs:
         clang -O2 recognize.c -lm -o recognize
         cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
 
-  # TODO: fix the torch backend and reenable
-  # torchbackend:
-  #   name: Torch Backend Tests
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       pydeps: "pillow torchvision expecttest"
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Lint with ruff
-  #     run: |
-  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
-  #       python3 -m ruff check extra/torch_backend/backend.py
-  #   - name: Test one op
-  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-  #   - name: Test ResNet-18
-  #     run: DEBUG=2 python3 extra/torch_backend/example.py
-  #   - name: My (custom) tests
-  #     run: python3 extra/torch_backend/test.py
-  #   - name: Test one op in torch tests
-  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-  #   - name: Test Ops with TINY_BACKEND
-  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-  #   - name: Test in-place operations on views
-  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-  #   - name: Test multi-gpu
-  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  torchbackend:
+    name: Torch Backend Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
+      run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: My (custom) tests
+      run: python3 extra/torch_backend/test.py
+    - name: Test one op in torch tests
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+    - name: Test kernel fusion
+      run: python3 extra/torch_backend/test_kernel_fusion.py
 
-  # torchbackendmore:
-  #   name: Torch Backend Tests More
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
-  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-  #   - name: Test some torch tests (expect failure)
-  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+
+  torchbackendmore:
+    name: Torch Backend Tests More
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+    - name: Test some torch tests (expect failure)
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
   bepython:
     name: Python Backend
@@ -261,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
@@ -306,6 +308,7 @@ jobs:
       with:
         key: spec-unit
         deps: testing_unit
+        python-version: '3.14'
     - name: Test SPEC=2
       run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
@@ -323,6 +326,8 @@ jobs:
         deps: testing_unit
     - name: Fuzz Test symbolic
       run: python test/external/fuzz_symbolic.py
+    - name: Fuzz Test symbolic (symbolic divisors)
+      run: python test/external/fuzz_symbolic_symbolic_div.py
     - name: Fuzz Test fast idiv
       run: python test/external/fuzz_fast_idiv.py
     - name: Fuzz Test shape ops
@@ -442,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)

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

docs/mnist.md

@@ -131,7 +131,7 @@ timeit.repeat(jit_step, repeat=5, number=1)
 
 1.0 ms is 75x faster! Note that we aren't syncing the GPU, so GPU time may be slower.
 
-The slowness the first two times is the JIT capturing the kernels. And this JIT will not run any Python in the function, it will just replay the tinygrad kernels that were run, so be aware that non tinygrad Python operations won't work. Randomness functions work as expected.
+The first two runs of the function execute normally, with the JIT capturing the kernels. Starting from the third run, only the tinygrad operations are replayed, removing the overhead by skipping Python code execution. So be aware that any non-tinygrad Python values affecting the kernels will be "frozen" from the second run. Note that `Tensor` randomness functions work as expected.
 
 Unlike other JITs, we JIT everything, including the optimizer. Think of it as a dumb replay on different data.
 

docs/ramp.py

@@ -1,293 +0,0 @@
-#!/usr/bin/env python3
-
-# this file is a "ramp" for people new to tinygrad to think about how to approach it
-# it is runnable and editable.
-# whenever you see stuff like DEBUG=2 or CPU=1 discussed, these are environment variables
-# in a unix shell like bash `DEBUG=2 CPU=1 python docs/ramp.py`
-
-# this pip installs tinygrad master for the system
-# the -e allows you to edit the tinygrad folder and update system tinygrad
-# tinygrad is pure Python, so you are encouraged to do this
-# git pull in the tinygrad directory will also get you the latest
-"""
-git clone https://github.com/tinygrad/tinygrad.git
-cd tinygrad
-python3 -m pip install -e .
-"""
-
-# %% ********
-print("******* PART 1 *******")
-
-# we start with a Device.
-# a Device is where Tensors are stored and compute is run
-# tinygrad autodetects the best device on your system and makes it the DEFAULT
-from tinygrad import Device
-print(Device.DEFAULT)  # on Mac, you can see this prints METAL
-
-# now, lets create a Tensor
-from tinygrad import Tensor, dtypes
-t = Tensor([1,2,3,4])
-
-# you can see this Tensor is on the DEFAULT device with int dtype and shape (4,)
-assert t.device == Device.DEFAULT
-assert t.dtype == dtypes.int
-assert t.shape == (4,)
-
-# unlike in torch, if we print it, it doesn't print the contents
-# this is because tinygrad is lazy
-# this Tensor has not been computed yet
-print(t)
-# <Tensor <UOp METAL (4,) int (<Ops.COPY: 7>, None)> on METAL with grad None>
-
-# the ".uop" property on Tensor contains the specification of how to compute it
-print(t.uop)
-"""
-UOp(Ops.COPY, dtypes.int, arg=None, src=(
-  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-    UOp(Ops.UNIQUE, dtypes.void, arg=0, src=()),
-    UOp(Ops.DEVICE, dtypes.void, arg='PYTHON', src=()),)),
-  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
-"""
-# as you can see, it's specifying a copy from PYTHON device
-# which is where the [1,2,3,4] array lives
-
-# UOps are the specification language in tinygrad
-# they are immutable and form a DAG
-# they have a "Ops", a "dtype", a tuple of srcs (parents), and an arg
-
-t.realize()
-# if we want to "realize" a tensor, we can with the "realize" method
-# now when we look at the uop, it's changed
-print(t.uop)
-"""
-UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-  UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
-  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
-"""
-# the copy was actually run, and now the "uop" of the Tensor is just a BUFFER
-# if you run this script with DEBUG=2 in the environment, you can see the copy happen
-# *** METAL      1 copy       16,   METAL <- PYTHON ...
-
-# now let's do some compute
-# we look at the uop to see the specification of the compute
-t_times_2 = t * 2
-print(t_times_2.uop)
-"""
-UOp(Ops.MUL, dtypes.int, arg=None, src=(
-  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
-    x2:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),)),
-  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
-    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
-      UOp(Ops.CONST, dtypes.int, arg=2, src=(
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
-           x2,)),)),)),)),))
-"""
-# the BUFFER from above is being multiplied by a CONST 2
-# it's RESHAPEd and EXPANDed to broadcast the CONST to the BUFFER
-
-# we can check the result with
-assert t_times_2.tolist() == [2, 4, 6, 8]
-
-# UOps are both immutable and globally unique
-# if i multiply the Tensor by 4 twice, these result Tensors will have the same uop specification
-t_times_4_try_1 = t * 4
-t_times_4_try_2 = t * 4
-assert t_times_4_try_1.uop is t_times_4_try_2.uop
-# the specification isn't just the same, it's the exact same Python object
-assert t_times_4_try_1 is not t_times_4_try_2
-# the Tensor is a different Python object
-
-# if we realize `t_times_4_try_1` ...
-t_times_4_try_1.realize()
-print(t_times_4_try_2.uop)
-"""
-UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-  UOp(Ops.UNIQUE, dtypes.void, arg=4, src=()),
-  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
-"""
-# ... `t_times_4_try_2` also becomes the same BUFFER
-assert t_times_4_try_1.uop is t_times_4_try_2.uop
-# so this print doesn't require any computation, just a copy back to the CPU so we can print it
-print("** only the copy start")
-print(t_times_4_try_2.tolist())  # [4, 8, 12, 16]
-print("** only the copy end")
-# you can confirm this with DEBUG=2, seeing what's printed in between the "**" prints
-
-# tinygrad has an auto differentiation engine that operates according to these same principles
-# the derivative of "log(x)" is "1/x", and you can see this on line 20 of gradient.py
-t_float = Tensor([3.0])
-t_log = t_float.log()
-t_log_grad, = t_log.sum().gradient(t_float)
-# due to how log is implemented, this gradient contains a lot of UOps
-print(t_log_grad.uop)
-# ...not shown here...
-# but if you run with DEBUG=4 (CPU=1 used here for simpler code), you can see the generated code
-"""
-void E_(float* restrict data0, float* restrict data1) {
-  float val0 = *(data1+0);
-  *(data0+0) = (1/val0);
-}
-"""
-# the derivative is close to 1/3
-assert (t_log_grad.item() - 1/3) < 1e-6
-
-# %% ********
-print("******* PART 2 *******")
-
-# we redefine the same t here so this cell can run on it's own
-from tinygrad import Tensor
-t = Tensor([1,2,3,4])
-
-# what's above gives you enough of an understanding to go use tinygrad as a library
-# however, a lot of the beauty of tinygrad is in how easy it is to interact with the internals
-# NOTE: the APIs here are subject to change
-
-t_plus_3_plus_4 = t + 3 + 4
-print(t_plus_3_plus_4.uop)
-"""
-UOp(Ops.ADD, dtypes.int, arg=None, src=(
-  UOp(Ops.ADD, dtypes.int, arg=None, src=(
-    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
-      x3:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
-    UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
-      UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
-        UOp(Ops.CONST, dtypes.int, arg=3, src=(
-          x7:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
-             x3,)),)),)),)),)),
-  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
-    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
-      UOp(Ops.CONST, dtypes.int, arg=4, src=(
-         x7,)),)),)),))
-"""
-# you can see it's adding both 3 and 4
-
-# but by the time we are actually running the code, it's adding 7
-# `kernelize` will simplify and group the operations in the graph into kernels
-t_plus_3_plus_4.kernelize()
-print(t_plus_3_plus_4.uop)
-"""
-UOp(Ops.ASSIGN, dtypes.int, arg=None, src=(
-  x0:=UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-    UOp(Ops.UNIQUE, dtypes.void, arg=7, src=()),
-    x2:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
-  UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 48>,) (__add__,)>, src=(
-     x0,
-    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
-      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
-       x2,)),)),))
-"""
-# ASSIGN has two srcs, src[0] is the BUFFER that's assigned to, and src[1] is the thing to assign
-# src[1] is the GPU Kernel that's going to be run
-# we can get the ast of the Kernel as follows
-kernel_ast = t_plus_3_plus_4.uop.src[1].arg.ast
-
-# almost everything in tinygrad functions as a rewrite of the UOps
-# the codegen rewrites the ast to a simplified form ready for "rendering"
-from tinygrad.codegen import full_rewrite_to_sink
-rewritten_ast = full_rewrite_to_sink(kernel_ast)
-print(rewritten_ast)
-"""
-UOp(Ops.SINK, dtypes.void, arg=None, src=(
-  UOp(Ops.STORE, dtypes.void, arg=None, src=(
-    UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
-      UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=0, src=()),
-      x3:=UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', 4), src=()),)),
-    UOp(Ops.ADD, dtypes.int, arg=None, src=(
-      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
-        UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=1, src=()),
-           x3,)),)),
-      UOp(Ops.CONST, dtypes.int, arg=7, src=()),)),)),))
-"""
-# you can see at this point we are adding 7, not 3 and 4
-
-# with DEBUG=4, we can see the code.
-# since optimizations are on, it UPCASTed the operation, explicitly writing out all 4 +7s
-t_plus_3_plus_4.realize()
-"""
-void E_4n2(int* restrict data0, int* restrict data1) {
-  int val0 = *(data1+0);
-  int val1 = *(data1+1);
-  int val2 = *(data1+2);
-  int val3 = *(data1+3);
-  *(data0+0) = (val0+7);
-  *(data0+1) = (val1+7);
-  *(data0+2) = (val2+7);
-  *(data0+3) = (val3+7);
-}
-"""
-# the function name E_4n2 is "E" for elementwise op (as opposed to "r" for reduce op)
-# "4" for the size, and "n2" for name deduping (it's the 3rd function with the same E and 4 in this session)
-# when you print the name with DEBUG=2, you'll see the 4 is yellow, meaning that it's upcasted
-# if you run with NOOPT=1 ...
-"""
-void E_4n2(int* restrict data0, int* restrict data1) {
-  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
-    int val0 = *(data1+ridx0);
-    *(data0+ridx0) = (val0+7);
-  }
-}
-"""
-# ... you get this unoptimized code with a loop and the 4 is blue (for global). the color code is in kernel.py
-
-# %% ********
-print("******* PART 3 *******")
-
-# now, we go even lower and understand UOps better and how the graph rewrite engine works.
-# it's much simpler than what's in LLVM or MLIR
-
-from tinygrad import dtypes
-from tinygrad.uop.ops import UOp, Ops
-
-# first, we'll construct some const UOps
-a = UOp(Ops.CONST, dtypes.int, arg=2)
-b = UOp(Ops.CONST, dtypes.int, arg=2)
-
-# if you have been paying attention, you should know these are the same Python object
-assert a is b
-
-# UOps support normal Python math operations, so a_plus_b expresses the spec for 2 + 2
-a_plus_b = a + b
-print(a_plus_b)
-"""
-UOp(Ops.ADD, dtypes.int, arg=None, src=(
-  x0:=UOp(Ops.CONST, dtypes.int, arg=2, src=()),
-   x0,))
-"""
-
-# we could actually render this 2+2 into a language like c and run it
-# or, we can use tinygrad's graph rewrite engine to "constant fold"
-
-from tinygrad.uop.ops import graph_rewrite, UPat, PatternMatcher
-
-# a `PatternMatcher` is a list of tuples. for each element in the list:
-# [0] is the pattern to match, and [1] is the function to run.
-# this function can return either a UOp to replace the pattern with, or None to not replace
-simple_pm = PatternMatcher([
-  (UPat(Ops.ADD, src=(UPat(Ops.CONST, name="c1"), UPat(Ops.CONST, name="c2"))),
-   lambda c1,c2: UOp(Ops.CONST, dtype=c1.dtype, arg=c1.arg+c2.arg)),
-])
-# this pattern matches the addition of two CONST and rewrites it into a single CONST UOp
-
-# to actually apply the pattern to a_plus_b, we use graph_rewrite
-a_plus_b_simplified = graph_rewrite(a_plus_b, simple_pm)
-print(a_plus_b_simplified)
-"""
-UOp(Ops.CONST, dtypes.int, arg=4, src=())
-"""
-# 2+2 is in fact, 4
-
-# we can also use syntactic sugar to write the pattern nicer
-simpler_pm = PatternMatcher([
-  (UPat.cvar("c1")+UPat.cvar("c2"), lambda c1,c2: c1.const_like(c1.arg+c2.arg))
-])
-assert graph_rewrite(a_plus_b, simple_pm) is graph_rewrite(a_plus_b, simpler_pm)
-# note again the use of is, UOps are immutable and globally unique
-
-# %% ********
-
-# that brings you to an understanding of the most core concepts in tinygrad
-# you can run this with VIZ=1 to use the web based graph rewrite explorer
-# hopefully now you understand it. the nodes in the graph are just UOps

docs/tinybox.md

@@ -41,7 +41,7 @@ The BMC also has a web interface you can use if you find that easier.
 It is recommended that you change the BMC password after setting up the box, as the password on the screen is only the initial password.
 
 If you do decide to change the BMC password and no longer want the initial password to be displayed, remove the `/root/.bmc_password` file.
-Reboot after making these changes or restart the `displayservice.service` service.
+Reboot after making these changes or restart the `tinybox-display.service` service.
 
 ## What do I use it for?
 

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/llama3.py

@@ -1,8 +1,6 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
-import tiktoken
-from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -12,6 +10,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
   pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
   def __init__(self, model_path: str):
+    import tiktoken
+    from tiktoken.load import load_tiktoken_bpe
     mergeable_ranks = load_tiktoken_bpe(model_path)
     self.num_base_tokens = len(mergeable_ranks)
     special_tokens = [

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

examples/stable_diffusion.py

@@ -9,7 +9,7 @@ from typing import Dict, Any
 from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
-from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
 from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
@@ -266,13 +266,16 @@ if __name__ == "__main__":
   parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
   args = parser.parse_args()
 
+  profile_marker("create model")
   model = StableDiffusion()
 
-  # load in weights
+  profile_marker("load in weights")
   with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
     if not args.fakeweights:
       model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
-      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
+      state_dict = torch_load(model_bin)['state_dict']
+      profile_marker("state dict loaded")
+      load_state_dict(model, state_dict, verbose=False, strict=False, realize=False)
 
     if args.fp16:
       for k,v in get_state_dict(model).items():
@@ -281,12 +284,13 @@ if __name__ == "__main__":
 
     Tensor.realize(*get_state_dict(model).values())
 
-  # run through CLIP to get context
+  profile_marker("run clip (conditional)")
   tokenizer = Tokenizer.ClipTokenizer()
   prompt = Tensor([tokenizer.encode(args.prompt)])
   context = model.cond_stage_model.transformer.text_model(prompt).realize()
   print("got CLIP context", context.shape)
 
+  profile_marker("run clip (unconditional)")
   prompt = Tensor([tokenizer.encode("")])
   unconditional_context = model.cond_stage_model.transformer.text_model(prompt).realize()
   print("got unconditional CLIP context", unconditional_context.shape)
@@ -310,6 +314,7 @@ if __name__ == "__main__":
   step_times = []
   with Context(BEAM=getenv("LATEBEAM")):
     for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
+      profile_marker(f"step {len(timesteps)-index-1}")
       GlobalCounters.reset()
       st = time.perf_counter_ns()
       t.set_description("%3d %3d" % (index, timestep))
@@ -319,24 +324,26 @@ if __name__ == "__main__":
           latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
           if args.timing: Device[Device.DEFAULT].synchronize()
       step_times.append((time.perf_counter_ns() - st)*1e-6)
+    # done with diffusion model
     del run
+    del model.model
 
   if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
     min_time = min(step_times)
     assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-  # upsample latent space to image with autoencoder
-  x = model.decode(latent)
+  profile_marker("run decoder") # upsample latent space to image with autoencoder
+  x = model.decode(latent).realize()
   print(x.shape)
 
-  # save image
+  profile_marker("save image")
   im = Image.fromarray(x.numpy())
   print(f"saving {args.out}")
   im.save(args.out)
   # Open image.
   if not args.noshow: im.show()
 
-  # validation!
   if args.prompt == default_prompt and args.steps == 6 and args.seed == 0 and args.guidance == 7.5:
+    profile_marker("validate")
     ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "stable_diffusion_seed0.png")))
     distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
     assert distance < 3e-3, colored(f"validation failed with {distance=}", "red")  # higher distance with WINO

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,12 +1,13 @@
+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
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
 
-N = 4096
+N = getenv("N", 4096)
 M = K = N
-run_count = 5
+run_count = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
@@ -140,29 +141,29 @@ 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):
       ets.append(ei.run(wait=True))
   print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
 
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    tc = (a @ b).realize()
-  with Context(DEBUG=0):
-    err = (hc - tc).square().mean().item()
-  print(f"mean squared error {err}")
-  if err > 1e-06:
-    raise RuntimeError("matmul is wrong!")
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2):
+      tc = (a @ b).realize()
+    with Context(DEBUG=0):
+      err = (hc - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > 1e-06:
+      raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
   test_matmul(hand_spec_kernel3(), N=N)

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

@@ -1,8 +1,10 @@
-import pathlib
+import os, pathlib
+
+# TODO: there is a timing bug without this
+os.environ["AMD_AQL"] = "1"
+
 from tinygrad.device import Device
 from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
-import time
-import os
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -44,9 +46,9 @@ if __name__=="__main__":
     raise RuntimeError("Error while initiating AMD device")
 
   COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103'}:
-    if DEV.arch == 'gfx1103':
-      NUM_WORKGROUPS = 8
+  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
+    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/nv_gpu_driver/nv_ioctl.py

@@ -64,14 +64,17 @@ nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(
           x.startswith("NV") and x[6:].startswith("_CTRL_") and isinstance(getattr(nv_gpu, x), int)}
 
 def get_classes():
-  hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_570.py").read_text()
-  clss = re.search(r'NV01_ROOT.*?NV_SEMAPHORE_SURFACE = \(0x000000da\) # macro', hdrpy, re.DOTALL).group()
-  pattern = r'([0-9a-zA-Z_]*) = +\((0x[0-9a-fA-F]+)\)'
-  matches = re.findall(pattern, clss, re.MULTILINE)
-  return {int(num, base=16):name for name, num in matches}
+  res = {}
+  known_classes = {"NV01_DEVICE_0", "NV01_ROOT", "NV1_MEMORY_SYSTEM", "NV01_MEMORY_VIRTUAL", "NV1_MEMORY_USER", "NV50_MEMORY_VIRTUAL", "NV_FERMI_VASPACE_A",
+                   "NV20_SUBDEVICE_0"}
+  for nm,val in nv_gpu.__dict__.items():
+    if not isinstance(val, int): continue
+    if 0x3000 < val < 0xffff: res[val] = nm
+    if nm in known_classes: res[val] = nm
+  return res
 nvclasses = get_classes()
 nvuvms = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("UVM_") and nv_gpu.__dict__.get(x+"_PARAMS")}
-nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
+nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC9B0_", "NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
 
 global_ioctl_id = 0
 gpus_user_modes = []

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/self_tokenize.py

@@ -8,10 +8,10 @@ from sz import NONCORE_DIRS
 # llama 3 tokenizer
 tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
 
-def read_code(base_path):
+def read_code(base_path, full=False):
   ret = []
   for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
-    if not getenv("CORE") and any(path.split("./")[1].startswith(x) for x in NONCORE_DIRS): continue
+    if not full and any(path.split("./")[1].startswith(x) for x in NONCORE_DIRS): continue
     for name in files:
       if not name.endswith(".py"): continue
       if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
@@ -23,9 +23,10 @@ def read_code(base_path):
 if __name__ == "__main__":
   parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
   parser.add_argument("--output", help="Output file to write the combined code to.")
+  parser.add_argument("--full", action="store_true", help="All directories")
   args = parser.parse_args()
 
-  ret = read_code(".")
+  ret = read_code(".", args.full)
 
   table = []
   for name,code in ret:

extra/sqtt/active_sqtt_parse.py

@@ -0,0 +1,151 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+os.environ["PROFILE"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+from dataclasses import replace
+import atexit, contextlib
+from tinygrad import Tensor
+from tinygrad.helpers import system, OSX
+from tinygrad.runtime.ops_amd import AMDProgram
+from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
+from tinygrad.device import Device, ProfileDeviceEvent
+
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+dev = Device["AMD"]
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[str, list[WaveExec]] = {}
+  yield sqtt
+  events = dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())]
+
+  #rctx = decode(events)
+  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  #sqtt.update(rctx.inst_execs)
+
+  for e in events:
+    if isinstance(e, ProfileSQTTEvent):
+      print(replace(e, blob=b''))
+      if e.se == 0:
+        parse_sqtt_print_packets(e.blob)
+
+template = """.text
+.globl matmul
+.p2align 8
+.type matmul,@function
+matmul:
+  INSTRUCTION
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 8
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+def run_asm(src, num_workgroups=1, num_waves=1):
+  WAVE_SIZE = 32
+  t = Tensor.empty(0x1000).realize()
+  buf = t.uop.buffer.ensure_allocated()
+  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
+  dev.compiler.disassemble(lib)
+  fxn = AMDProgram(dev, "matmul", lib)
+  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
+
+if __name__ == "__main__":
+  with save_sqtt() as sqtt:
+    run_asm([
+      "s_nop 100",
+      "s_nop 100",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+      "s_nop 100",
+      "s_nop 100",
+      "s_add_i32 s2, s2, 10",
+      "s_add_i32 s2, s2, 10",
+      "s_nop 100",
+      "s_nop 100",
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v0, 0",
+      "s_nop 100",
+      "s_nop 100",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "s_nop 100",
+      "s_nop 100",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "s_nop 100",
+      "s_nop 100",
+      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
+      "s_endpgm",
+    ], num_workgroups=1, num_waves=1)
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
+    #Tensor.empty(1, 64).sum(axis=1).realize()
+    Tensor.empty(1).log2().realize()
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    # what's in v0?
+    run_asm([
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v1, 0",
+      "s_clause 0x1",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+    ]+[
+      "global_load_b32 v1, v0, s[0:1]",
+    ]*10+[
+      "global_load_b32 v10, v1, s[0:1]",
+      "s_waitcnt vmcnt(0)",
+
+      #"v_rcp_f32 v1, v0"
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v5 v4 v4",
+      #"v_add_f32_e32 v7 v6 v6",
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v2 v1 v1",
+      #"s_nop 1"
+    ]*5+[
+      "v_add_f32_e32 v3 v2 v2",
+    ]*5+[
+      "v_mul_f32_e32 v3 v2 v2",
+    ]*7)

extra/sqtt/attempt_sqtt_parse.py

@@ -1,53 +1,168 @@
-import pickle
-from hexdump import hexdump
+import pickle, sys
+from tinygrad.helpers import getenv, Timing, colored
 from extra.sqtt.roc import decode, ProfileSQTTEvent
-from tinygrad.helpers import getenv
+
+# do these enums match fields in the packets?
+#from tinygrad.runtime.support.amd import import_soc
+#soc = import_soc([11])
+#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
 
 # Instruction packets (one per ISA op)
 # NOTE: these are bad guesses and may be wrong! feel free to update if you know better
+# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
+
+# we see 18 opcodes
+# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
+# if you exclude everything, you are left with 6
+# opcodes( 6): 10 11 14 15 16 17
+# sometimes we see a lot of B, but not repeatable
+
+# not seen
+# 7 A C
+
+# NOTE: INST runs before EXEC
+
+OPCODE_COLORS = {
+  # dispatches are BLACK
+  0x1: "BLACK",
+  0x18: "BLACK",
+
+  # execs are yellow
+  0x2: "yellow",
+  0x3: "yellow",
+  0x4: "YELLOW",
+  0x5: "YELLOW",
+
+  # waves are blue
+  0x8: "blue",
+  0x9: "blue",
+  0x6: "cyan",
+  0xb: "cyan",
+}
 
 OPCODE_NAMES = {
-  # Small metadata / structural packets (NOT ISA op kinds)
-  0x01: "META_SMALL_ID",          # 12-bit identifier / slot tag
-  0x02: "META_FLAG",              # 1-byte flag/mode (CF/AF/8F/DF...)
-  0x03: "META_SUBEVENT_CODE",     # 1-byte sub-event/classification code
-  0x04: "META_BASE_INDEX_TAG",    # 12-bit base index/tag (..D, 9D, 10D, 58D...)
+  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
+  0x01: "VALUINST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
+  0x02: "VMEMEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
+  0x03: "ALUEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
+  0x04: "IMMEDIATE",
+  0x05: "IMMEDIATE_MASK",
 
-  # Instruction / timing / timestamp packets
-  0x0F: "TIME_SHORT_DELTA_PLUS4", # short ts, raw_delta+4
-  0x11: "TIME_WAVE_STATE",        # compact wave timing/stall state record
-  0x14: "INST_EXEC_RECORD",       # per-instruction execution record
-  0x16: "TIME_LONG_OR_MARKER",    # long delta / marker with 6-byte payload
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
+  0x06: "WAVERDY",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
+  0x08: "WAVEEND",
+  0x09: "WAVESTART",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
+  0x0B: "WAVEALLOC",  # FFF00
 
-  # State / control / perf snapshots
-  0x09: "CONTROL_CONFIG_32B",     # 32-bit control/config word (bursts of FE88..., C488...)
-  0x15: "PERFCOUNTER_SNAPSHOT",   # perf / TT configuration snapshot (8-byte)
+  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
+  0x0D: "PERF",
+  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
+  0x12: "EVENT",
+  0x13: "EVENT_BIG",  # FFFFF800
+  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
+  0x14: "REG",
+  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+  0x18: "INST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
+  0x19: "UTILCTR",
 
-  # Extra descriptors / events / metrics
-  0x06: "META_DESCRIPTOR_24B",    # 24-bit descriptor (seen in complex kernels like GEMM)
-  0x08: "EVENT_SMALL",            # small in-stream event (5-nibble payload)
-  0x12: "TIME_SECONDARY_METRIC",  # 3-byte secondary timing/latency/perf metric
-  0x18: "EVENT_SMALL_PAYLOAD",    # generic small side-band payload (5 nibbles)
-  0x19: "EVENT_SUMMARY_48B",      # rare 6-byte summary/aggregate metric
+  # this is the first (8 byte) packet in the bitstream
+  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
 
-  # Pseudo / unknown / not yet observed
-  0x07: "UNK_DELTA",              # unknown
-  0x0A: "UNK_DELTA2",             # unknown
-  0x0B: "UNK_DELTA3",             # unknown
-  0x0C: "UNK_DELTA4",             # unknown
-  0x0D: "UNK_DELTA5",             # unknown
-  0x0E: "UNK_DELTA6",             # unknown
-  0x10: "UNK_PSEUDO",             # not seen; pseudo/placeholder
-  0x17: "UNK_NO_DELTA",           # unknown, likely non-timing event
+  # pure time (no extra bits)
+  0x0F: "TS_DELTA_SHORT",
+  0x10: "NOP",
+  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
+
+  # not a good name, but seen and understood mostly
+  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
+  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
+
+  # packets we haven't seen / rarely see 0x0b
+  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
+  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
+  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
 }
 
+#  SALU    =  0x0 / s_mov_b32
+#  SMEM    =  0x1 / s_load_b*
+#  JUMP    =  0x3 / s_cbranch_scc0
+#  NEXT    =  0x4 / s_cbranch_execz
+#  MESSAGE =  0x9 / s_sendmsg
+#  VALU    =  0xb / v_(exp,log)_f32_e32
+#  VALU    =  0xd / v_lshlrev_b64
+#  VALU    =  0xe / v_mad_u64_u32
+#  VMEM    = 0x21 / global_load_b32
+#  VMEM    = 0x22 / global_load_b32
+#  VMEM    = 0x24 / global_store_b32
+#  VMEM    = 0x25 / global_store_b64
+#  VMEM    = 0x27 / global_store
+#  VMEM    = 0x28 / global_store_b64
+#  LDS     = 0x29 / ds_load_b128
+#  LDS     = 0x2b / ds_store_b32
+#  LDS     = 0x2e / ds_store_b128
+#  ????    = 0x5a / hidden global_load  instruction
+#  ????    = 0x5b / hidden global_load  instruction
+#  ????    = 0x5c / hidden global_store instruction
+#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
+OPNAME = {
+  0x0: "SALU",
+  0x1: "SMEM",
+  0x3: "JUMP",
+  0x4: "NEXT",
+  0x9: "MESSAGE",
+  0xb: "VALU",
+  0xd: "VALU",
+  0xe: "VALU",
+  0x21: "VMEM_LOAD",
+  0x22: "VMEM_LOAD",
+  0x24: "VMEM_STORE",
+  0x25: "VMEM_STORE",
+  0x26: "VMEM_STORE",
+  0x27: "VMEM_STORE",
+  0x28: "VMEM_STORE",
+  0x29: "LDS_LOAD",
+  0x2b: "LDS_STORE",
+  0x2e: "LDS_STORE",
+  0x50: "__SIMD_LDS_LOAD",
+  0x51: "__SIMD_LDS_LOAD",
+  0x54: "__SIMD_LDS_STORE",
+  0x5a: "__SIMD_VMEM_LOAD",
+  0x5b: "__SIMD_VMEM_LOAD",
+  0x5c: "__SIMD_VMEM_STORE",
+  0x5d: "__SIMD_VMEM_STORE",
+  0x5e: "__SIMD_VMEM_STORE",
+  0x5f: "__SIMD_VMEM_STORE",
+  0x72: "SALU_OR",
+  0x73: "VALU_CMPX",
+}
+
+ALUSRC = {
+  1: "SALU",
+  2: "VALU",
+  3: "VALU_ALT",
+}
+
+MEMSRC = {
+  0: "LDS",
+  1: "__LDS",
+  2: "VMEM",
+  3: "__VMEM",
+}
+
+
 # these tables are from rocprof trace decoder
 # rocprof_trace_decoder_parse_data-0x11c6a0
 # parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
 
-# ---------- 1. local_138: 256-byte state->token table ----------
+# ---------- 1. local_138: 256-byte state->opcode table ----------
 
-STATE_TO_TOKEN: bytes = bytes([
+STATE_TO_OPCODE: bytes = bytes([
   0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
   0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
   0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
@@ -66,17 +181,47 @@ STATE_TO_TOKEN: bytes = bytes([
   0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
 ])
 
+# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
+
+opcode_mask = {
+  0x10: 0b1111,
+
+  0x16: 0b1111111,
+  0x17: 0b1111111,
+  0x07: 0b1111111,
+  0x19: 0b1111111,
+  0x11: 0b1111111,
+  0x12: 0b11111111,
+  0x13: 0b11111111,
+  0x15: 0b1111111,
+
+  0x18: 0b111,
+  0x1: 0b111,
+
+  0x5: 0b11111,
+  0x6: 0b11111,
+  0xb: 0b11111,
+  0x8: 0b11111,
+  0xc: 0b11111,
+  0xd: 0b11111,
+
+  0xf: 0b1111,
+  0x14: 0b1111,
+
+  0x9: 0b11111,
+  0xa: 0b11111,
+
+  0x4: 0b1111,
+  0x3: 0b1111,
+  0x2: 0b1111,
+}
 
 # ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
 
 NIBBLE_BUDGET = [
-  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40,
-  0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
-  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40,
-  0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
+  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
 ]
-assert len(NIBBLE_BUDGET) == 32
-
 
 # ---------- 3. delta_map from your hash nodes ----------
 
@@ -95,7 +240,8 @@ DELTA_MAP_DEFAULT = {
   0x0B: (5,  3),   # shift=5,  end=8
   0x0C: (5,  3),   # shift=5,  end=8
   0x0D: (5,  3),   # shift=5,  end=8
-  0x0E: (7,  2),   # shift=7,  end=9
+  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
+  #0x0E: (7,  2),   # shift=7,  end=9
   0x0F: (4,  4),   # shift=4,  end=8
   0x10: (0,  0),   # shift=0,  end=0  (no delta)
   0x11: (7,  9),   # shift=7,  end=16
@@ -111,189 +257,203 @@ DELTA_MAP_DEFAULT = {
 
 # ---------- 4. One-line-per-packet parser ----------
 
-def decode_packet_fields(opcode: int, reg: int, delta: int) -> str:
-  """
-  Conservative decoding of a few packet types.
+def reg_mask(opcode):
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
+  shift, width = DELTA_MAP_DEFAULT[opcode]
+  delta_mask = ((1 << width) - 1) << shift
+  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
+  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
 
-  Rules:
-    - We first mask the 64-bit shift register down to the actual packet
-      width using NIBBLE_BUDGET[opcode & 0x1F], so we never read bits
-      that aren't really part of the packet.
-    - Only layouts that are clearly visible from the decompiled C are
-      decoded, and names are kept generic (cfg_*, idx_*, id_*, etc).
+def decode_packet_fields(opcode: int, reg: int) -> str:
   """
-  # --- 0. Restrict to the real packet bits for this opcode -------------
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]  # this table is in bits
-  if nb_bits <= 0 or nb_bits >= 64:
-    pkt = reg & ((1 << 64) - 1)
-  else:
-    pkt = reg & ((1 << nb_bits) - 1)
-
+  Decode packet payloads conservatively, using:
+    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
+    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
+    - Per-opcode layouts derived from rocprof's decompiled consumers.
+  """
+  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
+  pkt = reg & reg_mask(opcode)
   fields: list[str] = []
 
-  # --- 1. Timestamp-ish opcodes ----------------------------------------
+  match opcode:
+    case 0x01: # VALUINST
+      # 6 bit field
+      flag = (pkt >> 6) & 1
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+      if flag: fields.append("flag")
+    case 0x02: # VMEMEXEC
+      # 2 bit field (pipe is a guess)
+      src = pkt>>6
+      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
+    case 0x03: # ALUEXEC
+      # 2 bit field
+      src = pkt>>6
+      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
+    case 0x04: # IMMEDIATE_4
+      # 5 bit field (actually 4)
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+    case 0x05: # IMMEDIATE_5
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x6:
+      # wave ready FFFF00
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x0d:
+      # 20 bit field
+      fields.append(f"arg = {pkt>>8:X}")
+    case 0x12:
+      fields.append(f"event = {pkt>>11:X}")
+    case 0x15:
+      fields.append(f"snap = {pkt>>10:X}")
+    case 0x19:
+      # wave end
+      fields.append(f"ctr = {pkt>>9:X}")
+    case 0xf:
+      extracted_delta = (reg >> 4) & 0xF
+      fields.append(f"strange_delta=0x{extracted_delta:x}")
+    case 0x11:
+      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
+      # FF0000 is the mask
+      coarse    = pkt >> 16
+      fields.append(f"coarse=0x{coarse:02x}")
+      # From decomp:
+      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
+      #  - when coarse&8, it marks all live waves as "terminated"
+      if coarse & 0x01:
+        fields.append("flag_wave_interest=1")
+      if coarse & 0x08:
+        fields.append("flag_terminate_all=1")
+    case 0x8:
+      # wave end, this is 20 bits (FFF00)
+      flag7   = (pkt >> 8) & 1
+      simd    = (pkt >> 9) & 3
+      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 15) & 0x1f
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+    case 0x9:
+      # From case 9 (WAVESTART) in multiple consumers:
+      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
+      #   cls2   = (w >> 8) & 3        (class / group)
+      #   slot4  = (w >> 10) & 0xf     (slot / group index)
+      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
+      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
+      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
+      flag7   = (pkt >> 7) & 1
+      simd    = (pkt >> 8) & 3
+      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 13) & 0x1F
+      id7     = (pkt >> 17)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+      fields.append(f"id7=0x{id7:x}")
+    case 0x18:
+      # FFF88 is the mask
+      # From case 0x18:
+      #   low3   = w & 7
+      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
+      #   flags  = bits 6 (B6) and 7 (B7)
+      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
+      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
+      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
+      flag1 = (pkt >> 3) & 1
+      flag2 = (pkt >> 7) & 1
+      wave = (pkt >> 8) & 0x1F
+      op = (pkt >> 13)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
+      if flag1: fields.append("flag1")
+      if flag2: fields.append("flag2")
+    case 0x14:
+      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
+      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
+      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
+      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
 
-  if opcode == 0x0F:  # TIME_SHORT_DELTA_PLUS4
-    # By the time we get here, `delta` is already raw_delta+4.
-    fields.append(f"ts_short_plus4={delta}")
-    return ", ".join(fields)
+      fields.append(f"subop=0x{subop:04x}")
+      fields.append(f"slot={slot}")
+      fields.append(f"val32=0x{val32:08x}")
 
-  if opcode == 0x11:  # TIME_WAVE_STATE (medium/large delta)
-    shift, width = DELTA_MAP_DEFAULT[opcode]
-    raw_delta = (pkt >> shift) & ((1 << width) - 1)
-    coarse    = (pkt >> (shift + width)) & 0xFF  # next byte above delta
-    fields.append(f"raw_delta={raw_delta}")
-    if coarse:
-      fields.append(f"raw_coarse=0x{coarse:02x}")
-    return ", ".join(fields)
+      if hi_byte & 0x80:
+        # Config flavour: writes config words into per-slot state arrays.
+        fields.append("kind=config")
+        if subop == 0x000C:
+          fields.append("slot=lo")
+        elif subop == 0x000D:
+          fields.append("slot=hi")
+      else:
+        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
+        if subop == 0xC342:
+          fields.append("kind=cor_stream")
+          if val32 == 0x434F5200:
+            fields.append("cor_magic='COR\\0'")
+    case 0x16:
+      # Bits:
+      #   bit8  -> 0x100
+      #   bit9  -> 0x200
+      #   bits 12..47 -> 36-bit field used as delta or marker
+      bit8 = bool(pkt & 0x100)
+      bit9 = bool(pkt & 0x200)
+      if not bit9:
+        mode = "delta"
+      elif not bit8:
+        mode = "marker"
+      else:
+        mode = "other"
+      # need to use reg here
+      val36 = (reg >> 12) & ((1 << 36) - 1)
+      fields.append(f"mode={mode}")
+      if mode != "delta":
+        fields.append(f"val36=0x{val36:x}")
+    case 0x17:
+      # From decomp (two sites with identical logic):
+      #   layout = (w >> 7) & 0x3f
+      #   mode   = (w >> 0xd) & 3
+      #   group  = (w >> 0xf) & 7
+      #   sel_a  = (w >> 0x1c) & 0xf
+      #   sel_b  = (w >> 0x21) & 7
+      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
+      layout = (pkt >> 7)  & 0x3F
+      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
+      group  = (pkt >> 15) & 0x7
+      sel_a  = (pkt >> 0x1C) & 0xF
+      sel_b  = (pkt >> 0x21) & 0x7
+      flag4  = (pkt >> 0x3B) & 0x1
 
-  if opcode == 0x16:  # TIME_LONG_OR_MARKER
-    bit8 = bool(pkt & 0x100)
-    bit9 = bool(pkt & 0x200)
-    if not bit9:
-      mode = "delta"
-    elif not bit8:
-      mode = "marker"
-    else:
-      mode = "other"
-    val36 = (pkt >> 12) & ((1 << 36) - 1)
-    fields.append(f"mode={mode}")
-    fields.append(f"val36=0x{val36:x}")
-    return ", ".join(fields)
+      fields.append(f"layout={layout}")
+      fields.append(f"group={group}")
+      fields.append(f"simd={simd}")
+      fields.append(f"sel_a={sel_a}")
+      fields.append(f"sel_b={sel_b}")
+      if layout == 4:
+        fields.append(f"layout4_flag={flag4}")
+    case _:
+      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
+  return ",".join(fields)
 
-  # --- 2. Opcode 0x14: exec/config record ------------------------------
+FILTER_LEVEL = getenv("FILTER", 1)
 
-  if opcode == 0x14:
-    subop   = (pkt >> 16) & 0xFFFF       # matches (short)(w >> 0x10)
-    val32   = (pkt >> 32) & 0xFFFFFFFF   # matches (uint)(w >> 0x20)
-    slot    = (pkt >> 7) & 0x7           # used as (idx & 4) + (idx & 3)
-    hi_byte = (pkt >> 8) & 0xFF
+DEFAULT_FILTER: tuple[int, ...] = tuple()
+# NOP + pure time + "sample"
+if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
+# reg + event + sample + marker
+# TODO: events are probably good
+if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
+# instruction runs + valuinst
+if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
+# instructions dispatch (inst, immed)
+if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
+# waves
+if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
 
-    fields.append(f"subop=0x{subop:04x}")
-    fields.append(f"slot={slot}")
-    fields.append(f"val32=0x{val32:08x}")
-
-    if hi_byte & 0x80:
-      # "config" flavour, writes into local_168[...] etc.
-      fields.append("kind=config")
-      if subop == 0x000C:
-        fields.append("cfg_target=local_168[slot].lo")
-      elif subop == 0x000D:
-        fields.append("cfg_target=local_168[slot].hi")
-    else:
-      # COR marker: subop 0xC342, val32==0x434F5200 ("COR\0")
-      if subop == 0xC342:
-        fields.append("kind=cor_stream")
-        if val32 == 0x434F5200:
-          fields.append("cor_magic='COR\\0'")
-
-    return ", ".join(fields)
-
-  # --- 3. Opcode 0x17: mode/layout header ------------------------------
-
-  if opcode == 0x17:
-    # From case 0x17:
-    #   layout = (w >> 7) & 0x3f
-    #   mode   = (w >> 0xd) & 3
-    #   group  = (w >> 0xf) & 7
-    #   sel_a  = (w >> 0x1c) & 0xf
-    #   sel_b  = (w >> 0x21) & 7
-    #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
-    layout = (pkt >> 7)  & 0x3F
-    mode   = (pkt >> 13) & 0x3
-    group  = (pkt >> 15) & 0x7
-    sel_a  = (pkt >> 0x1C) & 0xF
-    sel_b  = (pkt >> 0x21) & 0x7
-    flag4  = (pkt >> 0x3B) & 0x1
-
-    fields.append(f"layout={layout}")
-    fields.append(f"group={group}")
-    fields.append(f"mode={mode}")
-    fields.append(f"sel_a={sel_a}")
-    fields.append(f"sel_b={sel_b}")
-    if layout == 4:
-      fields.append(f"layout4_flag={flag4}")
-
-    return ", ".join(fields)
-
-  # --- 4. Opcode 0x09: state-ish / indirection record ------------------
-
-  if opcode == 0x09:
-    # From case 9 on puVar58[1] (here pkt):
-    #
-    #   uVar41 = (w & 0xffffffff) >> 7;  local_520 = uVar41 & 1
-    #   local_4a0 = (w >> 8) & 3;
-    #   local_4a8 = (w >> 10) & (7 or 0xf)   (depends on local_494)
-    #   uVar69    = (w >> 0xd) or (w >> 0xf) (depends on local_494)
-    #   local_518 = (w >> 0x19) & 0x7f;
-    #
-    # We *don’t* know local_494 here, so we just expose the raw slices.
-    flag7   = (pkt >> 7) & 0x1          # low bit of uVar41
-    cls2    = (pkt >> 8) & 0x3          # local_4a0
-    slot4   = (pkt >> 10) & 0xF         # superset of 3-bit local_4a8
-    idx_lo  = (pkt >> 13) & 0x1F        # matches uVar69&0x1F when layout<4
-    idx_hi  = (pkt >> 15) & 0x1F        # matches uVar69&0x1F when layout>=4
-    id7     = (pkt >> 0x19) & 0x7F      # local_518
-
-    fields.append(f"flag7={flag7}")
-    fields.append(f"cls2={cls2}")
-    fields.append(f"slot4=0x{slot4:x}")
-    fields.append(f"idx_lo5=0x{idx_lo:x}")
-    fields.append(f"idx_hi5=0x{idx_hi:x}")
-    fields.append(f"id7=0x{id7:x}")
-    return ", ".join(fields)
-
-  # --- 5. Opcode 0x18: perf/event trigger ------------------------------
-
-  if opcode == 0x18:
-    # From case 0x18:
-    #  - low 3 bits:    (w & 7)
-    #  - mid 3 bits:    (w >> 3) & 7   or (w >> 4) & 7  (layout–dependent)
-    #  - hi id:         (w >> 0xc) & 0xff   OR (w >> 0xd) & 0x7f
-    #  - flag bits at 6 / 7
-    #
-    # The *real* semantics depend on global local_494 and accumulated
-    # local_500, so we keep this as a raw view that’s still useful for
-    # debugging, but not layout-dependent.
-    low3     = pkt & 0x7
-    grp3_a   = (pkt >> 3) & 0x7
-    grp3_b   = (pkt >> 4) & 0x7
-    flag_b6  = (pkt >> 6) & 0x1
-    flag_b7  = (pkt >> 7) & 0x1
-    idx5_a   = (pkt >> 7) & 0x1F
-    idx5_b   = (pkt >> 8) & 0x1F
-    hi8      = (pkt >> 12) & 0xFF
-    hi7      = (pkt >> 13) & 0x7F
-
-    fields.append(f"low3=0x{low3:x}")
-    fields.append(f"grp3_a=0x{grp3_a:x}")
-    fields.append(f"grp3_b=0x{grp3_b:x}")
-    fields.append(f"flag_b6={flag_b6}")
-    fields.append(f"flag_b7={flag_b7}")
-    fields.append(f"idx5_a=0x{idx5_a:x}")
-    fields.append(f"idx5_b=0x{idx5_b:x}")
-    fields.append(f"hi8=0x{hi8:02x}")
-    fields.append(f"hi7=0x{hi7:02x}")
-    return ", ".join(fields)
-
-  # --- 6. Generic tiny event-ish packets -------------------------------
-
-  if opcode in (0x08, 0x12, 0x19):
-    # These are all "small event" style tokens. The exact layout depends
-    # on global state (local_494 etc), so we just show:
-    #   - low 8 bits as a kind/flag byte
-    #   - the rest as an opaque payload.
-    kind    = pkt & 0xFF
-    payload = pkt >> 8
-    fields.append(f"kind_byte=0x{kind:02x}")
-    if payload:
-      fields.append(f"payload=0x{payload:x}")
-    return ", ".join(fields)
-
-  # --- 7. Everything else: no extra decode -----------------------------
-  return ""
-
-def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000) -> None:
+def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
   """
   Minimal debug: print ONE LINE per decoded token (packet).
 
@@ -302,120 +462,87 @@ def parse_sqtt_print_packets(data: bytes, max_tokens: int = 100000) -> None:
   """
   n = len(data)
   time = 0
+  last_printed_time = 0
   reg = 0          # shift register
   offset = 0       # bit offset, in steps of 4 (one nibble)
   nib_budget = 0x40
   flags = 0
   token_index = 0
+  opcodes_seen = set()
 
-  while (offset >> 3) < n and token_index < max_tokens:
-    # Remember where we started refilling for this step (bit offset),
-    # but the *logical* start of the current packet is last_real_offset.
-    refill_start = offset
-
+  while (offset >> 3) < n:
     # 1) Fill register with nibbles according to nib_budget
     if nib_budget != 0:
-      target = refill_start + 4 + ((nib_budget - 1) & ~3)
-      cur = refill_start
-      while cur != target and (cur >> 3) < n:
-        byte_index = cur >> 3
-        byte = data[byte_index]
-        shift = 4 if (cur & 4) else 0  # low then high nibble
-        nib = (byte >> shift) & 0xF
+      target = offset + 4 + ((nib_budget - 1) & ~3)
+      while offset != target and (offset >> 3) < n:
+        byte = data[offset >> 3]
+        nib = (byte >> (offset & 4)) & 0xF
         reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
-        cur += 4
-      offset = cur
+        offset += 4
+      if offset != target: break  # don't parse past the end
 
     # 2) Decode token from low 8 bits
-    state = reg & 0xFF
-    opcode = STATE_TO_TOKEN[state]
+    opcode = STATE_TO_OPCODE[reg & 0xFF]
+    opcodes_seen.add(opcode)
 
-    # 3) Handle pseudo-token 0x10: need more bits, don't print. Looks like a NOP.
-    if opcode == 0x10:
-      # "need more bits" pseudo-token: adjust nibble budget and continue
-      nib_budget = 4
-      if (offset >> 3) >= n:
-        break
-      # Do NOT count this as a real packet; do not update last_real_offset.
-      continue
+    # 4) Set next nibble budget based on opcode
+    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
 
-    # 4) Set next nibble budget
-    nb_index = opcode & 0x1F
-    nib_budget = NIBBLE_BUDGET[nb_index]
-    time_before = time
-    note = ""
-    # 5) Special opcode 0x16 (timestamp / marker)
+    # 5) Get delta
+    shift, width = DELTA_MAP_DEFAULT[opcode]
+    delta = (reg >> shift) & ((1 << width) - 1)
+
+    # 6) Update time and handle special opcodes 0xF/0x16
     if opcode == 0x16:
       two_bits = (reg >> 8) & 0x3
       if two_bits == 1:
         flags |= 0x01
 
+      # Common 36-bit field at bits [12..47]
       if (reg & 0x200) == 0:
-        # delta mode: 36-bit delta at bits [12..47]
-        delta = (reg >> 12) & ((1 << 36) - 1)
-        time += delta
-        note = "0x16-delta"
+        # delta mode: add 36-bit delta to time
+        pass
+      elif (reg & 0x100) == 0:
+        # marker / other modes: no time advance
+        # real marker: bit9=1, bit8=0, non-zero payload
+        # "other" 0x16 variants, ignored for timing
+        delta = 0
       else:
-        # marker mode if bit9==1 and bit8==0
-        if (reg & 0x100) == 0:
-          val = (reg >> 12) & ((1 << 36) - 1)
-          delta = 0
-          note = f"0x16-marker val=0x{val:x}"
-        else:
-          delta = 0
-          note = "0x16-other"
-    else:
-      # 6) Generic opcode (including 0x0F)
-      shift, width = DELTA_MAP_DEFAULT[opcode]
-      mask = (1 << width) - 1
-      delta = (reg >> shift) & mask
-
-      # TODO: add more opcode parsers here that add notes to other opcodes
-      if opcode == 0x0F:
-        delta_with_fix = delta + 4
-        note = f"0x0f (+4) raw_delta={delta}"
-        time += delta_with_fix
-        delta = delta_with_fix
-      else:
-        time += delta
-
-    # ONE-LINE PRINT PER PACKET
-    #assert last_real_offset%8 == 0
-    #assert (offset)%8 == 0, f"misalign offset {offset}"
+        raise RuntimeError("unknown 0x16 delta")
+    elif opcode == 0x0F:
+      # opcode 0x0F has an offset of 4 to the delta
+      # update: it's actually computed to be 8 to match WAVESTART
+      delta = delta + 8
 
     # Append extra decoded fields into the note string
-    extra = decode_packet_fields(opcode, reg, delta)
-    if extra: note = (note + " ; " + extra) if note else extra
-
-    BORING_OPCODES = {0x11, 0x14}
-    if opcode not in BORING_OPCODES or getenv("BORING", 1):
-      my_reg = reg
-      my_reg &= (1 << nib_budget) - 1
-      print(
-        f"{token_index:4d}  "
-        f"off={offset//4:5d}  "
-        f"op=0x{opcode:02x} "
-        f"{OPCODE_NAMES[opcode]:24s} "
-        f" time={time_before:8d}+{delta:8d}  "
-        f"{my_reg:16X} "
-        f"{note}"
-      )
+    note = decode_packet_fields(opcode, reg)
 
+    # this delta happens before the instruction
+    time += delta
     token_index += 1
 
+    if verbose and (filter is None or opcode not in filter):
+      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
+      last_printed_time = time
+
   # Optional summary at the end
-  print(f"# done: tokens={token_index}, final_time={time}, flags=0x{flags:02x}")
+  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
+  if verbose:
+    print(f"opcodes({len(opcodes_seen):2d}):",
+          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
+
 
 def parse(fn:str):
-  dat = pickle.load(open(fn, "rb"))
-  ctx = decode(dat)
+  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
+  if getenv("ROCM", 0):
+    with Timing(f"decode {fn}: "): ctx = decode(dat)
   dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
   print(f"got {len(dat_sqtt)} SQTT events in {fn}")
   return dat_sqtt
 
 if __name__ == "__main__":
-  #dat_sqtt = parse("extra/sqtt/examples/profile_empty_run_0.pkl")
-  dat_sqtt = parse("extra/sqtt/examples/profile_plus_run_0.pkl")
-  #dat_sqtt = parse("extra/sqtt/examples/profile_gemm_run_0.pkl")
-  blob_0 = dat_sqtt[0].blob
-  parse_sqtt_print_packets(blob_0[8:])
+  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
+  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
+  for i,dat in enumerate(dat_sqtt):
+    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
+      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))

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

@@ -1,4 +1,5 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
+from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
@@ -31,54 +32,76 @@ def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
 @dataclasses.dataclass(frozen=True)
 class InstExec:
   typ:str
-  inst:str
+  pc:int
   stall:int
   dur:int
   time:int
 
 @dataclasses.dataclass(frozen=True)
-class WaveExec:
+class WaveSlot:
   wave_id:int
   cu:int
   simd:int
+  se:int
+  @property
+  def cu_loc(self) -> str: return f"SE:{self.se} CU:{self.cu}"
+  @property
+  def simd_loc(self) -> str: return f"{self.cu_loc} SIMD:{self.simd}"
+  @property
+  def wave_loc(self) -> str: return f"{self.simd_loc} W:{self.wave_id}"
+
+@dataclasses.dataclass(frozen=True)
+class WaveExec(WaveSlot):
   begin_time:int
   end_time:int
-  insts:list[InstExec]
+  insts:bytearray
+  def unpack_insts(self) -> Generator[InstExec, None, None]:
+    sz = ctypes.sizeof(struct:=rocprof.rocprofiler_thread_trace_decoder_inst_t)
+    insts_array = (struct*(len(self.insts)//sz)).from_buffer(self.insts)
+    for inst in insts_array:
+      inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst.category)
+      yield InstExec(inst_typ, inst.pc.address, inst.stall, inst.duration, inst.time)
+
+@dataclasses.dataclass(frozen=True)
+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[tuple[str, int], tuple[str, int]] = {}
-    self.inst_execs:dict[str, list[WaveExec]] = {}
+    self.disasms:dict[str, dict[int, tuple[str, int]]] = {}
+    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)
-      for addr, info in llvm_disasm(arch, unwrap(prog.lib)).items():
-        self.disasms[(prog.name, unwrap(prog.base) + addr)] = info
+      base = unwrap(prog.base)
+      self.disasms[prog.name] = asm = {base+addr:info for addr,info in llvm_disasm(arch, unwrap(prog.lib)).items()}
 
   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("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
+    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_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("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+    if DEBUG >= 5: print(f"WAVE {ev.wave_id=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.contexts=} {ev.begin_time=} {ev.end_time=}")
+    # Skip wave events without instruction timings, occupancy events give the start and duration.
+    if ev.instructions_size == 0: return
 
-    inst_execs:list[InstExec] = []
-    for j in range(ev.instructions_size):
-      inst_ev = ev.instructions_array[j]
-      inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst_ev.category)
-      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
-      inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
-      if DEBUG >= 8: print(inst_execs[-1])
+    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)
 
-    if ev.instructions_size > 0:
-      self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, ev.begin_time, ev.end_time, inst_execs))
+    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:
   dev_events:dict[str, ProfileDeviceEvent] = {}
@@ -92,27 +115,30 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
+  def copy_cb(buf, buf_size, _):
     if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
     buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
     buf_size[0] = len(prof_info)
     return len(prof_info)
 
   @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
+  def trace_cb(record_type, events_ptr, n, _):
     match record_type:
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME:
+        if DEBUG >= 5:
+          pairs = [(ev.shader_clock, ev.realtime_clock) for ev in (rocprof.rocprofiler_thread_trace_decoder_realtime_t * n).from_address(events_ptr)]
+          print(f"REALTIME {pairs}")
       case _:
         if DEBUG >= 5: print(rocprof.enum_rocprofiler_thread_trace_decoder_record_type_t.get(record_type), events_ptr, n)
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    if DEBUG >= 8: print(f"isa_cb {pc.address=} {pc.code_object_id=}")
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run)[0]][pc.address]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
     if instr == "s_endpgm": mem_size_ptr[0] = 0
@@ -125,11 +151,22 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  def worker():
+    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  (t:=threading.Thread(target=worker, daemon=True)).start()
+  t.join()
   return ROCParseCtx
 
+def print_pmc(ev:ProfilePMCEvent) -> None:
+  ptr = 0
+  view = memoryview(ev.blob).cast('Q')
+  for s in ev.sched:
+    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:<2} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
+      ptr += 1
+
 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)))
@@ -142,10 +179,4 @@ if __name__ == "__main__":
   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)

extra/sqtt/test_pmc.py

@@ -0,0 +1,49 @@
+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.uop.ops import UOp, KernelInfo, AxisType
+from tinygrad.runtime.ops_amd import ProfilePMCEvent
+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=()))
+
+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[0])
+    np.testing.assert_equal(a.numpy(), b.numpy())
+
+  def test_copy_uncoalesced(self): return self.test_copy(stride=17)
+
+if __name__ == "__main__":
+  unittest.main()

extra/sqtt/test_timing.py

@@ -2,16 +2,16 @@ 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
 import sys, contextlib
-from tinygrad import Tensor
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Device, ProfileDeviceEvent
 
 from extra.sqtt.roc import decode, WaveExec
@@ -75,7 +75,6 @@ class TestTiming(unittest.TestCase):
     inp = Tensor([-2.0]).realize()
     with save_sqtt() as sqtt:
       Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
-
     wave = list(sqtt.values())[0][0]
     for i in range(len(wave.insts)):
       if wave.insts[i].inst.startswith("global_store"):
@@ -102,7 +101,7 @@ class TestTiming(unittest.TestCase):
       assert data0.dtype.base == dtypes.ulong
       op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
       op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
+      op = custom("unsigned long long t1 = __builtin_readcyclecounter();", op)
       op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
       return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
     diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
@@ -113,5 +112,32 @@ class TestTiming(unittest.TestCase):
     # cycles = sleep dur + overhead of storing hi/lo REG_SHADER_CYCLES
     self.assertGreaterEqual(diff_hw_reg.item(), sleep.dur)
 
+  def test_nop(self):
+    with save_sqtt() as sqtt:
+      asm_kernel(["s_nop 1"]*10).realize()
+    wave = list(sqtt.values())[0][0]
+    for e in wave.insts:
+      print(f"{e.inst} {e.dur=} {e.stall=}")
+
+  def test_wave_sched(self):
+    num_waves = getenv("NUM_WAVES", 16)
+    num_wgps = getenv("NUM_WGPS", 2)
+    num_vgpr = getenv("NUM_VGPR", 256)
+    with save_sqtt() as sqtt:
+      # 1 cycle decode, no stall
+      asm_kernel([f"v_mov_b32_e32 v{i} {i}" for i in range(num_vgpr)], l=32*num_waves, g=num_wgps).realize()
+    waves = list(sqtt.values())[0]
+    print(len(waves), "waves decoded")
+    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/test_mi350.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+AMD=1 AMD_LLVM=1 python -m pytest -n=1 test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py --durations=20
+AMD=1 AMD_LLVM=0 python -m pytest -n=1 test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py --durations=20
+
+CNT=1 AMD_LLVM=0 DEBUG=2 FP8E4M3=1 HALF=0 BFLOAT16=0 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py
+CNT=1 AMD_LLVM=0 DEBUG=2 FP8E4M3=0 HALF=1 BFLOAT16=0 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py
+CNT=1 AMD_LLVM=0 DEBUG=2 FP8E4M3=0 HALF=0 BFLOAT16=1 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py
+
+CNT=1 AMD_LLVM=1 DEBUG=2 FP8E4M3=0 HALF=1 BFLOAT16=0 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py
+CNT=1 AMD_LLVM=1 DEBUG=2 FP8E4M3=0 HALF=0 BFLOAT16=1 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py
+CNT=1 AMD_LLVM=1 DEBUG=2 FP8E4M3=1 HALF=0 BFLOAT16=0 SHOULD_USE_TC=1 python extra/gemm/simple_matmul.py

extra/thunder/tiny/tk/__init__.py

@@ -1 +1,6 @@
-WARP_THREADS = 32
+from tinygrad.device import Device
+
+if Device.DEFAULT == "AMD":
+  WARP_THREADS = 64
+else:
+  WARP_THREADS = 32

extra/thunder/tiny/tk/group.py

@@ -7,22 +7,21 @@ 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, ST, RT, RV
+from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout
 
 class Group:
   def __init__(self, warps:int, ker):
     self.warps = warps
     self.group_threads = warps * WARP_THREADS
-    self.threadIdx_x = ker.threadIdx_x
     self.ker = ker
 
   # helpers
   @property
-  def laneid(self): return self.threadIdx_x % self.group_threads
+  def laneid(self): return self.ker.threadIdx_x % self.group_threads
   @property
   def warpid(self): return self.laneid // WARP_THREADS
   @property
-  def groupid(self): return self.threadIdx_x // self.group_threads
+  def groupid(self): return self.ker.threadIdx_x // self.group_threads
 
   # ops that only work on a single warp
 
@@ -40,6 +39,7 @@ class Group:
     return reg.after(reg_store).reshape(reg.shape)
 
   def zero(self, reg:ALL_TILES): return self.clear(reg, 0)
+  def ones(self, reg:ALL_TILES): return self.clear(reg, 1)
   def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
 
   copy_rid = 300
@@ -51,56 +51,145 @@ class Group:
     rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
     Group.copy_rid += len(dst.shape)
 
-    dst_store = dst[*rngs_for_shape].store(src[*rngs_for_shape].cast(dst.dtype.base)).end(*rngs_for_shape)
+    src_load = src[*rngs_for_shape]
+    if src.dtype.base != dst.dtype.base:
+      src_load = src_load.cast(dst.dtype.base)
+    dst_store = dst[*rngs_for_shape].store(src_load).end(*rngs_for_shape)
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
 
-  def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
+  def transpose(self, dst:UOp|RT, src:UOp|RT):
+    dst, src = cast(UOp, dst), cast(UOp, src)
+    assert self.warps == 1
+
+    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)
+
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store).reshape(dst.shape)
+
+  def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
     c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
     assert self.warps == 1
 
+    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))), ())
+    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))), ())
+    else:  raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
+
     for height in self.ker.range(c.shape[-3], track=False):
       for width in self.ker.range(c.shape[-2], track=False):
-        for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
-          wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+        for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
+          if a_base_shape.cols == 16:
+            a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
+            b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
+          elif a_base_shape.cols == 32:
+            a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
+            b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
+          else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
+          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
-          a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
-          b_in1 = UOp.vectorize(*([b[inner, width, i] for i in range(2)] + [b[inner, width, 4+i] for i in range(2)]))
-          c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
-          b_in2 = UOp.vectorize(*([b[inner, width, 2+i] for i in range(2)] + [b[inner, width, 6+i] for i in range(2)]))
-          c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
-
-          out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
-          out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
+          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape) if after else c_store
+    return c.after(c_store).reshape(c.shape)
 
-  def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
+  def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
     c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
     assert self.warps == 1
 
+    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))), ())
+    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))), ())
+    else:  raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
+
     for height in self.ker.range(c.shape[-3], track=False):
       for width in self.ker.range(c.shape[-2], track=False):
-        for inner in self.ker.range(a.shape[-2], AxisType.REDUCE, track=False):
-          wmma_arg = ("WMMA_8_16_16_bfloat16_float", (8, 16, 16), dtypes.bfloat16, dtypes.float, "CUDA", 32, (((4, 2), (3, 2), (8, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+        for inner in self.ker.range(a.shape[-2], axis_type=AxisType.REDUCE, track=False):
+          if a_base_shape.cols == 16:
+            a_in = UOp.vectorize(*[a[height, inner, i] for i in range(4)])
+            b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
+          elif a_base_shape.cols == 32:
+            a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
+            b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
+          else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
+          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
 
-          a_in = UOp.vectorize(*[a[height, inner, i] for i in range(8)])
-          b_in1 = UOp.vectorize(*([b[width, inner, i] for i in range(2)] + [b[width, inner, 4+i] for i in range(2)]))
-          c_out1 = UOp.vectorize(*[c[height, width, i] for i in range(4)])
-          b_in2 = UOp.vectorize(*([b[width, inner, 2+i] for i in range(2)] + [b[width, inner, 6+i] for i in range(2)]))
-          c_out2 = UOp.vectorize(*[c[height, width, 4+i] for i in range(4)])
-
-          out1 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in1, c_out1), arg=wmma_arg)
-          out2 = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in2, c_out2), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out1.gep(i)) for i in range(4)] + [c[height, width, 4+i].store(out2.gep(i)) for i in range(4)]
+          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape) if after else c_store
+    return c.after(c_store).reshape(c.shape)
+
+  def mma_AtB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
+    c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
+    assert self.warps == 1
+
+    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))), ())
+    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))), ())
+    else:  raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
+
+    for height in self.ker.range(c.shape[-3], track=False):
+      for width in self.ker.range(c.shape[-2], track=False):
+        for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
+          if a_base_shape.cols == 16:
+            a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
+            b_in = UOp.vectorize(*[b[inner, width, i] for i in range(4)])
+          elif a_base_shape.cols == 32:
+            a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
+            b_in = UOp.vectorize(*[b[inner, width, i] for i in range(8)])
+          else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
+          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
+
+          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_store = UOp.group(*c_i).end(height, width, inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape)
+
+  def mma_AtBt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
+    c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
+    assert self.warps == 1
+
+    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))), ())
+    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))), ())
+    else:  raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
+
+    for height in self.ker.range(c.shape[-3], track=False):
+      for width in self.ker.range(c.shape[-2], track=False):
+        for inner in self.ker.range(a.shape[-3], axis_type=AxisType.REDUCE, track=False):
+          if a_base_shape.cols == 16:
+            a_in = UOp.vectorize(*[a[inner, height, i] for i in range(4)])
+            b_in = UOp.vectorize(*[b[width, inner, i] for i in range(4)])
+          elif a_base_shape.cols == 32:
+            a_in = UOp.vectorize(*[a[inner, height, i] for i in range(8)])
+            b_in = UOp.vectorize(*[b[width, inner, i] for i in range(8)])
+          else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
+          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])
+
+          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
+          c_store = UOp.group(*c_i).end(height, width, inner)
+
+    self.ker.push_store(c_store, c)
+    return c.after(c_store).reshape(c.shape)
 
   map_rid = 400
   def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
@@ -120,171 +209,213 @@ class Group:
     self.ker.push_store(a_store, a)
     return a.after(a_store).reshape(a.shape)
 
-  def row_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp]):
+  def row_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp], init_value:float=0.0):
     vec, src = cast(UOp, vec), cast(UOp, src)
     assert self.warps == 1
 
-    red_local = self.ker.alloc((self.group_threads, 2), src.dtype.base, AddrSpace.LOCAL)
-    red_reg = self.ker.alloc((2,), src.dtype.base, AddrSpace.REG)
+    red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
+    red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
 
     for height in self.ker.range(src.shape[-3], track=False):
       i = UOp.range(red_reg.size, Group.clear_rid)
       Group.clear_rid += 1
       red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
-      reg_store = red_reg.flatten()[i].store(0.).end(i)
+      reg_store = red_reg.flatten()[i].store(init_value).end(i)
       red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
 
-      for outer in self.ker.range(2, track=False):
-        for width in self.ker.range(src.shape[-2], AxisType.REDUCE, track=False):
-          for inner in self.ker.range(4, AxisType.REDUCE, track=False):
-            elem_index = inner + 2 * (inner // 2) + outer * 2
-            reg_store = red_reg[outer].store(op(red_reg[outer], src[height, width, elem_index])).end(inner, width, outer)
-            red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
-
-      # store to shared memory
-      for outer in self.ker.range(2, track=False):
-        red_local_store = red_local[self.laneid, outer].store(red_reg[outer]).end(outer)
-        red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
-
-      # reduce from shared memory
-      for outer in self.ker.range(2, track=False):
-        for inner in self.ker.range(3, AxisType.REDUCE, track=False):
-          offset = (self.laneid // 4) * 4 + ((self.laneid + inner + 1) % 4)
-          reg_store = red_reg[outer].store(op(red_reg[outer], red_local[offset, outer])).end(inner, outer)
+      for width in self.ker.range(src.shape[-2], axis_type=AxisType.REDUCE, track=False):
+        for inner in self.ker.range(4, axis_type=AxisType.REDUCE, track=False):
+          reg_store = red_reg[0].store(op(red_reg[0], src[height, width, inner])).end(width, inner)
           red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
 
+      # store to shared memory
+      red_local_store = red_local[self.laneid].store(red_reg[0])
+      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+
+      # reduce from shared memory
+      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
+        offset = (self.laneid + (1 + inner) * 16) % self.group_threads
+        reg_store = red_reg[0].store(op(red_reg[0], red_local[offset])).end(inner)
+        red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
       # reduce with vec
-      for outer in self.ker.range(2, track=False):
-        vec_store = vec[height, 0, outer].store(op(vec[height, 0, outer], red_reg[outer])).end(outer, height)
+      vec_store = vec[height, 0].store(op(vec[height, 0], red_reg[0])).end(height)
+
+    self.ker.push_store(vec_store, vec)
+    return vec.after(vec_store).reshape(vec.shape)
+
+  def col_reduce(self, vec:UOp|RV, src:UOp|RT, op:Callable[[UOp, UOp], UOp], init_value:float=0.0):
+    vec, src = cast(UOp, vec), cast(UOp, src)
+    assert self.warps == 1
+
+    red_local = self.ker.alloc((self.group_threads,), src.dtype.base, AddrSpace.LOCAL)
+    red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
+
+    for width in self.ker.range(src.shape[-2], track=False):
+      i = UOp.range(red_reg.size, Group.clear_rid)
+      Group.clear_rid += 1
+      red_reg = red_reg.after(width, *[tkr._rng for tkr in self.ker.range_stack])
+      reg_store = red_reg.flatten()[i].store(init_value).end(i)
+      red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      for height in self.ker.range(src.shape[-3], axis_type=AxisType.REDUCE, track=False):
+        for inner in self.ker.range(4, axis_type=AxisType.REDUCE, track=False):
+          reg_store = red_reg[0].store(op(red_reg[0], src[height, width, inner])).end(height, inner)
+          red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      # store to shared memory
+      red_local_store = red_local[self.laneid].store(red_reg[0])
+      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+
+      # reduce from shared memory
+      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
+        offset = (self.laneid + (1 + inner) * 16) % self.group_threads
+        reg_store = red_reg[0].store(op(red_reg[0], red_local[offset])).end(inner)
+        red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
+
+      # reduce with vec
+      vec_store = vec[width, 0].store(op(vec[width, 0], red_reg[0])).end(width)
 
     self.ker.push_store(vec_store, vec)
     return vec.after(vec_store).reshape(vec.shape)
 
   # ops that can work across multiple warps
 
-  LOAD_INNER = 8
-  def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
+  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)
     if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
-      srcf = src.flatten(-2)
-
-      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
-      else: local_warpid = self.warpid
-      warp_laneid = self.threadIdx_x % WARP_THREADS
+      laneid = self.ker.laneid
+      rt, st = cast(RT, dst), cast(ST, src)
+      elements_per_thread = rt.base_shape.elements_per_thread
 
       for height in self.ker.range(dst.shape[-3], track=False):
         for width in self.ker.range(dst.shape[-2], track=False):
-          for inner in self.ker.range(RT.BASE_TILE_NEPT, track=False):
-            base_row = (local_warpid * dst.shape[-3] + height) * RT.BASE_TILE_ROWS
-            base_col = width * RT.BASE_TILE_COLS
-
-            if not transpose:
-              row = base_row + (warp_laneid // 4)
-              col = base_col + 2 * (warp_laneid % 4)
-
-              row_offset = ((inner % 4) // 2) * 8
-              col_offset = (inner % 2) + (inner // 4) * 8
+          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 = base_row + 2 * (warp_laneid % 4)
-              col = base_col + (warp_laneid // 4)
+              row = laneid % rt.base_shape.rows
+              col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
 
-              row_offset = (inner % 2) + (inner // 4) * 8
-              col_offset = ((inner % 4) // 2) * 8
+            srow, scol = cast(ST, src).swizzle(row, col)
 
-            src_i_last = (row + row_offset) * src.shape[-1] + col + col_offset
-
-            dst_store = dst[*dst_idxs, height, width, inner].store(srcf[*idxs[:-2], src_i_last])
+            src_load = src[*idxs[:-2], height, width, srow, scol]
+            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)
             dst_store = dst_store.end(height, width, inner)
     elif dst_dtype.addrspace == AddrSpace.LOCAL and src_dtype.addrspace == AddrSpace.GLOBAL:
-      dstf = dst.flatten(-2)
-
       srcf = src.flatten()
       row_stride = prod(src.shape[axis+1:])
 
-      idxs = tuple(idx * dst.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
-      idxs = tuple(idx * dst.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      st = cast(ST, dst)
+      idxs = tuple(idx * st.rows if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * st.cols if i == 3 else idx for i, idx in enumerate(idxs))
       src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
 
-      memcpy_per_row = dst.shape[-1] // Group.LOAD_INNER
-      total_calls = prod(dst.shape[-2:]) // (self.group_threads * Group.LOAD_INNER)
+      for height in self.ker.range(dst.shape[-4], track=False):
+        for width in self.ker.range(dst.shape[-3], track=False):
+          elements_per_thread = st.base_shape.elements_per_thread
+          memcpy_per_row = st.base_shape.cols // elements_per_thread
+          total_calls = st.base_shape.num_elements // (self.group_threads * elements_per_thread)
 
-      for outer in self.ker.range(total_calls, track=False):
-        for inner in self.ker.range(Group.LOAD_INNER, track=False):
-          load_idx = outer * self.group_threads + self.laneid
-          row = load_idx // memcpy_per_row
-          col = (load_idx * Group.LOAD_INNER) % dst.shape[-1]
+          for outer in self.ker.range(total_calls, track=False):
+            for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
+              load_idx = outer * self.group_threads + self.laneid
+              row = load_idx // memcpy_per_row
+              col = (load_idx * elements_per_thread) % st.base_shape.cols + inner
 
-          dst_i = row * dst.shape[-1] + col + inner
-          src_i += row * row_stride + col + inner
+              srow, scol = cast(ST, dst).swizzle(row, col)
 
-          dst_store = dstf[*dst_idxs, dst_i].store(srcf[src_i]).end(outer, inner)
+              src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
+              src_i += row * row_stride + col
+
+              src_load = srcf[src_i]
+              if src.dtype.base != dst.dtype.base:
+                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:
+      srcf = src.flatten()
+      row_stride = prod(src.shape[axis+1:])
+
+      laneid = self.ker.laneid
+      rt = cast(RT, dst)
+      elements_per_thread = rt.base_shape.elements_per_thread
+
+      idxs = tuple(idx * dst.shape[-3] * rt.base_shape.rows if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * dst.shape[-2] * rt.base_shape.cols 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 height in self.ker.range(dst.shape[-3], track=False):
+        for width in self.ker.range(dst.shape[-2], track=False):
+          for inner in self.ker.range(elements_per_thread, track=False):
+            base_row = height * rt.base_shape.rows
+            base_col = width * rt.base_shape.cols
+
+            if rt.layout == TileLayout.COL:
+              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 = base_row + row, base_col + col
+
+            src_i += srow * row_stride + scol
+
+            src_load = srcf[src_i]
+            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)
     else:
       raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
 
-    return dst.after(dst_store.barrier()).reshape(dst.shape)
+    self.ker.push_store(dst_store, dst)
+    return dst.after(dst_store).reshape(dst.shape)
 
-  STORE_INNER = 8
-  def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0, transpose:bool=False):
+  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.LOCAL:
-      dstf = dst.flatten(-2)
-
-      if self.warps % 4 == 0: local_warpid = (self.warpid // 4) + (self.warpid % 4) * (self.warps // 4)
-      else: local_warpid = self.warpid
-      warp_laneid = self.threadIdx_x % WARP_THREADS
-
-      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(RT.BASE_TILE_NEPT, track=False):
-            base_row = (local_warpid * src.shape[-3] + height) * RT.BASE_TILE_ROWS
-            base_col = width * RT.BASE_TILE_COLS
-
-            if not transpose:
-              row = base_row + (warp_laneid // 4)
-              col = base_col + 2 * (warp_laneid % 4)
-
-              row_offset = ((inner % 4) // 2) * 8
-              col_offset = (inner % 2) + (inner // 4) * 8
-            else:
-              row = base_row + 2 * (warp_laneid % 4)
-              col = base_col + (warp_laneid // 4)
-
-              row_offset = (inner % 2) + (inner // 4) * 8
-              col_offset = ((inner % 4) // 2) * 8
-
-            dst_i_last = (row + row_offset) * dst.shape[-1] + col + col_offset
-
-            dst_store = dstf[*idxs[:-2], dst_i_last].store(src[*src_idxs, height, width, inner])
-            dst_store = dst_store.end(height, width, inner)
-    elif src_dtype.addrspace == AddrSpace.LOCAL and dst_dtype.addrspace == AddrSpace.GLOBAL:
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL:
       dstf = dst.flatten()
       row_stride = prod(dst.shape[axis+1:])
 
-      idxs = tuple(idx * src.shape[-2] if i == axis else idx for i, idx in enumerate(idxs))
-      idxs = tuple(idx * src.shape[-1] if i == 3 else idx for i, idx in enumerate(idxs))
+      laneid = self.ker.laneid
+      rt = cast(RT, src)
+      elements_per_thread = rt.base_shape.elements_per_thread
+
+      idxs = tuple(idx * src.shape[-3] * rt.base_shape.rows if i == axis else idx for i, idx in enumerate(idxs))
+      idxs = tuple(idx * src.shape[-2] * rt.base_shape.cols if i == 3 else idx for i, idx in enumerate(idxs))
       dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
 
-      srcf = src.flatten(-2)
+      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):
+            base_row = height * rt.base_shape.rows
+            base_col = width * rt.base_shape.cols
 
-      memcpy_per_row = src.shape[-1] // Group.STORE_INNER
-      total_calls = prod(src.shape[-2:]) // (self.group_threads * Group.STORE_INNER)
+            if rt.layout == TileLayout.COL:
+              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
 
-      for outer in self.ker.range(total_calls, track=False):
-        for inner in self.ker.range(Group.STORE_INNER, track=False):
-          load_idx = outer * self.group_threads + self.laneid
-          row = load_idx // memcpy_per_row
-          col = (load_idx * Group.STORE_INNER) % src.shape[-1]
+            srow, scol = base_row + row, base_col + col
 
-          src_i = row * src.shape[-1] + col + inner
-          dst_i += row * row_stride + col + inner
+            dst_i += srow * row_stride + scol
 
-          dst_store = dstf[dst_i].store(srcf[*src_idxs, src_i]).end(outer, inner)
+            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 = dstf[dst_i].store(src_load).end(height, width, inner)
     else:
       raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
 
     self.ker.push_store(dst_store, dst)
-    return dst.after(dst_store.barrier()).reshape(dst.shape)
+    return dst.after(dst_store).reshape(dst.shape)

extra/thunder/tiny/tk/kernel.py

@@ -2,22 +2,24 @@ from contextlib import AbstractContextManager
 from tinygrad.uop.ops import UOp, KernelInfo, AxisType, AddrSpace
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.group import Group
-from extra.thunder.tiny.tk.tiles import GL, ST, RT, RV
+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, end:int, axis_type:AxisType): self.end, self.axis_type, self.done = end, axis_type, False
+  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.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, _tk_range.user_rid-1, axis_type=self.axis_type)
+      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")
@@ -29,10 +31,13 @@ class Kernel(AbstractContextManager):
     self.global_slot = 0
     self.shared_slot = 0
     self.register_slot = 0
+    self.range_id = 0
     self.allocs = {}
 
   @property
   def warpid(self): return self.threadIdx_x // WARP_THREADS
+  @property
+  def laneid(self): return self.threadIdx_x % WARP_THREADS
 
   def __enter__(self): return self
   def __exit__(self, exc_type, exc_value, traceback): pass
@@ -43,8 +48,10 @@ class Kernel(AbstractContextManager):
   @property
   def warpgroup(self): return self.group(4)
 
-  def range(self, end:int, axis_type:AxisType=AxisType.LOOP, track:bool=True):
-    rng = _tk_range(end, axis_type)
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+    if end == 0: start, end = 0, start
+    rng = _tk_range(start, end, step, axis_type, self.range_id)
+    self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
@@ -69,9 +76,9 @@ class Kernel(AbstractContextManager):
     return uop
 
   def gl(self, shape, dtype): return GL.create(shape, dtype, self)
-  def st(self, shape, dtype): return ST.create(shape, dtype, self)
-  def rt(self, shape, dtype): return RT.create(shape, dtype, self)
-  def rv(self, length, dtype, layout="naive"): return RV.create(length, dtype, layout, self)
+  def st(self, shape, dtype, layout=TileLayout.ROW, base_shape=ST_16X16): return ST.create(shape, dtype, layout, base_shape, self)
+  def rt(self, shape, dtype, layout=TileLayout.ROW, base_shape=RT_16X16): return RT.create(shape, dtype, layout, base_shape, self)
+  def rv(self, length, dtype, layout=VecLayout.ORTHO, rt_base_shape=RT_16X16): return RV.create(length, dtype, layout, rt_base_shape, self)
 
   def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
 
@@ -80,9 +87,13 @@ class Kernel(AbstractContextManager):
     rngs = []
     while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
 
-    return self.store_stack.pop()[0]._uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+    last_store = self.store_stack.pop()[0]
+    if hasattr(last_store, '_uop'): uop = last_store._uop
+    else: uop = last_store
+
+    return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
 
   def endrange(self):
     last_store = self.store_stack.pop()
     last_range = self.range_stack.pop()
-    return last_store[1].after(last_store[0].barrier().end(last_range._rng)).reshape(last_store[1].shape)
+    return last_store[1].after(last_store[0].end(last_range._rng)).reshape(last_store[1].shape)

extra/thunder/tiny/tk/tiles.py

@@ -1,5 +1,8 @@
+from enum import Enum, auto
 import functools
-from tinygrad.dtype import AddrSpace
+from typing import Callable
+from dataclasses import dataclass
+from tinygrad.dtype import AddrSpace, DType
 from tinygrad.mixin import MathMixin
 from tinygrad.uop.ops import UOp, Ops
 
@@ -11,9 +14,9 @@ def unwrap(x):
   if isinstance(x, dict): return {k: unwrap(v) for k,v in x.items()}
   return x
 
-def wrap(x, ker, cls):
-  if isinstance(x, UOp): return cls(x, ker)
-  if isinstance(x, (list, tuple)): return type(x)(wrap(y, ker, cls) for y in x)
+def wrap(x, s):
+  if isinstance(x, UOp): return s.ruop(x)
+  if isinstance(x, (list, tuple)): return type(x)(wrap(y, s) for y in x)
   return x
 
 def autowrap(source_cls, blacklist=None):
@@ -31,10 +34,10 @@ def autowrap(source_cls, blacklist=None):
       if callable(val):
         @functools.wraps(val)
         def proxy(*args, **kwargs):
-          return wrap(val(*unwrap(args), **unwrap(kwargs)), self.ker, cls)
+          return wrap(val(*unwrap(args), **unwrap(kwargs)), self)
         return proxy
       if name in UOp.__slots__: return val
-      return wrap(val, self.ker, cls)
+      return wrap(val, self)
     cls.__getattr__ = __getattr__
 
     for name in dir(source_cls):
@@ -46,9 +49,9 @@ def autowrap(source_cls, blacklist=None):
       else:
         original = getattr(source_cls, name)
         if callable(original):
-          def make_proxy(op_name, func):
+          def make_proxy(_, func):
             def proxy(self, *args, **kwargs):
-              return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self.ker, cls)
+              return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self)
             return proxy
           setattr(cls, name, make_proxy(name, original))
 
@@ -66,10 +69,13 @@ class TileMathMixin(MathMixin):
       elif isinstance(src[0], (int,float,bool)): uop = self.ker.warp.map(self._uop, lambda x: UOp.alu(x, op, inner_op(x.ufix(src[0]))))
       elif src[0]._shape is None: uop = UOp.alu(self._uop, op, inner_op(self._uop.ufix(src[0])))
       else:
-        if isinstance(self, RT) and isinstance(src[0], RV): uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0, (idx[2]%4)//2])))
+        if isinstance(self, RT) and isinstance(src[0], RV):
+          match self.layout:
+            case TileLayout.ROW: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0])))
+            case TileLayout.COL: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[1], 0])))
         else: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[*idx])))
     else: raise NotImplementedError
-    return type(self)(uop, self.ker)
+    return self.ruop(uop)
   def const_like(self, b): return b
 
   # override ops that do compute on the src uop
@@ -80,64 +86,186 @@ class TileMathMixin(MathMixin):
 
 @autowrap(UOp)
 class GL:
-  def __init__(self, uop, ker):
+  def __init__(self, uop:UOp, ker):
     self._uop, self.ker = uop, ker
 
+  def ruop(self, uop:UOp):
+    return GL(uop, self.ker)
+
   @classmethod
-  def create(cls, shape, dtype, ker):
+  def create(cls, shape, dtype:DType, ker):
     uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
     return cls(uop, ker)
 
+class TileLayout(Enum):
+  ROW = auto()
+  COL = auto()
+
+class VecLayout(Enum):
+  ORTHO = auto()
+
+@dataclass(frozen=True)
+class BaseShape:
+  rows: int
+  cols: int
+
+  @property
+  def num_elements(self): return self.rows * self.cols
+  @property
+  def elements_per_thread(self): return self.num_elements // WARP_THREADS
+
+@dataclass(frozen=True)
+class STBaseShape(BaseShape):
+  _swizzle: Callable[[UOp, DType], UOp]
+  bytes_per_thread: Callable[[DType], int]
+
+  def swizzle(self, row, col, dtype:DType):
+    offset = row * self.cols + col
+    offset *= dtype.itemsize
+    offset = self._swizzle(offset, dtype)
+    offset //= dtype.itemsize
+    return offset
+
+def st_16x16_swizzle(offset:UOp, _): return offset
+def st_16x16_bpt(dtype:DType):
+  if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
+  else: raise NotImplementedError
+ST_16X16 = STBaseShape(16, 16, st_16x16_swizzle, st_16x16_bpt)
+
+def st_16x16_swizzled_swizzle(offset:UOp, dtype:DType):
+  if dtype.itemsize == 2:
+    swizzle = ((offset % 512) >> 7) << 3
+    return offset ^ swizzle
+  elif dtype.itemsize == 4:
+    return offset
+  else: raise NotImplementedError
+def st_16x16_swizzled_bpt(dtype:DType):
+  if dtype.itemsize == 2: return 4
+  elif dtype.itemsize == 4: return 16
+  else: raise NotImplementedError
+ST_16X16_SWIZZLED = STBaseShape(16, 16, st_16x16_swizzled_swizzle, st_16x16_swizzled_bpt)
+
+def st_32x32_swizzle(offset:UOp, dtype:DType):
+  if dtype.itemsize == 2:
+    first_swizzle = ((offset % 1024) >> 9) << 5
+    second_swizzle = ((offset % 2048) >> 10) << 4
+    return offset ^ first_swizzle ^ second_swizzle
+  elif dtype.itemsize == 4:
+    return offset
+  else: raise NotImplementedError
+def st_32x32_bpt(dtype:DType):
+  if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
+  else: raise NotImplementedError
+ST_32X32 = STBaseShape(32, 32, st_32x32_swizzle, st_32x32_bpt)
+
+def st_16x32_swizzle(offset:UOp, dtype:DType):
+  if dtype.itemsize == 2:
+    swizzle = ((offset % 1024) >> 9) << 5
+    return offset ^ swizzle
+  elif dtype.itemsize == 4:
+    return offset
+  else: raise NotImplementedError
+def st_16x32_bpt(dtype:DType):
+  if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
+  else: raise NotImplementedError
+ST_16X32 = STBaseShape(16, 32, st_16x32_swizzle, st_16x32_bpt)
+
+def st_32x16_swizzle(offset:UOp, dtype:DType):
+  if dtype.itemsize == 2:
+    swizzle = ((offset % 1024) >> 9) << 4
+    return offset ^ swizzle
+  elif dtype.itemsize == 4:
+    return offset
+  else: raise NotImplementedError
+def st_32x16_bpt(dtype:DType):
+  if dtype.itemsize == 2 or dtype.itemsize == 4: return 16
+  else: raise NotImplementedError
+ST_32X16 = STBaseShape(32, 16, st_32x16_swizzle, st_32x16_bpt)
+
 @autowrap(UOp)
 class ST:
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
+  def __init__(self, uop:UOp, rows:int, cols:int, layout:TileLayout, base_shape:STBaseShape, ker):
+    self._uop, self.rows, self.cols, self.layout, self.base_shape, self.ker = uop, rows, cols, layout, base_shape, ker
+
+  def ruop(self, uop:UOp):
+    return ST(uop, self.rows, self.cols, self.layout, self.base_shape, self.ker)
 
   @classmethod
-  def create(cls, shape, dtype, ker):
-    uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
-    return cls(uop, ker)
+  def create(cls, shape, dtype:DType, layout:TileLayout, base_shape:STBaseShape, ker):
+    rows = shape[-2]
+    cols = shape[-1]
+    assert rows % base_shape.rows == 0
+    assert cols % base_shape.cols == 0
+    assert cols % base_shape.elements_per_thread == 0
+
+    height = rows // base_shape.rows
+    width = cols // base_shape.cols
+
+    uop = ker.alloc(shape[:-2] + (height, width, base_shape.rows, base_shape.cols), dtype, AddrSpace.LOCAL)
+    return cls(uop, rows, cols, layout, base_shape, ker)
+
+  def swizzle(self, row, col):
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.base.scalar())
+
+    row = swizzled_offset // self.base_shape.cols
+    col = swizzled_offset % self.base_shape.cols
+
+    return row, col
+
+@dataclass(frozen=True)
+class RTBaseShape(BaseShape):
+  stride: int
+
+  @property
+  def num_strides(self):
+    return self.elements_per_thread // self.stride
+
+RT_16X16 = RTBaseShape(rows=16, cols=16, stride=4)
+RT_32X32 = RTBaseShape(rows=32, cols=32, stride=4)
+RT_32X32_8 = RTBaseShape(rows=32, cols=32, stride=8)
+RT_16X32 = RTBaseShape(rows=16, cols=32, stride=8)
+RT_32X16 = RTBaseShape(rows=32, cols=16, stride=8)
+RT_32X16_4 = RTBaseShape(rows=32, cols=16, stride=4)
+RT_16X32_4 = RTBaseShape(rows=16, cols=32, stride=4)
 
 @autowrap(UOp)
 class RT(TileMathMixin):
-  BASE_TILE_ROWS, BASE_TILE_COLS = 16, 16
-  BASE_TILE_NE = BASE_TILE_ROWS * BASE_TILE_COLS
-  BASE_TILE_NEPT = BASE_TILE_NE // WARP_THREADS
+  def __init__(self, uop:UOp, layout:TileLayout, base_shape:RTBaseShape, ker):
+    self._uop, self.layout, self.base_shape, self.ker = uop, layout, base_shape, ker
 
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
+  def ruop(self, uop:UOp):
+    return RT(uop, self.layout, self.base_shape, self.ker)
 
   @classmethod
-  def create(cls, shape, dtype, ker):
+  def create(cls, shape, dtype:DType, layout:TileLayout, base_shape:RTBaseShape, ker):
     assert len(shape) == 2
-    assert shape[0] % RT.BASE_TILE_ROWS == 0
-    assert shape[1] % RT.BASE_TILE_COLS == 0
+    assert shape[0] % base_shape.rows == 0
+    assert shape[1] % base_shape.cols == 0
 
-    height = shape[0] // RT.BASE_TILE_ROWS
-    width = shape[1] // RT.BASE_TILE_COLS
+    height = shape[0] // base_shape.rows
+    width = shape[1] // base_shape.cols
 
-    uop = ker.alloc((height, width, RT.BASE_TILE_NEPT), dtype, AddrSpace.REG)
-    return cls(uop, ker)
+    uop = ker.alloc((height, width, base_shape.elements_per_thread), dtype, AddrSpace.REG)
+    return cls(uop, layout, base_shape, ker)
 
 @autowrap(UOp)
 class RV(TileMathMixin):
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
+  def __init__(self, uop:UOp, layout:VecLayout, ker):
+    self._uop, self.layout, self.ker = uop, layout, ker
+
+  def ruop(self, uop:UOp):
+    return RV(uop, self.layout, self.ker)
 
   @classmethod
-  def create(cls, length, dtype, layout, ker):
-    tiles = length // RT.BASE_TILE_ROWS
+  def create(cls, length, dtype:DType, layout:VecLayout, base_shape:RTBaseShape, ker):
+    tiles = length // base_shape.rows
 
     match layout:
-      case "naive":
-        inner_dim = 1
-        outer_dim = (tiles + 1) // 2
-      case "ortho":
+      case VecLayout.ORTHO:
         inner_dim = 1
         outer_dim = tiles
-      case _: raise NotImplementedError(f"rv layout {layout} not implemented")
 
-    uop = ker.alloc((outer_dim, inner_dim, 2), dtype, AddrSpace.REG)
-    return RV(uop, ker)
+    uop = ker.alloc((outer_dim, inner_dim), dtype, AddrSpace.REG)
+    return RV(uop, layout, ker)
 
 ALL_TILES = UOp | GL | ST | RT | RV

extra/thunder/tiny/visualize_tile.py

@@ -0,0 +1,156 @@
+from tinygrad.helpers import colored
+
+WARP_THREADS = 64
+BASE_TILE_ROWS = 16
+BASE_TILE_COLS = 16
+BASE_TILE_NEPT = (BASE_TILE_ROWS * BASE_TILE_COLS) // WARP_THREADS
+DTYPE_SIZE = 2
+INST = "ds_read_b64"
+
+def row_col(threadIdx_x):
+  local_warpid = threadIdx_x // WARP_THREADS
+  warp_laneid = threadIdx_x % WARP_THREADS
+
+  ret = []
+
+  for inner in range(BASE_TILE_NEPT):
+    if BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 16:
+      row = warp_laneid % 16
+      col = 4 * (warp_laneid // 16)
+    elif BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 32:
+      row = warp_laneid % 16
+      col = 8 * (warp_laneid // 16)
+
+    row_offset = 0
+    col_offset = inner
+
+    # swizzle then find row and col
+    offset = (row + row_offset) * BASE_TILE_COLS + (col + col_offset)
+    offset *= DTYPE_SIZE
+
+    if BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 16:
+      swizzle = ((offset % 512) >> 7) << 3
+      offset = offset ^ swizzle
+    elif BASE_TILE_ROWS == 16 and BASE_TILE_COLS == 32:
+      swizzle = ((offset % 1024) >> 9) << 5
+      offset = offset ^ swizzle
+
+    offset //= DTYPE_SIZE
+
+    row = offset // BASE_TILE_COLS
+    col = offset % BASE_TILE_COLS
+
+    ret.append((row, col))
+
+  return ret
+
+# ===
+
+def shm_phase(inst, threadIdx_x):
+  match inst:
+    case "ds_read_b128":
+      match threadIdx_x:
+        case 0 | 1 | 2 | 3 | 12 | 13 | 14 | 15 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27: return 0
+        case 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 16 | 17 | 18 | 19 | 28 | 29 | 30 | 31: return 1
+        case 32 | 33 | 34 | 35 | 44 | 45 | 46 | 47 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59: return 2
+        case 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 48 | 49 | 50 | 51 | 60 | 61 | 62 | 63: return 3
+    case "ds_read_b64":
+      if threadIdx_x < 32: return 0
+      else: return 1
+    case "ds_write_b64":
+      if threadIdx_x < 16: return 0
+      elif threadIdx_x < 32: return 1
+      elif threadIdx_x < 48: return 2
+      else: return 3
+
+def shm_bank(inst, row, col):
+  bank =  row * (BASE_TILE_COLS // 2) + (col // 2)
+
+  match inst:
+    case "ds_read_b128": bank = bank % 64
+    case "ds_read_b64": bank = bank % 64
+    case "ds_write_b64": bank = bank % 32
+
+  return bank
+
+def map_range(value, from_min, from_max, to_min, to_max):
+  ratio = (value - from_min) / (from_max - from_min)
+  return to_min + ratio * (to_max - to_min)
+
+def shm_bank_gradient(inst, bank):
+  # rgb color for each bank
+  # for 16 bit elements, two elements per bank row wise
+
+  # gradient from blue to red
+  amount = map_range(bank, 0, (64 if inst != "ds_write_b64" else 32) - 1, 0, 120)
+  amount = int(amount)
+  return (amount, amount // 2, 120 - amount)
+
+def color_code(phase):
+  match phase:
+    case 0: return "red"
+    case 1: return "green"
+    case 2: return "blue"
+    case 3: return "yellow"
+
+def rgb_bg(text, color):
+  return f"\033[48;2;{color[0]};{color[1]};{color[2]}m{text}\033[0m"
+
+def visualize_threads(inst=INST):
+  for threadIdx_x in range(WARP_THREADS):
+    row, col = zip(*row_col(threadIdx_x))
+    print(f"Thread {threadIdx_x:2}: ", end="")
+    for r, c in zip(row, col):
+      phase = shm_phase(inst, threadIdx_x)
+      color = color_code(phase)
+      print(f"{color}({r:3},{c:3})\033[0m ", end="")
+    print()
+
+  unique_pairs = set()
+  for threadIdx_x in range(WARP_THREADS):
+    rc_list = row_col(threadIdx_x)
+    for rc in rc_list:
+      unique_pairs.add(rc)
+  assert len(unique_pairs) == 64 * BASE_TILE_NEPT, f"Expected {64 * BASE_TILE_NEPT} unique pairs, got {len(unique_pairs)}"
+
+def visualize_tile(inst=INST):
+  tile = [[-1 for _ in range(BASE_TILE_COLS)] for _ in range(BASE_TILE_ROWS)]
+  for threadIdx_x in range(WARP_THREADS):
+    rc_list = row_col(threadIdx_x)
+    for r, c in rc_list:
+      try:
+        tile[r][c] = threadIdx_x
+      except:
+        pass
+
+  bank_conflicts = {}
+
+  print("\nTile layout (each number indicates the thread holding that position):")
+  for r in range(BASE_TILE_ROWS):
+    for c in range(BASE_TILE_COLS):
+      phase = shm_phase(inst, tile[r][c])
+      bank = shm_bank(inst, r, c)
+      color = color_code(phase)
+      bank_color = shm_bank_gradient(inst, bank)
+
+      if (bank, phase) not in bank_conflicts:
+        bank_conflicts[(bank, phase)] = []
+      bank_conflicts[(bank, phase)].append((r, c, tile[r][c]))
+
+      if phase == -1:
+        bank_color = (0, 0, 0)
+
+      text = colored(f"{tile[r][c]:2}", color)
+      text = rgb_bg(text, bank_color)
+      print(f"{text:2}", end=" ")
+    print()
+
+  for (bank, phase), positions in bank_conflicts.items():
+    if len(positions) > 1:
+      unique_threads = set(pos[2] for pos in positions)
+      if len(unique_threads) > 1:
+        print(f"{len(unique_threads)} way bank conflict: bank {bank}")
+
+if __name__ == "__main__":
+  visualize_tile()
+  # visualize_threads()

extra/tinyfs/fetch_file.py

@@ -8,4 +8,4 @@ if __name__ == "__main__":
   parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
   args = parser.parse_args()
 
-  Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()
+  Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()

extra/tinyfs/fetch_raid.py

@@ -1,4 +1,4 @@
-import json, multiprocessing
+import json, multiprocessing, functools
 from pathlib import Path
 
 from tinygrad.tensor import Tensor
@@ -14,23 +14,25 @@ def fetch_file(item):
   path.parent.mkdir(parents=True, exist_ok=True)
 
   try:
-    pt = Tensor(bytes.fromhex(h), device="CPU").load(size).to(f"disk:{path.as_posix()}").realize()
+    pt = Tensor(bytes.fromhex(h), device="CPU").fs_load(size).to(f"disk:{path.as_posix()}").realize()
   except Exception as e:
     print(f"error fetching {path}, {h}, {size}: {e}")
     raise
 
   pt.uop.buffer.deallocate()
 
-def fetch_mapping():
-  mapping_tensor = Tensor(bytes.fromhex("d734f5e3be9f1e9d863bfaa4fc6c1ef2")).load(175866113).realize()
+def fetch_mapping(h, l):
+  mapping_tensor = Tensor(bytes.fromhex(h)).fs_load(l).realize()
   mapping = mapping_tensor.data().tobytes().decode()
   mapping = json.loads(mapping)
   mapped_files = mapping.items()
   return list(mapped_files)
 
 if __name__ == "__main__":
+  h, l = getenv("HASH", "d734f5e3be9f1e9d863bfaa4fc6c1ef2"), getenv("LENGTH", 175866113)
+
   with multiprocessing.Pool(processes=1) as pool:
-    mapped_files = pool.apply(fetch_mapping)
+    mapped_files = pool.apply(functools.partial(fetch_mapping, h, l))
 
   print(f"fetched mapping for {len(mapped_files)} files")
 

extra/tinyfs/upload_raid.py

@@ -8,7 +8,7 @@ raid_root = Path("/raid")
 
 def upload_file(path: Path):
   pt = Tensor(path).realize()
-  h = pt.store().realize()
+  h = pt.fs_store().realize()
   pt.uop.realized.deallocate()
   return h.data().hex(), path, pt.nbytes()
 
@@ -26,6 +26,6 @@ if __name__ == "__main__":
 
   mapping = json.dumps(mapping).encode()
   mapping_tensor = Tensor(mapping, device="CPU")
-  h = mapping_tensor.store().realize()
+  h = mapping_tensor.fs_store().realize()
 
   print(f"final hash: {h.data().hex()}, size: {len(mapping)}")

extra/torch_backend/backend.py

@@ -4,10 +4,10 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import getenv, prod
+from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
-import torch, pathlib, math, operator, functools, inspect
+import torch, pathlib, math, operator, functools, weakref
 torch.autograd.grad_mode.set_multithreading_enabled(False)
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
 
@@ -18,7 +18,17 @@ def _to_torch_device(device: str): return torch.device("tiny", int(device.partit
 
 import torch.utils.cpp_extension
 mod = torch.utils.cpp_extension.load(name="custom_device_extension", sources=[str(pathlib.Path(__file__).parent / "wrapped_tensor.cpp")])
-def wrap(x:Tensor) -> torch.Tensor: return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
+def calculate_storage_offset(x: Tensor) -> int:
+  offset = 0
+  for u in x.uop.toposort():
+    if u.op == Ops.SHRINK:
+      u_strides = strides_for_shape(u.src[0].shape)
+      for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+  return offset
+def wrap(x: Tensor) -> torch.Tensor:
+  x._strides = strides_for_shape(x.shape) # always recalculate
+  if (not hasattr(x, '_storage_offset')) or (not x.uop.is_realized): x._storage_offset = calculate_storage_offset(x)
+  return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
 def unwrap(x:torch.Tensor) -> Tensor:
   assert isinstance(x, torch.Tensor), f"x isn't {type(x)}"
   return mod.unwrap(x)
@@ -35,17 +45,20 @@ torch.utils.generate_methods_for_privateuse1_backend()
 aten = torch.ops.aten
 
 # track view relationships for in place operations
-def is_view(tensor: Tensor): return hasattr(tensor, "_view_base")
 def canonical_base(view: Tensor): return getattr(view, "_view_base", view)
 def derived_views(base: Tensor): return [t for tref in getattr(base, "_views", set()) if (t:=tref()) is not None]
+def unwrap_args(args, kwargs):
+  return [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args], {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
 def wrap_view_op(fn):
-  def _wrap(*args,**kwargs):
-    args = [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args]
-    kwargs = {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
-    ret = fn(*args,**kwargs)
-    ret._view_base = base = canonical_base(args[0])
-    if not hasattr(base, "_views"): base._views = set()
+  @functools.wraps(fn)
+  def _wrap(*args, **kwargs):
+    args, kwargs = unwrap_args(args, kwargs)
+    ret = fn(*args, **kwargs)
+    base = canonical_base(args[0])
+    ret._view_base = base
+    base._views = getattr(base, "_views", set())
     base._views.add(weakref.ref(ret))
+    ret._view_ops = _get_view_ops(args[0]) + [(fn, args[1:], kwargs)]
     return wrap(ret)
   return _wrap
 
@@ -58,48 +71,83 @@ view_ops = {
   "aten.transpose.int": Tensor.transpose,
   "aten.squeeze.dim": Tensor.squeeze,
   "aten.unsqueeze": Tensor.unsqueeze,
-  "aten.detach": Tensor.detach,
   "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
-}
+  "aten.permute": Tensor.permute,
+  "aten.alias": lambda self: self,
+  }
+
+# torch 2.10 handles this natively
+if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
 
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
 
-# in place operations with views
-def realize_with_views(self: Tensor, views: Tensor):
-  if not self.uop.st.contiguous: self.replace(self.contiguous())
-  self.replace(self.clone().realize())
-  for v in views:
-    if v.uop.base.op is Ops.BUFFER_VIEW: continue # skip subbuffer, we just use the real buffer view
-    ret = self
-    st = ShapeTracker(self.uop.st.views + v.uop.st.views) # TODO: is this right?
-    for mo in cached_to_movement_ops(self.shape, st): ret = apply_mop(ret, mo)
-    v.replace(ret)
-def maybe_realize_storage(self: Tensor) -> bool:
-  if realize:=is_view(self): realize_with_views((base:=canonical_base(self)), derived_views(base))
-  return realize
-def inplace_fn(outvars: str|list[str]):
-  if type(outvars) is str: outvars = [outvars]
-  def decorator(fn):
-    sig = inspect.signature(fn)
-    def wrapper(*args, **kwargs):
-      bound = sig.bind(*args, **kwargs)
-      outs = [kwargs.get(v, bound.arguments.get(v)) for v in outvars]
-      outs = [unwrap(o) if isinstance(o, torch.Tensor) else o for o in outs]
-      realize = any(maybe_realize_storage(o) for o in outs)
-      ret = fn(*args, **kwargs)
-      if realize: Tensor.realize(*(o for o in outs))
-      return ret
-    return wrapper
-  return decorator
+def _get_view_ops(view): return getattr(view, "_view_ops", [])
+
+def _apply_view_ops(target, ops):
+  for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
+  return target
+
+# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
+def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
+  if not (req := argfix(*args)): return None
+  new_shape, infer_idx = [], -1
+  for i, s in enumerate(req):
+    if s is None: s = shape[i] if i < len(shape) else None
+    if not isinstance(s, int): return None
+    if s == -1:
+      if infer_idx != -1: return None
+      infer_idx = len(new_shape)
+    new_shape.append(s)
+  total = prod(shape)
+  if infer_idx != -1:
+    known = prod(x for x in new_shape if x != -1)
+    if known == 0:
+      if total != 0: return None
+      new_shape[infer_idx] = 0
+    else: new_shape[infer_idx] = total // known
+  return tuple(new_shape) if prod(new_shape) == total else None
+
+# TODO: can we get rid of this? only for test_flatten_reshape_add
+def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
+  if not (ops := _get_view_ops(view)): return False
+  shapes = [base.shape]
+  for fn, args, _ in ops:
+    if fn is Tensor.reshape:
+      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
+      shapes.append(next_shape)
+  if shapes[-1] != view.shape: return False
+  for s in reversed(shapes[:-1]): val = val.reshape(s)
+  base.assign(val)
+  return True
+
+def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
+  val = value if value.dtype == base.dtype else value.cast(base.dtype)
+  if view.shape == base.shape: return base.assign(val)
+  if _try_simple_reshape_view_write(base, view, val): return
+  idx_base = Tensor.arange(base.numel(), device=base.device, dtype=dtypes.int32).reshape(base.shape)
+  idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
+  flat_base = base.reshape(base.numel()).contiguous()
+  flat_base[idx_view] = val.reshape(-1)
+  base.assign(flat_base.reshape(base.shape))
+
+def _apply_inplace(target: Tensor, value: Tensor) -> None:
+  val = value if value.dtype == target.dtype else value.cast(target.dtype)
+  base = canonical_base(target)
+  views = derived_views(base)
+  if not views: return target.assign(val)
+  view_ops_map = {v: _get_view_ops(v) for v in views}
+  if target is base or target.uop is base.uop: base.assign(val)
+  else: _view_write(base, target, val)
+  for v in views: v.replace(_apply_view_ops(base, view_ops_map[v]))
 
 # *** bad functions on CPU ***
 
 @torch.library.impl("aten::_index_put_impl_", "privateuseone")
-@inplace_fn("self")
 def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
   # TODO: move to tinygrad
   ret = aten._index_put_impl_(self.cpu(), [x.cpu() if isinstance(x, torch.Tensor) else None for x in indices], values.cpu(), accumulate, unsafe).to(self.device)
-  return wrap(unwrap(self).assign(unwrap(ret)))
+  unwrap(self).assign(unwrap(ret))
+  return self
 
 @torch.library.impl("aten::index_put", "privateuseone")
 def index_put(self, indices, values, accumulate=False):
@@ -150,43 +198,23 @@ for i in [
 def index_tensor(x, y):
   return wrap(unwrap(x)[[unwrap(_y.to(x.device)) if _y is not None else slice(None) for _y in y]])
 
-@torch.library.impl("aten::zero_", "privateuseone")
-@inplace_fn("x")
-def zero_(x):
-  if TORCH_DEBUG: print(f"zero_ {x.shape}")
-  tt = unwrap(x)
-  tt.assign(tt.zeros_like())
-
-@torch.library.impl("aten::fill_.Scalar", "privateuseone")
-@inplace_fn("x")
-def fill_scalar(x, y):
-  if TORCH_DEBUG: print(f"fill_.Scalar {x.shape} {y}")
-  tt = unwrap(x)
-  tt.assign(tt.full_like(y))
 
 @torch.library.impl("aten::_local_scalar_dense", "privateuseone")
 def _local_scalar_dense(tensor): return unwrap(tensor).item()
 
-@functools.cache
-def cached_to_movement_ops(shape, st) -> list:
-  mops = to_movement_ops(st)
-  if mops[0] == (MovementOps.RESHAPE, shape): mops = mops[1:]
-  return mops
-
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps
-
 @wrap_view_op
-def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
-  # multiple as_strided do not compound
-  base = canonical_base(tensor)
-  # TODO: this is heavyweight
-  st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
-  ret = base
-  if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
-  if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
-  for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
-  return ret
+def _as_strided(tensor:Tensor, size, stride, storage_offset=0):
+  base = getattr(tensor, "_as_strided_base", canonical_base(tensor)).flatten()
+  if prod(size) == 1: return base[storage_offset].reshape(size)
+  indices = Tensor.zeros(size, dtype=dtypes.int32, device=base.device) + storage_offset
+  for dim, (sz, st) in enumerate(zip(size, stride)):
+    if st != 0:
+      dim_range = Tensor.arange(sz, device=base.device, dtype=dtypes.int32) * st
+      shape_for_broadcast = [1] * dim + [sz] + [1] * (len(size) - dim - 1)
+      indices = indices + dim_range.reshape(shape_for_broadcast)
+  result = base[indices.flatten()].reshape(size)
+  result._as_strided_base = base
+  return result
 
 @torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
@@ -245,15 +273,14 @@ def convolution_overrideable(input, weight, bias, stride, padding, dilation, tra
   if TORCH_DEBUG >= 1:
     print(f"convolution {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
   input, weight, bias = unwrap(input), unwrap(weight), unwrap(bias) if bias is not None else None
-  # TODO: fix test_biased_conv2d fails without realize()
-  if not transposed: return wrap(input.conv2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding).realize())
-  return wrap(input.conv_transpose2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding).realize())
+  if not transposed: return wrap(input.conv2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding))
+  return wrap(input.conv_transpose2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding))
 
 @torch.library.impl("aten::convolution_backward_overrideable", "privateuseone")
 def convolution_backward_overrideable(grad_out, input, weight, stride, padding, dilation, transposed, output_padding, groups, output_mask):
   if TORCH_DEBUG >= 1:
     print(f"convolution_backward {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
-  grad_out, input, weight, bias = unwrap(grad_out), unwrap(input), unwrap(weight), Tensor.zeros(weight.shape[0], device=_from_torch_device(weight.device))
+  grad_out, input, weight, bias = unwrap(grad_out).detach(), unwrap(input).detach(), unwrap(weight).detach(), Tensor.zeros(weight.shape[0], device=_from_torch_device(weight.device))
   if not transposed: out = Tensor.conv2d(input, weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding)
   else:
     bias = Tensor.zeros(weight.shape[1] * groups)
@@ -315,55 +342,57 @@ for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
 
 @torch.library.impl("aten::scatter_add.out", "privateuseone")
-@inplace_fn("out")
 def scatter_add(self, dim, index, src, out):
-  self, index, src, out = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
-  if self.shape == (): return wrap(out.assign(src))
-  return wrap(out.assign(Tensor.scatter_reduce(self, dim, index, src, reduce='sum')))
+  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
+  if self.shape == (): _apply_inplace(out_unwrapped, src)
+  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
+  return out
 
-@torch.library.impl("aten::_copy_from", "privateuseone")
-def _copy_from(src: torch.Tensor, dest, non_blocking=False):
-  realize = dest.is_tiny and maybe_realize_storage(unwrap(dest))
-  cast_dtype = _from_torch_dtype(dest.dtype)
+def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
-    to_device = _from_torch_device(dest.device)
-    src,dest = unwrap(src),unwrap(dest)
-    # TODO we need to properly match dest shape and strides, not blindly assign
-    if dest.uop.st.contiguous or dest.uop.is_realized: src = src.contiguous() # this only solves some cases
-    dest.assign(src.cast(cast_dtype).to(to_device))
-    if realize: Tensor.realize(dest)
+    src_t, dest_t = unwrap(src), unwrap(dest)
+    if dest_t.uop.is_contiguous() or dest_t.uop.is_realized: src_t = src_t.contiguous()
+    _apply_inplace(dest_t, src_t.cast(cast_dtype).to(to_device))
   elif src.is_tiny and dest.is_cpu:
-    # TODO: is there a better way?
     dest.resize_(src.numel()).resize_(src.shape)
     dest.copy_(torch.from_numpy(unwrap(src).cast(cast_dtype).numpy()))
   elif src.is_cpu and dest.is_tiny:
-    to_device = _from_torch_device(dest.device)
-    # TODO we need to properly match dest shape and strides, not blindly assign
     unwrap(dest).assign(Tensor(src.numpy()).cast(cast_dtype).to(to_device))
-    if realize: Tensor.realize(unwrap(dest))
   else:
     raise NotImplementedError(f"can't copy from {src.device} -> {dest.device}")
 
+@torch.library.impl("aten::_copy_from", "privateuseone")
+def _copy_from(src: torch.Tensor, dest, non_blocking=False):
+  cast_dtype = _from_torch_dtype(dest.dtype)
+  to_device = _from_torch_device(dest.device)
+  _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking)
+  return dest
+
+@torch.library.impl("aten::copy_", "privateuseone")
+def copy_(self, src, non_blocking=False):
+  cast_dtype = _from_torch_dtype(self.dtype)
+  to_device = _from_torch_device(self.device)
+  _copy_between_devices(src, self, cast_dtype, to_device, non_blocking)
+  return self
+
 @torch.library.impl("aten::cat.out", "privateuseone")
-@inplace_fn("out")
 def cat_out(tensors, dim=0, out=None):
-  unwrap(out).assign(Tensor.cat(*[unwrap(x) for x in tensors], dim=dim))
+  _apply_inplace(unwrap(out), Tensor.cat(*[unwrap(x) for x in tensors], dim=dim))
+  return out
 
 @torch.library.impl("aten::topk.values", "privateuseone")
-@inplace_fn(["values", "indices"])
 def topk_values(input, k, dim=None, largest=True, sorted=True, values=None, indices=None):
   out_values, out_indices = unwrap(input).topk(k, dim if dim is not None else -1, largest, sorted)
-  unwrap(values).assign(out_values)
-  unwrap(indices).assign(out_indices.cast(dtypes.int64))
-  return wrap(out_values), wrap(out_indices)
+  _apply_inplace(unwrap(values), out_values)
+  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
+  return values, indices
 
 @torch.library.impl("aten::sort.values_stable", "privateuseone")
-@inplace_fn(["values", "indices"])
 def sort_values(input, dim=-1, descending=False, stable=True, values=None, indices=None):
   out_values, out_indices = unwrap(input).sort(dim, descending)
-  unwrap(values).assign(out_values)
-  unwrap(indices).assign(out_indices.cast(dtypes.int64))
-  return wrap(out_values), wrap(out_indices)
+  _apply_inplace(unwrap(values), out_values)
+  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
+  return values, indices
 
 @torch.library.impl("aten::_linalg_svd", "privateuseone")
 def _linalg_svd(self, full_matrices=False):
@@ -373,7 +402,6 @@ def _linalg_svd(self, full_matrices=False):
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
-  aten.native_batch_norm, aten.native_batch_norm_backward,
   aten.native_layer_norm_backward,
   aten.linalg_cross,
   aten.addmm,
@@ -510,7 +538,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
 
 # we add the "out" here
 def wrap_out(f):
-  @inplace_fn("out")
   def _wrap_out(*args, **kwargs):
     out = kwargs.pop('out')
     assigned = f(*args, **kwargs)
@@ -518,22 +545,33 @@ def wrap_out(f):
     assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
     assert out.device == assigned.device, f"device mismatch: {assigned.device} -> {out.device}"
     assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    if out.uop.is_realized: assigned = assigned.contiguous() # TODO: how does this map to torch's semantics
     return out.assign(assigned)
   return _wrap_out
 
+def _inplace_op(t, new_value):
+  if not hasattr(t, "_view_base") and not getattr(canonical_base(t), "_views", set()): t.replace(new_value)
+  else: _apply_inplace(t, new_value)
+  return t
+
 tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
   "aten.floor_divide": lambda x,y: x//y,
-  "aten.floor_divide_.Tensor": inplace_fn("x")(lambda x,y: x.assign(x//y)),
+  "aten.floor_divide_.Tensor": lambda x,y: x//y,
   # TODO: use tinygrad methods, but they require x to be unsigned
   "aten.__lshift__.Scalar": lambda x,y: x*(2**y),
-  "aten.__ilshift__.Scalar": inplace_fn("x")(lambda x,y: x.assign(x*(2**y))),
+  "aten.__ilshift__.Scalar": lambda x,y: x*(2**y),
   "aten.__rshift__.Scalar": lambda x,y: x//(2**y),
-  "aten.__irshift__.Scalar": inplace_fn("x")(lambda x,y: x.assign(x//(2**y))),
+  "aten.__irshift__.Scalar": lambda x,y: x//(2**y),
+  # inplace ops using replace for fusion
+  "aten.zero_": lambda x: x.zeros_like(),
+  "aten.fill_.Scalar": lambda x, y: x.full_like(y),
+  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.mul_.Tensor": lambda self, other: self * other,
+  "aten.mul_.Scalar": lambda self, other: self * other,
   # relu doesn't have an out form?
   "aten.relu": Tensor.relu,
-  "aten.relu_": inplace_fn("x")(lambda x: x.assign(x.relu())),
+  "aten.relu_": lambda x: x.relu(),
   "aten.mean": Tensor.mean,
   "aten.mean.dim": Tensor.mean,
   "aten.min": Tensor.min,
@@ -554,19 +592,17 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
   "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
   "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": inplace_fn("self")(lambda self:
-    self.assign(Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype))),
-  "aten.random_.from": inplace_fn("self")(lambda self, from_, to:
-    self.assign(Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype))),
-  "aten.uniform_": inplace_fn("self")(lambda self, low=0, high=1: self.assign(Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype))),
-  "aten.normal_": inplace_fn("self")(lambda self, mean=0, std=1: self.assign(Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype))),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
+  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
+  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
+  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
   # these don't work in out form, they have size 0
   "aten.abs": Tensor.abs,
   "aten.logical_not": Tensor.logical_not,
-  "aten.logical_or_": inplace_fn("x")(lambda x, y: x.assign(x | y)),
+  "aten.logical_or_": lambda x, y: x | y,
   "aten.multinomial": Tensor.multinomial,
-  "aten.masked_fill_.Scalar": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
-  "aten.masked_fill_.Tensor": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
+  "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
+  "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
   "aten.masked_fill.Scalar": Tensor.masked_fill,
   "aten.masked_fill.Tensor": Tensor.masked_fill,
   "aten.masked_select": Tensor.masked_select,
@@ -580,7 +616,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.asinh": Tensor.asinh,
   "aten.mul": Tensor.mul,
   "aten.atanh": Tensor.atanh,
-  "aten.fill_.Tensor": Tensor.full, # TODO: looks wrong
+  "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
   "aten.squeeze_.dim": lambda self, dim: self.replace(self.squeeze(dim), allow_shape_mismatch=True), # TODO: inplace view op, here?
@@ -601,20 +637,51 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.unfold": Tensor.unfold,
 }}
 
+# operations that need inplace treatment (use _inplace_op instead of wrap_fxn) AKA return original tensor
+inplace_ops = {
+  "aten.zero_",
+  "aten.fill_.Scalar",
+  "aten.fill_.Tensor",
+  "aten.add_.Tensor",
+  "aten.add_.Scalar",
+  "aten.mul_.Tensor",
+  "aten.mul_.Scalar",
+  "aten.floor_divide_.Tensor",
+  "aten.__ilshift__.Scalar",
+  "aten.__irshift__.Scalar",
+  "aten.relu_",
+  "aten.random_",
+  "aten.random_.from",
+  "aten.uniform_",
+  "aten.normal_",
+  "aten.logical_or_",
+  "aten.masked_fill_.Scalar",
+  "aten.masked_fill_.Tensor",
+}
+
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
       print(k, len(args), [x.shape if isinstance(x, torch.Tensor) else x for x in args],
                           {k:v.shape if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()})
-    args = [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args]
-    kwargs = {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
+    args, kwargs = unwrap_args(args, kwargs)
     out = f(*args, **kwargs)
     if isinstance(out, Tensor): return wrap(out)
     elif isinstance(out, tuple): return tuple(wrap(x) for x in out)
     else: raise RuntimeError(f"unknown output type {type(out)}")
   return nf
 
-for k,v in tiny_backend.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_fxn(k,v))
+def wrap_inplace(k,f):
+  def nf(*args, **kwargs):
+    orig = args[0]
+    args, kwargs = unwrap_args(args, kwargs)
+    _inplace_op(args[0], f(*args, **kwargs))
+    return orig
+  return nf
+
+for k,v in tiny_backend.items():
+  wrapper = wrap_inplace if k in inplace_ops else wrap_fxn
+  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
 
 @torch.library.impl("aten::equal", "privateuseone")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -628,42 +695,72 @@ if TORCH_DEBUG:
       return func(*args, **(kwargs or {}))
   (_dispatch_log:=DispatchLog()).__enter__() # NOTE: must be kept alive
 
-# NOTE: patch torch optimizer step to avoid continously growing the computation graph
-import weakref
-_torch_modules_with_buffers: weakref.WeakSet[torch.nn.Module] = weakref.WeakSet()
-def register_torch_buffer(mod, _name, _buffer): _torch_modules_with_buffers.add(mod)
-def get_real_tinygrad_buffers():
-  res = set()
-  for mod in _torch_modules_with_buffers:
-    for _,b in mod.named_buffers(recurse=False):
-      if b is not None and b.is_tiny:
-        res.add(unwrap(b))
-  return res
-torch.nn.modules.module.register_module_buffer_registration_hook(register_torch_buffer)
+# this implementation is needed to allow the batchnorm kernels to fuse in e.g. mnist training
+# aten::native_batch_norm does more than Tensor.batchnorm
+@torch.library.impl("aten::native_batch_norm", "privateuseone")
+def native_batch_norm(input, weight, bias, running_mean, running_var, training, momentum, eps):
+  input_t, weight_t, bias_t = unwrap(input), unwrap(weight) if weight is not None else None, unwrap(bias) if bias is not None else None
+  running_mean_t, running_var_t = unwrap(running_mean) if running_mean is not None else None, unwrap(running_var) if running_var is not None else None
+  if training:
+    batch_var, batch_mean = input_t.var_mean(axis=tuple(x for x in range(input_t.ndim) if x != 1), correction=0)
+    batch_invstd = batch_var.add(eps).rsqrt()
+    out = input_t.batchnorm(weight_t, bias_t, batch_mean, batch_invstd)
+    if running_mean_t is not None and running_var_t is not None:
+      numel_ratio = input_t.numel() / (input_t.numel() - input_t.shape[1])
+      running_mean_t.assign((1 - momentum) * running_mean_t + momentum * batch_mean.detach())
+      running_var_t.assign((1 - momentum) * running_var_t + momentum * numel_ratio * batch_var.detach())
+    return wrap(out), wrap(batch_mean), wrap(batch_invstd)
+  else:
+    out = input_t.batchnorm(weight_t, bias_t, running_mean_t, running_var_t.add(eps).rsqrt())
+    return wrap(out), wrap(running_mean_t), wrap(running_var_t.add(eps).rsqrt())
 
-from torch.nn.modules import Module
-def param_hook(_grad):
-  if _grad is not None and _grad.is_tiny: Tensor.realize(unwrap(_grad))
-def module_hook(module:Module, _name, _submodule):
-  for param in _submodule.parameters(recurse=False):
-    if param.requires_grad: param.register_hook(param_hook)
-torch.nn.modules.module.register_module_module_registration_hook(module_hook)
+@torch.library.impl("aten::native_batch_norm_backward", "privateuseone")
+def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask):
+  grad_out_t, input_t = unwrap(grad_out), unwrap(input)
+  weight_t = unwrap(weight) if weight is not None else None
+  save_mean_t = unwrap(save_mean)
+  save_invstd_t = unwrap(save_invstd)
+  out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t)
+  targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m]
+  if targets:
+    grads = out.gradient(*targets, gradient=grad_out_t)
+    grad_input = grads.pop(0) if output_mask[0] else None
+    grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None
+  else:
+    grad_input, grad_weight = None, None
+  grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None
+  return (wrap(grad_input) if grad_input is not None else None,
+          wrap(grad_weight) if grad_weight is not None else None,
+          wrap(grad_bias) if grad_bias is not None else None)
 
-def realize_optimizer_step(optimizer: torch.optim.Optimizer, *args, **kwargs):
-  tinygrad_tensors = []
-  for param_group in optimizer.param_groups:
-    for param in param_group["params"]:
-      if param is None: continue
-      tinygrad_tensors.append(param.data)
-  for state_dict in optimizer.state.values():
-    for _, value in state_dict.items():
-      if torch.is_tensor(value): tinygrad_tensors.append(value)
-  real_tinygrad_tensors = [unwrap(x) for x in tinygrad_tensors if x.is_tiny]
-  real_tinygrad_tensors += get_real_tinygrad_buffers()
-  if len(real_tinygrad_tensors): Tensor.realize(*real_tinygrad_tensors)
+# _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
+# we need torch.autograd.Function with explicit AutogradPrivateUse1 registration
+class _PadCircular(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, input, padding):
+    ctx.save_for_backward(input)
+    ctx.padding = padding
+    return pad_forward(input, padding, mode="circular")
+  @staticmethod
+  def backward(ctx, grad_output):
+    input, = ctx.saved_tensors
+    return pad_backward(grad_output, input, ctx.padding, mode="circular"), None
 
-_optimizer_init = torch.optim.Optimizer.__init__
-def _optimizer_patched_init(self, *args, **kwargs):
-  _optimizer_init(self, *args, **kwargs)
-  self.register_step_post_hook(realize_optimizer_step)
-torch.optim.Optimizer.__init__ = _optimizer_patched_init
+@torch.library.impl("aten::_pad_circular", "privateuseone")
+def _pad_circular(self, padding): return _PadCircular.apply(self, padding)
+
+@torch.library.impl("aten::_pad_circular", "AutogradPrivateUse1")
+def _pad_circular_autograd(self, padding): return _PadCircular.apply(self, padding)
+
+# only needed for test_diag_backward_gradient_values
+# was going through torch before, but now we are using tinygrad directly and tracking views
+# Tensor.diagonal does not support all cases tests in the tests
+@torch.library.impl("aten::diagonal", "privateuseone")
+@wrap_view_op
+def diagonal(self, offset=0, dim1=0, dim2=1):
+  if offset != 0: raise NotImplementedError(f"diagonal with {offset=} not implemented")
+  dim1, dim2 = dim1 % self.ndim, dim2 % self.ndim
+  if dim1 != self.ndim - 2 or dim2 != self.ndim - 1: raise NotImplementedError(f"diagonal with {dim1=}, {dim2=} not implemented, only last two dims supported")
+  batch_shape, m, n = self.shape[:-2], self.shape[-2], self.shape[-1]
+  diag_len = min(m, n)
+  return self.reshape(*batch_shape, m*n).pad(tuple((0,0) for _ in batch_shape) + ((0, diag_len),)).reshape(*batch_shape, diag_len, n+1)[..., :, 0]

extra/torch_backend/example.py

@@ -1,12 +1,13 @@
 from PIL import Image
-from tinygrad.helpers import getenv
-import torch, torchvision, pathlib
+from tinygrad.helpers import getenv, GlobalCounters
+import torch, torchvision, pathlib, warnings
 import torchvision.transforms as transforms
 import extra.torch_backend.backend
 device = "tiny"
 torch.set_default_device(device)
 
 if __name__ == "__main__":
+  GlobalCounters.reset()
   img = Image.open(pathlib.Path(__file__).parent.parent.parent / "test/models/efficientnet/Chicken.jpg").convert('RGB')
   transform = transforms.Compose([
     transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(),
@@ -19,3 +20,10 @@ if __name__ == "__main__":
   out = model(img).detach().cpu().numpy()
   print("output:", out.shape, out.argmax())
   assert out.argmax() == 7  # cock
+
+  kernel_count = GlobalCounters.kernel_count
+  assert kernel_count > 0, "No kernels, test failed"
+  expected_kernels = 228
+  expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
+  if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
+  assert kernel_count <= expected_kernels, f"{expectation}"

extra/torch_backend/test.py

@@ -2,7 +2,7 @@
 import unittest
 import torch
 import numpy as np
-from tinygrad.helpers import getenv, Context, GlobalCounters
+from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
   import extra.torch_backend.backend2
   device = "cpu"
@@ -25,7 +25,7 @@ class TestTorchBackend(unittest.TestCase):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
 
-  def test_numpy_ones(self):
+  def test_numpy_ones_int32(self):
     a = torch.ones(4, dtype=torch.int32, device=device)
     assert a.dtype == torch.int32
     np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
@@ -219,7 +219,6 @@ class TestTorchBackend(unittest.TestCase):
     a = torch.ones(4, device=device)
     print(str(a))
 
-  @unittest.skip("failed")
   def test_floor_div(self):
     a = torch.tensor([10., 7., 5.], device=device)
     b = torch.tensor([3., 2., 2.], device=device)
@@ -248,5 +247,672 @@ class TestTorchBackend(unittest.TestCase):
   def test_diagonal_rectangular(self): self._test_diagonal(4, 5, 6)
   def test_diagonal_4d(self): self._test_diagonal(2, 3, 4, 5)
 
+  def test_pad_circular_simple(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2)
+    padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
+    expected = np.array([[[[3.,2.,3.,2.], [1.,0.,1.,0.], [3.,2.,3.,2.], [1.,0.,1.,0.]]]], dtype=np.float32)
+    np.testing.assert_allclose(padded.cpu().numpy(), expected)
+
+  def test_pad_circular_backward(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2).requires_grad_(True)
+    padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
+    loss = padded.sum()
+    loss.backward()
+    expected_grad = np.array([[[[4., 4.], [4., 4.]]]], dtype=np.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad)
+
+
+  def test_matmul_backward(self):
+    x = torch.randn(3, 4, device=device, dtype=torch.float32, requires_grad=True)
+    y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
+    z = (x @ y).sum()
+    z.backward()
+    assert x.grad is not None
+    assert y.grad is not None
+    assert x.grad.shape == x.shape
+    assert y.grad.shape == y.shape
+
+  def test_matmul_broadcast_backward(self):
+    x = torch.randn(2, 3, 4, device=device, dtype=torch.float32, requires_grad=True)
+    y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
+    z = (x @ y).sum()
+    z.backward()
+    assert x.grad is not None
+    assert y.grad is not None
+    assert x.grad.shape == x.shape
+    assert y.grad.shape == y.shape
+
+  def test_diag_vector_to_matrix(self):
+    vec = torch.tensor([1., 2., 3., 4., 5.], dtype=torch.float32, device=device)
+    mat = torch.diag(vec)
+    expected = np.diag([1., 2., 3., 4., 5.])
+    np.testing.assert_allclose(mat.cpu().numpy(), expected, rtol=1e-5)
+    assert mat.shape == (5, 5)
+
+  def test_diagonal_matrix_to_vector(self):
+    mat = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device)
+    vec = torch.linalg.diagonal(mat)
+    expected = np.array([1., 5., 9.])
+    np.testing.assert_allclose(vec.cpu().numpy(), expected, rtol=1e-5)
+    assert vec.shape == (3,)
+
+  def test_permute_2(self):
+    a = torch.randn(2, 3, 4, dtype=torch.float32, device=device)
+    b = a.permute(2, 0, 1)
+    assert b.shape == (4, 2, 3)
+    np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))
+
+  def test_batchnorm_unsqueeze(self):
+    bn = torch.nn.BatchNorm2d(4).to(device)
+    x = torch.randn(8, 4, 3, 3, device=device)
+    out = bn(x)
+    self.assertEqual(out.shape, x.shape)
+
+  def test_slice_inplace_zero(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b.zero_()
+    expected = np.array([[1., 1., 1.],
+                         [1., 0., 0.],
+                         [1., 0., 0.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_slice_inplace_fill(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b.fill_(5.0)
+    expected = np.array([[1., 1., 1.],
+                         [1., 5., 5.],
+                         [1., 5., 5.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_fill_tensor_value(self):
+    a = torch.zeros((2, 2), dtype=torch.float32, device=device)
+    value = torch.tensor(3, dtype=torch.int64, device=device)
+    a.fill_(value)
+    expected = np.full((2, 2), 3, dtype=np.float32)
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_slice_inplace_mul(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b *= 2
+    expected = np.array([[1., 1., 1.],
+                         [1., 2., 2.],
+                         [1., 2., 2.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_permute_slice_zero(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:].permute(1, 0)
+    b.zero_()
+    expected = np.array([[1., 1., 1.],
+                         [1., 0., 0.],
+                         [1., 0., 0.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_permute_slice_mul(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:].permute(1, 0)
+    b *= 2
+    expected = np.array([[1., 1., 1.],
+                         [1., 2., 2.],
+                         [1., 2., 2.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_simple_slice_setitem(self):
+    a = torch.tensor([10, 20, 30], device=device)
+    a[1] = 99
+    np.testing.assert_equal(a.cpu().numpy(), [10, 99, 30])
+
+  def test_2d_slice_setitem(self):
+    a = torch.zeros((3, 3), device=device)
+    a[1, 2] = 99
+    self.assertEqual(a[1, 2].item(), 99)
+    self.assertEqual(a.sum().item(), 99)
+
+  def test_view_copy(self):
+    a = torch.tensor([10, 20, 30], device=device)
+    view = a[1]
+    view.copy_(torch.tensor(88, device=device))
+    np.testing.assert_equal(a.cpu().numpy(), [10, 88, 30])
+
+  def test_diag_2d_input(self):
+    a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device)
+    d = torch.diag(a)
+    np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
+
+  def test_diag_1d_input(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    d = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(d.cpu().numpy(), expected)
+
+  def test_permute_view_tracking(self):
+    a = torch.ones((2, 3, 4), device=device)
+    b = a.permute(2, 0, 1)
+    self.assertEqual(b.shape, (4, 2, 3))
+
+  def test_detach_view_creation(self):
+    a = torch.tensor([1.0, 2.0, 3.0], device=device)
+    b = a.detach()
+    np.testing.assert_equal(b.cpu().numpy(), [1.0, 2.0, 3.0])
+
+  def test_view_zero_inplace(self):
+    a = torch.ones((4, 4), device=device)
+    view = a[1:3, 1:3]
+    view.zero_()
+    self.assertEqual(view.sum().item(), 0)
+
+  def test_view_fill_inplace(self):
+    a = torch.zeros((4, 4), device=device)
+    view = a[1:3, 1:3]
+    view.fill_(5)
+    self.assertEqual(view.sum().item(), 20)
+
+  def test_permute_contiguous(self):
+    a = torch.tensor([[1, 2], [3, 4]], device=device)
+    b = a.permute(1, 0)
+    c = b.contiguous()
+    expected = [[1, 3], [2, 4]]
+    np.testing.assert_equal(c.cpu().numpy(), expected)
+
+  def test_diag_2d_extract_diagonal(self):
+    a = torch.tensor([[1, 2], [3, 4]], device=device)
+    result = torch.diag(a)
+    np.testing.assert_equal(result.cpu().numpy(), [1, 4])
+
+  def test_slice_inplace_multiply_offset_preservation(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    a[1:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [1, 4, 6])
+
+  def test_slice_inplace_mul_pattern(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[:2] *= 3
+    a[2:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [3, 6, 6, 8])
+
+  def test_chained_slice_column(self):
+    a = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4)
+    torch_res = a[:, 1:2][:, 0:1].cpu().numpy()
+    cpu_res = torch.arange(16, dtype=torch.float32).reshape(4, 4)[:, 1:2][:, 0:1].numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_slice_with_step(self):
+    a = torch.arange(20, dtype=torch.float32, device=device)
+    torch_res = a[::2][1:4].cpu().numpy()
+    cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_slice_negative_dim(self):
+    a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
+    torch_chunks = a.chunk(3, -1)
+    cpu_chunks = torch.arange(13, dtype=torch.int32).repeat(8, 1).chunk(3, -1)
+    assert len(torch_chunks) == len(cpu_chunks)
+    for i in range(len(torch_chunks)):
+      np.testing.assert_equal(torch_chunks[i].cpu().numpy(), cpu_chunks[i].numpy())
+
+  def test_dot_vector_matrix(self):
+    a = torch.arange(65, dtype=torch.float32, device=device)
+    b = torch.arange(65*45, dtype=torch.float32, device=device).reshape(65, 45)
+    torch_res = a.matmul(b).reshape(-1).cpu().numpy()
+    cpu_res = torch.arange(65, dtype=torch.float32).matmul(torch.arange(65*45, dtype=torch.float32).reshape(65, 45)).numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_alias_passthrough(self):
+    a = torch.randn(3, 3, device=device)
+    alias_view = torch.ops.aten.alias(a)
+    alias_view += 1
+    np.testing.assert_equal(a.cpu().numpy(), alias_view.cpu().numpy())
+
+  def test_split_simple_vector(self):
+    a = torch.arange(10, dtype=torch.float32, device=device)
+    torch_chunks = a.split([1,4,5])
+    cpu_chunks = torch.arange(10, dtype=torch.float32).split([1,4,5])
+    for tc, cc in zip(torch_chunks, cpu_chunks):
+      np.testing.assert_equal(tc.cpu().numpy(), cc.cpu().numpy())
+
+  def test_split_matches_torch(self):
+    a = torch.arange(10, dtype=torch.float32, device=device)
+    torch_chunks = a.split([1,4,5])
+    tiny_chunks = [chunk.cpu().numpy() for chunk in torch_chunks]
+    cpu_chunks = [torch.arange(10, dtype=torch.float32).split([1,4,5])[i].numpy() for i in range(3)]
+    for tr, cr in zip(tiny_chunks, cpu_chunks): np.testing.assert_equal(tr, cr)
+
+  def test_sum_matches_torch(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2,3)
+    torch_res = a.sum().cpu().numpy()
+    cpu_res = torch.arange(6, dtype=torch.float32).reshape(2,3).sum().numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_view_matches_torch(self):
+    a = torch.arange(6, dtype=torch.float32, device=device)
+    torch_res = a.view(2, 3).cpu().numpy()
+    cpu_res = torch.arange(6, dtype=torch.float32).view(2, 3).numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_view_zero_with_indices(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[1:3].zero_()
+    np.testing.assert_equal(a.cpu().numpy(), [1, 0, 0, 4])
+
+  def test_view_fill_with_indices(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[::2].fill_(9)
+    np.testing.assert_equal(a.cpu().numpy(), [9, 2, 9, 4])
+
+  def test_nested_slice_inplace_ops(self):
+    a = torch.tensor([1, 2, 3, 4, 5, 6], device=device)
+    a[:3] += 10
+    a[3:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [11, 12, 13, 8, 10, 12])
+
+  def test_diag_1d(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    result = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(result.cpu().numpy(), expected)
+
+  def test_diag_backward(self):
+    a = torch.randn(5, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_diagonal(self):
+    a = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    expected = torch.tensor([1., 5., 9.], dtype=torch.float32)
+    self.assertEqual(b.shape, (3,))
+    np.testing.assert_allclose(b.detach().cpu().numpy(), expected.numpy(), rtol=1e-5)
+
+  def test_diagonal_backward(self):
+    a = torch.randn(5, 5, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_expand_backward(self):
+    a = torch.randn(4, 3, 1, 6, dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(4, 3, 2, 6)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_einsum_backward(self):
+    a = torch.randn(10, 10, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.einsum('ij->ji', a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_diag_backward_gradient_values(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.ones(3, dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_backward_gradient_values_2d_to_1d(self):
+    a = torch.tensor([[1.0, 2.0, 3.0],
+                      [4.0, 5.0, 6.0],
+                      [7.0, 8.0, 9.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[1.0, 0.0, 0.0],
+                                   [0.0, 1.0, 0.0],
+                                   [0.0, 0.0, 1.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_backward_gradient_values(self):
+    a = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(3, 4)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[4.0], [4.0], [4.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_backward_with_leading_dims(self):
+    a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(3, 1, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[3.0, 3.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_2d_to_1d_backward(self):
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_complex_backward(self):
+    a = torch.tensor([[[1.0, 2.0]]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(2, 3, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[[6.0, 6.0]]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_backward_with_scaling(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = (b * torch.tensor([[2.0, 0.0, 0.0],
+                               [0.0, 3.0, 0.0],
+                               [0.0, 0.0, 4.0]], device=device)).sum()
+    loss.backward()
+    expected_grad = torch.tensor([2.0, 3.0, 4.0], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_repeat_basic(self):
+    a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
+    b = a.repeat(2, 1)
+    expected = torch.tensor([[1, 2, 3], [1, 2, 3]], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_repeat_multidim(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
+    b = a.repeat(2, 3)
+    expected = torch.arange(6, dtype=torch.float32).reshape(2, 3).repeat(2, 3)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_repeat_backward(self):
+    a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.repeat(3, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[6.0, 6.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_cumsum_1d(self):
+    a = torch.tensor([1, 2, 3, 4], dtype=torch.float32, device=device)
+    b = torch.cumsum(a, dim=0)
+    expected = torch.tensor([1, 3, 6, 10], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_cumsum_2d(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    b = torch.cumsum(a, dim=0)
+    expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=0)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+    c = torch.cumsum(a, dim=1)
+    expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=1)
+    np.testing.assert_equal(c.cpu().numpy(), expected.numpy())
+
+  def test_cumsum_backward(self):
+    a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([4.0, 3.0, 2.0, 1.0], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_constant_pad_nd_1d(self):
+    a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
+    b = torch.nn.functional.pad(a, (1, 2), mode='constant', value=0)
+    expected = torch.tensor([0, 1, 2, 3, 0, 0], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_constant_pad_nd_2d(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
+    b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
+    expected = torch.nn.functional.pad(torch.arange(6, dtype=torch.float32).reshape(2, 3), (1, 1, 1, 1), mode='constant', value=0)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_constant_pad_nd_2d_backward(self):
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.ones((2, 2), dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_negative_strides_cumsum_backward(self):
+    a = torch.randn(5, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    b.sum().backward()
+    grad = a.grad.cpu().numpy()
+    self.assertEqual(len(grad), 5)
+
+  def test_cumsum_fix_gradient_values(self):
+    a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    loss = b.sum()
+    loss.backward()
+    expected = np.array([4.0, 3.0, 2.0, 1.0])
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected, rtol=1e-5)
+
+  def test_diag_1d_to_2d(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(b.detach().cpu().numpy(), expected)
+
+  def test_diag_2d_to_1d(self):
+    c = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float32, device=device)
+    d = torch.diag(c)
+    np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
+
+  def test_biased_conv2d(self):
+    # Test case for two sequential conv2d with same weights/bias and ReLU in between, this is as special case from test_ops.py
+    torch.manual_seed(0)
+    C = 8
+    x_cpu = torch.randn(1, C, 5, 5, requires_grad=True)
+    w_cpu = torch.randn(C, C, 1, 1, requires_grad=True)
+    b_cpu = torch.randn(C, requires_grad=True)
+    x_tiny = x_cpu.detach().to(device).requires_grad_(True)
+    w_tiny = w_cpu.detach().to(device).requires_grad_(True)
+    b_tiny = b_cpu.detach().to(device).requires_grad_(True)
+    out_cpu = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_cpu, w_cpu, b_cpu).relu(), w_cpu, b_cpu)
+    out_tiny = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_tiny, w_tiny, b_tiny).relu(), w_tiny, b_tiny)
+    grad_out = torch.randn_like(out_cpu)
+    out_cpu.backward(grad_out)
+    out_tiny.backward(grad_out.to(device))
+    np.testing.assert_allclose(x_tiny.grad.cpu().numpy(), x_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+    np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+    np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+
+
+from tinygrad import Tensor
+class TestBackendHelpers(unittest.TestCase):
+
+  def test_calculate_storage_offset_no_shrink(self):
+    t = Tensor.ones(3, 4)
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 0
+
+  def test_calculate_storage_offset_with_shrink(self):
+    t = Tensor.ones(10, 10)[2:5, 3:7]
+    # strides for (10, 10) are [10, 1]
+    # offset = 2*10 + 3*1 = 23
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 23
+
+  def test_calculate_storage_offset_multiple_shrinks(self):
+    t = Tensor.ones(5, 6, 7)[1:3, 2:4, 3:5]
+    # strides for (5, 6, 7) are [42, 7, 1]
+    # offset = 1*42 + 2*7 + 3*1 = 42 + 14 + 3 = 59
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 59
+
+  def test_calculate_storage_offset_with_reshape(self):
+    t = Tensor.ones(10, 10)
+    orig_offset = extra.torch_backend.backend.calculate_storage_offset(t)
+    assert orig_offset == 0
+    t = t.reshape(100)
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == orig_offset
+
+  def test_slice_values_match_torch(self):
+    torch_cpu = torch.arange(100, dtype=torch.float32).reshape(10, 10)
+    torch_tiny = torch_cpu.to(device)
+    sliced_cpu = torch_cpu[2:5, 3:7]
+    sliced_tiny = torch_tiny[2:5, 3:7]
+    np.testing.assert_equal(sliced_tiny.cpu().numpy(), sliced_cpu.numpy())
+
+  def test_slice_values_match_torch_3d(self):
+    torch_cpu_3d = torch.arange(210, dtype=torch.float32).reshape(5, 6, 7)
+    torch_tiny_3d = torch_cpu_3d.to(device)
+    sliced_cpu_3d = torch_cpu_3d[1:3, 2:4, 3:5]
+    sliced_tiny_3d = torch_tiny_3d[1:3, 2:4, 3:5]
+    np.testing.assert_equal(sliced_tiny_3d.cpu().numpy(), sliced_cpu_3d.numpy())
+
+  def test_topk_out(self):
+    a = torch.tensor([1, 3, 2, 4], device=device)
+    values = torch.empty(2, device=device)
+    indices = torch.empty(2, dtype=torch.int64, device=device)
+    ret_values, ret_indices = torch.topk(a, k=2, out=(values, indices))
+    np.testing.assert_equal(values.cpu().numpy(), [4, 3])
+    np.testing.assert_equal(indices.cpu().numpy(), [3, 1])
+    assert ret_values is values
+    assert ret_indices is indices
+
+  def test_sort_out(self):
+    a = torch.tensor([3, 1, 4, 2], device=device)
+    values = torch.empty(4, device=device)
+    indices = torch.empty(4, dtype=torch.int64, device=device)
+    ret_values, ret_indices = torch.sort(a, out=(values, indices))
+    np.testing.assert_equal(values.cpu().numpy(), [1, 2, 3, 4])
+    np.testing.assert_equal(indices.cpu().numpy(), [1, 3, 0, 2])
+    assert ret_values is values
+    assert ret_indices is indices
+
+  def test_cat_out(self):
+    a = torch.tensor([1, 2], device=device)
+    b = torch.tensor([3, 4], device=device)
+    out = torch.empty(4, device=device)
+    ret = torch.cat([a, b], out=out)
+    np.testing.assert_equal(out.cpu().numpy(), [1, 2, 3, 4])
+    assert ret is out
+
+  def test_scatter_add_out(self):
+    src = torch.tensor([[1, 2, 3], [4, 5, 6]], device=device, dtype=torch.float32)
+    index = torch.tensor([[0, 1, 2], [0, 1, 2]], device=device)
+    input = torch.zeros(3, 3, device=device, dtype=torch.float32)
+    out = torch.zeros(3, 3, device=device, dtype=torch.float32)
+    ret = torch.scatter_add(input, 0, index, src, out=out)
+    expected = torch.tensor([[5, 0, 0], [0, 7, 0], [0, 0, 9]], dtype=torch.float32)
+    np.testing.assert_allclose(out.cpu().numpy(), expected.cpu().numpy())
+    assert ret is out
+
+  def test_floor_divide_inplace_identity(self):
+    x = torch.tensor([10, 20, 30, 40], dtype=torch.int32, device=device)
+    y = torch.tensor([2, 4, 5, 8], dtype=torch.int32, device=device)
+    ret = x.floor_divide_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [5, 5, 6, 5])
+
+  def test_lshift_inplace_identity(self):
+    x = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
+    ret = x.__ilshift__(2)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
+
+  def test_rshift_inplace_identity(self):
+    x = torch.tensor([16, 32, 48, 64], dtype=torch.int32, device=device)
+    ret = x.__irshift__(2)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
+
+  def test_relu_inplace_identity(self):
+    x = torch.tensor([-1.0, 2.0, -3.0, 4.0], device=device)
+    ret = x.relu_()
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
+
+  def test_random_inplace_identity(self):
+    x = torch.zeros(10, dtype=torch.int32, device=device)
+    ret = x.random_()
+    assert ret is x
+    assert x.shape == (10,)
+
+  def test_random_from_inplace_identity(self):
+    x = torch.zeros(10, dtype=torch.int32, device=device)
+    ret = x.random_(5, 10)
+    assert ret is x
+    # values should be in range [5, 10)
+    assert torch.all(x >= 5).item() and torch.all(x < 10).item()
+
+  def test_uniform_inplace_identity(self):
+    x = torch.zeros(10, device=device)
+    ret = x.uniform_(0.0, 1.0)
+    assert ret is x
+    # values should be in range [0, 1)
+    assert torch.all(x >= 0.0).item() and torch.all(x < 1.0).item()
+
+  def test_normal_inplace_identity(self):
+    x = torch.zeros(100, device=device)
+    ret = x.normal_(0.0, 1.0)
+    assert ret is x
+    # just check that values changed from zeros
+    assert not torch.all(x == 0.0).item()
+
+  def test_logical_or_inplace_identity(self):
+    x = torch.tensor([True, False, True, False], device=device)
+    y = torch.tensor([False, False, True, True], device=device)
+    ret = x.logical_or_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [True, False, True, True])
+
+  def test_masked_fill_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    mask = torch.tensor([True, False, True, False], device=device)
+    ret = x.masked_fill_(mask, 0.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
+
+  def test_masked_fill_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    mask = torch.tensor([True, False, True, False], device=device)
+    value = torch.tensor(99.0, device=device)
+    ret = x.masked_fill_(mask, value)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [99.0, 2.0, 99.0, 4.0])
+
+  def test_zero_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.zero_()
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 0.0, 0.0, 0.0])
+
+  def test_fill_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.fill_(5.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [5.0, 5.0, 5.0, 5.0])
+
+  def test_fill_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    value = torch.tensor(7.0, device=device)
+    ret = x.fill_(value)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [7.0, 7.0, 7.0, 7.0])
+
+  def test_add_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    y = torch.tensor([10.0, 20.0, 30.0, 40.0], device=device)
+    ret = x.add_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [11.0, 22.0, 33.0, 44.0])
+
+  def test_add_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.add_(10.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [11.0, 12.0, 13.0, 14.0])
+
+  def test_mul_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    y = torch.tensor([2.0, 3.0, 4.0, 5.0], device=device)
+    ret = x.mul_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [2.0, 6.0, 12.0, 20.0])
+
+  def test_mul_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.mul_(2.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [2.0, 4.0, 6.0, 8.0])
+
 if __name__ == "__main__":
   unittest.main()

extra/torch_backend/test_kernel_fusion.py

@@ -0,0 +1,144 @@
+# simple tests
+import unittest
+import torch
+import warnings
+from tinygrad.helpers import getenv, GlobalCounters
+if getenv("TINY_BACKEND2"):
+  import extra.torch_backend.backend2
+  device = "cpu"
+else:
+  import extra.torch_backend.backend
+  device = "tiny"
+
+
+class TestKernelFusionRegression(unittest.TestCase):
+  def _realize(self, t): _ = t.detach().cpu().numpy()
+
+  def _check_kernel_count(self, fn, expected_kernels):
+    torch.manual_seed(42)
+    GlobalCounters.reset()
+    fn().detach().cpu().numpy()
+    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
+    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
+    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
+
+  def test_elementwise_fusion(self):
+    def fn():
+      x = torch.randn(128, 128, device=device)
+      return (x + 1.0) * 2.0 - 0.5
+    self._check_kernel_count(fn, 6)
+
+  def test_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 3, 32, 32, device=device)
+      conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
+      with torch.no_grad():
+        return torch.nn.functional.relu(conv(x))
+    self._check_kernel_count(fn, 8)
+
+  def test_batchnorm_fusion(self):
+    def fn():
+      x = torch.randn(2, 3, 16, 16, device=device)
+      conv = torch.nn.Conv2d(3, 8, 3, padding=1).to(device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.eval()
+      with torch.no_grad():
+        return torch.nn.functional.relu(bn(conv(x)))
+    self._check_kernel_count(fn, 16)
+
+  def test_reduce_fusion(self):
+    def fn():
+      x = torch.randn(64, 64, device=device)
+      return (x * 2.0).sum()
+    self._check_kernel_count(fn, 7)
+
+  def test_matmul_elementwise_fusion(self):
+    def fn():
+      x = torch.randn(32, 32, device=device)
+      w = torch.randn(32, 32, device=device)
+      return torch.nn.functional.relu(x @ w + 1.0)
+    self._check_kernel_count(fn, 6)
+
+  def test_pooling_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      return torch.nn.functional.max_pool2d(x * 2.0, 2)
+    self._check_kernel_count(fn, 5)
+
+  def test_residual_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      identity = torch.randn(1, 8, 16, 16, device=device)
+      out = x + identity
+      return torch.nn.functional.relu(out)
+    self._check_kernel_count(fn, 6)
+
+  def test_inplace_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 16, 32, 32, device=device)
+      y = torch.randn(1, 16, 32, 32, device=device)
+      x += y
+      return torch.nn.functional.relu(x)
+    self._check_kernel_count(fn, 6)
+
+  def test_conv_bn_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      identity = torch.randn(1, 8, 16, 16, device=device)
+      conv = torch.nn.Conv2d(8, 8, 3, padding=1, bias=False).to(device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.eval()
+      with torch.no_grad():
+        out = bn(conv(x))
+        out += identity
+        return torch.nn.functional.relu(out)
+    self._check_kernel_count(fn, 16)
+
+  def test_multiple_inplace_ops_fusion(self):
+    def fn():
+      x = torch.randn(64, 64, device=device)
+      x += 1.0
+      x *= 2.0
+      return torch.nn.functional.relu(x)
+    self._check_kernel_count(fn, 4)
+
+  def test_view_inplace_no_fusion_break(self):
+    def fn():
+      x = torch.randn(4, 64, device=device)
+      view = x[1:3]
+      view += 1.0
+      return x.sum()
+    self._check_kernel_count(fn, 8)
+
+  def test_batchnorm_running_stats_update(self):
+    def fn():
+      x = torch.randn(2, 8, 8, 8, device=device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.train()
+      with torch.no_grad():
+        return bn(x)
+    self._check_kernel_count(fn, 10)
+
+  # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
+  def test_mnist_training_fusion(self):
+    def fn():
+      model = torch.nn.Sequential(
+        torch.nn.Conv2d(1, 8, 3, padding=1),
+        torch.nn.ReLU(),
+        torch.nn.MaxPool2d(2),
+        torch.nn.Flatten(),
+        torch.nn.Linear(8*14*14, 10)
+      ).to(device)
+      optimizer = torch.optim.Adam(model.parameters(), 1e-3)
+      x = torch.randn(32, 1, 28, 28, device=device)
+      labels = torch.randint(0, 10, (32,), device=device)
+      out = model(x)
+      loss = torch.nn.functional.cross_entropy(out, labels)
+      optimizer.zero_grad()
+      loss.backward()
+      optimizer.step()
+      return loss
+    self._check_kernel_count(fn, 33)
+
+if __name__ == "__main__":
+  unittest.main()

extra/torch_backend/wrapped_tensor.cpp

@@ -113,16 +113,9 @@ int register_hook() {
 int temp_register_hook = register_hook();
 
 at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceIndex device_index) {
-  // TODO: we have to get the dtype and the shape from the tinygrad Tensor
   std::vector<int64_t> sizes = py_obj.attr("shape").cast<std::vector<int64_t>>();
-
-  py::list views = py_obj.attr("uop").attr("st").attr("views");
-  std::vector<int64_t> strides = views[views.size() - 1].attr("strides").cast<std::vector<int64_t>>();
-  int64_t storage_offset = 0;
-  for (auto& v: views) {
-    storage_offset += v.attr("offset").cast<int64_t>(); // TODO: is this correct?
-  }
-
+  std::vector<int64_t> strides = py_obj.attr("_strides").cast<std::vector<int64_t>>();
+  int64_t storage_offset = py_obj.attr("_storage_offset").cast<int64_t>();
   return at::detail::make_tensor<at::TinyOpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>>(
     at::DispatchKeySet(at::DispatchKey::PrivateUse1),
     c10::scalarTypeToTypeMeta(dtype),

mkdocs.yml

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

sz.py

@@ -4,6 +4,8 @@ import token
 import tokenize
 import itertools
 from tabulate import tabulate
+from tinygrad.uop import Ops
+from tinygrad.helpers import ContextVar
 
 TOKEN_WHITELIST = [token.OP, token.NAME, token.NUMBER, token.STRING]
 
@@ -79,11 +81,15 @@ if __name__ == "__main__":
       print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", floatfmt=".1f")+"\n")
       groups = sorted([('/'.join(x[0].rsplit("/", 1)[0].split("/")[0:2]), x[1], x[2]) for x in table])
       dir_sizes = {}
-      for dir_name, group in itertools.groupby(groups, key=lambda x:x[0]):
+      for dir_name, _group in itertools.groupby(groups, key=lambda x:x[0]):
+        group = list(_group)
         dir_sizes[dir_name] = sum([x[1] for x in group])
-        print(f"{dir_name:30s} : {dir_sizes[dir_name]:6d}")
-      print(f"\n core line count: {sum([v for k,v in dir_sizes.items() if k not in NONCORE_DIRS])}")
+        print(f"{dir_name:30s} : {dir_sizes[dir_name]:6d} in {len(group):2d} files")
+      print()
+      print(f"        ops: {len(Ops)}")
+      print(f"      flags: {len(ContextVar._cache)}")
+      print(f" core lines: {sum([v for k,v in dir_sizes.items() if k not in NONCORE_DIRS])}")
       total_lines = sum([x[1] for x in table])
-      print(f"total line count: {total_lines}")
+      print(f"total lines: {total_lines}")
       max_line_count = int(os.getenv("MAX_LINE_COUNT", "-1"))
       assert max_line_count == -1 or total_lines <= max_line_count, f"OVER {max_line_count} LINES"

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_benchmark_pyrender.py

@@ -0,0 +1,34 @@
+# benchmark speed of pyrender for all created UOps saved with TRACK_MATCH_STATS=2
+import functools, pickle
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.helpers import tqdm, temp, time_to_str, cpu_profile
+
+BENCHMARK_OPS = {Ops.INDEX, Ops.BUFFERIZE}
+
+@functools.cache
+def create_uop(a:int) -> UOp:
+  op, dtype, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
+
+if __name__ == "__main__":
+  # load rewrite trace
+  with open(temp("rewrites.pkl", append_user=True), "rb") as f:
+    trace = pickle.load(f)
+
+  # benchmark
+  result:list[tuple[str, int]] = []
+  try:
+    for steps in tqdm(trace.rewrites):
+      for r in steps:
+        for _,yn,_,__ in r.matches:
+          y = create_uop(yn)
+          if y.op in BENCHMARK_OPS:
+            with cpu_profile("pyrender") as e:
+              try: ren = y.render()
+              except Exception: ren = "PYRENDER_ERR"
+            result.append((ren, float(e.en-e.st)/1e6))
+  finally:
+    N = 10
+    print(f"Slowst {N} renders from {len(result)} samples:")
+    for ren,tm in sorted(result, key=lambda x:x[1], reverse=True)[:N]:
+      print(f"{time_to_str(tm).strip():<10s} {ren}")

test/external/external_test_example.py

@@ -1,13 +1,14 @@
 import unittest
 from tinygrad import Device
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import getenv, CI, OSX
 
 def multidevice_test(fxn):
   exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
   def ret(self):
     for device in Device._devices:
-      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"]: continue
+      # broken on OSX USB AMD, why?
+      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
       if not CI: print(device)
       if device in exclude_devices:
         if not CI: print(f"WARNING: {device} test is excluded")

test/external/external_test_onnx_backend.py

@@ -184,9 +184,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,7 +1,9 @@
 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
+from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
+from tinygrad.uop.divandmod import fold_divmod_general
 from test.test_tiny import TestTiny
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
@@ -67,8 +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/fuzz_symbolic_symbolic_div.py

@@ -0,0 +1,66 @@
+import random
+import z3
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.validate import uops_to_z3
+from tinygrad.helpers import DEBUG, Context, colored
+
+seed = random.randint(0, 100)
+print(f"Seed: {seed}")
+random.seed(seed)
+
+def get_random_term(ranges, factors):
+  # 10% chance of nesting
+  if random.randint(0,9) == 0: return get_random_expr(ranges, factors)
+  return random.choice(ranges)*random.choice(factors)*random.choice([1, 1, 1, -1])
+
+def get_random_expr(ranges, factors):
+  num_terms = random.randint(2,4)
+  x = UOp.sum(*[get_random_term(ranges, factors) for _ in range(num_terms)])
+  return x.alu(random.choice([Ops.IDIV, Ops.MOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
+
+if __name__ == "__main__":
+  skipped = 0
+  for i in range(700):
+    if i % 100 == 0:
+      print(f"Running test {i}")
+    upper_bounds = [*list(range(1, 4)), 16, 33, 53, 64, 256]
+    variable_names = ["i", "j", "k"]
+    variables = [UOp.variable(s, 1, random.choice(upper_bounds)) for s in variable_names]
+    factors = variables+upper_bounds
+    # add some products
+    for _ in range(2): factors.append(random.choice(variables)*random.choice(variables))
+    # add some adds
+    for _ in range(2): factors.append(random.choice(variables)+random.choice(factors))
+    num_ranges = 4
+    ranges = [UOp.range(random.choice(factors), i) for i in range(num_ranges)]
+    variable_names += [f"r{i}" for i in range(num_ranges)]
+    expr = get_random_expr(ranges, factors)
+
+    with Context(CORRECT_DIVMOD_FOLDING=1):
+      simplified_expr = expr.simplify()
+
+    if DEBUG>=1:
+      print(expr.render(simplify=False), "  -->  ", simplified_expr.render(simplify=False))
+
+    solver = z3.Solver()
+    solver.set(timeout=3000)  # some expressions take very long verify, but its very unlikely they actually return sat
+    z3_expr, z3_simplified_expr, *z3_vars = uops_to_z3(solver, expr, simplified_expr, *variables, *ranges)
+    check = solver.check(z3_simplified_expr != z3_expr)
+    if check == z3.unknown and DEBUG>=1:
+      skipped += 1
+      print("skipped z3 verification due to timeout")
+    elif check == z3.sat:
+      print(colored("simplify INCORRECT!", "red"))
+      print(solver.model())
+      var_vals = {s:solver.model()[z] for s,z in zip(variable_names, z3_vars)}
+      print("reproduce with:")
+      print("var_vals = ", var_vals)
+      print("globals = var_vals|{'cdiv':cdiv,'cmod':cmod}")
+      print("expr = ast.simplify()")
+      print("assert eval(ast.render(pm=renderer_infer, simplify=False),globals) == eval(expr.render(pm=renderer_infer, simplify=False),globals)")
+      print()
+
+      assert False
+
+    if DEBUG >= 2: print(f"validated {expr.render()}")
+  print(f"Skipped {skipped} expressions due to timeout")

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))
@@ -37,6 +36,8 @@ def trunc_log(x):
 
 # user config
 SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+# uncomment this to disable by default
+#SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
 if REF == "master": SKIP_PROCESS_REPLAY = True
 class ProcessReplayWarning(Warning): pass
 
@@ -50,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
@@ -65,7 +64,10 @@ def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)
 
-replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_rangeify_map":replay_get_rangeify_map, "get_program":replay_get_program}
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
+replayers["get_program"] = replay_get_program
+# disable this for speed, does it ever find things?
+#replayers["get_rangeify_map"] = replay_get_rangeify_map
 
 # *** run replayers on captured rows and print diffs
 

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/amd/amdgpu.py

@@ -124,6 +124,7 @@ class PM4Executor(AMDQueue):
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
       else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")
+    elif mem_data_sel == 0: pass # no write
     else: raise RuntimeError(f"Unknown {mem_data_sel=}")
 
   def _exec_copy_data(self, n):

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_onnx.py

@@ -58,6 +58,18 @@ class TestOnnxModel(unittest.TestCase):
     print(cls, _LABELS[cls])
     assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"
 
+  def test_pad_list_value(self):
+    from tinygrad.nn.onnx import onnx_ops
+    from tinygrad import Tensor
+    Pad = onnx_ops['Pad']
+    x = Tensor([1, 2, 3])
+    out = Pad(x, pads=[0, 1], value=[-float('inf')])
+    assert out.shape == (4,)
+    assert out.numpy()[-1] == -float('inf')
+    out2 = Pad(x, pads=[1, 0], constant_value=[5.0])
+    assert out2.shape == (4,)
+    assert out2.numpy()[0] == 5.0
+
 @unittest.skipUnless(Device.DEFAULT == "METAL", "only run on METAL")
 class TestHuggingFaceOnnxModels(unittest.TestCase):
   @classmethod

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

@@ -1,6 +1,7 @@
 import unittest
 import pathlib
 from examples.whisper import init_whisper, load_file_waveform, transcribe_file, transcribe_waveform
+import examples.mlperf.metrics as metrics
 from tinygrad.helpers import CI, fetch, CPU_LLVM
 from tinygrad import Device, dtypes
 from tinygrad.device import is_dtype_supported
@@ -14,7 +15,39 @@ TEST_FILE_2 = str(pathlib.Path(__file__).parent / "whisper/test2.wav")
 TRANSCRIPTION_2 = "a slightly longer audio file so that we can test batch transcriptions of varying length."
 # TODO this file will possibly not survive long. find another 1-2 minute sound file online to transcribe
 TEST_FILE_3_URL = 'https://homepage.ntu.edu.tw/~karchung/miniconversations/mc45.mp3'
-TRANSCRIPTION_3 = "Just lie back and relax. Is the level of pressure about right? Yes, it's fine, and I'd like conditioner please. Sure. I'm going to start the second lathering now. Would you like some Q-tips? How'd you like it cut? I'd like my bangs and the back trimmed, and I'd like the rest thinned out a bit and layered. Where would you like the part? On the left, right about here. Here, have a look. What do you think? It's fine. Here's a thousand anti-dollars. It's 30-ant extra for the rants. Here's your change and receipt. Thank you, and please come again. So how do you like it? It could have been worse, but you'll notice that I didn't ask her for her card. Hmm, yeah. Maybe you can try that place over there next time."   # noqa: E501
+TRANSCRIPTION_3 = """Just lie back and relax.
+Is the level of pressure about right?
+Yes, it's fine. And I'd like conditioner, please.
+Sure. I'm going to start the second lathering now.
+Would you like some Q-tips?
+How'd you like it cut?
+I'd like my bangs and the back trimmed,
+and I'd like the rest thinned out a bit and layered.
+Where would you like the part?
+On the left, right about here.
+Here, have a look. What do you think?
+It's fine. Here's thousand NT dollars.
+It's 30 NT extra for the rinse. Here's your change and receipt.
+Thank you, and please come again!
+So, how do you like it?
+It could have been worse. But you'll notice that I didn't ask her for her card.
+Hmm, yeah.
+Mm, maybe you can try that place over there next time."""
+
+TRANSCRIPTION_3_ALT = "Just lie back and relax. Is the level of pressure about right? Yes, it's fine. And I'd like conditioner please. Sure. I'm going to start the second lathering now. Would you like some Q-tips? How'd you like it cut? I'd like my bangs on the back trimmed, and I'd like the rest to stand out a bit and layered. Where would you like the part? On the left, right about here. Here. Have a look. What do you think? It's fine. Here's a thousand and eighty dollars. It's thirty and t extra for the rants. Here's your change and receipt. Thank you, and please come again. So how do you like it? It could have been worse, but you'll notice that I didn't ask her for her card. Hmm, yeah. Maybe you can try that place over there next time." #noqa: E501
+# NOTE: same as TRANSCRIPTION_3 but with minor changes that should only amount to ~0.079 WER difference (see test_wer_same)
+# 'and'     --> 'on'
+# 'thinned' --> 'to stand'
+# 'nt'      --> 'and eighty'
+# '30 nt'   --> 'thirty and t'
+# 'rinse'   --> 'rants'
+# 'mm'      --> ''
+
+def wer_helper(result: str, reference: str)->float:
+  result = metrics.normalize_string(result)
+  reference = metrics.normalize_string(reference)
+  wer, _, _ = metrics.word_error_rate([result], [reference])
+  return wer
 
 @unittest.skipIf(Device.DEFAULT in ["CPU"], "slow")
 @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support")
@@ -30,6 +63,15 @@ class TestWhisper(unittest.TestCase):
     del cls.model
     del cls.enc
 
+  def assertWER(self, actual: str, expected: str, threshold: float):
+    __tracebackhide__ = True  # Hide traceback for py.test
+    wer = wer_helper(actual, expected)
+    if wer > threshold:
+      err = f"WER={wer:.3f} > {threshold}"
+      raise AssertionError(
+        err
+      )
+
   def test_transcribe_file1(self):
     self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)
 
@@ -52,20 +94,38 @@ 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.assertEqual(TRANSCRIPTION_3, transcription)
+    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)]
 
     trancriptions = transcribe_waveform(self.model, self.enc, waveforms)
     self.assertEqual(2, len(trancriptions))
-    self.assertEqual(TRANSCRIPTION_3, trancriptions[0])
+    self.assertWER(trancriptions[0], TRANSCRIPTION_3, 0.085)
     self.assertEqual(TRANSCRIPTION_1, trancriptions[1])
 
+  def test_wer_same(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER(TRANSCRIPTION_3_ALT, reference, 0.079)
+
+  def test_wer_different(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER("[no speech]", reference, 1.0)
+
+  def test_wer_different_2(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER("", reference, 1.0)
+
+  def test_wer_different_3(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER(reference[:len(reference)//2], reference, 0.524)
+
 if __name__ == '__main__':
   unittest.main()

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_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_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"

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)
@@ -765,6 +768,16 @@ class TestMultiTensor(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       Tensor.rand_like(t, device=(d3, d4))
 
+  def test_full_like_on_shard(self, axis=None):
+    t = Tensor.empty((16, 16)).shard(devices_2, axis=axis)
+    t2 = Tensor.full_like(t, 1.0)
+    self.assertEqual(t.shape, t2.shape)
+    self.assertEqual(t.device, t2.device)
+    self.assertEqual(t.dtype, t2.dtype)
+    self.assertEqual(t.uop.axis, t2.uop.axis)
+    t2.realize()
+  def test_full_like_on_shard_axis(self): self.test_full_like_on_shard(0)
+
   def test_dropout_on_shard(self):
     with Tensor.train():
       X = Tensor.ones(256).to(devices_2)
@@ -817,6 +830,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()
@@ -841,6 +855,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()
@@ -962,6 +979,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
@@ -1116,9 +1136,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))
 
@@ -1161,6 +1183,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
@@ -1170,6 +1193,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()
@@ -1232,6 +1258,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

@@ -2699,6 +2699,9 @@ class TestOps(unittest.TestCase):
     a = Tensor(3.14)
     np.testing.assert_allclose(Tensor.stack(a, a).numpy(), Tensor([3.14, 3.14]).numpy())
 
+  def test_stack_max(self):
+    helper_test_op(None, lambda x, y: torch.stack((x, y)).max(axis=0)[0], lambda x, y: Tensor.stack(x, y).max(axis=0), vals=[[1.], [2.]])
+
   def test_repeat(self):
     x = Tensor.randn(4, 6, 3)
     base_repeats = [2, 4, 3]
@@ -2726,6 +2729,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))

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_outerworld.py

@@ -1,5 +1,6 @@
 import unittest
-from tinygrad import Tensor, UOp
+import numpy as np
+from tinygrad import Tensor, UOp, nn
 from tinygrad.uop.ops import AxisType, Ops
 
 class TestOuterworldReduce(unittest.TestCase):
@@ -50,8 +51,37 @@ class TestOuterRange(unittest.TestCase):
     # 3 matmuls with outer world range
     i = UOp.range(3, -100, AxisType.OUTER)
     vec_i = Tensor(vec.uop.after(i))
-    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
-    out = Tensor(vec.uop.after(vec_i.uop.store((vec_i.contiguous() @ mats[vi]).uop).end(i)))
+    comp = vec_i.contiguous() @ mats[i]
+    store = vec_i.uop.store(comp.uop).end(i)
+    out = Tensor(vec.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+class TestOuterScan(unittest.TestCase):
+  def _test_scan(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
+
+    # 3 matmuls in "scan"
+    vec1 = vec @ mats[0]
+    vec2 = vec1 @ mats[1]
+    vec3 = vec2 @ mats[2]
+    ref = Tensor.stack(vec1, vec2, vec3)
+    ref.realize()
+    return vec, mats, ref
+
+  def test_uop_scan_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with SCAN
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(3, 1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
+    comp = phi @ mats[i]
+    store = out[i].uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
     out.realize()
 
     # TODO: testing allclose
@@ -116,5 +146,85 @@ class TestOuterworld(unittest.TestCase):
     out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
     self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
+class TestVmap(unittest.TestCase):
+  def test_vmap_inner(self, axis_type=AxisType.LOOP, fuse=False, grad=False):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
+
+    ref = x @ mats
+    if fuse: ref = ref * 2
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, axis_type)
+    out = x @ mats[a]
+    out = out.reshape(1, 10).pad(((a,(3-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    if fuse: out = out * 2
+    if grad:
+      out.mean().backward()
+      np.testing.assert_allclose(mats.grad.numpy(), (2./30) if fuse else (1./30))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
+  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
+  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
+
+  def test_vmap_inner_grad(self): self.test_vmap_inner(grad=True)
+  def test_vmap_inner_fuse_grad(self): self.test_vmap_inner(fuse=True, grad=True)
+  def test_vmap_outer_grad(self): self.test_vmap_inner(AxisType.OUTER, grad=True)
+
+  def test_vmap_convs(self):
+    layers = [
+      nn.Conv2d(1, 8, 3), Tensor.relu,
+      nn.Conv2d(8, 8, 3), Tensor.relu]
+    img = Tensor.randn(4, 1, 16, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None, None, None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  def test_vmap_gemm(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    img = Tensor.randn(4, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  @unittest.skip("this is broken, we need to lower the outer reduce in the outer graph")
+  def test_vmap_gemm_grad(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    layer_tensors = nn.state.get_parameters(layers)
+    img = Tensor.randn(4, 16).realize(*layer_tensors)
+    for l in layer_tensors: l.requires_grad_()
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    out_grads = [x.numpy() for x in grads]
+
+    # compute reference grads
+    for l in layer_tensors: l.grad = None
+    img.sequential(layers).mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    ref_grads = [x.numpy() for x in grads]
+
+    # compare
+    for o,r in zip(out_grads, ref_grads): np.testing.assert_allclose(o, r, atol=1e-6)
+
 if __name__ == '__main__':
   unittest.main()

test/test_pickle.py

@@ -20,7 +20,7 @@ class TestPickle(unittest.TestCase):
     self.assertEqual(pm2.rewrite(sink).key, tt.key)
 
   def test_pickle_main_pattern_matcher(self):
-    from tinygrad.codegen.late.devectorizer import sym
+    from tinygrad.uop.symbolic import sym
     ssym = pickle.dumps(sym)
     dsym = pickle.loads(ssym)
     self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

test/test_profiler.py

@@ -17,7 +17,7 @@ def helper_collect_profile(*devs):
   cpu_events.clear()
 
   profile_list = []
-  with Context(VIZ=1):
+  with Context(VIZ=1, PROFILE=1):
     yield profile_list
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()
@@ -199,7 +199,7 @@ class TestProfiler(unittest.TestCase):
     #self.assertLess(e1.st, e2.st)
     #self.assertGreater(e1.en-e1.st, e2.en-e2.st)
 
-  @unittest.skipIf(not CI, "this test is flaky locally")
+  @unittest.skip("this test is flaky")
   @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
   def test_graph(self):
     from test.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs

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)
@@ -672,33 +669,6 @@ class TestSchedule(unittest.TestCase):
     c = (a.sum(2).contiguous() + b).contiguous()
     check_schedule(c, 2)
 
-  def test_kernelize(self):
-    a = Tensor.empty(10)
-    b = Tensor.empty(10)
-    c = (a+b).kernelize()
-    d = c+2
-    check_schedule(d, 2)
-
-  def test_kernelize_view(self):
-    a = Tensor.empty(4,1)
-    b = a*2
-    c = b.kernelize()+Tensor.empty(4,4)
-    check_schedule(c, 2)
-
-  def test_kernelize_diamond(self):
-    a = Tensor([0]).realize()
-    prev_a = (a+1).contiguous()
-    a.assign(Tensor([2]))
-    a.kernelize(prev_a)
-    self.assertEqual((prev_a+a*3).item(), 1+2*3)
-
-  def test_kernelize_sym(self):
-    a = Tensor([1])+Tensor([2])
-    a.kernelize()
-    b = a/a
-    check_schedule(b, 0)
-    self.assertEqual(b.item(), 1)
-
   # TODO: this requires supporting multiple stores in the AST
   @unittest.expectedFailure
   def test_multioutput_ast(self):
@@ -710,35 +680,6 @@ class TestSchedule(unittest.TestCase):
     self.assertEqual(a.buffer.numpy(), [7])
     self.assertEqual(b.buffer.numpy(), [12])
 
-  # unlike schedule, kernelize can be called multiple times on a Tensor
-  def test_double_kernelize(self):
-    a = Tensor.empty(10)
-    b = Tensor.empty(10)
-    c = (a+b)
-    d = c.kernelize()+2
-    e = c.kernelize()+d.kernelize()
-    check_schedule(e, 3)
-
-  def test_kernelize_bw(self):
-    a = Tensor.full((3,), 2.0, requires_grad=True).contiguous()
-    b = Tensor.full((3,), 3.0, requires_grad=True).contiguous()
-    x = (a*b).kernelize()
-    y = Tensor.eye(3, requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-    self.assertEqual(z.item(), 18.0)
-    self.assertEqual(z.grad.item(), 1.0)
-
-  def test_kernelize_bw_view(self):
-    a = Tensor.full((3,1), 2.0, requires_grad=True).contiguous()
-    b = Tensor.full((3,1), 3.0, requires_grad=True).contiguous()
-    x = (a*b).kernelize()
-    y = Tensor.eye(6, requires_grad=True)
-    z = y.matmul(x.expand(3,2).reshape(6)).sum()
-    z.backward()
-    self.assertEqual(z.item(), 36.0)
-    self.assertEqual(z.grad.item(), 1.0)
-
   @unittest.skip("no longer supported")
   def test_double_from(self):
     x = Tensor([1,2,3,4])
@@ -1915,18 +1856,6 @@ class TestSchedule(unittest.TestCase):
     for X in range(1,N): root = root + bufs[X][vi] + bufs[X][vj]
     self.assertEqual(root.item(), N * 2)
 
-  def test_limit_bufs_kernelize(self):
-    N = 31
-    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
-      bufs = [Tensor(i).contiguous().realize() for i in range(N)]
-    x = bufs[0]
-    for y in bufs[1:]: x = x+y
-    x.kernelize()
-    kcount = len([s for s in x.uop.toposort() if s.op is Ops.KERNEL])
-    z = x+Tensor.empty(1) # z only loads 2 buffers
-    sched = z.schedule()
-    self.assertEqual(len(sched), kcount+1)
-
 class TestSwizzle(unittest.TestCase):
   def test_swizzle_simple(self):
     Tensor.manual_seed(0)
@@ -2118,7 +2047,7 @@ class TestCopyFolding(unittest.TestCase):
     b = Tensor.empty(4, device="CPU")
     add = a+b
     assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
-    add.kernelize()
+    add.schedule()
 
   def test_alu_before_copy(self):
     buf = Tensor.ones(1).contiguous().realize()
@@ -2438,5 +2367,12 @@ class TestUOpBecome(unittest.TestCase):
     b.shrink(((0,4),)).assign(a_view).realize()
     self.assertListEqual(b.tolist(), [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
+class TestSimpleSchedule(unittest.TestCase):
+  def test_reduce_doesnt_split(self):
+    a = Tensor.empty(16,16).sum(axis=1)
+    a1 = a.reshape(4,4)
+    a2 = a.reshape(16,1,1)
+    self.assertEqual(len(Tensor.schedule(a1, a2)), 1)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

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

@@ -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_tiny.py

@@ -35,9 +35,9 @@ class TestTiny(unittest.TestCase):
     out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.zeros(8).contiguous())
     self.assertListEqual(out.tolist(), [1]*8+[0]*8)
 
-  def test_sum(self):
-    out = Tensor.ones(256).contiguous().sum()
-    self.assertEqual(out.item(), 256)
+  def test_sum(self, N=getenv("SUM_N", 256)):
+    out = Tensor.ones(N).contiguous().sum()
+    self.assertEqual(out.item(), N)
 
   def test_gemm(self, N=getenv("GEMM_N", 64), out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()

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)
@@ -517,7 +518,7 @@ class TestUOpStr(unittest.TestCase):
 
 class TestUPatHelpers(unittest.TestCase):
   def test_location(self):
-    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "math.py")
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
     self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])

test/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,22 +1,27 @@
 import unittest, math
 
 from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.helpers import CI
-from tinygrad.renderer.ptx import PTXRenderer
 import numpy as np
 
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
+from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout
 
-@unittest.skipUnless(Device.DEFAULT in ["CUDA", "NV"], "only cuda")
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "no ptx")
+@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 = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    N = 8192
+    BLOCK_SIZE = 64
+    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,11 +30,10 @@ 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)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
+      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
 
       col, row = ker.blockIdx_x, ker.blockIdx_y
 
@@ -39,13 +43,12 @@ class TestTK(unittest.TestCase):
         b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
 
         a_reg = warp.load(a_reg, a_smem)
-        b_reg = warp.load(b_reg, b_smem, transpose=True)
+        b_reg = warp.load(b_reg, b_smem)
 
         c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
       c_reg = ker.endrange()
 
-      c_smem = warp.store(c_smem, c_reg)
-      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+      c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
 
       sink = ker.finish()
 
@@ -65,27 +68,26 @@ class TestTK(unittest.TestCase):
 
   @unittest.skipIf(CI, "no wmma in ci")
   def test_simple_matmul_transposed(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    N = 8192
+    BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128
+    with Kernel("simple_matmul_transposed", (N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       c = ker.gl((1, 1, N, N), dtypes.float32)
       a = ker.gl((1, 1, N, N), dtypes.bfloat16)
       b = ker.gl((1, 1, N, N), dtypes.bfloat16)
 
-      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      a_smem = ker.st((BLOCK_N, BLOCK_K), dtypes.bfloat16, base_shape=ST_16X32)
+      b_smem = ker.st((BLOCK_M, BLOCK_K), dtypes.bfloat16, base_shape=ST_16X32)
 
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      a_reg = ker.rt((BLOCK_N, BLOCK_K), dtypes.bfloat16, base_shape=RT_16X32)
+      b_reg = ker.rt((BLOCK_M, BLOCK_K), dtypes.bfloat16, base_shape=RT_16X32)
+      c_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL, base_shape=RT_16X16)
 
       col, row = ker.blockIdx_x, ker.blockIdx_y
 
       c_reg = warp.zero(c_reg)
-      for tile in ker.range(N // BLOCK_SIZE):
+      for tile in ker.range(N // BLOCK_K):
         a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
         b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
 
@@ -95,8 +97,7 @@ class TestTK(unittest.TestCase):
         c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
       c_reg = ker.endrange()
 
-      c_smem = warp.store(c_smem, c_reg)
-      c = warp.store(c, c_smem, (0, 0, row, col), (), axis=2)
+      c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
 
       sink = ker.finish()
 
@@ -115,16 +116,15 @@ class TestTK(unittest.TestCase):
     np.testing.assert_allclose(c.numpy(), ref.numpy())
 
   def test_load_store(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    N = 64
+    BLOCK_SIZE = 32
+    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)
       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)
@@ -134,8 +134,45 @@ class TestTK(unittest.TestCase):
       a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
       a_reg = warp.load(a_reg, a_smem)
       b_reg = warp.copy(b_reg, a_reg)
-      b_smem = warp.store(b_smem, b_reg)
-      b = warp.store(b, b_smem, (0, 0, row, col), (), axis=2)
+      b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float()
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+
+  @unittest.skip("TODO")
+  def test_load_store_group(self):
+    N = 256
+    BLOCK_SIZE = 64
+    with Kernel("load_store_group", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
+      warp = ker.warp
+      group = ker.group(2)
+
+      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)
+
+      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 = group.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 = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
 
       sink = ker.finish()
 
@@ -153,9 +190,9 @@ class TestTK(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), ref.numpy())
 
   def test_add(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("add", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -172,8 +209,7 @@ class TestTK(unittest.TestCase):
 
           a_reg += 1
 
-          a_smem = warp.store(a_smem, a_reg)
-          b = warp.store(b, a_smem, (0, 0, tile_row, tile_col), (), axis=2)
+          b = warp.store(b, a_reg, (0, 0, tile_row, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -191,36 +227,34 @@ class TestTK(unittest.TestCase):
       np.testing.assert_allclose(b.numpy(), ref.numpy())
 
   def test_max(self):
-    N = 16
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("max", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
       a = ker.gl((1, 1, N, N), dtypes.float32)
 
       a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
 
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
 
-      max_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+      max_reg = ker.rv(BLOCK_SIZE, dtypes.float32)
 
-      for tile_row in ker.range(N // BLOCK_SIZE):
-        max_reg = warp.neg_inf(max_reg.after(tile_row))
+      for tile_col in ker.range(N // BLOCK_SIZE):
+        max_reg = warp.neg_inf(max_reg.after(tile_col))
 
-        for tile_col in ker.range(N // BLOCK_SIZE):
+        for tile_row in ker.range(N // BLOCK_SIZE):
           a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
           a_reg = warp.load(a_reg, a_smem)
-          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+          max_reg = warp.col_reduce(max_reg, a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
         max_reg = ker.endrange()
 
-        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
+        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[1], 0])
 
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+        for tile_row in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -233,41 +267,39 @@ class TestTK(unittest.TestCase):
     for _ in range(5): ei.run(wait=True)
     b = b.float()
 
-    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+    ref = a.float().max(axis=2, keepdim=True).expand(a.shape)
 
     np.testing.assert_allclose(b.numpy(), ref.numpy())
 
   def test_max_nonsquare(self):
-    N, M = 16, 64
+    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)
       a = ker.gl((1, 1, N, M), dtypes.float32)
 
       a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
 
-      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
+      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
 
-      max_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
+      max_reg = ker.rv(BLOCK_M, dtypes.float32)
 
-      for tile_row in ker.range(N // BLOCK_N):
-        max_reg = warp.neg_inf(max_reg.after(tile_row))
+      for tile_col in ker.range(M // BLOCK_M):
+        max_reg = warp.neg_inf(max_reg.after(tile_col))
 
-        for tile_col in ker.range(M // BLOCK_M):
+        for tile_row in ker.range(N // BLOCK_N):
           a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
           a_reg = warp.load(a_reg, a_smem)
-          max_reg = warp.row_reduce(max_reg, a_reg, lambda a, b: a.maximum(b))
+          max_reg = warp.col_reduce(max_reg, a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
         max_reg = ker.endrange()
 
-        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
+        b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[1], 0])
 
-        for tile_col in ker.range(M // BLOCK_M):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+        for tile_row in ker.range(N // BLOCK_N):
+          b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -280,41 +312,39 @@ class TestTK(unittest.TestCase):
     for _ in range(5): ei.run(wait=True)
     b = b.float()
 
-    ref = a.float().max(axis=3, keepdim=True).expand(a.shape)
+    ref = a.float().max(axis=2, keepdim=True).expand(a.shape)
 
     np.testing.assert_allclose(b.numpy(), ref.numpy())
 
   def test_sum(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("sum", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
       a = ker.gl((1, 1, N, N), dtypes.float32)
 
       a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
 
-      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
-      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
 
-      sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+      sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32)
 
-      for tile_row in ker.range(N // BLOCK_SIZE):
-        sum_reg = warp.zero(sum_reg.after(tile_row))
+      for tile_col in ker.range(N // BLOCK_SIZE):
+        sum_reg = warp.zero(sum_reg.after(tile_col))
 
-        for tile_col in ker.range(N // BLOCK_SIZE):
+        for tile_row in ker.range(N // BLOCK_SIZE):
           a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
           a_reg = warp.load(a_reg, a_smem)
-          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+          sum_reg = warp.col_reduce(sum_reg, a_reg, lambda a, b: a + b)
         sum_reg = ker.endrange()
 
-        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[1], 0])
 
-        for tile_col in ker.range(N // BLOCK_SIZE):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+        for tile_row in ker.range(N // BLOCK_SIZE):
+          b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -327,41 +357,39 @@ class TestTK(unittest.TestCase):
     for _ in range(5): ei.run(wait=True)
     b = b.float()
 
-    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+    ref = a.float().sum(axis=2, keepdim=True).expand(a.shape)
 
     np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
 
   def test_sum_nonsquare(self):
-    N, M = 16, 64
+    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)
       a = ker.gl((1, 1, N, M), dtypes.float32)
 
       a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
 
-      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
-      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32)
+      a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
+      b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
 
-      sum_reg = ker.rv(BLOCK_N, dtypes.float32, "ortho")
+      sum_reg = ker.rv(BLOCK_M, dtypes.float32)
 
-      for tile_row in ker.range(N // BLOCK_N):
-        sum_reg = warp.zero(sum_reg.after(tile_row))
+      for tile_col in ker.range(M // BLOCK_M):
+        sum_reg = warp.zero(sum_reg.after(tile_col))
 
-        for tile_col in ker.range(M // BLOCK_M):
+        for tile_row in ker.range(N // BLOCK_N):
           a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
           a_reg = warp.load(a_reg, a_smem)
-          sum_reg = warp.row_reduce(sum_reg, a_reg, lambda a, b: a + b)
+          sum_reg = warp.col_reduce(sum_reg, a_reg, lambda a, b: a + b)
         sum_reg = ker.endrange()
 
-        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[0], 0, (idx[2]%4)//2])
-        b_smem = warp.store(b_smem, b_reg)
+        b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[1], 0])
 
-        for tile_col in ker.range(M // BLOCK_M):
-          b = warp.store(b, b_smem, (0, 0, tile_row, tile_col), (), axis=2)
+        for tile_row in ker.range(N // BLOCK_N):
+          b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -374,15 +402,14 @@ class TestTK(unittest.TestCase):
     for _ in range(5): ei.run(wait=True)
     b = b.float()
 
-    ref = a.float().sum(axis=3, keepdim=True).expand(a.shape)
+    ref = a.float().sum(axis=2, keepdim=True).expand(a.shape)
 
     np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
 
-  @unittest.skip("fake range not ended")
   def test_softmax(self):
-    N = 32
-    BLOCK_SIZE = 16
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("softmax", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
@@ -392,37 +419,37 @@ class TestTK(unittest.TestCase):
 
       a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
 
-      max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-      max_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
-      norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32, "ortho")
+      max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32)
+      max_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
+      norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
 
       max_vec = warp.neg_inf(max_vec)
       norm_vec = warp.zero(norm_vec)
 
       for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
+        a_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, a_reg, lambda a, b: a.maximum(b))
-        a_reg = (a_reg - max_vec).exp2()
+        max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
+        a_reg_ = (a_reg_ - max_vec).exp2()
         max_vec_last = (max_vec_last - max_vec).exp2()
         norm_vec *= max_vec_last
-        norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
+        norm_vec = warp.row_reduce(norm_vec, a_reg_, lambda a, b: a + b)
       norm_vec = ker.endrange()
+      max_vec = max_vec.after(norm_vec)
 
       for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
+        a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg_ = warp.load(a_reg, 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
 
-        a_smem = warp.store(a_smem, a_reg)
-        b = warp.store(b, a_smem, (0, 0, 0, tile_col), (), axis=2)
+        b = warp.store(b, a_reg_, (0, 0, 0, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -439,5 +466,179 @@ class TestTK(unittest.TestCase):
 
     np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
 
+  def test_softmax_col(self):
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("softmax_col", (1, 1, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
+      a = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+
+      max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32)
+      max_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
+      norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
+
+      max_vec = warp.neg_inf(max_vec)
+      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_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_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 = 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_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)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, BLOCK_SIZE, dtype="float32")
+      b = Tensor.empty(1, 1, N, BLOCK_SIZE, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float().softmax(axis=2)
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy(), atol=1e-5, rtol=1e-5)
+
+  def test_fa(self):
+    NUM_WORKERS = 1
+    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("fa", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
+      warp = ker.warp
+
+      # kernel
+      o = ker.gl((B, N, H, D), dtypes.bfloat16)
+      q = ker.gl((B, N, H, D), dtypes.bfloat16)
+      k = ker.gl((B, N, H_KV, D), dtypes.bfloat16)
+      v = ker.gl((B, N, H_KV, D), dtypes.bfloat16)
+
+      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)
+      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)
+
+        # mask for causal
+        q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % 16)
+        kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // 16) * 4
+        att_block = warp.map(att_block,
+                             lambda x, idx: ((kv_base + idx[0]*16 + idx[2]) > (q_base + idx[1]*16)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
+
+        # 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)
+
+      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)
+
+      sink = ker.finish()
+
+    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()
+      out = Tensor.empty(B, N, H, D, dtype=dtypes.bfloat16)
+      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)
+    out = out.float()
+
+    q_permuted = q.permute(0, 2, 1, 3)
+    k_permuted = k.permute(0, 2, 1, 3)
+    v_permuted = v.permute(0, 2, 1, 3)
+    ref = q_permuted.scaled_dot_product_attention(k_permuted, v_permuted, is_causal=True, enable_gqa=True).float()
+    ref = ref.permute(0, 2, 1, 3)
+
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_autogen.py

@@ -5,16 +5,16 @@ from tinygrad.runtime.support.c import Struct
 class TestAutogen(unittest.TestCase):
   def test_packed_struct_sizeof(self):
     layout = [('a', ctypes.c_char), ('b', ctypes.c_int, 5), ('c', ctypes.c_char)]
-    class X(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Y(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Z(Struct): _packed_, _fields_ = True, layout
-    self.assertNotEqual(ctypes.sizeof(X), 4) # ctypes bug! gcc-13.3.0 says this should have size 4
+    class Z(Struct): pass
+    Z._packed_, Z._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Y), 6)
     self.assertEqual(ctypes.sizeof(Z), 3)
     layout = [('a', ctypes.c_int, 31), ('b', ctypes.c_int, 31), ('c', ctypes.c_int, 1), ('d', ctypes.c_int, 1)]
     class Foo(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Bar(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Baz(Struct): _fields_, _packed_ = layout, True
+    class Baz(Struct): pass
+    Baz._packed_, Baz._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Foo), 12)
     self.assertEqual(ctypes.sizeof(Bar), 12)
     self.assertEqual(ctypes.sizeof(Baz), 8)
@@ -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

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_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])