fixed and test is real

* viz: tabulate pmc

* linter

* enable nesting

* pmc comes before waves

.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
@@ -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/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
@@ -263,7 +261,9 @@ jobs:
     - name: Check Device.DEFAULT
       run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
     - name: Run unit tests
-      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+      run: |
+        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
     - name: Run targetted tests on NULL backend
       run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
     # TODO: too slow
@@ -447,7 +447,7 @@ jobs:
         with:
           key: onnxoptl
           deps: testing
-          pydeps: "tensorflow==2.15.1 tensorflow_addons"
+          pydeps: "tensorflow==2.19"
           python-version: '3.11'
           opencl: 'true'
       - name: Test ONNX (CL)

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.
 

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

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

extra/assembly/rdna3/generate.py

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

extra/assembly/rdna3/parse.py

@@ -0,0 +1,15 @@
+from tinygrad import Tensor, nn
+import xml.etree.ElementTree as ET
+
+if __name__ == "__main__":
+  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
+  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
+  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
+  root = ET.fromstring(xml_str)
+
+  for op_el in root.findall("./ISA/OperandTypes/OperandType"):
+    op_name = op_el.findtext("OperandTypeName")
+    val_dict = {}
+    for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
+      val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
+    print(op_name, val_dict)

extra/gemm/amd_uop_matmul.py

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

extra/hevc/.gitignore

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

extra/hevc/decode.py

@@ -0,0 +1,83 @@
+import argparse, os, hashlib
+from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
+from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
+from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--input_file", type=str, default="")
+  parser.add_argument("--output_dir", type=str, default="extra/hevc/out")
+  args = parser.parse_args()
+
+  if args.input_file == "":
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    hevc_tensor = Tensor.from_url(url, device="CPU")
+  else:
+    hevc_tensor = Tensor.empty(os.stat(args.input_file).st_size, dtype=dtypes.uint8, device=f"disk:{args.input_file}").to("CPU")
+
+  dat = bytes(hevc_tensor.data())
+  dat_hash = hashlib.md5(dat).hexdigest()
+
+  with Timing("prep infos: "):
+    dat_nv = hevc_tensor.to("NV")
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
+
+  frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
+
+  # move all needed data to gpu
+  #all_slices = []
+  with Timing("copy to gpu: "):
+    opaque_nv = opaque.to("NV").contiguous().realize()
+    hevc_tensor = hevc_tensor.to("NV")
+
+  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+  max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)
+
+  # define variables
+  v_pos = Variable("pos", 0, max_hist + 1)
+  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
+  v_sz = Variable("sz", 0, hevc_tensor.numel())
+  v_i = Variable("i", 0, len(frame_info)-1)
+
+  @TinyJit
+  def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
+    return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
+
+  # warm up
+  history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
+  for i in range(3):
+    hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
+    decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
+
+  out_images = []
+  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
+    for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
+      history = history[-max_hist:] if max_hist > 0 else []
+      # TODO: this shrink should work as a slice
+      hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
+
+      outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
+      out_images.append(outimg)
+      if is_hist: history.append(outimg)
+
+    Device.default.synchronize()
+
+  if getenv("VALIDATE", 0):
+    import pickle
+    if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
+      url = "https://github.com/nimlgen/hevc_validate_set/raw/refs/heads/main/decoded_frames_b813bfdbec194fd17fdf0e3ceb8cea1c.pkl"
+      decoded_frames = pickle.load(fetch(url).open("rb"))
+    else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
+  else: import cv2
+
+  for i, img in tqdm(enumerate(out_images)):
+    if getenv("VALIDATE", 0):
+      if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
+        img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
+        assert img.data() == decoded_frames[i], f"Frame {i} does not match reference decoder!"
+        print(f"Frame {i} matches reference decoder!")
+    else:
+      if len(args.output_dir):
+        os.makedirs(args.output_dir, exist_ok=True)
+        img = to_bgr(img, h, w, luma_w, chroma_off).realize()
+        cv2.imwrite(f"{args.output_dir}/out_frame_{i:04d}.png", img.numpy())

extra/hevc/hevc.py

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

extra/hip_gpu_driver/hip_ioctl.py

@@ -66,7 +66,7 @@ def ioctl(fd, request, argp):
     print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
     if name == "AMDKFD_IOC_SVM":
       out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
-      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.kfd_ioctl_svm_attr_type__enumvalues[out[i].type]:40s}: {out[i].value:#x}")
+      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
   else:
     print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
           f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")

extra/mmapeak/mmapeak.py

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

extra/nv_gpu_driver/clc9b0.h

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

extra/qcom_gpu_driver/qcom_opencl_interop.py

@@ -29,8 +29,9 @@ rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw
 
 # create QCOM tensor with the externally managed buffer
 x = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
-y = (x + 1).numpy()
-print(y)
+y = (x + 1).reshape(-1).tolist()
+print(y[:10])
+assert y == [i + 1 for i in range(64)]
 
 # all calculations are done, save to free the object
 cl.clReleaseMemObject(cl_buf)
@@ -49,7 +50,7 @@ for i in range(4):
   cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
   rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20]
 
-  y = calc(x = Tensor.from_blob(rawbuf_ptr, (2, 2), dtype=dtypes.int, device='QCOM')).numpy()
+  y = calc(x = Tensor.from_blob(rawbuf_ptr, (2, 2), dtype=dtypes.int, device='QCOM')).tolist()
   print(f'jit {i}\n', y)
 
   # all calculations are done, save to free the object
@@ -80,8 +81,19 @@ rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw
 # dtypes.imageh = cl.cl_image_format(cl.CL_RGBA, cl.CL_HALF_FLOAT)
 # dtypes.imagef = cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT)
 x = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
-y = (x + 1).numpy()
-print(y)
+y = (x + 1).tolist()
+print(y[:10])
 
 # all calculations are done, save to free the object
 cl.clReleaseMemObject(cl_img)
+
+# from numpy
+import numpy as np
+
+YUV_SIZE = 50
+a_np = (32*np.random.randn(YUV_SIZE).astype(np.float32) + 128).clip(0,255).astype(np.uint8)
+a = Tensor.from_blob(a_np.ctypes.data, (YUV_SIZE,), dtype=dtypes.uint8, device='QCOM').realize()
+
+print(a.numpy()[:10], a_np[:10])
+assert np.all(a.numpy() == a_np)
+assert np.all((a - 1).numpy() == a_np - 1)

extra/remu/src/helpers.rs

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

extra/remu/src/thread.rs

@@ -1024,7 +1024,7 @@ impl<'a> Thread<'a> {
                     let vdst = (instr & 0xff) as usize;
                     let abs = ((instr >> 8) & 0x7) as usize;
                     let opsel = ((instr >> 11) & 0xf) as usize;
-                    let cm = (instr >> 15) & 0x1;
+                    let cm = ((instr >> 15) & 0x1) != 0;
 
                     let s = |n: usize| ((instr >> n) & 0x1ff) as usize;
                     let src = (s(32), s(41), s(50));
@@ -1032,7 +1032,9 @@ impl<'a> Thread<'a> {
                     let omod = (instr >> 59) & 0x3;
                     let neg = ((instr >> 61) & 0x7) as usize;
                     assert_eq!(omod, 0);
-                    assert_eq!(cm, 0);
+                    if op != 272 && cm {
+                        return todo_instr!(op); // TODO: add VOP3 clamp for all ops
+                    }
                     assert_eq!(opsel, 0);
 
                     match op {
@@ -1266,7 +1268,7 @@ impl<'a> Thread<'a> {
                             }
 
                             let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 543 | 551 | 567 | 606 | 796 => {
                                     let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                     let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                     let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1275,12 +1277,26 @@ impl<'a> Thread<'a> {
                                         260 => s0 - s1,
                                         261 => s1 - s0,
                                         264 => s0 * s1,
-                                        272 => f32::max(s0, s1),
+                                        272 => f32::max(s0, s1).clmp(cm),
                                         299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                         426 => s0.recip(),
                                         430 => 1.0 / f32::sqrt(s0),
                                         531 => f32::mul_add(s0, s1, s2),
                                         537 => f32::min(f32::min(s0, s1), s2),
+                                        543 => {
+                                            if s0.is_nan() || s1.is_nan() || s2.is_nan() {
+                                                f32::min(f32::min(s0, s1), s2)
+                                            } else {
+                                                let max = f32::max(f32::max(s0, s1), s2);
+                                                if max == s0 {
+                                                    f32::max(s1, s2)
+                                                } else if max == s1 {
+                                                    f32::max(s0, s2)
+                                                } else {
+                                                    f32::max(s0, s1)
+                                                }
+                                            }
+                                        },
                                         540 => f32::max(f32::max(s0, s1), s2),
                                         551 => s2 / s1,
                                         567 => {
@@ -1290,6 +1306,7 @@ impl<'a> Thread<'a> {
                                                 false => ret,
                                             }
                                         }
+                                        606 => f32::min(f32::max(s0, s1), s2),
                                         796 => s0 * 2f32.powi(s1.to_bits() as i32),
                                         // cnd_mask isn't a float only ALU but supports neg
                                         257 => {

extra/sqtt/attempt_sqtt_parse.py

@@ -119,7 +119,6 @@ OPNAME = {
   0xb: "VALU",
   0xd: "VALU",
   0xe: "VALU",
-  0x10: "__END",
   0x21: "VMEM_LOAD",
   0x22: "VMEM_LOAD",
   0x24: "VMEM_STORE",
@@ -480,6 +479,7 @@ def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -
         nib = (byte >> (offset & 4)) & 0xF
         reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
         offset += 4
+      if offset != target: break  # don't parse past the end
 
     # 2) Decode token from low 8 bits
     opcode = STATE_TO_OPCODE[reg & 0xFF]

extra/sqtt/rgptool.py

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

extra/sqtt/roc.py

@@ -44,9 +44,11 @@ class WaveSlot:
   simd:int
   se:int
   @property
-  def simd_loc(self) -> str: return f"SE:{self.se} CU:{self.cu} SIMD:{self.simd}"
+  def cu_loc(self) -> str: return f"SE:{self.se} CU:{self.cu}"
   @property
-  def wave_loc(self) -> str: return f"{self.simd_loc} WAVE:{self.wave_id}"
+  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):
@@ -65,12 +67,14 @@ class OccEvent(WaveSlot):
   time:int
   start:int
 
+RunKey = tuple[str, int]
+
 class _ROCParseCtx:
   def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
     self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
     self.disasms:dict[str, dict[int, tuple[str, int]]] = {}
-    self.inst_execs:dict[str, list[WaveExec]] = {}
-    self.occ_events:dict[str, list[OccEvent]] = {}
+    self.inst_execs:dict[RunKey, list[WaveExec]] = {}
+    self.occ_events:dict[RunKey, list[OccEvent]] = {}
 
     for prog in prog_evs:
       arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
@@ -79,14 +83,14 @@ class _ROCParseCtx:
 
   def next_sqtt(self):
     x = next(self.sqtt_evs, None)
-    self.active_kern = x.kern if x is not None else None
+    self.active_run = (x.kern, x.exec_tag) if x is not None else None
     self.active_se = x.se if x is not None else None
     self.active_blob = (ctypes.c_ubyte * len(x.blob)).from_buffer_copy(x.blob) if x is not None else None
     return self.active_blob
 
   def on_occupancy_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_occupancy_t):
     if DEBUG >= 5: print(f"OCC {ev.time=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.wave_id=} {ev.start=}")
-    self.occ_events.setdefault(unwrap(self.active_kern), []).append(OccEvent(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.time, ev.start))
+    self.occ_events.setdefault(unwrap(self.active_run), []).append(OccEvent(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.time, ev.start))
 
   def on_wave_ev(self, ev:rocprof.rocprofiler_thread_trace_decoder_wave_t):
     if DEBUG >= 5: print(f"WAVE {ev.wave_id=} {self.active_se=} {ev.cu=} {ev.simd=} {ev.contexts=} {ev.begin_time=} {ev.end_time=}")
@@ -96,7 +100,7 @@ class _ROCParseCtx:
     insts_blob = bytearray(sz:=ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t))
     ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
 
-    self.inst_execs.setdefault(unwrap(self.active_kern), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
+    self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
                                                                              ev.end_time, insts_blob))
 
 def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
@@ -134,7 +138,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_kern)][pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run)[0]][pc.address]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
     if instr == "s_endpgm": mem_size_ptr[0] = 0
@@ -154,6 +158,15 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   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)))
@@ -166,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,7 +2,9 @@ import os
 os.environ["PYTHONPATH"] = "."
 os.environ["SQTT"] = "1"
 if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["VIZ"] = "1"
+os.environ["PROFILE"] = "1"
+# VIZ=1 to launch server
+# os.environ["VIZ"] = "1"
 os.environ["AMD_LLVM"] = "0"
 
 import unittest
@@ -129,5 +131,13 @@ class TestTiming(unittest.TestCase):
     for w in waves:
       print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
 
+  def test_ones(self):
+    N = getenv("N", 4096)
+    CNT = getenv("CNT", 2)
+    with save_sqtt() as sqtt:
+      for _ in range(CNT):
+        Tensor.ones(N, N).contiguous().realize()
+    self.assertEqual(len(sqtt), CNT)
+
 if __name__ == "__main__":
   unittest.main()

extra/thunder/tiny/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,7 +51,22 @@ 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 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)
@@ -60,20 +75,27 @@ class Group:
     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], axis_type=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))), ())
+          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)
@@ -83,20 +105,87 @@ class Group:
     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], axis_type=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))), ())
+          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)])
+          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)
 
-          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)]
+    self.ker.push_store(c_store, c)
+    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)
@@ -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], axis_type=AxisType.REDUCE, track=False):
-          for inner in self.ker.range(4, axis_type=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, axis_type=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 = 4
-  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 = 4
-  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,24 +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, start:int, end:int, step:int, axis_type:AxisType):
+  def __init__(self, start:int, end:int, step:int, axis_type:AxisType, rid:int):
     self.start, self.end, self.step = start, end, step
-    self.axis_type, self.done = axis_type, False
+    self.axis_type, self.rid, self.done = axis_type, rid, False
   def __iter__(self): return self
   def __next__(self):
     if not self.done:
       self.done = True
-      _tk_range.user_rid += 1
-      self._rng = UOp.range(self.end // self.step, _tk_range.user_rid-1, axis_type=self.axis_type) * self.step + self.start
+      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
       return self._rng
     raise StopIteration
 
 class Kernel(AbstractContextManager):
-  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+  def __init__(self, name:str, grid_size:tuple[int, int, int], block_size:int):
+    self.name = name
+
     self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
     self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
     self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
@@ -31,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
@@ -47,7 +50,8 @@ class Kernel(AbstractContextManager):
 
   def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
     if end == 0: start, end = 0, start
-    rng = _tk_range(start, end, step, axis_type)
+    rng = _tk_range(start, end, step, axis_type, self.range_id)
+    self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
@@ -72,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))
 
@@ -87,9 +91,9 @@ class Kernel(AbstractContextManager):
     if hasattr(last_store, '_uop'): uop = last_store._uop
     else: uop = last_store
 
-    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+    return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
 
   def endrange(self):
     last_store = self.store_stack.pop()
     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
 
@@ -66,7 +69,10 @@ 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 self.ruop(uop)
@@ -80,76 +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):
+  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):
-    return ST(uop, self.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):
-    return RT(uop, self.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, layout, ker):
+  def __init__(self, uop:UOp, layout:VecLayout, ker):
     self._uop, self.layout, self.ker = uop, layout, ker
 
-  def ruop(self, uop):
+  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)
+    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()

mkdocs.yml

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

test/device/test_amd_llvm.py

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

test/device/test_hcq.py

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

test/external/external_test_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,8 @@
 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
 
@@ -68,9 +69,12 @@ if __name__ == "__main__":
     t()
 
     # these caches will keep uops alive
+    schedule_cache.clear()
     method_cache.clear()
     apply_movement_op.cache_clear()
+    _apply_reshape.cache_clear()
     fold_divmod_general.cache_clear()
+    UOp.const.cache_clear()
     Tensor._device_seeds.clear()
     Tensor._device_rng_counters.clear()
 

test/external/mlperf_resnet/lars_optimizer.py

@@ -29,7 +29,6 @@ from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import linalg_ops
 from tensorflow.python.ops import math_ops
-from tensorflow.python.training import training_ops
 from tensorflow.python.ops import state_ops
 
 
@@ -147,20 +146,7 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
     return scaled_lr, grad
 
   def _apply_dense(self, grad, var, apply_state=None):
-    var_device, var_dtype = var.device, var.dtype.base_dtype
-    coefficients = ((apply_state or {}).get((var_device, var_dtype))
-                    or self._fallback_apply_state(var_device, var_dtype))
-
-    scaled_lr, grad = self.compute_lr(grad, var, coefficients)
-    mom = self.get_slot(var, "momentum")
-    return training_ops.apply_momentum(
-        var,
-        mom,
-        math_ops.cast(1.0, var.dtype.base_dtype),
-        grad * scaled_lr,
-        self.momentum,
-        use_locking=False,
-        use_nesterov=self.use_nesterov)
+    return self._resource_apply_dense(grad, var, apply_state)
 
   def _resource_apply_dense(self, grad, var, apply_state=None):
     var_device, var_dtype = var.device, var.dtype.base_dtype
@@ -194,13 +180,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                     or self._fallback_apply_state(var_device, var_dtype))
 
     mom = self.get_slot(var, "momentum")
-    return training_ops.sparse_apply_momentum(
-        var,
-        mom,
-        coefficients["learning_rate"],
-        grad.values,
-        grad.indices,
-        self.momentum,
+    return tf.raw_ops.SparseApplyMomentum(
+        var=var,
+        accum=mom,
+        lr=coefficients["learning_rate"],
+        grad=grad.values,
+        indices=grad.indices,
+        momentum=self.momentum,
         use_locking=False,
         use_nesterov=self.use_nesterov)
 
@@ -210,13 +196,13 @@ class LARSOptimizer(optimizer_v2.OptimizerV2):
                     or self._fallback_apply_state(var_device, var_dtype))
 
     mom = self.get_slot(var, "momentum")
-    return training_ops.resource_sparse_apply_keras_momentum(
-        var.handle,
-        mom.handle,
-        coefficients["learning_rate"],
-        grad,
-        indices,
-        self.momentum,
+    return tf.raw_ops.ResourceSparseApplyKerasMomentum(
+        var=var.handle,
+        accum=mom.handle,
+        lr=coefficients["learning_rate"],
+        grad=grad,
+        indices=indices,
+        momentum=self.momentum,
         use_locking=False,
         use_nesterov=self.use_nesterov)
 

test/external/process_replay/process_replay.py

@@ -14,7 +14,6 @@ try:
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen.opt import Opt
   from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
-  from tinygrad.device import Device
 except ImportError as e:
   print(repr(e))
   exit(int(ASSERT_DIFF))
@@ -36,9 +35,9 @@ def trunc_log(x):
   logging.info("\n".join(lines))
 
 # user config
-# NOTE: process replay is slow so it's now disabled by default. add [pr] to enable it
-#SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
-SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
+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
 
@@ -52,12 +51,10 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
   sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
   input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
-  # if no renderer was provided, open the device to get it
-  if renderer is None: renderer = Device[p.device].renderer
   p2 = get_program(input_ast, renderer=renderer)
   def to_str(ret:ProgramSpec) -> str:
     # PYTHON renderer pickles UOps, first unpickle and decode here
@@ -67,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/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

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

test/opt/test_gen_float4.py

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

test/test_arange.py

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

test/test_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)
@@ -827,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()
@@ -851,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()
@@ -972,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
@@ -1126,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))
 
@@ -1171,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
@@ -1180,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()
@@ -1242,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

@@ -2729,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_profiler.py

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

@@ -660,6 +660,18 @@ class TestUOpGraph(unittest.TestCase):
     bad_gate = UOp.const(dtypes.int, 1)
     with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, dtypes.void, (glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])
 
+  def test_after_end(self):
+    r = UOp.range(10, 0)
+
+    c = r + 1
+    self.assertIn(r, c.ranges)
+
+    e = UOp.const(dtypes.void, None).end(r)
+    self.assertNotIn(r, e.ranges)
+
+    a = c.after(e)
+    self.assertNotIn(r, a.ranges)
+
 @track_rewrites()
 def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)
 

test/test_uops.py

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

test/test_uops_stats.py

@@ -175,13 +175,13 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
-    p = get_program(self.ast_gemm, opts=[])
+    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
     except KernelOptError:
       raise unittest.SkipTest("no tensor cores")
     print(p.src)
@@ -190,18 +190,19 @@ class TestStatsOptimized(unittest.TestCase):
   # this is a good lesson about why UPCASTing is a good idea
 
   def test_gemm_one_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
-    p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
+                    opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
                                            Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
     except KernelOptError:
       raise unittest.SkipTest("no locals")
@@ -210,7 +211,7 @@ class TestStatsOptimized(unittest.TestCase):
 
   def test_gemm_group(self):
     try:
-      p = get_program(self.ast_gemm, opts=[Opt(OptOps.GROUP, 0, 4)])
+      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 4)])
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
@@ -219,14 +220,14 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
-    p = get_program(self.ast_reduce, opts=[])
+    p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
     print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
     self.assertEqual(p.estimates.ops, N*N)
     self.assertEqual(p.estimates.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
     try:
-      p = get_program(self.ast_reduce, opts=[Opt(OptOps.GROUP, 0, 50)])
+      p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 50)])
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     # NOTE: these are wrong, they don't respect the if statement

test/testextra/test_hevc.py

@@ -0,0 +1,65 @@
+import unittest
+
+from tinygrad import Device
+from tinygrad.helpers import fetch
+from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
+
+class TestHevc(unittest.TestCase):
+  def test_hevc_parser(self):
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat, device=Device.DEFAULT)
+
+    def _test_common(frame, bts):
+      self.assertEqual(frame0.pic_width_in_luma_samples, 1952)
+      self.assertEqual(frame0.pic_height_in_luma_samples, 1216)
+      self.assertEqual(frame0.chroma_format_idc, 1)
+      self.assertEqual(frame0.bit_depth_luma, 8)
+      self.assertEqual(frame0.bit_depth_chroma, 8)
+      self.assertEqual(frame0.log2_min_luma_coding_block_size, 3)
+      self.assertEqual(frame0.log2_max_luma_coding_block_size, 5)
+      self.assertEqual(frame0.log2_min_transform_block_size, 2)
+      self.assertEqual(frame0.log2_max_transform_block_size, 5)
+      self.assertEqual(frame0.num_tile_columns, 3)
+      self.assertEqual(frame0.num_tile_rows, 1)
+      self.assertEqual(frame0.colMvBuffersize, 589)
+      self.assertEqual(frame0.HevcSaoBufferOffset, 2888)
+      self.assertEqual(frame0.HevcBsdCtrlOffset, 25992)
+      self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcFltAboveOffset, 26714)
+      self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcSaoAboveOffset, 36214)
+
+      # tiles
+      self.assertEqual(bytes(bts[0x200:0x210]), b'\x18\x00&\x00\x18\x00&\x00\r\x00&\x00\x00\x00\x00\x00')
+
+    frame0 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[0].data())
+    _test_common(frame0, opaque[0].data())
+    self.assertEqual(frame0.stream_len, 148063)
+    self.assertEqual(frame0.IDR_picture_flag, 1)
+    self.assertEqual(frame0.RAP_picture_flag, 1)
+    self.assertEqual(frame0.sw_hdr_skip_length, 0)
+    self.assertEqual(frame0.num_ref_frames, 0)
+
+    frame1 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[1].data())
+    _test_common(frame1, opaque[1].data())
+    self.assertEqual(frame1.stream_len, 57110)
+    self.assertEqual(frame1.IDR_picture_flag, 0)
+    self.assertEqual(frame1.RAP_picture_flag, 0)
+    self.assertEqual(frame1.sw_hdr_skip_length, 9)
+    self.assertEqual(frame1.num_ref_frames, 1)
+    self.assertEqual(list(frame1.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame1.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame1.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+
+    frame3 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[3].data())
+    _test_common(frame3, opaque[3].data())
+    self.assertEqual(frame3.stream_len, 47036)
+    self.assertEqual(frame3.IDR_picture_flag, 0)
+    self.assertEqual(frame3.RAP_picture_flag, 0)
+    self.assertEqual(frame3.sw_hdr_skip_length, 9)
+    self.assertEqual(frame3.num_ref_frames, 1)
+    self.assertEqual(list(frame3.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+    self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+if __name__ == "__main__":
+  unittest.main()

test/testextra/test_tk.py

@@ -1,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.skipIf(CI and Device.DEFAULT not 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

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

test/unit/test_indexing.py

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

test/unit/test_kernelize.py

@@ -1,37 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.uop import Ops
-
-class TestKernelize(unittest.TestCase):
-  def test_add_reshaped(self):
-    a = Tensor.ones(16,16).contiguous()
-    b = Tensor.zeros(16,16).contiguous()
-    ret = (a+b).sum(axis=1)
-    ret_reshaped_1 = ret.reshape(4,4)
-    ret_reshaped_2 = ret.reshape(2,8)
-    ret.kernelize()
-    self.assertIs(ret_reshaped_1.uop.src[0], ret_reshaped_2.uop.src[0])
-
-  def test_two_reduce(self):
-    a = Tensor.ones(16,16).contiguous()
-    a1 = a.sum(axis=1)
-    a0 = a1.sum(axis=0)
-    a0.kernelize()
-    self.assertEqual(len([s for s in a0.uop.toposort() if s.op is Ops.KERNEL]), 2)
-    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
-    # input Tensor and user contiguous kernelize
-    self.assertIs(a0.uop.base.op, Ops.AFTER)
-    self.assertIs(a.uop.base.op, Ops.AFTER)
-
-  def test_two_reduce_w_add(self):
-    a = Tensor.ones(16,16).contiguous()
-    a1 = a.sum(axis=1)
-    a0 = (a1+1).sum(axis=0)
-    a0.kernelize()
-    # NOTE: the +1 is fused with a1, so a1 is not kernelized
-    self.assertIs(a1.uop.base.op, Ops.REDUCE_AXIS)
-    # the input to the REDUCE_AXIS is an ASSIGN though
-    self.assertIs(a1.uop.base.src[0].base.op, Ops.AFTER)
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_microbenchmarks.py

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

test/unit/test_schedule_cache.py

@@ -0,0 +1,24 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.engine.schedule import schedule_cache
+
+class TestScheduleCache(unittest.TestCase):
+  def test_simple(self):
+    a = Tensor.ones(10).contiguous()
+    b = Tensor.ones(10).contiguous()
+    Tensor.realize(a, b)
+
+    # warm up
+    for _ in range(2):
+      num = (a.sum().contiguous()+b.sum().contiguous()).item()
+      print(num)
+
+    # confirm schedule cache doesn't grow
+    start_len_schedule_cache = len(schedule_cache)
+    for _ in range(3):
+      num = (a.sum().contiguous()+b.sum().contiguous()).item()
+      print(num)
+    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_shm_tensor.py

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

test/unit/test_simple_schedule.py

@@ -1,15 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.uop.ops import Ops
-
-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)
-    Tensor.kernelize(a1, a2)
-    kernels = [x for x in a1.uop.sink(a2.uop).toposort() if x.op is Ops.KERNEL]
-    self.assertEqual(len(kernels), 1)
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_tinyfs.py

@@ -3,19 +3,19 @@ from tinygrad import Tensor
 
 class TestLoadStore(unittest.TestCase):
   def test_load_shape(self):
-    t = Tensor(bytes(16)).fs_load(1024).kernelize()
+    t = Tensor(bytes(16)).fs_load(1024)
     assert t.shape == (1024,), t.shape
 
   def test_store_shape(self):
-    t = Tensor.zeros(1024).fs_store().kernelize()
+    t = Tensor.zeros(1024).fs_store()
     assert t.shape == (16,), t.shape
 
   def test_load_large_shape(self):
-    t = Tensor(bytes(16)).fs_load(10_000_000).kernelize()
+    t = Tensor(bytes(16)).fs_load(10_000_000)
     assert t.shape == (10_000_000,), t.shape
 
   def test_store_large_shape(self):
-    t = Tensor.zeros(10_000_000).fs_store().kernelize()
+    t = Tensor.zeros(10_000_000).fs_store()
     assert t.shape == (16,), t.shape
 
 if __name__ == "__main__":

test/unit/test_uop_symbolic.py

@@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
-from tinygrad.uop.symbolic import sym, commutative
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -112,8 +112,9 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 1, 8)
     self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
     self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
+    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2))
     self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
-    self.assertEqual((a*a*3).divide_exact(a*a.const_like(-3)).simplify(), a*-1)
+    self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1)
     self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
     self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
     self.assertEqual((a*b*3+a*b*b).divide_exact(a*b).simplify(), b+3)
@@ -322,12 +323,12 @@ class TestSymbolic(unittest.TestCase):
   def test_mod_mod_wrong_sign(self):
     v1=Variable("v1", 0, 128)
     v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -3, 4, "(v1%2*2+(v3+-1)%5+-2)")
 
   def test_mod_mod_wrong_sign2(self):
     v2=Variable("v2", 0, 8)
     v3=Variable("v3", 0, 4)
-    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")
+    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -2, 6, "(((v2+((v3+3)%7))+-2)%7)")
 
   def test_mul_mul(self):
     self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
@@ -377,9 +378,9 @@ class TestSymbolic(unittest.TestCase):
   def test_big_mod(self):
     self.helper_test_variable(Variable("a", -20, 20)%10, -9, 9, "(a%10)")
     self.helper_test_variable(Variable("a", -20, 0)%10, -9, 0, "(((a*-1)%10)*-1)")
-    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 9, "(a%10)")  # TODO: tighter max
+    self.helper_test_variable(Variable("a", -20, 1)%10, -9, 1, "(a%10)")
     self.helper_test_variable(Variable("a", 0, 20)%10, 0, 9, "(a%10)")
-    self.helper_test_variable(Variable("a", -1, 20)%10, -9, 9, "(a%10)")  # TODO: tighter min
+    self.helper_test_variable(Variable("a", -1, 20)%10, -1, 9, "(a%10)")
 
   def test_ge_remove(self):
     self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "False")
@@ -580,12 +581,12 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")
     self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")
 
-  @unittest.expectedFailure  # TODO: improve nest_div_by_smallest_factor
   def test_sum_div_complex4(self):
     gidx0 = Variable("gidx0", 0, 2)
     lidx2 = Variable("lidx2", 0, 12)
     lidx3 = Variable("lidx3", 0, 12)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")
+    # TODO: improve nest_div_by_smallest_factor to get ((lidx2+(lidx3*2))//3)
+    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((gidx0+(lidx2*19+lidx3*38)//3)//19)")
 
   def test_sum_mul_distribute(self):
     gidx0 = Variable("gidx0", 0, 7)
@@ -1029,6 +1030,25 @@ class TestSymbolicRealWorld(unittest.TestCase):
     self.assertIn(idx.render(),
       ("(lidx3+((lidx5+1)//16*802816+(lidx5+1)%16*49+gidx0*3211264+gidx1*784+gidx2*8+lidx4*100352)+2207744)",))
 
+class TestGatedUopGivenValid(unittest.TestCase):
+  def test_invalid_gate_simplifies_index(self):
+    r0 = Variable("r0", 0, 2)
+
+    idx:UOp = (r0 < 3).where((r0 + uconst(-1)) // uconst(3), UOp.invalid())
+    idx = graph_rewrite(idx, pm_simplify_valid)
+    self.assertEqual(idx, (r0 < 3).where(uconst(0), UOp.invalid()))
+
+  def test_invalid_gate_simplifies_vectorize(self):
+    r0 = Variable("r0", 0, 2)
+
+    idx0 = (r0 + uconst(-1)) // uconst(3)
+    idx1 = r0 % uconst(3)
+    idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.index.vec(2), (idx0, idx1)), UOp.invalid())
+    idx = graph_rewrite(idx, pm_simplify_valid)
+    # NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
+    expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
+    self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
+
 class TestBounds(unittest.TestCase):
   def test_unrolled_arange(self):
     # #include <metal_stdlib>

test/unit/test_viz.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatch
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
-from tinygrad.helpers import VIZ
+from tinygrad.helpers import VIZ, cpu_profile
 from tinygrad.device import Buffer
 
 @track_rewrites(name=True)
@@ -262,14 +262,6 @@ from tinygrad import Tensor, Device
 from tinygrad.engine.realize import get_program
 
 class TestVizIntegration(BaseTestViz):
-  # kernelize has a custom name function in VIZ
-  def test_kernelize_tracing(self):
-    a = Tensor.empty(4, 4)
-    Tensor.kernelize(a+1, a+2)
-    lst = get_viz_list()
-    self.assertEqual(len(lst), 1)
-    self.assertEqual(lst[0]["name"], "Schedule 2 Kernels n1")
-
   # codegen supports rendering of code blocks
   def test_codegen_tracing(self):
     ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
@@ -284,7 +276,7 @@ class TestVizIntegration(BaseTestViz):
       a = Tensor.empty(1)
       b = Tensor.empty(1)
       metadata = (alu:=a+b).uop.metadata
-      alu.kernelize()
+      alu.schedule()
       graph = next(get_viz_details(0, 0))["graph"]
     self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
 
@@ -367,7 +359,7 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
       for _ in range(event_count):
         alloc, ts, key = u("<BII")
         if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIBB") for _ in range(u("<I")[0])]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 class TestVizProfiler(BaseTestViz):
@@ -415,8 +407,8 @@ class TestVizProfiler(BaseTestViz):
 
     tracks = list(j['layout'])
     self.assertEqual(tracks[0], 'NV')
-    self.assertEqual(tracks[1], 'NV:1')
-    self.assertEqual(tracks[2], 'NV Graph')
+    self.assertEqual(tracks[1], 'NV Graph')
+    self.assertEqual(tracks[2], 'NV:1')
 
     nv_events = j['layout']['NV']['events']
     self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
@@ -470,6 +462,14 @@ class TestVizProfiler(BaseTestViz):
     assert kernels[0]["st"] <= markers[0]["ts"] <= kernels[1]["st"]
     assert markers[1]["ts"] >= kernels[1]["st"]+kernels[1]["dur"]
 
+  def test_layout_order(self):
+    def fn(): return
+    for dname in ["TINY", "USER", "TEST:1 N1", "TEST:2 N1", "TEST:1 N2"]:
+      with cpu_profile("fn", dname): fn()
+    layout = list(load_profile(cpu_events)["layout"])
+    self.assertListEqual(layout[:2], ["USER","TINY"])
+    self.assertListEqual(layout[2:], ["TEST:1 N1","TEST:1 N2", "TEST:2 N1"])
+
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
   a.uop.buffer.allocate()

tinygrad/codegen/late/devectorizer.py

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

tinygrad/codegen/opt/postrange.py

@@ -7,6 +7,7 @@ from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
+from tinygrad.helpers import ALLOW_TF32
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
@@ -228,6 +229,7 @@ class Scheduler:
       except IndexError:
         raise KernelOptError(f"invalid tensor core choice {tc_select}")
       for tc in tensor_cores:
+        if self.ren.device in ("CUDA", "NV") and tc.dtype_in == dtypes.float and not ALLOW_TF32: continue
         if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
           # tensor cores have three ranges. X, Y, and REDUCE
           in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)

tinygrad/codegen/opt/tc.py

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

tinygrad/codegen/simplify.py

@@ -88,7 +88,8 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
   # lift x+y out of reduce on lt
   ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
-  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
+  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
+   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
   # fold the range
   # bound from below
   ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),

tinygrad/device.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator
+from typing import Any, Generic, TypeVar, Iterator, Generator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
-from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM
-from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
-from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ
+from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
+from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
+from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -238,6 +238,7 @@ class Allocator(Generic[DeviceType]):
   # def _as_buffer(self, src) -> memoryview:
   # def _offset(self, buf, size:int, offset:int):
   # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
+  def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional
 
 class LRUAllocator(Allocator, Generic[DeviceType]):
   """
@@ -276,29 +277,49 @@ class Compiler:
     return lib
   def disassemble(self, lib:bytes): pass
 
-CompilerPairT = tuple[functools.partial|type[Renderer], functools.partial|type[Compiler]]
+@dataclass(frozen=True)
+class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial; ctrl_var:ContextVar|None = None # noqa: E702
+
+@dataclass(frozen=True)
+class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
+
 class Compiled:
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
-  def __init__(self, device:str, allocator:Allocator, compilers:Sequence[CompilerPairT]|None, runtime, graph=None, group_id=None):
+  def __init__(self, device:str, allocator:Allocator, compilers:CompilerSet|None, runtime, graph=None, group_id=None):
     self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id
-    self.compilers = cast(list[CompilerPairT], compilers or [(Renderer, Compiler)])
 
-    envnames = [self._get_compiler_envvar(c) for r,c in self.compilers]
-    enable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 1)
-    disable_comps = set((en, comp_pair) for en, comp_pair in zip(envnames, self.compilers) if en is not None and getenv(en, -1) == 0)
+    self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
+    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial]]] = {}
+    self.cached_pair:dict[Any, tuple[Renderer, Compiler]] = {}
+    for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
+      self.comp_sets[self._compiler_name(cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
 
-    if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}")
-    for _, comp_pair in disable_comps: self.compilers.remove(comp_pair)
+  @property
+  def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
 
-    self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers,
-                                                       f"No compiler for {self.device} is available")
+  @property
+  def compiler(self) -> Compiler: return self._select_compiler_pair()[1]
 
-    if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}")
+  def _compiler_name(self, c:type[Compiler]|functools.partial) -> str:
+    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 
-  def _get_compiler_envvar(self, c):
-    compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}"
-    return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}"
+  def _select_compiler_pair(self) -> tuple[Renderer, Compiler]:
+    # select forced compiler from global env var.
+    forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])
+
+    # add forced compilers from individual env vars.
+    forced_comps |= set(rc for en, rc in self.comp_sets.values() if en is not None and en.value == 1)
+    if len(forced_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {forced_comps}")
+
+    # select remaining compilers (all or forced only)
+    comps = list(rc for en, rc in self.comp_sets.values())
+
+    # remove disabled compilers
+    for en, rc in self.comp_sets.values():
+      if en is not None and en.value == 0 and rc in comps: comps.remove(rc)
+
+    return select_first_inited(list(forced_comps) if len(forced_comps)>0 else comps, f"No compiler for {self.device} is available", self.cached_pair)
 
   def synchronize(self):
     """
@@ -324,11 +345,13 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
   if device is None: device = Device.DEFAULT
   if dtype == dtypes.bfloat16:
     if device == "METAL": return not CI
-    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
-    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not getenv("CPU_LVP")
+    if device == "CUDA": return not CI and not CUDA_PTX
+    if device == "NV": return not CI and not NV_PTX and not NV_NAK
+    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
     return device in {"AMD", "PYTHON", "NULL"}
   if dtype in dtypes.fp8s:
-    if device in {"CUDA", "NV"}: return not CI and not getenv(f"{device}_PTX") and not getenv("NV_NAK")
+    if device == "CUDA": return not CI and not CUDA_PTX
+    if device == "NV": return not CI and not NV_PTX and not NV_NAK
     if device == "AMD": return not CI and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
     return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
@@ -365,16 +388,24 @@ def enumerate_devices_str() -> Generator[str, None, None]:
   for device in ALL_DEVICES:
     compilers_results, any_works = [], False
     try:
-      default_compiler = (d:=Device[device]).compiler
-      for i,(r,c) in enumerate(d.compilers):
-        try:
-          d.renderer, d.compiler = r(), c()
-          with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
-          if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-          default_text = '(default)' if type(default_compiler) is type(d.compiler) else f'({d._get_compiler_envvar(c)}=1 to make default)'
-          compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
-          any_works = True
-        except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
+      d = Device[device]
+      default_comp_pairs, default_compiler, cc_ctrl_var = d.comp_sets, d.compiler, d.comps_ctrl_var
+      try:
+        for k,(en,(r,c)) in default_comp_pairs.items():
+          d.comp_sets = {k:(None,(r,c))} # env var set to None, so it doesn't interfere
+          d.comps_ctrl_var = None
+          try:
+            # d.renderer, d.compiler = r(), c()
+            with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
+            if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(c)} to make default)' if cc_ctrl_var is not None else ''
+            default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
+            compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
+            any_works = True
+          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
+      finally:
+        # put the defaults back!
+        d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var
       result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results])
     except Exception as e:
       result = f"{colored('FAIL', 'red')} {e}"

tinygrad/dtype.py

@@ -108,7 +108,7 @@ class dtypes:
   def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats or isinstance(x, ImageDType)
   @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
   @functools.cache
-  def is_int(x: DType) -> bool: return x.scalar() in dtypes.ints + (dtypes.index,)
+  def is_int(x: DType) -> bool: return x.scalar() in dtypes.index_like
   @staticmethod
   @functools.cache
   def is_unsigned(x: DType) -> bool: return x.scalar() in dtypes.uints
@@ -128,6 +128,8 @@ class dtypes:
       assert len(val) == dtype.count, f"mismatch {val} {dtype}"
       return tuple(dtypes.as_const(x, dtype) for x in val)
     if isinstance(val, InvalidType): return val
+    # NOTE: float('nan') != float('nan'), so we canonicalize here
+    if isinstance(val, float) and math.isnan(val): val = math.nan
     return int(val) if dtypes.is_int(dtype) else float(val) if dtypes.is_float(dtype) else bool(val)
   @staticmethod
   @functools.cache
@@ -185,6 +187,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
+  index_like = ints + (index,)
   all = floats + ints + (bool, index) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):

tinygrad/engine/jit.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv,
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops
-from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, CompiledRunner, Runner, Estimates
+from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
 from tinygrad.engine.memory import _internal_memory_planner
 from tinygrad.nn.state import get_parameters
 from tinygrad.schedule.rangeify import mop_cleanup
@@ -143,7 +143,7 @@ class MultiGraphRunner(GraphRunner):
 
 def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
   if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]
-  if isinstance(ei.prg, (BufferCopy, BufferXfer)): return [cast(Buffer, ei.bufs[0])]
+  if isinstance(ei.prg, (BufferCopy, BufferXfer, EncDec)): return [cast(Buffer, ei.bufs[0])]
   return []
 
 def update_depends(depends:set[Buffer|None], jit_cache:list[ExecItem]):
@@ -318,7 +318,7 @@ class TinyJit(Generic[ReturnType]):
 
       # memory planning (optional)
       # Exclude buffers involved in transfer ops to preserve parallelism.
-      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy)) for b in ji.bufs}
+      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec)) for b in ji.bufs}
       assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], noopt_buffers, debug_prefix="JIT ")
       jit_cache = [ExecItem(item.prg, [assigned.get(b,b).ensure_allocated() for b in item.bufs if b is not None],
                             item.metadata, item.fixedvars) for item in jit_cache]

tinygrad/engine/realize.py

@@ -3,7 +3,7 @@ import time, pprint, random, itertools, math
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
-from tinygrad.helpers import unwrap, disable_gc
+from tinygrad.helpers import unwrap
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, print_uops, track_rewrites, KernelInfo, pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -13,9 +13,8 @@ from tinygrad.codegen.opt import Opt
 
 # **************** Program Creation ****************
 
-@disable_gc()
 @track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
@@ -31,7 +30,6 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
   if DEBUG >= 5: print(pyrender(ast))
 
   # linearize
-  if renderer is None: renderer = Device.default.renderer
   if opts is not None:
     assert ast.arg is None, "can't apply opts if sink has an arg"
     ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
@@ -142,6 +140,19 @@ class BufferCopy(Runner):
 class BufferXfer(BufferCopy):
   def copy(self, dest, src): dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
 
+class EncDec(Runner):
+  def __init__(self, encdec:UOp, total_sz:int, device:str):
+    self.shape, self.pos_var = encdec.arg[0], encdec.variables()[0].expr
+    name = f"enc/dec {total_sz/1e6:7.2f}M, HEVC" if total_sz >= 1e6 else f"enc/dec {total_sz:8d}, HEVC"
+    super().__init__(colored(name, "yellow"), device, Estimates(lds=total_sz, mem=total_sz))
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
+    st = time.perf_counter()
+    rawbufs[0].allocator._encode_decode(rawbufs[0]._buf, rawbufs[1]._buf, rawbufs[2]._buf,
+                                        [x._buf for x in rawbufs[3:]], self.shape, var_vals[self.pos_var])
+    if wait:
+      Device[rawbufs[0].device].synchronize()
+      return time.perf_counter() - st
+
 # **************** method cache ****************
 
 method_cache: dict[tuple[str, type, bytes, tuple[int, ...], bool], CompiledRunner] = {}
@@ -202,6 +213,7 @@ si_lowerer = PatternMatcher([
   (UPat(Ops.COPY, name="copy"), lambda ctx,copy: ((BufferXfer(ctx[0].nbytes, ctx[0].device, ctx[1].device) \
       if hasattr(Device[ctx[0].device].allocator, '_transfer') and all_same([x.device.split(":")[0] for x in ctx]) \
       else BufferCopy(ctx[0].nbytes, ctx[0].device, ctx[1].device)), list(ctx))),
+  (UPat(Ops.ENCDEC, name="encdec"), lambda ctx,encdec: ((EncDec(encdec, ctx[0].nbytes, ctx[1].device)), list(ctx))),
 ])
 def lower_schedule_item(si:ScheduleItem) -> ExecItem:
   return ExecItem(*cast(tuple[Runner,list], si_lowerer.rewrite(si.ast, si.bufs)), si.metadata, si.fixedvars)

tinygrad/engine/schedule.py

@@ -1,9 +1,12 @@
+import time
 from typing import cast
 from dataclasses import dataclass, field, replace
-from collections import deque, defaultdict
-from tinygrad.uop.ops import UOp, Ops, buffers
-from tinygrad.device import Device, Buffer, MultiBuffer
-from tinygrad.helpers import Metadata, all_same
+from collections import deque
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites
+from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
+from tinygrad.uop.spec import type_verify, tensor_spec
+from tinygrad.device import Buffer, MultiBuffer
+from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
 
 # **** ScheduleItem return type
 
@@ -18,99 +21,180 @@ class ScheduleItem:
 # **** schedule linearizer
 
 def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[str, int]]:
-  # construct the KERNEL children graph based on assigns
-  children: defaultdict[UOp, list[UOp]] = defaultdict(list)
-  in_degree: dict[UOp, int] = {}
-  var_vals: dict[str, int] = {}
-  for u in sched_sink.toposort():
-    if u.op is Ops.RANGE:
-      in_degree.setdefault(u, 0)
-      continue
-    if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
-    k = u.src[1]
-    in_degree.setdefault(k, 0)
-    for s in k.src[0].src if k.op is Ops.END else k.src:
-      if s.op is Ops.AFTER:
-        children[s.src[1]].append(k)
-        in_degree[k] += 1
-      elif s.op in {Ops.MSELECT, Ops.MSTACK}:
-        for ss in s.src:
-          if ss.op is Ops.MSELECT: ss = ss.src[0]
-          if ss.op is not Ops.BUFFER:
-            assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
-            children[ss.src[1]].append(k)
-            in_degree[k] += 1
-      elif s.op is Ops.BUFFER:
-        pass  # a BUFFER is already realized, nothing to do here
-      elif s.op is Ops.BIND:
-        # for RANGE this is in fixedvars
-        if s.src[1].op is not Ops.RANGE:
-          var, val = s.unbind()
-          assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
-          var_vals[var.expr] = val
+  with cpu_profile(TracingKey("toposort sched_sink")):
+    # construct the KERNEL children graph based on assigns
+    children: dict[UOp, list[UOp]] = {}
+    in_degree: dict[UOp, int] = {}
+    var_vals: dict[str, int] = {}
+    for u in sched_sink.toposort():
+      if u.op is Ops.RANGE:
+        in_degree.setdefault(u, 0)
+        continue
+      if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
+      k = u.src[1]
+      in_degree.setdefault(k, 0)
+      for s in k.src[0].src if k.op is Ops.END else k.src:
+        if s.op is Ops.AFTER:
+          children.setdefault(s.src[1], []).append(k)
+          in_degree[k] += 1
+        elif s.op in {Ops.MSELECT, Ops.MSTACK}:
+          for ss in s.src:
+            if ss.op is Ops.MSELECT: ss = ss.src[0]
+            if ss.op is not Ops.BUFFER:
+              assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
+              children.setdefault(ss.src[1], []).append(k)
+              in_degree[k] += 1
+        elif s.op is Ops.BUFFER:
+          pass  # a BUFFER is already realized, nothing to do here
+        elif s.op is Ops.BIND:
+          # for RANGE this is in fixedvars
+          if s.src[1].op is not Ops.RANGE:
+            var, val = s.unbind()
+            assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+            var_vals[var.expr] = val
+        else:
+          raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")
+
+  with cpu_profile(TracingKey("linearize to ScheduleItem")):
+    queue: deque[UOp] = deque()
+    for k,v in in_degree.items():
+      if v == 0: queue.append(k)
+
+    schedule: list[ScheduleItem|UOp] = []
+    while len(queue):
+      k = rk = queue.popleft()
+      if k.op is Ops.END: k = k.src[0]
+      if k.op is Ops.RANGE: schedule.append(k)
+      elif k.op is Ops.KERNEL:
+        ast = k.arg.ast
+        # create subbuffers if needed
+        if ast.op is Ops.BUFFER_VIEW:
+          base = k.src[1].buf_uop.buffer
+          assert isinstance(base, Buffer), "base can't be MultiBuffer"
+          buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
+        ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
+        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and s.src[1].op is Ops.RANGE)
+        if any(isinstance(x, MultiBuffer) for x in ubufs):
+          assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
+          dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
+          for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
+            schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}, bound_ranges=bound_ranges))
+        else:
+          # ONE -> ONE
+          schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata, bound_ranges=bound_ranges))
+        if rk.op is Ops.END: schedule.append(rk)
       else:
-        raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")
+        raise RuntimeError(f"can't schedule {k.op}")
+      for x in children.get(rk, []):
+        in_degree[x] -= 1
+        if in_degree[x] == 0: queue.append(x)
 
-  # linearize KERNEL UOps into ScheduleItems in BFS order
-
-  def _heuristic(k: UOp):
-    if k.op is Ops.KERNEL and k.arg.ast.op is Ops.COPY and not all_same([Device[cast(Buffer, s.buf_uop.buffer).device].group_id for s in k.src]):
-      return 1000
-    return 0
-
-  last_heuristic: int = 0
-  queues: defaultdict[int, deque[UOp]] = defaultdict(deque)
-  last_queue: deque[UOp] = deque()
-  for k,v in in_degree.items():
-    if v == 0: queues[_heuristic(k)].append(k)
-
-  schedule: list[ScheduleItem|UOp] = []
-  while last_queue or any(queues.values()):
-    if not last_queue: last_heuristic, last_queue = min((it for it in queues.items() if it[1]), key=lambda x: abs(x[0]-last_heuristic))
-    k = rk = last_queue.popleft()
-    if k.op is Ops.END: k = k.src[0]
-    if k.op is Ops.RANGE: schedule.append(k)
-    elif k.op is Ops.KERNEL:
-      ast = k.arg.ast
-      # create subbuffers if needed
-      if ast.op is Ops.BUFFER_VIEW:
-        base = k.src[1].buf_uop.buffer
-        assert isinstance(base, Buffer), "base can't be MultiBuffer"
-        buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-      ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
-      bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and s.src[1].op is Ops.RANGE)
-      if any(isinstance(x, MultiBuffer) for x in ubufs):
-        assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
-        dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
-        for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-          schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}, bound_ranges=bound_ranges))
+  with cpu_profile(TracingKey("expand ranges")):
+    real_schedule: list[ScheduleItem] = []
+    sched_ptr = 0
+    in_ranges = {}
+    range_ptrs = {}
+    while sched_ptr < len(schedule):
+      si = schedule[sched_ptr]
+      if isinstance(si, UOp):
+        if si.op is Ops.RANGE:
+          in_ranges[si] = 0
+          range_ptrs[si] = sched_ptr + 1
+        elif si.op is Ops.END:
+          if in_ranges[si.src[1]] < si.src[1].vmax:
+            in_ranges[si.src[1]] += 1
+            sched_ptr = range_ptrs[si.src[1]]
+            continue
       else:
-        # ONE -> ONE
-        schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata, bound_ranges=bound_ranges))
-      if rk.op is Ops.END: schedule.append(rk)
-    else:
-      raise RuntimeError(f"can't schedule {k.op}")
-    for x in children[rk]:
-      in_degree[x] -= 1
-      if in_degree[x] == 0: queues[_heuristic(x)].append(x)
-
-  # expand the ranges in the schedule
-  real_schedule: list[ScheduleItem] = []
-  sched_ptr = 0
-  in_ranges = {}
-  range_ptrs = {}
-  while sched_ptr < len(schedule):
-    si = schedule[sched_ptr]
-    if isinstance(si, UOp):
-      if si.op is Ops.RANGE:
-        in_ranges[si] = 0
-        range_ptrs[si] = sched_ptr + 1
-      elif si.op is Ops.END:
-        if in_ranges[si.src[1]] < si.src[1].vmax:
-          in_ranges[si.src[1]] += 1
-          sched_ptr = range_ptrs[si.src[1]]
-          continue
-    else:
-      real_schedule.append(replace(si, fixedvars=si.fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in si.bound_ranges}, bound_ranges=()))
-    sched_ptr += 1
+        real_schedule.append(replace(si, fixedvars=si.fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in si.bound_ranges}, bound_ranges=()))
+      sched_ptr += 1
   return real_schedule, var_vals
+
+from tinygrad.engine.memory import memory_planner
+from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.schedule.multi import get_multi_map
+
+def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
+    if b.op is Ops.BUFFER:
+      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
+    else:
+      # TODO: flip args in CONST
+      assert b.op is Ops.CONST
+      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
+  return ret
+
+pm_pre_sched_cache = PatternMatcher([
+  # replace input buffers
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
+  # remove unique consts
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
+])
+
+def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
+    assert b.op is Ops.BUFFER
+    # if it's not in the cache, create a new buffer
+    ctx[b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
+  return ret
+
+pm_post_sched_cache = PatternMatcher([
+  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer_back),
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.LUNIQUE)), name="b"), replace_input_buffer_back),
+])
+
+schedule_cache: dict[bytes, tuple[UOp, UOp]] = {}
+@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
+def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ScheduleItem], dict[str, int]]:
+  # big_sink srcs are all the Tensors
+  st = time.perf_counter()
+
+  # replace all UNIQUE buffers with LUNIQUE
+  input_buffers: dict[UOp, UOp] = {}
+  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
+  sched_cache_key = big_sink_cache.key
+
+  if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
+    # verify Tensors match the spec (on big_sink, we only need to do this if cache misses)
+    if SPEC: type_verify(big_sink, tensor_spec)
+
+    # hack to preserve metadata
+    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
+
+    # tensor map is what we return
+    tensor_map: dict[UOp, UOp] = {}
+
+    if any(isinstance(x._device, tuple) for x in big_sink_cache.toposort()):
+      tensor_map |= get_multi_map(big_sink_cache)
+      big_sink_cache = big_sink_cache.substitute(tensor_map, name="Apply Multi Map")
+      big_sink_cache = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink_cache.src]))
+
+    tensor_map |= get_rangeify_map(big_sink_cache)
+    big_sink = big_sink_cache.substitute(tensor_map, name="Apply Kernelize Map")
+
+    # save in schedule cache
+    tensor_map_sink = UOp.sink(*flatten([(k,v) for k,v in tensor_map.items()]))
+    schedule_cache[sched_cache_key] = (big_sink, tensor_map_sink)
+  else:
+    # schedule cache hit
+    del big_sink_cache
+    big_sink, tensor_map_sink = sc_ret
+
+  # replace all the LUNIQUEs with UNIQUEs
+  input_buffers_reverse = {v:k for k,v in input_buffers.items()}
+  big_sink = graph_rewrite(big_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for sched cache")
+  tm_src = graph_rewrite(tensor_map_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for tensor map").src
+  tensor_map = {tm_src[i]:tm_src[i+1] for i in range(0, len(tm_src), 2)}
+
+  # create the schedule
+  schedule, var_vals = create_schedule_with_vars(big_sink)
+  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
+
+  # remove all AFTERs, after scheduling, the tensors are just buffers
+  tensor_map |= {u:u.buf_uop for u in big_sink.toposort() if u.op is Ops.AFTER}
+
+  if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
+    print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
+          f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
+          f" | {len(UOpMetaClass.ucache)} uops in cache")
+  return tensor_map, schedule, var_vals

tinygrad/helpers.py

@@ -114,10 +114,14 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T:
+def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg:str, cache:dict|None=None) -> tuple[T,...]|T:
   excs = []
   for typ in candidates:
-    try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+    if cache is not None and typ in cache: return cache[typ]
+    try:
+      x = tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)()
+      if cache is not None: cache[typ] = x
+      return x
     except Exception as e: excs.append(e)
   raise ExceptionGroup(err_msg, excs)
 
@@ -173,14 +177,17 @@ WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1),
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
-PICKLE_BUFFERS, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("LRU", 1)
+LRU = ContextVar("LRU", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
+# Compilers
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
+NV_PTX, CUDA_PTX, NV_NAK = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0)
+AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", VIZ.value)
@@ -193,6 +200,8 @@ DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
 TUPLE_ORDER = ContextVar("TUPLE_ORDER", 1)
 # set to 0 to disable the compiler cache
 CCACHE = ContextVar("CCACHE", 1)
+# allow tf32 to be used on NVIDIA GPUs
+ALLOW_TF32 = ContextVar("ALLOW_TF32", 0)
 
 @dataclass(frozen=True)
 class Metadata:
@@ -291,6 +300,15 @@ def cpu_profile(name:str|TracingKey, device="TINY", is_copy=False, display=True)
 def profile_marker(name:str, color="gray") -> None:
   cpu_events.append(ProfilePointEvent("TINY", "marker", None, {"name":name, "color":color}))
 
+if getenv("DEBUG_GC"):
+  gc_start: decimal.Decimal = perf_counter_us()
+  def my_gc_callback(phase, info):
+    global gc_start
+    if phase == 'start': gc_start = perf_counter_us()
+    elif phase == "stop":
+      cpu_events.append(ProfileRangeEvent("GC", f"collected: {info['collected']} (gen {info['generation']})", gc_start, perf_counter_us()))
+  if PROFILE: gc.callbacks.append(my_gc_callback)
+
 # *** universal database cache ***
 
 cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches" if OSX else "~/.cache")), "tinygrad")
@@ -363,14 +381,16 @@ def _ensure_downloads_dir() -> pathlib.Path:
   return pathlib.Path(cache_dir) / "downloads"
 
 def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False,
-          allow_caching=not getenv("DISABLE_HTTP_CACHE")) -> pathlib.Path:
+          allow_caching=not getenv("DISABLE_HTTP_CACHE"), headers:dict[str, str]={}) -> pathlib.Path:
   if url.startswith(("/", ".")): return pathlib.Path(url)
   if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name)
-  else: fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + (".gunzip" if gunzip else ""))
+  else:
+    hh = "_"+hashlib.md5(("\n".join(f"{k.strip()}:{v.strip()}" for k,v in sorted(headers.items()))).encode("utf-8")).hexdigest() if headers else ""
+    fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else ""))
   if not fp.is_file() or not allow_caching:
     (_dir := fp.parent).mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.11.0"}), timeout=10) as r:
-      assert r.status == 200, r.status
+    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.11.0", **headers}), timeout=10) as r:
+      assert r.status in {200, 206}, r.status
       length = int(r.headers.get('content-length', 0)) if not gunzip else None
       readfile = gzip.GzipFile(fileobj=r) if gunzip else r
       progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)

tinygrad/nn/onnx.py

@@ -710,7 +710,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
           mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
-    value = constant_value or value
+    value = _resolve_const(constant_value or value)
     axes = axes or list(range(x.ndim))
     real_pads = [0] * (x.ndim*2)
     for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
@@ -1158,7 +1158,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     b_idx = Tensor.arange(b, device=x.device).reshape(b, *(1,)*(indices.ndim - 2)).expand(*indices.shape[:-1])
     ret = x[(b_idx,) + tuple(i.squeeze(-1) for i in indices.split(1, -1))]
     return ret.reshape(*x_shape[:batch_dims], *i_shape[batch_dims:-1], *ret.shape[indices.ndim-1:])
-  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul"]='none'):
+  def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul", "max", "min"]='none'):
     assert updates.shape == indices.shape[:-1] + x.shape[cast(int, indices.shape[-1]):]
     x = x.contiguous()
     for index, u in zip(indices.split(1, 0), updates.split(1, 0)):
@@ -1167,7 +1167,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
       if reduction == "none": x[i] = u
       elif reduction == "add": x[i] += u
       elif reduction == "mul": x[i] *= u
-      else: raise NotImplementedError("reduction doesn't support max or min")
+      elif reduction == "max": x[i] = x[i].maximum(u)
+      elif reduction == "min": x[i] = x[i].minimum(u)
     return x
 
   def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Literal["none", "add", "mul", "min", "max"]="none"):

tinygrad/nn/optim.py

@@ -9,13 +9,12 @@ class Optimizer:
   Base class for all optimizers.
   """
   def __init__(self, params: list[Tensor], lr: float, fused=FUSE_OPTIM):
-    # if it's None, but being put into an optimizer, set it to True
+    # if requires_grad is None, but being put into an optimizer, set it to True
     for x in params:
       if x.requires_grad is None: x.requires_grad = True
 
     self.params: list[Tensor] = dedup([x for x in params if x.requires_grad])
     assert len(self.params) != 0, "optimizer must have at least one param"
-    self.device = self.params[0].device
     self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized
     self.fused = fused
     # store lr in at least float32 precision
@@ -23,6 +22,9 @@ class Optimizer:
                      dtype=least_upper_dtype(dtypes.default_float, dtypes.float32))
     if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))
 
+  @property
+  def device(self): return self.params[0].device
+
   def _new_optim_param(self) -> list[Tensor]:
     param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
     if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False).contiguous()]
@@ -49,6 +51,7 @@ class Optimizer:
                 - help: Consider setting Tensor.training=True before calling Optimizer.step().""")
     if self.fused:
       # optimizer fusion just concatenates all the buffers, runs the _step, then splits them back up
+      # NOTE: contiguous is for speed
       out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
                               [Tensor.cat(*[unwrap(t.grad).contiguous().flatten() for t in self.params], dim=0)])
       updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]

tinygrad/nn/state.py

@@ -1,4 +1,4 @@
-import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
+import json, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -151,7 +151,8 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
         if DEBUG >= 1: print(f"WARNING: not loading {k}")
         continue
       if v.shape != state_dict[k].shape:
-        raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
+        if {(), (1,)} == {state_dict[k].shape, v.shape}: state_dict[k] = state_dict[k].reshape(v.shape)
+        else: raise ValueError(f'Shape mismatch in layer `{k}`: Expected shape {v.shape}, but found {state_dict[k].shape} in state dict.')
       if isinstance(v.device, tuple):
         if isinstance(state_dict[k].device, tuple): v.replace(state_dict[k])
         else: v.replace(state_dict[k].shard(v.device, v.uop.axis))
@@ -161,6 +162,27 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
       ret.append(v)
   return ret
 
+@accept_filename
+def zip_extract(t: Tensor) -> dict[str, Tensor]:
+  files: dict[str, Tensor] = {}
+  file_offsets: dict[str, tuple[Tensor, int, int]] = {}
+  with zipfile.ZipFile(TensorIO(t), "r") as myzip:
+    for zi in myzip.filelist:
+      file_offset = zi.header_offset+30+t[zi.header_offset+26:zi.header_offset+30].bitcast(dtypes.uint16).to("CPU").sum()
+      file_offsets[zi.filename] = (file_offset, zi.compress_size, zi.compress_type)
+  # sadly, the extra length needs to be read from the local header of each file. this is a limitation of the zip file format
+  Tensor.realize(*[x[0] for x in file_offsets.values()])
+  for filename, (file_offset, compress_size, compress_type) in file_offsets.items():
+    # possible to remove this realize/item? it's slow
+    file_offset_int = int(file_offset.item())
+    files[filename] = t[file_offset_int:file_offset_int+compress_size]
+    match compress_type:
+      case zipfile.ZIP_STORED: pass
+      # TODO: we need a zlib UOp so this can be lazy
+      case zipfile.ZIP_DEFLATED: files[filename] = Tensor(zlib.decompress(files[filename].data(), -15))
+      case _: raise NotImplementedError(f"compression {compress_type} not supported")
+  return files
+
 @accept_filename
 def tar_extract(t: Tensor) -> dict[str, Tensor]:
   """
@@ -179,6 +201,7 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
 
 # torch support!
 
+# TODO: this should use tar_extract and zip_extract
 @accept_filename
 def torch_load(t:Tensor) -> dict[str, Tensor]:
   """

tinygrad/renderer/cstyle.py

@@ -373,6 +373,7 @@ class CUDARenderer(CStyleLanguage):
   code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+int(x))}", "l": lambda x: f"threadIdx.{chr(120+int(x))}",
                        "i": lambda x: f"(blockIdx.{chr(120+int(x))}*blockDim.{chr(120+int(x))}+threadIdx.{chr(120+int(x))})"}
   code_for_op = { **CStyleLanguage.code_for_op,
+    Ops.TRUNC: lambda x,dtype: f"htrunc({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"trunc({x})",
     Ops.SIN: lambda x,dtype: f"hsin({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sin({x})",
     Ops.LOG2: lambda x,dtype: f"hlog2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"log2({x})",
     Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})",

tinygrad/runtime/autogen/__init__.py

@@ -4,6 +4,8 @@ from tinygrad.helpers import fetch, flatten, system, getenv
 root = (here:=pathlib.Path(__file__).parent).parents[2]
 nv_src = {"nv_570": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/81fe4fb417c8ac3b9bdcc1d56827d116743892a5.tar.gz",
           "nv_580": "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/2af9f1f0f7de4988432d4ae875b5858ffdb09cc2.tar.gz"}
+ffmpeg_src = "https://ffmpeg.org/releases/ffmpeg-8.0.1.tar.gz"
+rocr_src = "https://github.com/ROCm/rocm-systems/archive/refs/tags/rocm-7.1.1.tar.gz"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"
 
 def load(name, dll, files, **kwargs):
@@ -27,6 +29,7 @@ def __getattr__(nm):
     case "libc": return load("libc", ["find_library('c')"], lambda: (
       [i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
       ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), use_errno=True)
+    case "avcodec": return load("avcodec", [], ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], tarball=ffmpeg_src)
     case "opencl": return load("opencl", ["find_library('OpenCL')"], ["/usr/include/CL/cl.h"])
     case "cuda": return load("cuda", ["find_library('cuda')"], ["/usr/include/cuda.h"], args=["-D__CUDA_API_VERSION_INTERNAL"], parse_macros=False)
     case "nvrtc": return load("nvrtc", ["find_library('nvrtc')"], ["/usr/include/nvrtc.h"])
@@ -34,14 +37,14 @@ def __getattr__(nm):
     case "kfd": return load("kfd", [], ["/usr/include/linux/kfd_ioctl.h"])
     case "nv_570" | "nv_580":
       return load(nm, [], [
-        *[root/"extra/nv_gpu_driver"/s for s in ["clc6c0qmd.h","clcec0qmd.h"]], "{}/kernel-open/common/inc/nvmisc.h",
-        *[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
+        *[root/"extra/nv_gpu_driver"/s for s in ["clc9b0.h", "clc6c0qmd.h","clcec0qmd.h", "nvdec_drv.h"]], "{}/kernel-open/common/inc/nvmisc.h",
+        *[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0070", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
                                                                     "83de", "c6c0", "cdc0"]],
-        *[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "uvm_ioctl", "uvm_linux_ioctl", "hwref/ampere/ga100/dev_fault"]],
+        *[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "clcfb0", "uvm_ioctl", "uvm_linux_ioctl", "hwref/ampere/ga100/dev_fault"]],
         *[f"{{}}/src/nvidia/arch/nvalloc/unix/include/nv{s}.h" for s in ["_escape", "-ioctl", "-ioctl-numbers",
                                                                          "-ioctl-numa", "-unix-nvos-params-wrappers"]],
         *[f"{{}}/src/common/sdk/nvidia/inc/{s}.h" for s in ["alloc/alloc_channel", "nvos", "ctrl/ctrlc36f", "ctrl/ctrlcb33",
-                                                            "ctrl/ctrla06c", "ctrl/ctrl90f1"]],
+                                                            "ctrl/ctrla06c", "ctrl/ctrl90f1", "ctrl/ctrla06f/ctrla06fgpfifo"]],
         *[f"{{}}/src/common/sdk/nvidia/inc/ctrl/ctrl{s}/*.h" for s in ["0000", "0080", "2080", "83de"]],
         "{}/kernel-open/common/inc/nvstatus.h", "{}/src/nvidia/generated/g_allclasses.h"
       ], args=[
@@ -87,8 +90,11 @@ def __getattr__(nm):
         "os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so'", "'/usr/local/lib/libamd_comgr.dylib'", "'/opt/homebrew/lib/libamd_comgr.dylib'"
       ], ["/opt/rocm/include/amd_comgr/amd_comgr.h"], args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"])
     case "hsa": return load("hsa", ["os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so'", "find_library('hsa-runtime64')"], [
-      f"/opt/rocm/include/hsa/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code", "hsa_ext_finalize",
-                                               "hsa_ext_image", "hsa_ven_amd_aqlprofile"] ], args=["-I/opt/rocm/include"])
+      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
+      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue",
+                                                                              "amd_hsa_kernel_code", "hsa_ext_finalize",
+                                                                              "hsa_ext_image", "hsa_ven_amd_aqlprofile"]]],
+      tarball=rocr_src, args=["-DLITTLEENDIAN_CPU"])
     case "amd_gpu": return load("amd_gpu", [], [root/f"extra/hip_gpu_driver/{s}.h" for s in ["sdma_registers", "nvd", "gc_11_0_0_offset",
                                                                                              "sienna_cichlid_ip_offset"]],
                                 args=["-I/opt/rocm/include", "-x", "c++"])
@@ -129,4 +135,4 @@ python3 src/compiler/builtin_types_h.py gen/builtin_types.h""", cwd=path, shell=
       return load("metal", ["find_library('Metal')"],[f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
                   ["ComputeCommandEncoder", "ComputePipeline", "CommandQueue", "Device", "IndirectCommandBuffer", "Resource", "CommandEncoder"]],
                   args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
-    case _: raise AttributeError(f"no such autogen: {nm}")
+    case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/autogen/avcodec.py

@@ -0,0 +1,543 @@
+# mypy: ignore-errors
+import ctypes
+from tinygrad.helpers import unwrap
+from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+enum_HEVCNALUnitType = CEnum(ctypes.c_uint32)
+HEVC_NAL_TRAIL_N = enum_HEVCNALUnitType.define('HEVC_NAL_TRAIL_N', 0)
+HEVC_NAL_TRAIL_R = enum_HEVCNALUnitType.define('HEVC_NAL_TRAIL_R', 1)
+HEVC_NAL_TSA_N = enum_HEVCNALUnitType.define('HEVC_NAL_TSA_N', 2)
+HEVC_NAL_TSA_R = enum_HEVCNALUnitType.define('HEVC_NAL_TSA_R', 3)
+HEVC_NAL_STSA_N = enum_HEVCNALUnitType.define('HEVC_NAL_STSA_N', 4)
+HEVC_NAL_STSA_R = enum_HEVCNALUnitType.define('HEVC_NAL_STSA_R', 5)
+HEVC_NAL_RADL_N = enum_HEVCNALUnitType.define('HEVC_NAL_RADL_N', 6)
+HEVC_NAL_RADL_R = enum_HEVCNALUnitType.define('HEVC_NAL_RADL_R', 7)
+HEVC_NAL_RASL_N = enum_HEVCNALUnitType.define('HEVC_NAL_RASL_N', 8)
+HEVC_NAL_RASL_R = enum_HEVCNALUnitType.define('HEVC_NAL_RASL_R', 9)
+HEVC_NAL_VCL_N10 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_N10', 10)
+HEVC_NAL_VCL_R11 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_R11', 11)
+HEVC_NAL_VCL_N12 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_N12', 12)
+HEVC_NAL_VCL_R13 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_R13', 13)
+HEVC_NAL_VCL_N14 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_N14', 14)
+HEVC_NAL_VCL_R15 = enum_HEVCNALUnitType.define('HEVC_NAL_VCL_R15', 15)
+HEVC_NAL_BLA_W_LP = enum_HEVCNALUnitType.define('HEVC_NAL_BLA_W_LP', 16)
+HEVC_NAL_BLA_W_RADL = enum_HEVCNALUnitType.define('HEVC_NAL_BLA_W_RADL', 17)
+HEVC_NAL_BLA_N_LP = enum_HEVCNALUnitType.define('HEVC_NAL_BLA_N_LP', 18)
+HEVC_NAL_IDR_W_RADL = enum_HEVCNALUnitType.define('HEVC_NAL_IDR_W_RADL', 19)
+HEVC_NAL_IDR_N_LP = enum_HEVCNALUnitType.define('HEVC_NAL_IDR_N_LP', 20)
+HEVC_NAL_CRA_NUT = enum_HEVCNALUnitType.define('HEVC_NAL_CRA_NUT', 21)
+HEVC_NAL_RSV_IRAP_VCL22 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_IRAP_VCL22', 22)
+HEVC_NAL_RSV_IRAP_VCL23 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_IRAP_VCL23', 23)
+HEVC_NAL_RSV_VCL24 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL24', 24)
+HEVC_NAL_RSV_VCL25 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL25', 25)
+HEVC_NAL_RSV_VCL26 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL26', 26)
+HEVC_NAL_RSV_VCL27 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL27', 27)
+HEVC_NAL_RSV_VCL28 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL28', 28)
+HEVC_NAL_RSV_VCL29 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL29', 29)
+HEVC_NAL_RSV_VCL30 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL30', 30)
+HEVC_NAL_RSV_VCL31 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_VCL31', 31)
+HEVC_NAL_VPS = enum_HEVCNALUnitType.define('HEVC_NAL_VPS', 32)
+HEVC_NAL_SPS = enum_HEVCNALUnitType.define('HEVC_NAL_SPS', 33)
+HEVC_NAL_PPS = enum_HEVCNALUnitType.define('HEVC_NAL_PPS', 34)
+HEVC_NAL_AUD = enum_HEVCNALUnitType.define('HEVC_NAL_AUD', 35)
+HEVC_NAL_EOS_NUT = enum_HEVCNALUnitType.define('HEVC_NAL_EOS_NUT', 36)
+HEVC_NAL_EOB_NUT = enum_HEVCNALUnitType.define('HEVC_NAL_EOB_NUT', 37)
+HEVC_NAL_FD_NUT = enum_HEVCNALUnitType.define('HEVC_NAL_FD_NUT', 38)
+HEVC_NAL_SEI_PREFIX = enum_HEVCNALUnitType.define('HEVC_NAL_SEI_PREFIX', 39)
+HEVC_NAL_SEI_SUFFIX = enum_HEVCNALUnitType.define('HEVC_NAL_SEI_SUFFIX', 40)
+HEVC_NAL_RSV_NVCL41 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL41', 41)
+HEVC_NAL_RSV_NVCL42 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL42', 42)
+HEVC_NAL_RSV_NVCL43 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL43', 43)
+HEVC_NAL_RSV_NVCL44 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL44', 44)
+HEVC_NAL_RSV_NVCL45 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL45', 45)
+HEVC_NAL_RSV_NVCL46 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL46', 46)
+HEVC_NAL_RSV_NVCL47 = enum_HEVCNALUnitType.define('HEVC_NAL_RSV_NVCL47', 47)
+HEVC_NAL_UNSPEC48 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC48', 48)
+HEVC_NAL_UNSPEC49 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC49', 49)
+HEVC_NAL_UNSPEC50 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC50', 50)
+HEVC_NAL_UNSPEC51 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC51', 51)
+HEVC_NAL_UNSPEC52 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC52', 52)
+HEVC_NAL_UNSPEC53 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC53', 53)
+HEVC_NAL_UNSPEC54 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC54', 54)
+HEVC_NAL_UNSPEC55 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC55', 55)
+HEVC_NAL_UNSPEC56 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC56', 56)
+HEVC_NAL_UNSPEC57 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC57', 57)
+HEVC_NAL_UNSPEC58 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC58', 58)
+HEVC_NAL_UNSPEC59 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC59', 59)
+HEVC_NAL_UNSPEC60 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC60', 60)
+HEVC_NAL_UNSPEC61 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC61', 61)
+HEVC_NAL_UNSPEC62 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC62', 62)
+HEVC_NAL_UNSPEC63 = enum_HEVCNALUnitType.define('HEVC_NAL_UNSPEC63', 63)
+
+enum_HEVCSliceType = CEnum(ctypes.c_uint32)
+HEVC_SLICE_B = enum_HEVCSliceType.define('HEVC_SLICE_B', 0)
+HEVC_SLICE_P = enum_HEVCSliceType.define('HEVC_SLICE_P', 1)
+HEVC_SLICE_I = enum_HEVCSliceType.define('HEVC_SLICE_I', 2)
+
+_anonenum0 = CEnum(ctypes.c_uint32)
+HEVC_MAX_LAYERS = _anonenum0.define('HEVC_MAX_LAYERS', 63)
+HEVC_MAX_SUB_LAYERS = _anonenum0.define('HEVC_MAX_SUB_LAYERS', 7)
+HEVC_MAX_LAYER_SETS = _anonenum0.define('HEVC_MAX_LAYER_SETS', 1024)
+HEVC_MAX_LAYER_ID = _anonenum0.define('HEVC_MAX_LAYER_ID', 63)
+HEVC_MAX_NUH_LAYER_ID = _anonenum0.define('HEVC_MAX_NUH_LAYER_ID', 62)
+HEVC_MAX_VPS_COUNT = _anonenum0.define('HEVC_MAX_VPS_COUNT', 16)
+HEVC_MAX_SPS_COUNT = _anonenum0.define('HEVC_MAX_SPS_COUNT', 16)
+HEVC_MAX_PPS_COUNT = _anonenum0.define('HEVC_MAX_PPS_COUNT', 64)
+HEVC_MAX_DPB_SIZE = _anonenum0.define('HEVC_MAX_DPB_SIZE', 16)
+HEVC_MAX_REFS = _anonenum0.define('HEVC_MAX_REFS', 16)
+HEVC_MAX_SHORT_TERM_REF_PIC_SETS = _anonenum0.define('HEVC_MAX_SHORT_TERM_REF_PIC_SETS', 64)
+HEVC_MAX_LONG_TERM_REF_PICS = _anonenum0.define('HEVC_MAX_LONG_TERM_REF_PICS', 32)
+HEVC_MIN_LOG2_CTB_SIZE = _anonenum0.define('HEVC_MIN_LOG2_CTB_SIZE', 4)
+HEVC_MAX_LOG2_CTB_SIZE = _anonenum0.define('HEVC_MAX_LOG2_CTB_SIZE', 6)
+HEVC_MAX_CPB_CNT = _anonenum0.define('HEVC_MAX_CPB_CNT', 32)
+HEVC_MAX_LUMA_PS = _anonenum0.define('HEVC_MAX_LUMA_PS', 35651584)
+HEVC_MAX_WIDTH = _anonenum0.define('HEVC_MAX_WIDTH', 16888)
+HEVC_MAX_HEIGHT = _anonenum0.define('HEVC_MAX_HEIGHT', 16888)
+HEVC_MAX_TILE_ROWS = _anonenum0.define('HEVC_MAX_TILE_ROWS', 22)
+HEVC_MAX_TILE_COLUMNS = _anonenum0.define('HEVC_MAX_TILE_COLUMNS', 20)
+HEVC_MAX_SLICE_SEGMENTS = _anonenum0.define('HEVC_MAX_SLICE_SEGMENTS', 600)
+HEVC_MAX_ENTRY_POINT_OFFSETS = _anonenum0.define('HEVC_MAX_ENTRY_POINT_OFFSETS', 2700)
+HEVC_MAX_PALETTE_PREDICTOR_SIZE = _anonenum0.define('HEVC_MAX_PALETTE_PREDICTOR_SIZE', 128)
+
+enum_HEVCScalabilityMask = CEnum(ctypes.c_uint32)
+HEVC_SCALABILITY_DEPTH = enum_HEVCScalabilityMask.define('HEVC_SCALABILITY_DEPTH', 32768)
+HEVC_SCALABILITY_MULTIVIEW = enum_HEVCScalabilityMask.define('HEVC_SCALABILITY_MULTIVIEW', 16384)
+HEVC_SCALABILITY_SPATIAL = enum_HEVCScalabilityMask.define('HEVC_SCALABILITY_SPATIAL', 8192)
+HEVC_SCALABILITY_AUXILIARY = enum_HEVCScalabilityMask.define('HEVC_SCALABILITY_AUXILIARY', 4096)
+HEVC_SCALABILITY_MASK_MAX = enum_HEVCScalabilityMask.define('HEVC_SCALABILITY_MASK_MAX', 65535)
+
+enum_HEVCAuxId = CEnum(ctypes.c_uint32)
+HEVC_AUX_ALPHA = enum_HEVCAuxId.define('HEVC_AUX_ALPHA', 1)
+HEVC_AUX_DEPTH = enum_HEVCAuxId.define('HEVC_AUX_DEPTH', 2)
+
+class struct_H265RawNALUnitHeader(Struct): pass
+uint8_t = ctypes.c_ubyte
+struct_H265RawNALUnitHeader._fields_ = [
+  ('nal_unit_type', uint8_t),
+  ('nuh_layer_id', uint8_t),
+  ('nuh_temporal_id_plus1', uint8_t),
+]
+H265RawNALUnitHeader = struct_H265RawNALUnitHeader
+class struct_H265RawProfileTierLevel(Struct): pass
+struct_H265RawProfileTierLevel._fields_ = [
+  ('general_profile_space', uint8_t),
+  ('general_tier_flag', uint8_t),
+  ('general_profile_idc', uint8_t),
+  ('general_profile_compatibility_flag', (uint8_t * 32)),
+  ('general_progressive_source_flag', uint8_t),
+  ('general_interlaced_source_flag', uint8_t),
+  ('general_non_packed_constraint_flag', uint8_t),
+  ('general_frame_only_constraint_flag', uint8_t),
+  ('general_max_12bit_constraint_flag', uint8_t),
+  ('general_max_10bit_constraint_flag', uint8_t),
+  ('general_max_8bit_constraint_flag', uint8_t),
+  ('general_max_422chroma_constraint_flag', uint8_t),
+  ('general_max_420chroma_constraint_flag', uint8_t),
+  ('general_max_monochrome_constraint_flag', uint8_t),
+  ('general_intra_constraint_flag', uint8_t),
+  ('general_one_picture_only_constraint_flag', uint8_t),
+  ('general_lower_bit_rate_constraint_flag', uint8_t),
+  ('general_max_14bit_constraint_flag', uint8_t),
+  ('general_inbld_flag', uint8_t),
+  ('general_level_idc', uint8_t),
+  ('sub_layer_profile_present_flag', (uint8_t * 7)),
+  ('sub_layer_level_present_flag', (uint8_t * 7)),
+  ('sub_layer_profile_space', (uint8_t * 7)),
+  ('sub_layer_tier_flag', (uint8_t * 7)),
+  ('sub_layer_profile_idc', (uint8_t * 7)),
+  ('sub_layer_profile_compatibility_flag', ((uint8_t * 32) * 7)),
+  ('sub_layer_progressive_source_flag', (uint8_t * 7)),
+  ('sub_layer_interlaced_source_flag', (uint8_t * 7)),
+  ('sub_layer_non_packed_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_frame_only_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_12bit_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_10bit_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_8bit_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_422chroma_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_420chroma_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_monochrome_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_intra_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_one_picture_only_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_lower_bit_rate_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_max_14bit_constraint_flag', (uint8_t * 7)),
+  ('sub_layer_inbld_flag', (uint8_t * 7)),
+  ('sub_layer_level_idc', (uint8_t * 7)),
+]
+H265RawProfileTierLevel = struct_H265RawProfileTierLevel
+class struct_H265RawSubLayerHRDParameters(Struct): pass
+uint32_t = ctypes.c_uint32
+struct_H265RawSubLayerHRDParameters._fields_ = [
+  ('bit_rate_value_minus1', (uint32_t * 32)),
+  ('cpb_size_value_minus1', (uint32_t * 32)),
+  ('cpb_size_du_value_minus1', (uint32_t * 32)),
+  ('bit_rate_du_value_minus1', (uint32_t * 32)),
+  ('cbr_flag', (uint8_t * 32)),
+]
+H265RawSubLayerHRDParameters = struct_H265RawSubLayerHRDParameters
+class struct_H265RawHRDParameters(Struct): pass
+uint16_t = ctypes.c_uint16
+struct_H265RawHRDParameters._fields_ = [
+  ('nal_hrd_parameters_present_flag', uint8_t),
+  ('vcl_hrd_parameters_present_flag', uint8_t),
+  ('sub_pic_hrd_params_present_flag', uint8_t),
+  ('tick_divisor_minus2', uint8_t),
+  ('du_cpb_removal_delay_increment_length_minus1', uint8_t),
+  ('sub_pic_cpb_params_in_pic_timing_sei_flag', uint8_t),
+  ('dpb_output_delay_du_length_minus1', uint8_t),
+  ('bit_rate_scale', uint8_t),
+  ('cpb_size_scale', uint8_t),
+  ('cpb_size_du_scale', uint8_t),
+  ('initial_cpb_removal_delay_length_minus1', uint8_t),
+  ('au_cpb_removal_delay_length_minus1', uint8_t),
+  ('dpb_output_delay_length_minus1', uint8_t),
+  ('fixed_pic_rate_general_flag', (uint8_t * 7)),
+  ('fixed_pic_rate_within_cvs_flag', (uint8_t * 7)),
+  ('elemental_duration_in_tc_minus1', (uint16_t * 7)),
+  ('low_delay_hrd_flag', (uint8_t * 7)),
+  ('cpb_cnt_minus1', (uint8_t * 7)),
+  ('nal_sub_layer_hrd_parameters', (H265RawSubLayerHRDParameters * 7)),
+  ('vcl_sub_layer_hrd_parameters', (H265RawSubLayerHRDParameters * 7)),
+]
+H265RawHRDParameters = struct_H265RawHRDParameters
+class struct_H265RawVUI(Struct): pass
+struct_H265RawVUI._fields_ = [
+  ('aspect_ratio_info_present_flag', uint8_t),
+  ('aspect_ratio_idc', uint8_t),
+  ('sar_width', uint16_t),
+  ('sar_height', uint16_t),
+  ('overscan_info_present_flag', uint8_t),
+  ('overscan_appropriate_flag', uint8_t),
+  ('video_signal_type_present_flag', uint8_t),
+  ('video_format', uint8_t),
+  ('video_full_range_flag', uint8_t),
+  ('colour_description_present_flag', uint8_t),
+  ('colour_primaries', uint8_t),
+  ('transfer_characteristics', uint8_t),
+  ('matrix_coefficients', uint8_t),
+  ('chroma_loc_info_present_flag', uint8_t),
+  ('chroma_sample_loc_type_top_field', uint8_t),
+  ('chroma_sample_loc_type_bottom_field', uint8_t),
+  ('neutral_chroma_indication_flag', uint8_t),
+  ('field_seq_flag', uint8_t),
+  ('frame_field_info_present_flag', uint8_t),
+  ('default_display_window_flag', uint8_t),
+  ('def_disp_win_left_offset', uint16_t),
+  ('def_disp_win_right_offset', uint16_t),
+  ('def_disp_win_top_offset', uint16_t),
+  ('def_disp_win_bottom_offset', uint16_t),
+  ('vui_timing_info_present_flag', uint8_t),
+  ('vui_num_units_in_tick', uint32_t),
+  ('vui_time_scale', uint32_t),
+  ('vui_poc_proportional_to_timing_flag', uint8_t),
+  ('vui_num_ticks_poc_diff_one_minus1', uint32_t),
+  ('vui_hrd_parameters_present_flag', uint8_t),
+  ('hrd_parameters', H265RawHRDParameters),
+  ('bitstream_restriction_flag', uint8_t),
+  ('tiles_fixed_structure_flag', uint8_t),
+  ('motion_vectors_over_pic_boundaries_flag', uint8_t),
+  ('restricted_ref_pic_lists_flag', uint8_t),
+  ('min_spatial_segmentation_idc', uint16_t),
+  ('max_bytes_per_pic_denom', uint8_t),
+  ('max_bits_per_min_cu_denom', uint8_t),
+  ('log2_max_mv_length_horizontal', uint8_t),
+  ('log2_max_mv_length_vertical', uint8_t),
+]
+H265RawVUI = struct_H265RawVUI
+class struct_H265RawExtensionData(Struct): pass
+H265RawExtensionData = struct_H265RawExtensionData
+class struct_H265RawVPS(Struct): pass
+H265RawVPS = struct_H265RawVPS
+class struct_H265RawSTRefPicSet(Struct): pass
+struct_H265RawSTRefPicSet._fields_ = [
+  ('inter_ref_pic_set_prediction_flag', uint8_t),
+  ('delta_idx_minus1', uint8_t),
+  ('delta_rps_sign', uint8_t),
+  ('abs_delta_rps_minus1', uint16_t),
+  ('used_by_curr_pic_flag', (uint8_t * 16)),
+  ('use_delta_flag', (uint8_t * 16)),
+  ('num_negative_pics', uint8_t),
+  ('num_positive_pics', uint8_t),
+  ('delta_poc_s0_minus1', (uint16_t * 16)),
+  ('used_by_curr_pic_s0_flag', (uint8_t * 16)),
+  ('delta_poc_s1_minus1', (uint16_t * 16)),
+  ('used_by_curr_pic_s1_flag', (uint8_t * 16)),
+]
+H265RawSTRefPicSet = struct_H265RawSTRefPicSet
+class struct_H265RawScalingList(Struct): pass
+int16_t = ctypes.c_int16
+int8_t = ctypes.c_byte
+struct_H265RawScalingList._fields_ = [
+  ('scaling_list_pred_mode_flag', ((uint8_t * 6) * 4)),
+  ('scaling_list_pred_matrix_id_delta', ((uint8_t * 6) * 4)),
+  ('scaling_list_dc_coef_minus8', ((int16_t * 6) * 4)),
+  ('scaling_list_delta_coeff', (((int8_t * 64) * 6) * 4)),
+]
+H265RawScalingList = struct_H265RawScalingList
+class struct_H265RawSPS(Struct): pass
+H265RawSPS = struct_H265RawSPS
+class struct_H265RawPPS(Struct): pass
+H265RawPPS = struct_H265RawPPS
+class struct_H265RawAUD(Struct): pass
+struct_H265RawAUD._fields_ = [
+  ('nal_unit_header', H265RawNALUnitHeader),
+  ('pic_type', uint8_t),
+]
+H265RawAUD = struct_H265RawAUD
+class struct_H265RawSliceHeader(Struct): pass
+struct_H265RawSliceHeader._fields_ = [
+  ('nal_unit_header', H265RawNALUnitHeader),
+  ('first_slice_segment_in_pic_flag', uint8_t),
+  ('no_output_of_prior_pics_flag', uint8_t),
+  ('slice_pic_parameter_set_id', uint8_t),
+  ('dependent_slice_segment_flag', uint8_t),
+  ('slice_segment_address', uint16_t),
+  ('slice_reserved_flag', (uint8_t * 8)),
+  ('slice_type', uint8_t),
+  ('pic_output_flag', uint8_t),
+  ('colour_plane_id', uint8_t),
+  ('slice_pic_order_cnt_lsb', uint16_t),
+  ('short_term_ref_pic_set_sps_flag', uint8_t),
+  ('short_term_ref_pic_set', H265RawSTRefPicSet),
+  ('short_term_ref_pic_set_idx', uint8_t),
+  ('num_long_term_sps', uint8_t),
+  ('num_long_term_pics', uint8_t),
+  ('lt_idx_sps', (uint8_t * 16)),
+  ('poc_lsb_lt', (uint8_t * 16)),
+  ('used_by_curr_pic_lt_flag', (uint8_t * 16)),
+  ('delta_poc_msb_present_flag', (uint8_t * 16)),
+  ('delta_poc_msb_cycle_lt', (uint32_t * 16)),
+  ('slice_temporal_mvp_enabled_flag', uint8_t),
+  ('slice_sao_luma_flag', uint8_t),
+  ('slice_sao_chroma_flag', uint8_t),
+  ('num_ref_idx_active_override_flag', uint8_t),
+  ('num_ref_idx_l0_active_minus1', uint8_t),
+  ('num_ref_idx_l1_active_minus1', uint8_t),
+  ('ref_pic_list_modification_flag_l0', uint8_t),
+  ('list_entry_l0', (uint8_t * 16)),
+  ('ref_pic_list_modification_flag_l1', uint8_t),
+  ('list_entry_l1', (uint8_t * 16)),
+  ('mvd_l1_zero_flag', uint8_t),
+  ('cabac_init_flag', uint8_t),
+  ('collocated_from_l0_flag', uint8_t),
+  ('collocated_ref_idx', uint8_t),
+  ('luma_log2_weight_denom', uint8_t),
+  ('delta_chroma_log2_weight_denom', int8_t),
+  ('luma_weight_l0_flag', (uint8_t * 16)),
+  ('chroma_weight_l0_flag', (uint8_t * 16)),
+  ('delta_luma_weight_l0', (int8_t * 16)),
+  ('luma_offset_l0', (int16_t * 16)),
+  ('delta_chroma_weight_l0', ((int8_t * 2) * 16)),
+  ('chroma_offset_l0', ((int16_t * 2) * 16)),
+  ('luma_weight_l1_flag', (uint8_t * 16)),
+  ('chroma_weight_l1_flag', (uint8_t * 16)),
+  ('delta_luma_weight_l1', (int8_t * 16)),
+  ('luma_offset_l1', (int16_t * 16)),
+  ('delta_chroma_weight_l1', ((int8_t * 2) * 16)),
+  ('chroma_offset_l1', ((int16_t * 2) * 16)),
+  ('five_minus_max_num_merge_cand', uint8_t),
+  ('use_integer_mv_flag', uint8_t),
+  ('slice_qp_delta', int8_t),
+  ('slice_cb_qp_offset', int8_t),
+  ('slice_cr_qp_offset', int8_t),
+  ('slice_act_y_qp_offset', int8_t),
+  ('slice_act_cb_qp_offset', int8_t),
+  ('slice_act_cr_qp_offset', int8_t),
+  ('cu_chroma_qp_offset_enabled_flag', uint8_t),
+  ('deblocking_filter_override_flag', uint8_t),
+  ('slice_deblocking_filter_disabled_flag', uint8_t),
+  ('slice_beta_offset_div2', int8_t),
+  ('slice_tc_offset_div2', int8_t),
+  ('slice_loop_filter_across_slices_enabled_flag', uint8_t),
+  ('num_entry_point_offsets', uint16_t),
+  ('offset_len_minus1', uint8_t),
+  ('entry_point_offset_minus1', (uint32_t * 2700)),
+  ('slice_segment_header_extension_length', uint16_t),
+  ('slice_segment_header_extension_data_byte', (uint8_t * 256)),
+]
+H265RawSliceHeader = struct_H265RawSliceHeader
+class struct_H265RawSlice(Struct): pass
+H265RawSlice = struct_H265RawSlice
+class struct_H265RawSEIBufferingPeriod(Struct): pass
+struct_H265RawSEIBufferingPeriod._fields_ = [
+  ('bp_seq_parameter_set_id', uint8_t),
+  ('irap_cpb_params_present_flag', uint8_t),
+  ('cpb_delay_offset', uint32_t),
+  ('dpb_delay_offset', uint32_t),
+  ('concatenation_flag', uint8_t),
+  ('au_cpb_removal_delay_delta_minus1', uint32_t),
+  ('nal_initial_cpb_removal_delay', (uint32_t * 32)),
+  ('nal_initial_cpb_removal_offset', (uint32_t * 32)),
+  ('nal_initial_alt_cpb_removal_delay', (uint32_t * 32)),
+  ('nal_initial_alt_cpb_removal_offset', (uint32_t * 32)),
+  ('vcl_initial_cpb_removal_delay', (uint32_t * 32)),
+  ('vcl_initial_cpb_removal_offset', (uint32_t * 32)),
+  ('vcl_initial_alt_cpb_removal_delay', (uint32_t * 32)),
+  ('vcl_initial_alt_cpb_removal_offset', (uint32_t * 32)),
+  ('use_alt_cpb_params_flag', uint8_t),
+]
+H265RawSEIBufferingPeriod = struct_H265RawSEIBufferingPeriod
+class struct_H265RawSEIPicTiming(Struct): pass
+struct_H265RawSEIPicTiming._fields_ = [
+  ('pic_struct', uint8_t),
+  ('source_scan_type', uint8_t),
+  ('duplicate_flag', uint8_t),
+  ('au_cpb_removal_delay_minus1', uint32_t),
+  ('pic_dpb_output_delay', uint32_t),
+  ('pic_dpb_output_du_delay', uint32_t),
+  ('num_decoding_units_minus1', uint16_t),
+  ('du_common_cpb_removal_delay_flag', uint8_t),
+  ('du_common_cpb_removal_delay_increment_minus1', uint32_t),
+  ('num_nalus_in_du_minus1', (uint16_t * 600)),
+  ('du_cpb_removal_delay_increment_minus1', (uint32_t * 600)),
+]
+H265RawSEIPicTiming = struct_H265RawSEIPicTiming
+class struct_H265RawSEIPanScanRect(Struct): pass
+int32_t = ctypes.c_int32
+struct_H265RawSEIPanScanRect._fields_ = [
+  ('pan_scan_rect_id', uint32_t),
+  ('pan_scan_rect_cancel_flag', uint8_t),
+  ('pan_scan_cnt_minus1', uint8_t),
+  ('pan_scan_rect_left_offset', (int32_t * 3)),
+  ('pan_scan_rect_right_offset', (int32_t * 3)),
+  ('pan_scan_rect_top_offset', (int32_t * 3)),
+  ('pan_scan_rect_bottom_offset', (int32_t * 3)),
+  ('pan_scan_rect_persistence_flag', uint16_t),
+]
+H265RawSEIPanScanRect = struct_H265RawSEIPanScanRect
+class struct_H265RawSEIRecoveryPoint(Struct): pass
+struct_H265RawSEIRecoveryPoint._fields_ = [
+  ('recovery_poc_cnt', int16_t),
+  ('exact_match_flag', uint8_t),
+  ('broken_link_flag', uint8_t),
+]
+H265RawSEIRecoveryPoint = struct_H265RawSEIRecoveryPoint
+class struct_H265RawFilmGrainCharacteristics(Struct): pass
+struct_H265RawFilmGrainCharacteristics._fields_ = [
+  ('film_grain_characteristics_cancel_flag', uint8_t),
+  ('film_grain_model_id', uint8_t),
+  ('separate_colour_description_present_flag', uint8_t),
+  ('film_grain_bit_depth_luma_minus8', uint8_t),
+  ('film_grain_bit_depth_chroma_minus8', uint8_t),
+  ('film_grain_full_range_flag', uint8_t),
+  ('film_grain_colour_primaries', uint8_t),
+  ('film_grain_transfer_characteristics', uint8_t),
+  ('film_grain_matrix_coeffs', uint8_t),
+  ('blending_mode_id', uint8_t),
+  ('log2_scale_factor', uint8_t),
+  ('comp_model_present_flag', (uint8_t * 3)),
+  ('num_intensity_intervals_minus1', (uint8_t * 3)),
+  ('num_model_values_minus1', (uint8_t * 3)),
+  ('intensity_interval_lower_bound', ((uint8_t * 256) * 3)),
+  ('intensity_interval_upper_bound', ((uint8_t * 256) * 3)),
+  ('comp_model_value', (((int16_t * 6) * 256) * 3)),
+  ('film_grain_characteristics_persistence_flag', uint8_t),
+]
+H265RawFilmGrainCharacteristics = struct_H265RawFilmGrainCharacteristics
+class struct_H265RawSEIDisplayOrientation(Struct): pass
+struct_H265RawSEIDisplayOrientation._fields_ = [
+  ('display_orientation_cancel_flag', uint8_t),
+  ('hor_flip', uint8_t),
+  ('ver_flip', uint8_t),
+  ('anticlockwise_rotation', uint16_t),
+  ('display_orientation_repetition_period', uint16_t),
+  ('display_orientation_persistence_flag', uint8_t),
+]
+H265RawSEIDisplayOrientation = struct_H265RawSEIDisplayOrientation
+class struct_H265RawSEIActiveParameterSets(Struct): pass
+struct_H265RawSEIActiveParameterSets._fields_ = [
+  ('active_video_parameter_set_id', uint8_t),
+  ('self_contained_cvs_flag', uint8_t),
+  ('no_parameter_set_update_flag', uint8_t),
+  ('num_sps_ids_minus1', uint8_t),
+  ('active_seq_parameter_set_id', (uint8_t * 16)),
+  ('layer_sps_idx', (uint8_t * 63)),
+]
+H265RawSEIActiveParameterSets = struct_H265RawSEIActiveParameterSets
+class struct_H265RawSEIDecodedPictureHash(Struct): pass
+struct_H265RawSEIDecodedPictureHash._fields_ = [
+  ('hash_type', uint8_t),
+  ('picture_md5', ((uint8_t * 16) * 3)),
+  ('picture_crc', (uint16_t * 3)),
+  ('picture_checksum', (uint32_t * 3)),
+]
+H265RawSEIDecodedPictureHash = struct_H265RawSEIDecodedPictureHash
+class struct_H265RawSEITimeCode(Struct): pass
+struct_H265RawSEITimeCode._fields_ = [
+  ('num_clock_ts', uint8_t),
+  ('clock_timestamp_flag', (uint8_t * 3)),
+  ('units_field_based_flag', (uint8_t * 3)),
+  ('counting_type', (uint8_t * 3)),
+  ('full_timestamp_flag', (uint8_t * 3)),
+  ('discontinuity_flag', (uint8_t * 3)),
+  ('cnt_dropped_flag', (uint8_t * 3)),
+  ('n_frames', (uint16_t * 3)),
+  ('seconds_value', (uint8_t * 3)),
+  ('minutes_value', (uint8_t * 3)),
+  ('hours_value', (uint8_t * 3)),
+  ('seconds_flag', (uint8_t * 3)),
+  ('minutes_flag', (uint8_t * 3)),
+  ('hours_flag', (uint8_t * 3)),
+  ('time_offset_length', (uint8_t * 3)),
+  ('time_offset_value', (int32_t * 3)),
+]
+H265RawSEITimeCode = struct_H265RawSEITimeCode
+class struct_H265RawSEIAlphaChannelInfo(Struct): pass
+struct_H265RawSEIAlphaChannelInfo._fields_ = [
+  ('alpha_channel_cancel_flag', uint8_t),
+  ('alpha_channel_use_idc', uint8_t),
+  ('alpha_channel_bit_depth_minus8', uint8_t),
+  ('alpha_transparent_value', uint16_t),
+  ('alpha_opaque_value', uint16_t),
+  ('alpha_channel_incr_flag', uint8_t),
+  ('alpha_channel_clip_flag', uint8_t),
+  ('alpha_channel_clip_type_flag', uint8_t),
+]
+H265RawSEIAlphaChannelInfo = struct_H265RawSEIAlphaChannelInfo
+class struct_H265RawSEI3DReferenceDisplaysInfo(Struct): pass
+struct_H265RawSEI3DReferenceDisplaysInfo._fields_ = [
+  ('prec_ref_display_width', uint8_t),
+  ('ref_viewing_distance_flag', uint8_t),
+  ('prec_ref_viewing_dist', uint8_t),
+  ('num_ref_displays_minus1', uint8_t),
+  ('left_view_id', (uint16_t * 32)),
+  ('right_view_id', (uint16_t * 32)),
+  ('exponent_ref_display_width', (uint8_t * 32)),
+  ('mantissa_ref_display_width', (uint8_t * 32)),
+  ('exponent_ref_viewing_distance', (uint8_t * 32)),
+  ('mantissa_ref_viewing_distance', (uint8_t * 32)),
+  ('additional_shift_present_flag', (uint8_t * 32)),
+  ('num_sample_shift_plus512', (uint16_t * 32)),
+  ('three_dimensional_reference_displays_extension_flag', uint8_t),
+]
+H265RawSEI3DReferenceDisplaysInfo = struct_H265RawSEI3DReferenceDisplaysInfo
+class struct_H265RawSEI(Struct): pass
+class struct_SEIRawMessageList(Struct): pass
+SEIRawMessageList = struct_SEIRawMessageList
+class struct_SEIRawMessage(Struct): pass
+SEIRawMessage = struct_SEIRawMessage
+size_t = ctypes.c_uint64
+struct_SEIRawMessage._fields_ = [
+  ('payload_type', uint32_t),
+  ('payload_size', uint32_t),
+  ('payload', ctypes.c_void_p),
+  ('payload_ref', ctypes.c_void_p),
+  ('extension_data', ctypes.POINTER(uint8_t)),
+  ('extension_bit_length', size_t),
+]
+struct_SEIRawMessageList._fields_ = [
+  ('messages', ctypes.POINTER(SEIRawMessage)),
+  ('nb_messages', ctypes.c_int32),
+  ('nb_messages_allocated', ctypes.c_int32),
+]
+struct_H265RawSEI._fields_ = [
+  ('nal_unit_header', H265RawNALUnitHeader),
+  ('message_list', SEIRawMessageList),
+]
+H265RawSEI = struct_H265RawSEI
+class struct_H265RawFiller(Struct): pass
+struct_H265RawFiller._fields_ = [
+  ('nal_unit_header', H265RawNALUnitHeader),
+  ('filler_size', uint32_t),
+]
+H265RawFiller = struct_H265RawFiller
+class struct_CodedBitstreamH265Context(Struct): pass
+CodedBitstreamH265Context = struct_CodedBitstreamH265Context

tinygrad/runtime/autogen/hsa.py

@@ -11,6 +11,248 @@ def dll():
   return None
 dll = dll()
 
+enum_SQ_RSRC_BUF_TYPE = CEnum(ctypes.c_uint32)
+SQ_RSRC_BUF = enum_SQ_RSRC_BUF_TYPE.define('SQ_RSRC_BUF', 0)
+SQ_RSRC_BUF_RSVD_1 = enum_SQ_RSRC_BUF_TYPE.define('SQ_RSRC_BUF_RSVD_1', 1)
+SQ_RSRC_BUF_RSVD_2 = enum_SQ_RSRC_BUF_TYPE.define('SQ_RSRC_BUF_RSVD_2', 2)
+SQ_RSRC_BUF_RSVD_3 = enum_SQ_RSRC_BUF_TYPE.define('SQ_RSRC_BUF_RSVD_3', 3)
+
+SQ_RSRC_BUF_TYPE = enum_SQ_RSRC_BUF_TYPE
+enum_BUF_DATA_FORMAT = CEnum(ctypes.c_uint32)
+BUF_DATA_FORMAT_INVALID = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_INVALID', 0)
+BUF_DATA_FORMAT_8 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_8', 1)
+BUF_DATA_FORMAT_16 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_16', 2)
+BUF_DATA_FORMAT_8_8 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_8_8', 3)
+BUF_DATA_FORMAT_32 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_32', 4)
+BUF_DATA_FORMAT_16_16 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_16_16', 5)
+BUF_DATA_FORMAT_10_11_11 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_10_11_11', 6)
+BUF_DATA_FORMAT_11_11_10 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_11_11_10', 7)
+BUF_DATA_FORMAT_10_10_10_2 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_10_10_10_2', 8)
+BUF_DATA_FORMAT_2_10_10_10 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_2_10_10_10', 9)
+BUF_DATA_FORMAT_8_8_8_8 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_8_8_8_8', 10)
+BUF_DATA_FORMAT_32_32 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_32_32', 11)
+BUF_DATA_FORMAT_16_16_16_16 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_16_16_16_16', 12)
+BUF_DATA_FORMAT_32_32_32 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_32_32_32', 13)
+BUF_DATA_FORMAT_32_32_32_32 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_32_32_32_32', 14)
+BUF_DATA_FORMAT_RESERVED_15 = enum_BUF_DATA_FORMAT.define('BUF_DATA_FORMAT_RESERVED_15', 15)
+
+BUF_DATA_FORMAT = enum_BUF_DATA_FORMAT
+enum_BUF_NUM_FORMAT = CEnum(ctypes.c_uint32)
+BUF_NUM_FORMAT_UNORM = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_UNORM', 0)
+BUF_NUM_FORMAT_SNORM = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_SNORM', 1)
+BUF_NUM_FORMAT_USCALED = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_USCALED', 2)
+BUF_NUM_FORMAT_SSCALED = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_SSCALED', 3)
+BUF_NUM_FORMAT_UINT = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_UINT', 4)
+BUF_NUM_FORMAT_SINT = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_SINT', 5)
+BUF_NUM_FORMAT_SNORM_OGL__SI__CI = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_SNORM_OGL__SI__CI', 6)
+BUF_NUM_FORMAT_RESERVED_6__VI = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_RESERVED_6__VI', 6)
+BUF_NUM_FORMAT_FLOAT = enum_BUF_NUM_FORMAT.define('BUF_NUM_FORMAT_FLOAT', 7)
+
+BUF_NUM_FORMAT = enum_BUF_NUM_FORMAT
+enum_BUF_FORMAT = CEnum(ctypes.c_uint32)
+BUF_FORMAT_32_UINT = enum_BUF_FORMAT.define('BUF_FORMAT_32_UINT', 20)
+
+BUF_FORMAT = enum_BUF_FORMAT
+enum_SQ_SEL_XYZW01 = CEnum(ctypes.c_uint32)
+SQ_SEL_0 = enum_SQ_SEL_XYZW01.define('SQ_SEL_0', 0)
+SQ_SEL_1 = enum_SQ_SEL_XYZW01.define('SQ_SEL_1', 1)
+SQ_SEL_RESERVED_0 = enum_SQ_SEL_XYZW01.define('SQ_SEL_RESERVED_0', 2)
+SQ_SEL_RESERVED_1 = enum_SQ_SEL_XYZW01.define('SQ_SEL_RESERVED_1', 3)
+SQ_SEL_X = enum_SQ_SEL_XYZW01.define('SQ_SEL_X', 4)
+SQ_SEL_Y = enum_SQ_SEL_XYZW01.define('SQ_SEL_Y', 5)
+SQ_SEL_Z = enum_SQ_SEL_XYZW01.define('SQ_SEL_Z', 6)
+SQ_SEL_W = enum_SQ_SEL_XYZW01.define('SQ_SEL_W', 7)
+
+SQ_SEL_XYZW01 = enum_SQ_SEL_XYZW01
+class union_COMPUTE_TMPRING_SIZE(ctypes.Union): pass
+class union_COMPUTE_TMPRING_SIZE_bitfields(Struct): pass
+union_COMPUTE_TMPRING_SIZE_bitfields._fields_ = [
+  ('WAVES', ctypes.c_uint32,12),
+  ('WAVESIZE', ctypes.c_uint32,13),
+  ('', ctypes.c_uint32,7),
+]
+union_COMPUTE_TMPRING_SIZE._fields_ = [
+  ('bitfields', union_COMPUTE_TMPRING_SIZE_bitfields),
+  ('bits', union_COMPUTE_TMPRING_SIZE_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_COMPUTE_TMPRING_SIZE_GFX11(ctypes.Union): pass
+class union_COMPUTE_TMPRING_SIZE_GFX11_bitfields(Struct): pass
+union_COMPUTE_TMPRING_SIZE_GFX11_bitfields._fields_ = [
+  ('WAVES', ctypes.c_uint32,12),
+  ('WAVESIZE', ctypes.c_uint32,15),
+  ('', ctypes.c_uint32,5),
+]
+union_COMPUTE_TMPRING_SIZE_GFX11._fields_ = [
+  ('bitfields', union_COMPUTE_TMPRING_SIZE_GFX11_bitfields),
+  ('bits', union_COMPUTE_TMPRING_SIZE_GFX11_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_COMPUTE_TMPRING_SIZE_GFX12(ctypes.Union): pass
+class union_COMPUTE_TMPRING_SIZE_GFX12_bitfields(Struct): pass
+union_COMPUTE_TMPRING_SIZE_GFX12_bitfields._fields_ = [
+  ('WAVES', ctypes.c_uint32,12),
+  ('WAVESIZE', ctypes.c_uint32,18),
+  ('', ctypes.c_uint32,2),
+]
+union_COMPUTE_TMPRING_SIZE_GFX12._fields_ = [
+  ('bitfields', union_COMPUTE_TMPRING_SIZE_GFX12_bitfields),
+  ('bits', union_COMPUTE_TMPRING_SIZE_GFX12_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD0(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD0_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD0_bitfields._fields_ = [
+  ('BASE_ADDRESS', ctypes.c_uint32,32),
+]
+union_SQ_BUF_RSRC_WORD0._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD0_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD0_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD1(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD1_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD1_bitfields._fields_ = [
+  ('BASE_ADDRESS_HI', ctypes.c_uint32,16),
+  ('STRIDE', ctypes.c_uint32,14),
+  ('CACHE_SWIZZLE', ctypes.c_uint32,1),
+  ('SWIZZLE_ENABLE', ctypes.c_uint32,1),
+]
+union_SQ_BUF_RSRC_WORD1._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD1_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD1_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD1_GFX11(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD1_GFX11_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD1_GFX11_bitfields._fields_ = [
+  ('BASE_ADDRESS_HI', ctypes.c_uint32,16),
+  ('STRIDE', ctypes.c_uint32,14),
+  ('SWIZZLE_ENABLE', ctypes.c_uint32,2),
+]
+union_SQ_BUF_RSRC_WORD1_GFX11._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD1_GFX11_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD1_GFX11_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD2(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD2_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD2_bitfields._fields_ = [
+  ('NUM_RECORDS', ctypes.c_uint32,32),
+]
+union_SQ_BUF_RSRC_WORD2._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD2_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD2_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD3(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD3_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD3_bitfields._fields_ = [
+  ('DST_SEL_X', ctypes.c_uint32,3),
+  ('DST_SEL_Y', ctypes.c_uint32,3),
+  ('DST_SEL_Z', ctypes.c_uint32,3),
+  ('DST_SEL_W', ctypes.c_uint32,3),
+  ('NUM_FORMAT', ctypes.c_uint32,3),
+  ('DATA_FORMAT', ctypes.c_uint32,4),
+  ('ELEMENT_SIZE', ctypes.c_uint32,2),
+  ('INDEX_STRIDE', ctypes.c_uint32,2),
+  ('ADD_TID_ENABLE', ctypes.c_uint32,1),
+  ('ATC__CI__VI', ctypes.c_uint32,1),
+  ('HASH_ENABLE', ctypes.c_uint32,1),
+  ('HEAP', ctypes.c_uint32,1),
+  ('MTYPE__CI__VI', ctypes.c_uint32,3),
+  ('TYPE', ctypes.c_uint32,2),
+]
+union_SQ_BUF_RSRC_WORD3._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD3_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD3_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD3_GFX10(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD3_GFX10_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD3_GFX10_bitfields._fields_ = [
+  ('DST_SEL_X', ctypes.c_uint32,3),
+  ('DST_SEL_Y', ctypes.c_uint32,3),
+  ('DST_SEL_Z', ctypes.c_uint32,3),
+  ('DST_SEL_W', ctypes.c_uint32,3),
+  ('FORMAT', ctypes.c_uint32,7),
+  ('RESERVED1', ctypes.c_uint32,2),
+  ('INDEX_STRIDE', ctypes.c_uint32,2),
+  ('ADD_TID_ENABLE', ctypes.c_uint32,1),
+  ('RESOURCE_LEVEL', ctypes.c_uint32,1),
+  ('RESERVED2', ctypes.c_uint32,3),
+  ('OOB_SELECT', ctypes.c_uint32,2),
+  ('TYPE', ctypes.c_uint32,2),
+]
+union_SQ_BUF_RSRC_WORD3_GFX10._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD3_GFX10_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD3_GFX10_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD3_GFX11(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD3_GFX11_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD3_GFX11_bitfields._fields_ = [
+  ('DST_SEL_X', ctypes.c_uint32,3),
+  ('DST_SEL_Y', ctypes.c_uint32,3),
+  ('DST_SEL_Z', ctypes.c_uint32,3),
+  ('DST_SEL_W', ctypes.c_uint32,3),
+  ('FORMAT', ctypes.c_uint32,6),
+  ('RESERVED1', ctypes.c_uint32,3),
+  ('INDEX_STRIDE', ctypes.c_uint32,2),
+  ('ADD_TID_ENABLE', ctypes.c_uint32,1),
+  ('RESERVED2', ctypes.c_uint32,4),
+  ('OOB_SELECT', ctypes.c_uint32,2),
+  ('TYPE', ctypes.c_uint32,2),
+]
+union_SQ_BUF_RSRC_WORD3_GFX11._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD3_GFX11_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD3_GFX11_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
+class union_SQ_BUF_RSRC_WORD3_GFX12(ctypes.Union): pass
+class union_SQ_BUF_RSRC_WORD3_GFX12_bitfields(Struct): pass
+union_SQ_BUF_RSRC_WORD3_GFX12_bitfields._fields_ = [
+  ('DST_SEL_X', ctypes.c_uint32,3),
+  ('DST_SEL_Y', ctypes.c_uint32,3),
+  ('DST_SEL_Z', ctypes.c_uint32,3),
+  ('DST_SEL_W', ctypes.c_uint32,3),
+  ('FORMAT', ctypes.c_uint32,6),
+  ('RESERVED1', ctypes.c_uint32,3),
+  ('INDEX_STRIDE', ctypes.c_uint32,2),
+  ('ADD_TID_ENABLE', ctypes.c_uint32,1),
+  ('WRITE_COMPRESS_ENABLE', ctypes.c_uint32,1),
+  ('COMPRESSION_EN', ctypes.c_uint32,1),
+  ('COMPRESSION_ACCESS_MODE', ctypes.c_uint32,2),
+  ('OOB_SELECT', ctypes.c_uint32,2),
+  ('TYPE', ctypes.c_uint32,2),
+]
+union_SQ_BUF_RSRC_WORD3_GFX12._fields_ = [
+  ('bitfields', union_SQ_BUF_RSRC_WORD3_GFX12_bitfields),
+  ('bits', union_SQ_BUF_RSRC_WORD3_GFX12_bitfields),
+  ('u32All', ctypes.c_uint32),
+  ('i32All', ctypes.c_int32),
+  ('f32All', ctypes.c_float),
+]
 hsa_status_t = CEnum(ctypes.c_uint32)
 HSA_STATUS_SUCCESS = hsa_status_t.define('HSA_STATUS_SUCCESS', 0)
 HSA_STATUS_INFO_BREAK = hsa_status_t.define('HSA_STATUS_INFO_BREAK', 1)
@@ -150,6 +392,7 @@ hsa_device_type_t = CEnum(ctypes.c_uint32)
 HSA_DEVICE_TYPE_CPU = hsa_device_type_t.define('HSA_DEVICE_TYPE_CPU', 0)
 HSA_DEVICE_TYPE_GPU = hsa_device_type_t.define('HSA_DEVICE_TYPE_GPU', 1)
 HSA_DEVICE_TYPE_DSP = hsa_device_type_t.define('HSA_DEVICE_TYPE_DSP', 2)
+HSA_DEVICE_TYPE_AIE = hsa_device_type_t.define('HSA_DEVICE_TYPE_AIE', 3)
 
 hsa_default_float_rounding_mode_t = CEnum(ctypes.c_uint32)
 HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT = hsa_default_float_rounding_mode_t.define('HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT', 0)
@@ -591,9 +834,20 @@ hsa_kernel_dispatch_packet_setup_width_t = CEnum(ctypes.c_uint32)
 HSA_KERNEL_DISPATCH_PACKET_SETUP_WIDTH_DIMENSIONS = hsa_kernel_dispatch_packet_setup_width_t.define('HSA_KERNEL_DISPATCH_PACKET_SETUP_WIDTH_DIMENSIONS', 2)
 
 class struct_hsa_kernel_dispatch_packet_s(Struct): pass
-struct_hsa_kernel_dispatch_packet_s._fields_ = [
+class struct_hsa_kernel_dispatch_packet_s_0(ctypes.Union): pass
+class struct_hsa_kernel_dispatch_packet_s_0_0(Struct): pass
+struct_hsa_kernel_dispatch_packet_s_0_0._fields_ = [
   ('header', uint16_t),
   ('setup', uint16_t),
+]
+struct_hsa_kernel_dispatch_packet_s_0._anonymous_ = ['_0']
+struct_hsa_kernel_dispatch_packet_s_0._fields_ = [
+  ('_0', struct_hsa_kernel_dispatch_packet_s_0_0),
+  ('full_header', uint32_t),
+]
+struct_hsa_kernel_dispatch_packet_s._anonymous_ = ['_0']
+struct_hsa_kernel_dispatch_packet_s._fields_ = [
+  ('_0', struct_hsa_kernel_dispatch_packet_s_0),
   ('workgroup_size_x', uint16_t),
   ('workgroup_size_y', uint16_t),
   ('workgroup_size_z', uint16_t),
@@ -970,6 +1224,7 @@ except AttributeError: pass
 hsa_signal_condition32_t = ctypes.c_uint32
 hsa_amd_packet_type_t = CEnum(ctypes.c_uint32)
 HSA_AMD_PACKET_TYPE_BARRIER_VALUE = hsa_amd_packet_type_t.define('HSA_AMD_PACKET_TYPE_BARRIER_VALUE', 2)
+HSA_AMD_PACKET_TYPE_AIE_ERT = hsa_amd_packet_type_t.define('HSA_AMD_PACKET_TYPE_AIE_ERT', 3)
 
 hsa_amd_packet_type8_t = ctypes.c_ubyte
 class struct_hsa_amd_packet_header_s(Struct): pass
@@ -994,6 +1249,72 @@ struct_hsa_amd_barrier_value_packet_s._fields_ = [
   ('completion_signal', hsa_signal_t),
 ]
 hsa_amd_barrier_value_packet_t = struct_hsa_amd_barrier_value_packet_s
+hsa_amd_aie_ert_state = CEnum(ctypes.c_uint32)
+HSA_AMD_AIE_ERT_STATE_NEW = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_NEW', 1)
+HSA_AMD_AIE_ERT_STATE_QUEUED = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_QUEUED', 2)
+HSA_AMD_AIE_ERT_STATE_RUNNING = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_RUNNING', 3)
+HSA_AMD_AIE_ERT_STATE_COMPLETED = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_COMPLETED', 4)
+HSA_AMD_AIE_ERT_STATE_ERROR = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_ERROR', 5)
+HSA_AMD_AIE_ERT_STATE_ABORT = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_ABORT', 6)
+HSA_AMD_AIE_ERT_STATE_SUBMITTED = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_SUBMITTED', 7)
+HSA_AMD_AIE_ERT_STATE_TIMEOUT = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_TIMEOUT', 8)
+HSA_AMD_AIE_ERT_STATE_NORESPONSE = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_NORESPONSE', 9)
+HSA_AMD_AIE_ERT_STATE_SKERROR = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_SKERROR', 10)
+HSA_AMD_AIE_ERT_STATE_SKCRASHED = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_SKCRASHED', 11)
+HSA_AMD_AIE_ERT_STATE_MAX = hsa_amd_aie_ert_state.define('HSA_AMD_AIE_ERT_STATE_MAX', 12)
+
+hsa_amd_aie_ert_cmd_opcode_t = CEnum(ctypes.c_uint32)
+HSA_AMD_AIE_ERT_START_CU = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_CU', 0)
+HSA_AMD_AIE_ERT_START_KERNEL = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_KERNEL', 0)
+HSA_AMD_AIE_ERT_CONFIGURE = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_CONFIGURE', 2)
+HSA_AMD_AIE_ERT_EXIT = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_EXIT', 3)
+HSA_AMD_AIE_ERT_ABORT = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_ABORT', 4)
+HSA_AMD_AIE_ERT_EXEC_WRITE = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_EXEC_WRITE', 5)
+HSA_AMD_AIE_ERT_CU_STAT = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_CU_STAT', 6)
+HSA_AMD_AIE_ERT_START_COPYBO = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_COPYBO', 7)
+HSA_AMD_AIE_ERT_SK_CONFIG = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_SK_CONFIG', 8)
+HSA_AMD_AIE_ERT_SK_START = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_SK_START', 9)
+HSA_AMD_AIE_ERT_SK_UNCONFIG = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_SK_UNCONFIG', 10)
+HSA_AMD_AIE_ERT_INIT_CU = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_INIT_CU', 11)
+HSA_AMD_AIE_ERT_START_FA = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_FA', 12)
+HSA_AMD_AIE_ERT_CLK_CALIB = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_CLK_CALIB', 13)
+HSA_AMD_AIE_ERT_MB_VALIDATE = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_MB_VALIDATE', 14)
+HSA_AMD_AIE_ERT_START_KEY_VAL = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_KEY_VAL', 15)
+HSA_AMD_AIE_ERT_ACCESS_TEST_C = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_ACCESS_TEST_C', 16)
+HSA_AMD_AIE_ERT_ACCESS_TEST = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_ACCESS_TEST', 17)
+HSA_AMD_AIE_ERT_START_DPU = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_DPU', 18)
+HSA_AMD_AIE_ERT_CMD_CHAIN = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_CMD_CHAIN', 19)
+HSA_AMD_AIE_ERT_START_NPU = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_NPU', 20)
+HSA_AMD_AIE_ERT_START_NPU_PREEMPT = hsa_amd_aie_ert_cmd_opcode_t.define('HSA_AMD_AIE_ERT_START_NPU_PREEMPT', 21)
+
+class struct_hsa_amd_aie_ert_start_kernel_data_s(Struct): pass
+struct_hsa_amd_aie_ert_start_kernel_data_s._fields_ = [
+  ('pdi_addr', ctypes.c_void_p),
+  ('data', (uint32_t * 0)),
+]
+hsa_amd_aie_ert_start_kernel_data_t = struct_hsa_amd_aie_ert_start_kernel_data_s
+class struct_hsa_amd_aie_ert_packet_s(Struct): pass
+class struct_hsa_amd_aie_ert_packet_s_0(Struct): pass
+struct_hsa_amd_aie_ert_packet_s_0._fields_ = [
+  ('state', uint32_t,4),
+  ('custom', uint32_t,8),
+  ('count', uint32_t,11),
+  ('opcode', uint32_t,5),
+  ('type', uint32_t,4),
+]
+struct_hsa_amd_aie_ert_packet_s._anonymous_ = ['_0']
+struct_hsa_amd_aie_ert_packet_s._fields_ = [
+  ('header', hsa_amd_vendor_packet_header_t),
+  ('_0', struct_hsa_amd_aie_ert_packet_s_0),
+  ('reserved0', uint64_t),
+  ('reserved1', uint64_t),
+  ('reserved2', uint64_t),
+  ('reserved3', uint64_t),
+  ('reserved4', uint64_t),
+  ('reserved5', uint64_t),
+  ('payload_data', uint64_t),
+]
+hsa_amd_aie_ert_packet_t = struct_hsa_amd_aie_ert_packet_s
 _anonenum0 = CEnum(ctypes.c_uint32)
 HSA_STATUS_ERROR_INVALID_MEMORY_POOL = _anonenum0.define('HSA_STATUS_ERROR_INVALID_MEMORY_POOL', 40)
 HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION = _anonenum0.define('HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION', 41)
@@ -1002,11 +1323,20 @@ HSA_STATUS_ERROR_MEMORY_FAULT = _anonenum0.define('HSA_STATUS_ERROR_MEMORY_FAULT
 HSA_STATUS_CU_MASK_REDUCED = _anonenum0.define('HSA_STATUS_CU_MASK_REDUCED', 44)
 HSA_STATUS_ERROR_OUT_OF_REGISTERS = _anonenum0.define('HSA_STATUS_ERROR_OUT_OF_REGISTERS', 45)
 HSA_STATUS_ERROR_RESOURCE_BUSY = _anonenum0.define('HSA_STATUS_ERROR_RESOURCE_BUSY', 46)
+HSA_STATUS_ERROR_NOT_SUPPORTED = _anonenum0.define('HSA_STATUS_ERROR_NOT_SUPPORTED', 47)
 
 hsa_amd_iommu_version_t = CEnum(ctypes.c_uint32)
 HSA_IOMMU_SUPPORT_NONE = hsa_amd_iommu_version_t.define('HSA_IOMMU_SUPPORT_NONE', 0)
 HSA_IOMMU_SUPPORT_V2 = hsa_amd_iommu_version_t.define('HSA_IOMMU_SUPPORT_V2', 1)
 
+class struct_hsa_amd_clock_counters_s(Struct): pass
+struct_hsa_amd_clock_counters_s._fields_ = [
+  ('gpu_clock_counter', uint64_t),
+  ('cpu_clock_counter', uint64_t),
+  ('system_clock_counter', uint64_t),
+  ('system_clock_frequency', uint64_t),
+]
+hsa_amd_clock_counters_t = struct_hsa_amd_clock_counters_s
 enum_hsa_amd_agent_info_s = CEnum(ctypes.c_uint32)
 HSA_AMD_AGENT_INFO_CHIP_ID = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_CHIP_ID', 40960)
 HSA_AMD_AGENT_INFO_CACHELINE_SIZE = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_CACHELINE_SIZE', 40961)
@@ -1042,6 +1372,9 @@ HSA_AMD_AGENT_INFO_DRIVER_UID = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_
 HSA_AMD_AGENT_INFO_NEAREST_CPU = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_NEAREST_CPU', 41235)
 HSA_AMD_AGENT_INFO_MEMORY_PROPERTIES = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_MEMORY_PROPERTIES', 41236)
 HSA_AMD_AGENT_INFO_AQL_EXTENSIONS = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_AQL_EXTENSIONS', 41237)
+HSA_AMD_AGENT_INFO_SCRATCH_LIMIT_MAX = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_SCRATCH_LIMIT_MAX', 41238)
+HSA_AMD_AGENT_INFO_SCRATCH_LIMIT_CURRENT = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_SCRATCH_LIMIT_CURRENT', 41239)
+HSA_AMD_AGENT_INFO_CLOCK_COUNTERS = enum_hsa_amd_agent_info_s.define('HSA_AMD_AGENT_INFO_CLOCK_COUNTERS', 41240)
 
 hsa_amd_agent_info_t = enum_hsa_amd_agent_info_s
 enum_hsa_amd_agent_memory_properties_s = CEnum(ctypes.c_uint32)
@@ -1085,6 +1418,11 @@ HSA_AMD_COHERENCY_TYPE_COHERENT = enum_hsa_amd_coherency_type_s.define('HSA_AMD_
 HSA_AMD_COHERENCY_TYPE_NONCOHERENT = enum_hsa_amd_coherency_type_s.define('HSA_AMD_COHERENCY_TYPE_NONCOHERENT', 1)
 
 hsa_amd_coherency_type_t = enum_hsa_amd_coherency_type_s
+enum_hsa_amd_dma_buf_mapping_type_s = CEnum(ctypes.c_uint32)
+HSA_AMD_DMABUF_MAPPING_TYPE_NONE = enum_hsa_amd_dma_buf_mapping_type_s.define('HSA_AMD_DMABUF_MAPPING_TYPE_NONE', 0)
+HSA_AMD_DMABUF_MAPPING_TYPE_PCIE = enum_hsa_amd_dma_buf_mapping_type_s.define('HSA_AMD_DMABUF_MAPPING_TYPE_PCIE', 1)
+
+hsa_amd_dma_buf_mapping_type_t = enum_hsa_amd_dma_buf_mapping_type_s
 try: (hsa_amd_coherency_get_type:=dll.hsa_amd_coherency_get_type).restype, hsa_amd_coherency_get_type.argtypes = hsa_status_t, [hsa_agent_t, ctypes.POINTER(hsa_amd_coherency_type_t)]
 except AttributeError: pass
 
@@ -1132,12 +1470,58 @@ hsa_amd_signal_handler = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_int64, ctypes.
 try: (hsa_amd_signal_async_handler:=dll.hsa_amd_signal_async_handler).restype, hsa_amd_signal_async_handler.argtypes = hsa_status_t, [hsa_signal_t, hsa_signal_condition_t, hsa_signal_value_t, hsa_amd_signal_handler, ctypes.c_void_p]
 except AttributeError: pass
 
-try: (hsa_amd_async_function:=dll.hsa_amd_async_function).restype, hsa_amd_async_function.argtypes = hsa_status_t, [ctypes.CFUNCTYPE(None, ctypes.c_void_p), ctypes.c_void_p]
+try: (hsa_amd_signal_wait_all:=dll.hsa_amd_signal_wait_all).restype, hsa_amd_signal_wait_all.argtypes = uint32_t, [uint32_t, ctypes.POINTER(hsa_signal_t), ctypes.POINTER(hsa_signal_condition_t), ctypes.POINTER(hsa_signal_value_t), uint64_t, hsa_wait_state_t, ctypes.POINTER(hsa_signal_value_t)]
 except AttributeError: pass
 
 try: (hsa_amd_signal_wait_any:=dll.hsa_amd_signal_wait_any).restype, hsa_amd_signal_wait_any.argtypes = uint32_t, [uint32_t, ctypes.POINTER(hsa_signal_t), ctypes.POINTER(hsa_signal_condition_t), ctypes.POINTER(hsa_signal_value_t), uint64_t, hsa_wait_state_t, ctypes.POINTER(hsa_signal_value_t)]
 except AttributeError: pass
 
+try: (hsa_amd_async_function:=dll.hsa_amd_async_function).restype, hsa_amd_async_function.argtypes = hsa_status_t, [ctypes.CFUNCTYPE(None, ctypes.c_void_p), ctypes.c_void_p]
+except AttributeError: pass
+
+class struct_hsa_amd_image_descriptor_s(Struct): pass
+struct_hsa_amd_image_descriptor_s._fields_ = [
+  ('version', uint32_t),
+  ('deviceID', uint32_t),
+  ('data', (uint32_t * 1)),
+]
+hsa_amd_image_descriptor_t = struct_hsa_amd_image_descriptor_s
+class struct_hsa_ext_image_descriptor_s(Struct): pass
+hsa_ext_image_descriptor_t = struct_hsa_ext_image_descriptor_s
+hsa_ext_image_geometry_t = CEnum(ctypes.c_uint32)
+HSA_EXT_IMAGE_GEOMETRY_1D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1D', 0)
+HSA_EXT_IMAGE_GEOMETRY_2D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2D', 1)
+HSA_EXT_IMAGE_GEOMETRY_3D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_3D', 2)
+HSA_EXT_IMAGE_GEOMETRY_1DA = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1DA', 3)
+HSA_EXT_IMAGE_GEOMETRY_2DA = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DA', 4)
+HSA_EXT_IMAGE_GEOMETRY_1DB = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1DB', 5)
+HSA_EXT_IMAGE_GEOMETRY_2DDEPTH = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DDEPTH', 6)
+HSA_EXT_IMAGE_GEOMETRY_2DADEPTH = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DADEPTH', 7)
+
+class struct_hsa_ext_image_format_s(Struct): pass
+hsa_ext_image_format_t = struct_hsa_ext_image_format_s
+hsa_ext_image_channel_type32_t = ctypes.c_uint32
+hsa_ext_image_channel_order32_t = ctypes.c_uint32
+struct_hsa_ext_image_format_s._fields_ = [
+  ('channel_type', hsa_ext_image_channel_type32_t),
+  ('channel_order', hsa_ext_image_channel_order32_t),
+]
+struct_hsa_ext_image_descriptor_s._fields_ = [
+  ('geometry', hsa_ext_image_geometry_t),
+  ('width', size_t),
+  ('height', size_t),
+  ('depth', size_t),
+  ('array_size', size_t),
+  ('format', hsa_ext_image_format_t),
+]
+class struct_hsa_ext_image_s(Struct): pass
+hsa_ext_image_t = struct_hsa_ext_image_s
+struct_hsa_ext_image_s._fields_ = [
+  ('handle', uint64_t),
+]
+try: (hsa_amd_image_create:=dll.hsa_amd_image_create).restype, hsa_amd_image_create.argtypes = hsa_status_t, [hsa_agent_t, ctypes.POINTER(hsa_ext_image_descriptor_t), ctypes.POINTER(hsa_amd_image_descriptor_t), ctypes.c_void_p, hsa_access_permission_t, ctypes.POINTER(hsa_ext_image_t)]
+except AttributeError: pass
+
 try: (hsa_amd_image_get_info_max_dim:=dll.hsa_amd_image_get_info_max_dim).restype, hsa_amd_image_get_info_max_dim.argtypes = hsa_status_t, [hsa_agent_t, hsa_agent_info_t, ctypes.c_void_p]
 except AttributeError: pass
 
@@ -1186,6 +1570,8 @@ enum_hsa_amd_memory_pool_flag_s = CEnum(ctypes.c_uint32)
 HSA_AMD_MEMORY_POOL_STANDARD_FLAG = enum_hsa_amd_memory_pool_flag_s.define('HSA_AMD_MEMORY_POOL_STANDARD_FLAG', 0)
 HSA_AMD_MEMORY_POOL_PCIE_FLAG = enum_hsa_amd_memory_pool_flag_s.define('HSA_AMD_MEMORY_POOL_PCIE_FLAG', 1)
 HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG = enum_hsa_amd_memory_pool_flag_s.define('HSA_AMD_MEMORY_POOL_CONTIGUOUS_FLAG', 2)
+HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG = enum_hsa_amd_memory_pool_flag_s.define('HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG', 4)
+HSA_AMD_MEMORY_POOL_UNCACHED_FLAG = enum_hsa_amd_memory_pool_flag_s.define('HSA_AMD_MEMORY_POOL_UNCACHED_FLAG', 8)
 
 hsa_amd_memory_pool_flag_t = enum_hsa_amd_memory_pool_flag_s
 try: (hsa_amd_memory_pool_get_info:=dll.hsa_amd_memory_pool_get_info).restype, hsa_amd_memory_pool_get_info.argtypes = hsa_status_t, [hsa_amd_memory_pool_t, hsa_amd_memory_pool_info_t, ctypes.c_void_p]
@@ -1209,6 +1595,9 @@ except AttributeError: pass
 try: (hsa_amd_memory_copy_engine_status:=dll.hsa_amd_memory_copy_engine_status).restype, hsa_amd_memory_copy_engine_status.argtypes = hsa_status_t, [hsa_agent_t, hsa_agent_t, ctypes.POINTER(uint32_t)]
 except AttributeError: pass
 
+try: (hsa_amd_memory_get_preferred_copy_engine:=dll.hsa_amd_memory_get_preferred_copy_engine).restype, hsa_amd_memory_get_preferred_copy_engine.argtypes = hsa_status_t, [hsa_agent_t, hsa_agent_t, ctypes.POINTER(uint32_t)]
+except AttributeError: pass
+
 class struct_hsa_pitched_ptr_s(Struct): pass
 struct_hsa_pitched_ptr_s._fields_ = [
   ('base', ctypes.c_void_p),
@@ -1285,55 +1674,14 @@ except AttributeError: pass
 try: (hsa_amd_interop_unmap_buffer:=dll.hsa_amd_interop_unmap_buffer).restype, hsa_amd_interop_unmap_buffer.argtypes = hsa_status_t, [ctypes.c_void_p]
 except AttributeError: pass
 
-class struct_hsa_amd_image_descriptor_s(Struct): pass
-struct_hsa_amd_image_descriptor_s._fields_ = [
-  ('version', uint32_t),
-  ('deviceID', uint32_t),
-  ('data', (uint32_t * 1)),
-]
-hsa_amd_image_descriptor_t = struct_hsa_amd_image_descriptor_s
-class struct_hsa_ext_image_descriptor_s(Struct): pass
-hsa_ext_image_descriptor_t = struct_hsa_ext_image_descriptor_s
-hsa_ext_image_geometry_t = CEnum(ctypes.c_uint32)
-HSA_EXT_IMAGE_GEOMETRY_1D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1D', 0)
-HSA_EXT_IMAGE_GEOMETRY_2D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2D', 1)
-HSA_EXT_IMAGE_GEOMETRY_3D = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_3D', 2)
-HSA_EXT_IMAGE_GEOMETRY_1DA = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1DA', 3)
-HSA_EXT_IMAGE_GEOMETRY_2DA = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DA', 4)
-HSA_EXT_IMAGE_GEOMETRY_1DB = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_1DB', 5)
-HSA_EXT_IMAGE_GEOMETRY_2DDEPTH = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DDEPTH', 6)
-HSA_EXT_IMAGE_GEOMETRY_2DADEPTH = hsa_ext_image_geometry_t.define('HSA_EXT_IMAGE_GEOMETRY_2DADEPTH', 7)
-
-class struct_hsa_ext_image_format_s(Struct): pass
-hsa_ext_image_format_t = struct_hsa_ext_image_format_s
-hsa_ext_image_channel_type32_t = ctypes.c_uint32
-hsa_ext_image_channel_order32_t = ctypes.c_uint32
-struct_hsa_ext_image_format_s._fields_ = [
-  ('channel_type', hsa_ext_image_channel_type32_t),
-  ('channel_order', hsa_ext_image_channel_order32_t),
-]
-struct_hsa_ext_image_descriptor_s._fields_ = [
-  ('geometry', hsa_ext_image_geometry_t),
-  ('width', size_t),
-  ('height', size_t),
-  ('depth', size_t),
-  ('array_size', size_t),
-  ('format', hsa_ext_image_format_t),
-]
-class struct_hsa_ext_image_s(Struct): pass
-hsa_ext_image_t = struct_hsa_ext_image_s
-struct_hsa_ext_image_s._fields_ = [
-  ('handle', uint64_t),
-]
-try: (hsa_amd_image_create:=dll.hsa_amd_image_create).restype, hsa_amd_image_create.argtypes = hsa_status_t, [hsa_agent_t, ctypes.POINTER(hsa_ext_image_descriptor_t), ctypes.POINTER(hsa_amd_image_descriptor_t), ctypes.c_void_p, hsa_access_permission_t, ctypes.POINTER(hsa_ext_image_t)]
-except AttributeError: pass
-
 hsa_amd_pointer_type_t = CEnum(ctypes.c_uint32)
 HSA_EXT_POINTER_TYPE_UNKNOWN = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_UNKNOWN', 0)
 HSA_EXT_POINTER_TYPE_HSA = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_HSA', 1)
 HSA_EXT_POINTER_TYPE_LOCKED = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_LOCKED', 2)
 HSA_EXT_POINTER_TYPE_GRAPHICS = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_GRAPHICS', 3)
 HSA_EXT_POINTER_TYPE_IPC = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_IPC', 4)
+HSA_EXT_POINTER_TYPE_RESERVED_ADDR = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_RESERVED_ADDR', 5)
+HSA_EXT_POINTER_TYPE_HSA_VMEM = hsa_amd_pointer_type_t.define('HSA_EXT_POINTER_TYPE_HSA_VMEM', 6)
 
 class struct_hsa_amd_pointer_info_s(Struct): pass
 struct_hsa_amd_pointer_info_s._fields_ = [
@@ -1345,6 +1693,7 @@ struct_hsa_amd_pointer_info_s._fields_ = [
   ('userData', ctypes.c_void_p),
   ('agentOwner', hsa_agent_t),
   ('global_flags', uint32_t),
+  ('registered', ctypes.c_bool),
 ]
 hsa_amd_pointer_info_t = struct_hsa_amd_pointer_info_s
 try: (hsa_amd_pointer_info:=dll.hsa_amd_pointer_info).restype, hsa_amd_pointer_info.argtypes = hsa_status_t, [ctypes.c_void_p, ctypes.POINTER(hsa_amd_pointer_info_t), ctypes.CFUNCTYPE(ctypes.c_void_p, size_t), ctypes.POINTER(uint32_t), ctypes.POINTER(ctypes.POINTER(hsa_agent_t))]
@@ -1377,6 +1726,7 @@ except AttributeError: pass
 enum_hsa_amd_event_type_s = CEnum(ctypes.c_uint32)
 HSA_AMD_GPU_MEMORY_FAULT_EVENT = enum_hsa_amd_event_type_s.define('HSA_AMD_GPU_MEMORY_FAULT_EVENT', 0)
 HSA_AMD_GPU_HW_EXCEPTION_EVENT = enum_hsa_amd_event_type_s.define('HSA_AMD_GPU_HW_EXCEPTION_EVENT', 1)
+HSA_AMD_GPU_MEMORY_ERROR_EVENT = enum_hsa_amd_event_type_s.define('HSA_AMD_GPU_MEMORY_ERROR_EVENT', 2)
 
 hsa_amd_event_type_t = enum_hsa_amd_event_type_s
 hsa_amd_memory_fault_reason_t = CEnum(ctypes.c_uint32)
@@ -1396,6 +1746,16 @@ struct_hsa_amd_gpu_memory_fault_info_s._fields_ = [
   ('fault_reason_mask', uint32_t),
 ]
 hsa_amd_gpu_memory_fault_info_t = struct_hsa_amd_gpu_memory_fault_info_s
+hsa_amd_memory_error_reason_t = CEnum(ctypes.c_uint32)
+HSA_AMD_MEMORY_ERROR_MEMORY_IN_USE = hsa_amd_memory_error_reason_t.define('HSA_AMD_MEMORY_ERROR_MEMORY_IN_USE', 1)
+
+class struct_hsa_amd_gpu_memory_error_info_s(Struct): pass
+struct_hsa_amd_gpu_memory_error_info_s._fields_ = [
+  ('agent', hsa_agent_t),
+  ('virtual_address', uint64_t),
+  ('error_reason_mask', uint32_t),
+]
+hsa_amd_gpu_memory_error_info_t = struct_hsa_amd_gpu_memory_error_info_s
 hsa_amd_hw_exception_reset_type_t = CEnum(ctypes.c_uint32)
 HSA_AMD_HW_EXCEPTION_RESET_TYPE_OTHER = hsa_amd_hw_exception_reset_type_t.define('HSA_AMD_HW_EXCEPTION_RESET_TYPE_OTHER', 1)
 
@@ -1415,6 +1775,7 @@ class struct_hsa_amd_event_s_0(ctypes.Union): pass
 struct_hsa_amd_event_s_0._fields_ = [
   ('memory_fault', hsa_amd_gpu_memory_fault_info_t),
   ('hw_exception', hsa_amd_gpu_hw_exception_info_t),
+  ('memory_error', hsa_amd_gpu_memory_error_info_t),
 ]
 struct_hsa_amd_event_s._anonymous_ = ['_0']
 struct_hsa_amd_event_s._fields_ = [
@@ -1441,6 +1802,11 @@ hsa_amd_queue_priority_t = enum_hsa_amd_queue_priority_s
 try: (hsa_amd_queue_set_priority:=dll.hsa_amd_queue_set_priority).restype, hsa_amd_queue_set_priority.argtypes = hsa_status_t, [ctypes.POINTER(hsa_queue_t), hsa_amd_queue_priority_t]
 except AttributeError: pass
 
+hsa_amd_queue_create_flag_t = CEnum(ctypes.c_uint32)
+HSA_AMD_QUEUE_CREATE_SYSTEM_MEM = hsa_amd_queue_create_flag_t.define('HSA_AMD_QUEUE_CREATE_SYSTEM_MEM', 0)
+HSA_AMD_QUEUE_CREATE_DEVICE_MEM_RING_BUF = hsa_amd_queue_create_flag_t.define('HSA_AMD_QUEUE_CREATE_DEVICE_MEM_RING_BUF', 1)
+HSA_AMD_QUEUE_CREATE_DEVICE_MEM_QUEUE_DESCRIPTOR = hsa_amd_queue_create_flag_t.define('HSA_AMD_QUEUE_CREATE_DEVICE_MEM_QUEUE_DESCRIPTOR', 2)
+
 hsa_amd_deallocation_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p)
 try: (hsa_amd_register_deallocation_callback:=dll.hsa_amd_register_deallocation_callback).restype, hsa_amd_register_deallocation_callback.argtypes = hsa_status_t, [ctypes.c_void_p, hsa_amd_deallocation_callback_t, ctypes.c_void_p]
 except AttributeError: pass
@@ -1496,9 +1862,16 @@ except AttributeError: pass
 try: (hsa_amd_portable_export_dmabuf:=dll.hsa_amd_portable_export_dmabuf).restype, hsa_amd_portable_export_dmabuf.argtypes = hsa_status_t, [ctypes.c_void_p, size_t, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(uint64_t)]
 except AttributeError: pass
 
+try: (hsa_amd_portable_export_dmabuf_v2:=dll.hsa_amd_portable_export_dmabuf_v2).restype, hsa_amd_portable_export_dmabuf_v2.argtypes = hsa_status_t, [ctypes.c_void_p, size_t, ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(uint64_t), uint64_t]
+except AttributeError: pass
+
 try: (hsa_amd_portable_close_dmabuf:=dll.hsa_amd_portable_close_dmabuf).restype, hsa_amd_portable_close_dmabuf.argtypes = hsa_status_t, [ctypes.c_int32]
 except AttributeError: pass
 
+enum_hsa_amd_vmem_address_reserve_flag_s = CEnum(ctypes.c_uint32)
+HSA_AMD_VMEM_ADDRESS_NO_REGISTER = enum_hsa_amd_vmem_address_reserve_flag_s.define('HSA_AMD_VMEM_ADDRESS_NO_REGISTER', 1)
+
+hsa_amd_vmem_address_reserve_flag_t = enum_hsa_amd_vmem_address_reserve_flag_s
 try: (hsa_amd_vmem_address_reserve:=dll.hsa_amd_vmem_address_reserve).restype, hsa_amd_vmem_address_reserve.argtypes = hsa_status_t, [ctypes.POINTER(ctypes.c_void_p), size_t, uint64_t, uint64_t]
 except AttributeError: pass
 
@@ -1563,6 +1936,35 @@ HSA_AMD_QUEUE_INFO_DOORBELL_ID = hsa_queue_info_attribute_t.define('HSA_AMD_QUEU
 try: (hsa_amd_queue_get_info:=dll.hsa_amd_queue_get_info).restype, hsa_amd_queue_get_info.argtypes = hsa_status_t, [ctypes.POINTER(hsa_queue_t), hsa_queue_info_attribute_t, ctypes.c_void_p]
 except AttributeError: pass
 
+class struct_hsa_amd_ais_file_handle_s(Struct): pass
+class struct_hsa_amd_ais_file_handle_s_0(ctypes.Union): pass
+struct_hsa_amd_ais_file_handle_s_0._fields_ = [
+  ('handle', ctypes.c_void_p),
+  ('fd', ctypes.c_int32),
+  ('pad', (uint8_t * 8)),
+]
+struct_hsa_amd_ais_file_handle_s._anonymous_ = ['_0']
+struct_hsa_amd_ais_file_handle_s._fields_ = [
+  ('_0', struct_hsa_amd_ais_file_handle_s_0),
+]
+hsa_amd_ais_file_handle_t = struct_hsa_amd_ais_file_handle_s
+int64_t = ctypes.c_int64
+try: (hsa_amd_ais_file_write:=dll.hsa_amd_ais_file_write).restype, hsa_amd_ais_file_write.argtypes = hsa_status_t, [hsa_amd_ais_file_handle_t, ctypes.c_void_p, uint64_t, int64_t, ctypes.POINTER(uint64_t), ctypes.POINTER(int32_t)]
+except AttributeError: pass
+
+try: (hsa_amd_ais_file_read:=dll.hsa_amd_ais_file_read).restype, hsa_amd_ais_file_read.argtypes = hsa_status_t, [hsa_amd_ais_file_handle_t, ctypes.c_void_p, uint64_t, int64_t, ctypes.POINTER(uint64_t), ctypes.POINTER(int32_t)]
+except AttributeError: pass
+
+enum_hsa_amd_log_flag_s = CEnum(ctypes.c_uint32)
+HSA_AMD_LOG_FLAG_BLIT_KERNEL_PKTS = enum_hsa_amd_log_flag_s.define('HSA_AMD_LOG_FLAG_BLIT_KERNEL_PKTS', 0)
+HSA_AMD_LOG_FLAG_AQL = enum_hsa_amd_log_flag_s.define('HSA_AMD_LOG_FLAG_AQL', 0)
+HSA_AMD_LOG_FLAG_SDMA = enum_hsa_amd_log_flag_s.define('HSA_AMD_LOG_FLAG_SDMA', 1)
+HSA_AMD_LOG_FLAG_INFO = enum_hsa_amd_log_flag_s.define('HSA_AMD_LOG_FLAG_INFO', 2)
+
+hsa_amd_log_flag_t = enum_hsa_amd_log_flag_s
+try: (hsa_amd_enable_logging:=dll.hsa_amd_enable_logging).restype, hsa_amd_enable_logging.argtypes = hsa_status_t, [ctypes.POINTER(uint8_t), ctypes.c_void_p]
+except AttributeError: pass
+
 amd_signal_kind64_t = ctypes.c_int64
 enum_amd_signal_kind_t = CEnum(ctypes.c_int32)
 AMD_SIGNAL_KIND_INVALID = enum_amd_signal_kind_t.define('AMD_SIGNAL_KIND_INVALID', 0)
@@ -1572,17 +1974,21 @@ AMD_SIGNAL_KIND_LEGACY_DOORBELL = enum_amd_signal_kind_t.define('AMD_SIGNAL_KIND
 
 class struct_amd_signal_s(Struct): pass
 class struct_amd_signal_s_0(ctypes.Union): pass
-int64_t = ctypes.c_int64
 struct_amd_signal_s_0._fields_ = [
   ('value', int64_t),
-  ('legacy_hardware_doorbell_ptr', ctypes.POINTER(uint32_t)),
   ('hardware_doorbell_ptr', ctypes.POINTER(uint64_t)),
 ]
 class struct_amd_signal_s_1(ctypes.Union): pass
-class struct_amd_queue_s(Struct): pass
-amd_queue_t = struct_amd_queue_s
+class struct_amd_queue_v2_s(Struct): pass
+amd_queue_v2_t = struct_amd_queue_v2_s
 amd_queue_properties32_t = ctypes.c_uint32
-struct_amd_queue_s._fields_ = [
+class struct_scratch_last_used_index_xcc_s(Struct): pass
+scratch_last_used_index_xcc_t = struct_scratch_last_used_index_xcc_s
+struct_scratch_last_used_index_xcc_s._fields_ = [
+  ('main', uint64_t),
+  ('alt', uint64_t),
+]
+struct_amd_queue_v2_s._fields_ = [
   ('hsa_queue', hsa_queue_t),
   ('caps', uint32_t),
   ('reserved1', (uint32_t * 3)),
@@ -1602,21 +2008,21 @@ struct_amd_queue_s._fields_ = [
   ('scratch_backing_memory_byte_size', uint64_t),
   ('scratch_wave64_lane_byte_size', uint32_t),
   ('queue_properties', amd_queue_properties32_t),
-  ('scratch_last_used_index', uint64_t),
+  ('scratch_max_use_index', uint64_t),
   ('queue_inactive_signal', hsa_signal_t),
-  ('reserved4', (uint32_t * 2)),
-  ('alt_scratch_last_used_index', uint64_t),
+  ('alt_scratch_max_use_index', uint64_t),
+  ('alt_scratch_resource_descriptor', (uint32_t * 4)),
   ('alt_scratch_backing_memory_location', uint64_t),
-  ('alt_scratch_backing_memory_byte_size', uint64_t),
   ('alt_scratch_dispatch_limit_x', uint32_t),
   ('alt_scratch_dispatch_limit_y', uint32_t),
   ('alt_scratch_dispatch_limit_z', uint32_t),
   ('alt_scratch_wave64_lane_byte_size', uint32_t),
   ('alt_compute_tmpring_size', uint32_t),
   ('reserved5', uint32_t),
+  ('scratch_last_used_index', (scratch_last_used_index_xcc_t * 128)),
 ]
 struct_amd_signal_s_1._fields_ = [
-  ('queue_ptr', ctypes.POINTER(amd_queue_t)),
+  ('queue_ptr', ctypes.POINTER(amd_queue_v2_t)),
   ('reserved2', uint64_t),
 ]
 struct_amd_signal_s._anonymous_ = ['_0', '_1']
@@ -1654,10 +2060,39 @@ AMD_QUEUE_PROPERTIES_RESERVED1 = enum_amd_queue_properties_t.define('AMD_QUEUE_P
 
 amd_queue_capabilities32_t = ctypes.c_uint32
 enum_amd_queue_capabilities_t = CEnum(ctypes.c_uint32)
-AMD_QUEUE_CAPS_ASYNC_RECLAIM_SHIFT = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_ASYNC_RECLAIM_SHIFT', 0)
-AMD_QUEUE_CAPS_ASYNC_RECLAIM_WIDTH = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_ASYNC_RECLAIM_WIDTH', 1)
-AMD_QUEUE_CAPS_ASYNC_RECLAIM = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_ASYNC_RECLAIM', 1)
+AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM_SHIFT = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM_SHIFT', 0)
+AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM_WIDTH = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM_WIDTH', 1)
+AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_CP_ASYNC_RECLAIM', 1)
+AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM_SHIFT = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM_SHIFT', 1)
+AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM_WIDTH = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM_WIDTH', 1)
+AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM = enum_amd_queue_capabilities_t.define('AMD_QUEUE_CAPS_SW_ASYNC_RECLAIM', 2)
 
+class struct_amd_queue_s(Struct): pass
+struct_amd_queue_s._fields_ = [
+  ('hsa_queue', hsa_queue_t),
+  ('caps', uint32_t),
+  ('reserved1', (uint32_t * 3)),
+  ('write_dispatch_id', uint64_t),
+  ('group_segment_aperture_base_hi', uint32_t),
+  ('private_segment_aperture_base_hi', uint32_t),
+  ('max_cu_id', uint32_t),
+  ('max_wave_id', uint32_t),
+  ('max_legacy_doorbell_dispatch_id_plus_1', uint64_t),
+  ('legacy_doorbell_lock', uint32_t),
+  ('reserved2', (uint32_t * 9)),
+  ('read_dispatch_id', uint64_t),
+  ('read_dispatch_id_field_base_byte_offset', uint32_t),
+  ('compute_tmpring_size', uint32_t),
+  ('scratch_resource_descriptor', (uint32_t * 4)),
+  ('scratch_backing_memory_location', uint64_t),
+  ('reserved3', (uint32_t * 2)),
+  ('scratch_wave64_lane_byte_size', uint32_t),
+  ('queue_properties', amd_queue_properties32_t),
+  ('reserved4', (uint32_t * 2)),
+  ('queue_inactive_signal', hsa_signal_t),
+  ('reserved5', (uint32_t * 14)),
+]
+amd_queue_t = struct_amd_queue_s
 amd_kernel_code_version32_t = ctypes.c_uint32
 enum_amd_kernel_code_version_t = CEnum(ctypes.c_uint32)
 AMD_KERNEL_CODE_VERSION_MAJOR = enum_amd_kernel_code_version_t.define('AMD_KERNEL_CODE_VERSION_MAJOR', 1)
@@ -1830,9 +2265,12 @@ AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Y = enum_amd_kernel_
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z_SHIFT = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z_SHIFT', 9)
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z_WIDTH = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z_WIDTH', 1)
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_GRID_WORKGROUP_COUNT_Z', 512)
-AMD_KERNEL_CODE_PROPERTIES_RESERVED1_SHIFT = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1_SHIFT', 10)
-AMD_KERNEL_CODE_PROPERTIES_RESERVED1_WIDTH = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1_WIDTH', 6)
-AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1', 64512)
+AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32_SHIFT = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32_SHIFT', 10)
+AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32_WIDTH = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32_WIDTH', 1)
+AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32 = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_WAVEFRONT_SIZE32', 1024)
+AMD_KERNEL_CODE_PROPERTIES_RESERVED1_SHIFT = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1_SHIFT', 11)
+AMD_KERNEL_CODE_PROPERTIES_RESERVED1_WIDTH = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1_WIDTH', 5)
+AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_RESERVED1', 63488)
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS_SHIFT = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS_SHIFT', 16)
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS_WIDTH = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS_WIDTH', 1)
 AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS = enum_amd_kernel_code_properties_t.define('AMD_KERNEL_CODE_PROPERTIES_ENABLE_ORDERED_APPEND_GDS', 65536)
@@ -2164,9 +2602,19 @@ struct_hsa_ext_sampler_descriptor_s._fields_ = [
   ('address_mode', hsa_ext_sampler_addressing_mode32_t),
 ]
 hsa_ext_sampler_descriptor_t = struct_hsa_ext_sampler_descriptor_s
+class struct_hsa_ext_sampler_descriptor_v2_s(Struct): pass
+struct_hsa_ext_sampler_descriptor_v2_s._fields_ = [
+  ('coordinate_mode', hsa_ext_sampler_coordinate_mode32_t),
+  ('filter_mode', hsa_ext_sampler_filter_mode32_t),
+  ('address_modes', (hsa_ext_sampler_addressing_mode32_t * 3)),
+]
+hsa_ext_sampler_descriptor_v2_t = struct_hsa_ext_sampler_descriptor_v2_s
 try: (hsa_ext_sampler_create:=dll.hsa_ext_sampler_create).restype, hsa_ext_sampler_create.argtypes = hsa_status_t, [hsa_agent_t, ctypes.POINTER(hsa_ext_sampler_descriptor_t), ctypes.POINTER(hsa_ext_sampler_t)]
 except AttributeError: pass
 
+try: (hsa_ext_sampler_create_v2:=dll.hsa_ext_sampler_create_v2).restype, hsa_ext_sampler_create_v2.argtypes = hsa_status_t, [hsa_agent_t, ctypes.POINTER(hsa_ext_sampler_descriptor_v2_t), ctypes.POINTER(hsa_ext_sampler_t)]
+except AttributeError: pass
+
 try: (hsa_ext_sampler_destroy:=dll.hsa_ext_sampler_destroy).restype, hsa_ext_sampler_destroy.argtypes = hsa_status_t, [hsa_agent_t, hsa_ext_sampler_t]
 except AttributeError: pass
 
@@ -2199,6 +2647,7 @@ struct_hsa_ext_images_1_pfn_s._fields_ = [
   ('hsa_ext_image_get_capability_with_layout', ctypes.CFUNCTYPE(hsa_status_t, hsa_agent_t, hsa_ext_image_geometry_t, ctypes.POINTER(hsa_ext_image_format_t), hsa_ext_image_data_layout_t, ctypes.POINTER(uint32_t))),
   ('hsa_ext_image_data_get_info_with_layout', ctypes.CFUNCTYPE(hsa_status_t, hsa_agent_t, ctypes.POINTER(hsa_ext_image_descriptor_t), hsa_access_permission_t, hsa_ext_image_data_layout_t, size_t, size_t, ctypes.POINTER(hsa_ext_image_data_info_t))),
   ('hsa_ext_image_create_with_layout', ctypes.CFUNCTYPE(hsa_status_t, hsa_agent_t, ctypes.POINTER(hsa_ext_image_descriptor_t), ctypes.c_void_p, hsa_access_permission_t, hsa_ext_image_data_layout_t, size_t, size_t, ctypes.POINTER(hsa_ext_image_t))),
+  ('hsa_ext_sampler_create_v2', ctypes.CFUNCTYPE(hsa_status_t, hsa_agent_t, ctypes.POINTER(hsa_ext_sampler_descriptor_v2_t), ctypes.POINTER(hsa_ext_sampler_t))),
 ]
 hsa_ext_images_1_pfn_t = struct_hsa_ext_images_1_pfn_s
 try: (hsa_ven_amd_aqlprofile_version_major:=dll.hsa_ven_amd_aqlprofile_version_major).restype, hsa_ven_amd_aqlprofile_version_major.argtypes = uint32_t, []
@@ -2393,9 +2842,10 @@ hsa_ven_amd_aqlprofile_1_00_pfn_t = struct_hsa_ven_amd_aqlprofile_1_00_pfn_s
 hsa_ven_amd_aqlprofile_pfn_t = struct_hsa_ven_amd_aqlprofile_1_00_pfn_s
 HSA_VERSION_1_0 = 1
 HSA_AMD_INTERFACE_VERSION_MAJOR = 1
-HSA_AMD_INTERFACE_VERSION_MINOR = 6
+HSA_AMD_INTERFACE_VERSION_MINOR = 14
 AMD_SIGNAL_ALIGN_BYTES = 64
 AMD_QUEUE_ALIGN_BYTES = 64
+MAX_NUM_XCC = 128
 AMD_CONTROL_DIRECTIVES_ALIGN_BYTES = 64
 AMD_ISA_ALIGN_BYTES = 256
 AMD_KERNEL_CODE_ALIGN_BYTES = 64

tinygrad/runtime/graph/hcq.py

@@ -1,6 +1,6 @@
 import collections, time
 from typing import Any, cast
-from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv, dedup
+from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv, dedup, suppress_finalizing
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
@@ -221,6 +221,7 @@ class HCQGraph(MultiGraphRunner):
 
   def dev_name(self, dev) -> str: return dev.device.replace(":", "_")
 
+  @suppress_finalizing
   def __del__(self):
     for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])