qwen work

* llm: add created/model fields, non-streaming support, and tests

- Add `created` timestamp and `model` fields to response (required by OpenAI spec)
- Add non-streaming mode support for /v1/chat/completions
- Add `send_data` helper to HTTPRequestHandler for responses with Content-Length
- Refactor viz/serve.py to use send_data
- Add integration tests using real OpenAI client

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* add openai to testing

* toml

* Remove 'openai' from dependencies

Removed 'openai' from the dependencies list.

* bump cache

---------

Co-authored-by: Claude Opus 4.5 <[email protected]>

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

@@ -298,7 +298,7 @@ runs:
     - name: Install mesa (linux)
       if: inputs.mesa == 'true' && runner.os == 'Linux'
       shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
     - name: Install mesa (macOS)
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash

.github/workflows/autogen.yml

@@ -116,11 +116,9 @@ jobs:
     - name: Verify Qualcomm autogen
       run: |
         mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
         mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
-        python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
+        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
         diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.py
         diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
     - name: Verify libusb autogen
       run: |

.github/workflows/benchmark.yml

@@ -54,9 +54,9 @@ jobs:
     - name: Print macOS version
       run: sw_vers
     - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
     - name: Run Stable Diffusion v2
       # TODO: very slow step time
       run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
@@ -64,7 +64,7 @@ jobs:
     - name: Run SDXL
       run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run model inference benchmark
-      run: METAL=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
     - name: Test speed vs torch
       run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
@@ -132,6 +132,10 @@ jobs:
       run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     #- name: UsbGPU openpilot test
     #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    - name: UsbGPU (USB4/TB) boot time
+      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU (USB4/TB) tiny tests
+      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -322,29 +326,27 @@ jobs:
       run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Train MNIST
       run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-    # TODO: too slow
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
-    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     # - name: Run 10 CIFAR training steps w winograd
     #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    # - name: Run full CIFAR training steps w 6 GPUS
-    #   run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
     - name: Run MLPerf resnet eval on training data
       run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (NVIDIA Training)
@@ -435,9 +437,8 @@ jobs:
       run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Run Stable Diffusion
       run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run LLaMA 7B
       run: |
         BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -527,22 +528,21 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: Train MNIST
       run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-    # TODO: too slow
     - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
-    # - name: Run 10 CIFAR training steps w HALF
-    #   run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
     # - name: Run 10 CIFAR training steps w BF16
     #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     # TODO: too slow
     # - name: Run 10 CIFAR training steps w winograd
     #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    #- name: Run full CIFAR training steps w 6 GPUS
-    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD Training)
@@ -592,13 +592,13 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: Run MLPerf resnet eval
       run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - name: Run 10 MLPerf Bert training steps (6 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD MLPerf)
@@ -627,18 +627,14 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    # - name: openpilot compile3 0.9.9 driving_vision
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
-    # - name: openpilot compile3 0.9.9 driving_policy
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
-    # - name: openpilot compile3 0.9.9 dmonitoring
-    #   run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.10.0 driving_policy
       run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.0 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
       run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_vision
       run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
@@ -708,10 +704,8 @@ jobs:
       run: |
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
-    # TODO: enable
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
     # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
     #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
     - name: Run 10 MLPerf Bert training steps (1 gpu)
@@ -772,11 +766,10 @@ jobs:
         NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
     - name: Test LLAMA-3
       run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
-    # TODO: too slow
-    # - name: Run full CIFAR training w 1 GPU
-    #   run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
+  CACHE_VERSION: '15'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -71,9 +71,7 @@ jobs:
     - name: Test Docs Build
       run: python -m mkdocs build --strict
     - name: Test Docs
-      run: |
-        python docs/abstractions2.py
-        python docs/abstractions3.py
+      run: python docs/abstractions3.py
     - name: Test README
       run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
     - name: Test Quickstart
@@ -263,7 +261,9 @@ jobs:
     - name: Check Device.DEFAULT
       run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
     - name: Run unit tests
-      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+      run: |
+        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
     - name: Run targetted tests on NULL backend
       run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
     # TODO: too slow
@@ -289,8 +289,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 19000 lines
-      run: MAX_LINE_COUNT=19000 python sz.py
+    - name: Repo line count < 20000 lines
+      run: MAX_LINE_COUNT=20000 python sz.py
 
   spec:
     strategy:
@@ -447,7 +447,7 @@ jobs:
         with:
           key: onnxoptl
           deps: testing
-          pydeps: "tensorflow==2.15.1 tensorflow_addons"
+          pydeps: "tensorflow==2.19"
           python-version: '3.11'
           opencl: 'true'
       - name: Test ONNX (CL)
@@ -465,7 +465,7 @@ jobs:
       - name: Test Bert training
         run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -972,3 +972,33 @@ jobs:
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
+
+# ****** Compile-only Tests ******
+
+  compiletests:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [ir3, nak]
+    name: Compile-only (${{ matrix.backend }})
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-${{ matrix.backend }}
+          deps: testing_minimal
+          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
+          python-version: '3.14'
+      - name: Set env
+        shell: bash
+        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/test_ops.py --durations=20

docs/abstractions2.py

@@ -1,135 +0,0 @@
-# tinygrad is a tensor library, and as a tensor library it has multiple parts
-# 1. a "runtime". this allows buffer management, compilation, and running programs
-# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
-# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
-# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
-
-
-print("******** first, the runtime ***********")
-
-from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
-
-cpu = CPUDevice()
-
-# allocate some buffers
-out = cpu.allocator.alloc(4)
-a = cpu.allocator.alloc(4)
-b = cpu.allocator.alloc(4)
-
-# load in some values (little endian)
-cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
-cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
-
-# compile a program to a binary
-lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
-
-# create a runtime for the program
-fxn = cpu.runtime("add", lib)
-
-# run the program
-fxn(out, a, b)
-
-# check the data out
-print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
-assert val == 5
-
-
-print("******** second, the Device ***********")
-
-DEVICE = "CPU"   # NOTE: you can change this!
-
-import struct
-from tinygrad.dtype import dtypes
-from tinygrad.device import Buffer, Device
-from tinygrad.uop.ops import UOp, Ops
-
-# allocate some buffers + load in values
-out = Buffer(DEVICE, 1, dtypes.int32).allocate()
-a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
-b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-# NOTE: a._buf is the same as the return from cpu.allocator.alloc
-
-# describe the computation
-idx = UOp.const(dtypes.index, 0)
-buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
-buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-alu = buf_1.index(idx) + buf_2.index(idx)
-output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
-st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
-s = UOp(Ops.SINK, dtypes.void, (st_0,))
-
-# convert the computation to a "linearized" format (print the format)
-from tinygrad.engine.realize import get_program, CompiledRunner
-program = get_program(s, Device[DEVICE].renderer)
-
-# compile a program (and print the source)
-fxn = CompiledRunner(program)
-print(fxn.p.src)
-# NOTE: fxn.clprg is the CPUProgram
-
-# run the program
-fxn.exec([out, a, b])
-
-# check the data out
-assert out.as_buffer().cast('I')[0] == 5
-
-
-print("******** third, the UOp ***********")
-
-from tinygrad.engine.realize import run_schedule
-from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.schedule.rangeify import get_rangeify_map
-
-# allocate some values + load in values
-a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
-b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
-a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
-b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-
-# describe the computation
-out = a + b
-s = UOp(Ops.SINK, dtypes.void, (out,))
-
-# group the computation into kernels
-becomes_map = get_rangeify_map(s)
-
-# the compute maps to an assign
-assign = becomes_map[a+b].base
-
-# the first source is the output buffer (data)
-assert assign.src[0].op is Ops.BUFFER
-# the second source is the kernel (compute)
-assert assign.src[1].op is Ops.KERNEL
-
-# schedule the kernel graph in a linear list
-s = UOp(Ops.SINK, dtypes.void, (assign,))
-sched, _ = create_schedule_with_vars(s)
-assert len(sched) == 1
-
-# DEBUGGING: print the compute ast
-print(sched[-1].ast)
-# NOTE: sched[-1].ast is the same as st_0 above
-
-# the output will be stored in a new buffer
-out = assign.buf_uop
-assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
-print(out)
-
-# run that schedule
-run_schedule(sched)
-
-# check the data out
-assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
-
-
-print("******** fourth, the Tensor ***********")
-
-from tinygrad import Tensor
-
-a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
-b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
-out = a + b
-
-# check the data out
-print(val:=out.item())
-assert val == 5

examples/beautiful_mnist.py

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

examples/gradaccum_mnist.py

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

examples/mlperf/initializers.py

@@ -59,9 +59,7 @@ class EmbeddingBert(nn.Embedding):
     arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
     if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
     arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    # TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
-    # Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
-    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()
+    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
 
 class LayerNormBert:
   def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):

examples/mlperf/model_train.py

@@ -918,40 +918,6 @@ def train_rnnt():
   # TODO: RNN-T
   pass
 
-@TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
-  optimizer.zero_grad()
-
-  for i in range(grad_acc):
-    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
-    # NOTE: these two have different names
-    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
-    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
-
-    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-      else: t.to_(GPUS[0])
-
-    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-    (loss * loss_scaler).backward()
-    # TODO: OOM without this realize with large grad_acc
-    Tensor.realize(*[p.grad for p in optimizer.params])
-
-  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
-  for p in optimizer.params:
-    p.grad = p.grad / loss_scaler
-    global_norm += p.grad.float().square().sum()
-  global_norm = global_norm.sqrt().contiguous()
-  for p in optimizer.params:
-    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
-
-  optimizer.step()
-  scheduler.step()
-  # TODO: no to("CPU") here because it blocks and messes the python time
-  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
-  return loss, global_norm, optimizer.optimizers[0].lr
-
 @TinyJit
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                    masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
@@ -1014,7 +980,8 @@ def train_bert():
   # ** hyperparameters **
   BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  # TODO: mlperf logging
+  # TODO: implement grad accumulation + mlperf logging
+  assert grad_acc == 1
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
   max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
@@ -1073,8 +1040,8 @@ def train_bert():
 
   # ** Optimizer **
   parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
-  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  parameters_wd = [x for x in parameters if x not in set(parameters_no_wd)]
+  optimizer_wd = LAMB(parameters_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
   optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
   optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)
 
@@ -1131,12 +1098,38 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
+  i, train_data = start_step, next(train_it)
 
   if RUNMLPERF:
     if MLLOGGER:
       MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
 
+  @TinyJit
+  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
+                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+    optimizer_group.zero_grad()
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+
+    global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer_group[0].device)
+    for p in optimizer_group.params:
+      p.grad = p.grad / loss_scaler
+      global_norm += p.grad.float().square().sum()
+    global_norm = global_norm.sqrt().contiguous()
+    for p in optimizer_group.params:
+      p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
+
+    optimizer_group.step()
+    scheduler_group.step()
+    # TODO: no to("CPU") here because it blocks and messes the python time
+    Tensor.realize(loss, global_norm, optimizer_group.optimizers[0].lr)
+    return loss, global_norm, optimizer_group.optimizers[0].lr
+
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
       Tensor.training = True
@@ -1144,16 +1137,12 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
-        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)
+        loss, global_norm, lr = train_step_bert(
+          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
+          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
 
         pt = time.perf_counter()
-
-        try:
-          next_data = [next(train_it) for _ in range(grad_acc)]
-        except StopIteration:
-          next_data = None
-
+        next_data = next(train_it)
         dt = time.perf_counter()
 
         device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
@@ -1188,8 +1177,7 @@ def train_bert():
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
       if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None:
         train_step_bert.captured.free_intermediates()
       eval_lm_losses = []
       eval_clsf_losses = []
@@ -1224,7 +1212,7 @@ def train_bert():
           return
 
       if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
 
       del eval_data
       avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1300,6 +1288,7 @@ def train_llama3():
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  assert grad_acc == 1, f"{grad_acc=} is not supported"
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1374,20 +1363,17 @@ def train_llama3():
 
   @TinyJit
   @Tensor.train()
-  def train_step(model, tokens:Tensor, grad_acc:int):
+  def train_step(model, tokens:Tensor):
     optim.zero_grad()
-    # grad acc
-    for batch in tokens.split(tokens.shape[0]//grad_acc):
-      if (DP := getenv("DP", 1)) > 1:
-        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-        batch = batch.shard(device, 0)
-      if (MP := getenv("MP", 1)) > 1:
-        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-        batch = batch.shard(device)
-      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
-      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
-      loss.backward()
-      Tensor.realize(*[p.grad for p in optim.params])
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss.backward()
     # L2 norm grad clip
     # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
     # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1426,14 +1412,14 @@ def train_llama3():
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
-      return fake_data(GBS, SAMPLES)
+      return fake_data(BS, SAMPLES)
     else:
       if SMALL:
         from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
       else:
         from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
 
   def get_eval_iter():
     if getenv("FAKEDATA", 0):
@@ -1451,7 +1437,7 @@ def train_llama3():
   for tokens in tqdm(iter, total=SAMPLES//GBS):
     t = time.perf_counter()
     GlobalCounters.reset()
-    loss, lr = train_step(model, tokens, grad_acc)
+    loss, lr = train_step(model, tokens)
     loss = loss.float().item()
 
     i += 1

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

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

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

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

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

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

examples/sdv2.py

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

extra/amdpci/proclogs.py

@@ -1,48 +1,59 @@
 import re, ctypes, sys, importlib
+from tinygrad.helpers import getenv
 
 from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
 class AMDFake(AMDev):
-  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
-    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+  def __init__(self, pci_dev, dma_regions=None):
+    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
     self._run_discovery()
     self._build_regs()
 
 amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
-
 from tinygrad.runtime.ops_amd import PCIIface
 
-def parse_amdgpu_logs(log_content, register_names=None):
-  register_map = register_names
+def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+  register_map = register_names or {}
 
-  final = ""
   def replace_register(match):
-    register = match.group(1)
-    return f"Reading register {register_map.get(int(register, base=16), register)}"
+    reg = match.group(1)
+    return f"Reading register {register_map.get(int(reg, 16), reg)}"
 
-  pattern = r'Reading register (0x[0-9a-fA-F]+)'
-
-  processed_log = re.sub(pattern, replace_register, log_content)
+  processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
 
   def replace_register_2(match):
-    register = match.group(1)
-    return f"Writing register {register_map.get(int(register, base=16), register)}"
+    reg = match.group(1)
+    return f"Writing register {register_map.get(int(reg, 16), reg)}"
+
+  processed_log = re.sub(r'Writing register (0x[0-9a-fA-F]+)', replace_register_2, processed_log)
+
+  # remove timing prefix
+  processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
+
+  # keep only xcc=0 lines (but keep lines with no xcc at all)
+  if only_xcc0:
+    kept = []
+    for line in processed_log.splitlines(True):
+      if "xcc=" not in line or re.search(r'\bxcc=0\b', line): kept.append(line)
+    processed_log = "".join(kept)
 
-  pattern = r'Writing register (0x[0-9a-fA-F]+)'
-  processed_log = re.sub(pattern, replace_register_2, processed_log)
   return processed_log
 
 def main():
+  only_xcc0 = bool(getenv("ONLY_XCC0", 0))
+
   reg_names = {}
   dev = PCIIface(None, 0)
   for x, y in dev.dev_impl.__dict__.items():
     if isinstance(y, AMRegister):
-      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+      for xcc, addr in y.addr.items():
+        reg_names[addr] = f"{x}, xcc={xcc}"
 
   with open(sys.argv[1], 'r') as f:
-    log_content = log_content_them = f.read()
+    log_content = f.read()
 
-  processed_log = parse_amdgpu_logs(log_content, reg_names)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)
@@ -51,5 +62,4 @@ if __name__ == '__main__':
   if len(sys.argv) != 3:
     print("Usage: <input_file_path> <output_file_path>")
     sys.exit(1)
-
-  main()
+  main()

extra/assembly/rdna3/generate.py

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

extra/assembly/rdna3/parse.py

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

extra/hevc/decode.py

@@ -1,7 +1,7 @@
 import argparse, os, hashlib
 from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
 from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
-from tinygrad import Tensor, dtypes, Device, Variable
+from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
@@ -9,8 +9,6 @@ if __name__ == "__main__":
   parser.add_argument("--output_dir", type=str, default="extra/hevc/out")
   args = parser.parse_args()
 
-  os.makedirs(args.output_dir, exist_ok=True)
-
   if args.input_file == "":
     url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
     hevc_tensor = Tensor.from_url(url, device="CPU")
@@ -27,26 +25,38 @@ if __name__ == "__main__":
   frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
 
   # move all needed data to gpu
-  all_slices = []
-  with Timing("prep slices to gpu: "):
+  #all_slices = []
+  with Timing("copy to gpu: "):
     opaque_nv = opaque.to("NV").contiguous().realize()
-
-    for i, (offset, sz, frame_pos, history_sz, _) in enumerate(frame_info):
-      all_slices.append(hevc_tensor[offset:offset+sz].to("NV").contiguous().realize())
-
-    Device.default.synchronize()
+    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)
-  pos = Variable("pos", 0, max_hist + 1)
 
-  history = []
+  # 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[-history_sz:] if history_sz > 0 else []
+      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 = all_slices[i].decode_hevc_frame(pos.bind(frame_pos), out_image_size, opaque_nv[i], history).realize()
+      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)
 
@@ -67,5 +77,7 @@ if __name__ == "__main__":
         assert img.data() == decoded_frames[i], f"Frame {i} does not match reference decoder!"
         print(f"Frame {i} matches reference decoder!")
     else:
-      img = to_bgr(img, h, w, luma_w, chroma_off).realize()
-      cv2.imwrite(f"{args.output_dir}/out_frame_{i:04d}.png", img.numpy())
+      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

@@ -322,6 +322,7 @@ def parse_hevc_file_headers(dat:bytes, device="NV"):
     if hdr.slice_type == avcodec.HEVC_SLICE_B: device_ctx.initreflistidxl1 = (ctypes.c_uint8 * 16)(*[idx for _,idx in after_list + before_list])
 
     locl_ctx_bytes = bytes(device_ctx)
+    locl_ctx_bytes += b'\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x00\x00' # blackwell extension
     locl_ctx_bytes += bytes(0x200 - len(locl_ctx_bytes)) # pad to 512 bytes
 
     pic_width_in_ctbs = ceildiv(sps.pic_width_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size))

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

@@ -1,4 +1,5 @@
 import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
+from tabulate import tabulate
 from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
@@ -67,10 +68,7 @@ class OccEvent(WaveSlot):
   time:int
   start:int
 
-@dataclasses.dataclass(frozen=True)
-class RunKey:
-  prg:str
-  tag:int
+RunKey = tuple[str, int]
 
 class _ROCParseCtx:
   def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
@@ -86,7 +84,7 @@ class _ROCParseCtx:
 
   def next_sqtt(self):
     x = next(self.sqtt_evs, None)
-    self.active_run = RunKey(x.kern, x.exec_tag) 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
@@ -141,7 +139,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[unwrap(ROCParseCtx.active_run).prg][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
@@ -161,6 +159,13 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   t.join()
   return ROCParseCtx
 
+def print_pmc(events:list[ProfilePMCEvent]) -> None:
+  from tinygrad.viz.serve import unpack_pmc
+  for e in events:
+    print("**", e.kern)
+    data = unpack_pmc(e)
+    print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
+
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
@@ -170,13 +175,4 @@ if __name__ == "__main__":
   rctx = decode(profile)
   print('SQTT:', rctx.inst_execs.keys())
 
-  for ev in profile:
-    if not isinstance(ev, ProfilePMCEvent): continue
-    print(f"PMC Event: dev={ev.device} kern={ev.kern}")
-    ptr = 0
-    for s in ev.sched:
-      view = memoryview(ev.blob).cast('Q')
-      print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
-        ptr += 1
+  print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])

extra/sqtt/test_pmc.py

@@ -0,0 +1,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)
+    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/thunder/tiny/fa.py

@@ -0,0 +1,166 @@
+import math
+
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+
+from extra.thunder.tiny.tk import WARP_THREADS
+from extra.thunder.tiny.tk.kernel import Kernel
+from extra.thunder.tiny.tk.tiles import GL, TileLayout
+
+NUM_WORKERS = 1
+Q_BLOCK_SIZE = 16
+KV_BLOCK_SIZE = 16
+
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
+  if len(xq.shape) == 3: xq, xk, xv = xq.unsqueeze(0), xk.unsqueeze(0), xv.unsqueeze(0)
+
+  odtype = xq.dtype
+  xq, xk, xv = xq.transpose(1, 2).cast(dtypes.bfloat16), xk.transpose(1, 2).cast(dtypes.bfloat16), xv.transpose(1, 2).cast(dtypes.bfloat16)
+
+  _, N_, _, D_ = xq.shape
+  block_size = max(Q_BLOCK_SIZE, KV_BLOCK_SIZE)
+  assert D_ % block_size == 0, f"embedding dimension must be multiple of block size, got {D_=} {block_size=}"
+
+  # pad to multiple of block size
+  xq = xq.pad(((0, 0), (0, (block_size - (xq.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
+  xk = xk.pad(((0, 0), (0, (block_size - (xk.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
+  xv = xv.pad(((0, 0), (0, (block_size - (xv.shape[1] % block_size)) % block_size), (0, 0), (0, 0)))
+
+  B, N, H, D = xq.shape
+  H_KV = xk.shape[2]
+  GROUP_SIZE = H // H_KV
+  print(f"Flash Attention {B=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
+
+  def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, mu:UOp) -> UOp:
+    with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
+      warp = ker.warp
+
+      o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(mu, ker), GL(l_vecu, ker)
+
+      head = ker.blockIdx_x
+      head_kv = head // GROUP_SIZE
+      batch = ker.blockIdx_z
+      q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
+
+      k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
+      v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
+
+      q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
+      q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
+      q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
+      k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
+      k_reg_transposed = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
+      v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
+      o_reg = ker.rt((D, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      o_reg_transposed = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
+      att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
+      mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
+      mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+
+      max_vec_last = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
+      max_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
+      norm_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
+      scale_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
+
+      max_vec = warp.neg_inf(max_vec)
+      norm_vec = warp.zero(norm_vec)
+      o_reg = warp.zero(o_reg)
+      scale_vec = warp.ones(scale_vec)
+
+      # load q tile
+      q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
+      q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
+      q_reg = warp.copy(q_reg, q_reg_fl)
+      q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
+
+      for kv_idx in ker.range(N // KV_BLOCK_SIZE):
+        k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
+        v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
+
+        k_reg = warp.load(k_reg, k_smem)
+        v_reg = warp.load(v_reg, v_smem)
+
+        # mma qk^t
+        att_block = warp.zero(att_block.after(kv_idx))
+        k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
+        att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
+
+        # apply attention mask
+        mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
+        mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
+        att_block += mask_reg_transposed
+
+        # softmax
+        max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
+        max_vec = warp.row_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)
+
+        scale_vec = warp.map(scale_vec.after(max_vec_last, max_vec), lambda _, idx: max_vec_last[*idx] - max_vec[*idx])
+        scale_vec = scale_vec.exp2()
+
+        o_reg *= scale_vec
+        norm_vec *= scale_vec
+
+        att_block -= max_vec
+        att_block = att_block.exp2()
+
+        norm_vec = warp.row_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)
+
+        # mma av
+        att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
+        o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
+      o_reg = ker.endrange()
+      norm_vec = norm_vec.after(o_reg)
+      max_vec = max_vec.after(o_reg)
+
+      o_reg /= norm_vec
+
+      o_reg_transposed = warp.transpose(o_reg_transposed, o_reg)
+      o = warp.store(o, o_reg_transposed, (batch, q_seq, head, 0), (), axis=1)
+
+      norm_vec = norm_vec.after(o)
+      max_vec = max_vec.after(o)
+
+      max_vec *= math.log(2)
+      norm_vec = norm_vec.log2() * math.log(2)
+      norm_vec += max_vec
+      l_vec = warp.store(l_vec, norm_vec, (batch, head, 0, q_seq), (), axis=2)
+      o = o.after(l_vec)
+
+      return ker.finish()
+
+  def custom_backward_q(out_qu:UOp, gradu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return UOp.sink(arg=KernelInfo(name="fa_custom_backward_q"))
+
+  def custom_backward_kv(out_ku:UOp, out_vu:UOp, gradu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
+    return UOp.sink(arg=KernelInfo(name="fa_custom_backward_kv"))
+
+  if is_causal:
+    if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
+    attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.bool).tril()
+  if attn_mask is not None:
+    if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
+  else:
+    attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.float32)
+
+  attn = Tensor.empty_like(xq)
+  l_vec = Tensor.empty(B, H, 1, N, requires_grad=False, device=xq.device, dtype=dtypes.float32).detach()
+
+  def grad(grad:UOp, kernel:UOp) -> tuple[None, None, UOp, UOp, UOp, None]:
+    grad_q = Tensor.empty_like(q := Tensor(kernel.src[2]))
+    grad_k = Tensor.empty_like(k := Tensor(kernel.src[3]))
+    grad_v = Tensor.empty_like(v := Tensor(kernel.src[4]))
+    mask = Tensor(kernel.src[5])
+
+    delta_vec = (Tensor(grad) * attn).sum(-1).unsqueeze(-2).detach()
+
+    print(l_vec.numpy())
+
+    grad_q = Tensor.custom_kernel(grad_q, Tensor(grad), q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
+    grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, Tensor(grad), q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
+    return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
+
+  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
+  attn = attn[:, :N_, :, :D_]
+
+  return attn.transpose(1, 2).cast(odtype)

extra/thunder/tiny/tk/group.py

@@ -7,7 +7,7 @@ from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.helpers import getenv, prod
 
 from extra.thunder.tiny.tk import WARP_THREADS
-from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout
+from extra.thunder.tiny.tk.tiles import ALL_TILES, GL, RT_16X16, RT_16X32, ST, RT, RV, TileLayout, VecLayout
 
 class Group:
   def __init__(self, warps:int, ker):
@@ -66,7 +66,10 @@ class Group:
     for height in self.ker.range(src.shape[-3], track=False):
       for width in self.ker.range(src.shape[-2], track=False):
         for inner in self.ker.range(src.shape[-1], track=False):
-          dst_store = dst[width, height, inner].store(src[height, width, inner]).end(height, width, inner)
+          src_load = src[height, width, inner]
+          if src.dtype.base != dst.dtype.base:
+            src_load = src_load.cast(dst.dtype.base)
+          dst_store = dst[width, height, inner].store(src_load).end(height, width, inner)
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
@@ -77,9 +80,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     else:  raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -107,9 +110,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     else:  raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -137,9 +140,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     else:  raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -167,9 +170,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ())
+      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
     else:  raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -200,9 +203,9 @@ class Group:
     Group.map_rid += len(a.shape)
 
     if op.__code__.co_argcount == 1:
-      to_store = op(a[*rngs_for_shape])
+      to_store = op(a[*rngs_for_shape]) # type: ignore
     else:
-      to_store = op(a[*rngs_for_shape], rngs_for_shape)
+      to_store = op(a[*rngs_for_shape], rngs_for_shape) # type: ignore
 
     a_store = a[*rngs_for_shape].store(to_store).end(*rngs_for_shape)
 
@@ -284,7 +287,7 @@ class Group:
   def load(self, dst:ALL_TILES, src:ALL_TILES, dst_idxs:tuple[UOp|int,...]=(), idxs:tuple[UOp|int,...]=(), axis:int=0):
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
+    dst_dtype, src_dtype = dst.dtype, src.dtype
     if dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.LOCAL:
       laneid = self.ker.laneid
       rt, st = cast(RT, dst), cast(ST, src)
@@ -338,7 +341,7 @@ class Group:
                 src_load = src_load.cast(dst.dtype.base)
               dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
               dst_store = dst_store.end(height, width, outer, inner).barrier()
-    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace ==AddrSpace.GLOBAL:
+    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
       srcf = src.flatten()
       row_stride = prod(src.shape[axis+1:])
 
@@ -371,8 +374,28 @@ class Group:
             if src.dtype.base != dst.dtype.base:
               src_load = src_load.cast(dst.dtype.base)
             dst_store = dst[*dst_idxs, height, width, inner].store(src_load).end(height, width, inner)
+    elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RV):
+      srcf = src.flatten()
+      row_stride = prod(src.shape[axis+1:])
+
+      laneid = self.ker.laneid
+      rv = cast(RV, dst)
+      reductions = rv.base_shape.rows
+
+      assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
+
+      idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
+      src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
+
+      for outer in self.ker.range(dst.shape[-2]):
+        src_i += outer * reductions + (laneid % reductions)
+
+        src_load = srcf[src_i]
+        if src.dtype.base != dst.dtype.base:
+          src_load = src_load.cast(dst.dtype.base)
+        dst_store = dst[outer, 0].store(src_load).end(outer)
     else:
-      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
@@ -380,8 +403,8 @@ class Group:
   def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0):
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL:
+    dst_dtype, src_dtype = dst.dtype, src.dtype
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
       dstf = dst.flatten()
       row_stride = prod(dst.shape[axis+1:])
 
@@ -414,8 +437,28 @@ class Group:
             if src.dtype.base != dst.dtype.base:
               src_load = src_load.cast(dst.dtype.base)
             dst_store = dstf[dst_i].store(src_load).end(height, width, inner)
+    elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RV):
+      dstf = dst.flatten()
+      row_stride = prod(dst.shape[axis+1:])
+
+      laneid = self.ker.laneid
+      rv = cast(RV, src)
+      reductions = rv.base_shape.rows
+
+      assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
+
+      idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
+      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
+
+      for outer in self.ker.range(src.shape[-2]):
+        dst_i += outer * reductions + (laneid % reductions)
+
+        src_load = src[outer, 0]
+        if src.dtype.base != dst.dtype.base:
+          src_load = src_load.cast(dst.dtype.base)
+        dst_store = dstf[dst_i].store(src_load).end(outer)
     else:
-      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)

extra/thunder/tiny/tk/kernel.py

@@ -5,33 +5,34 @@ from extra.thunder.tiny.tk.group import Group
 from extra.thunder.tiny.tk.tiles import GL, ST_16X16, ST_16X16_SWIZZLED, ST, RT_16X16, RT, RV, TileLayout, VecLayout
 
 class _tk_range:
-  user_rid = 0
-  def __init__(self, start:int, end:int, step:int, axis_type:AxisType):
+  def __init__(self, start:int, end:int, step:int, axis_type:AxisType, rid:int):
     self.start, self.end, self.step = start, end, step
-    self.axis_type, self.done = axis_type, False
+    self.axis_type, self.rid, self.done = axis_type, rid, False
   def __iter__(self): return self
   def __next__(self):
     if not self.done:
       self.done = True
-      _tk_range.user_rid += 1
-      self._rng = UOp.range(self.end // self.step, _tk_range.user_rid-1, axis_type=self.axis_type) * self.step + self.start
+      self._rng = UOp.range(self.end // self.step, self.rid, axis_type=self.axis_type) * self.step + self.start
       return self._rng
     raise StopIteration
 
 class Kernel(AbstractContextManager):
-  def __init__(self, grid_size:tuple[int, int, int], block_size:int):
+  def __init__(self, name:str, grid_size:tuple[int, int, int], block_size:int):
+    self.name = name
+
     self.blockIdx_x = UOp.special(grid_size[0], "gidx0")
     self.blockIdx_y = UOp.special(grid_size[1], "gidx1")
     self.blockIdx_z = UOp.special(grid_size[2], "gidx2")
     self.threadIdx_x = UOp.special(block_size, "lidx0")
 
-    self.range_stack = []
-    self.store_stack = []
+    self.range_stack: list[_tk_range] = []
+    self.store_stack: list[tuple[UOp, UOp]] = []
 
     self.global_slot = 0
     self.shared_slot = 0
     self.register_slot = 0
-    self.allocs = {}
+    self.range_id = 0
+    self.allocs: dict[tuple[str, tuple], UOp] = {}
 
   @property
   def warpid(self): return self.threadIdx_x // WARP_THREADS
@@ -49,7 +50,8 @@ class Kernel(AbstractContextManager):
 
   def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
     if end == 0: start, end = 0, start
-    rng = _tk_range(start, end, step, axis_type)
+    rng = _tk_range(start, end, step, axis_type, self.range_id)
+    self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
@@ -89,7 +91,7 @@ class Kernel(AbstractContextManager):
     if hasattr(last_store, '_uop'): uop = last_store._uop
     else: uop = last_store
 
-    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+    return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
 
   def endrange(self):
     last_store = self.store_stack.pop()

extra/thunder/tiny/tk/tiles.py

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

pyproject.toml

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

test/external/external_test_am.py

@@ -32,6 +32,8 @@ class FakeAM:
     self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
   def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
   def paddr2mc(self, paddr:int) -> int: return paddr
+  def paddr2xgmi(self, paddr:int) -> int: return paddr
+  def xgmi2paddr(self, xgmi_paddr:int) -> int: return xgmi_paddr
 
 #  * PTE format:
 #  * 63:59 reserved

test/external/external_test_onnx_backend.py

@@ -170,6 +170,19 @@ backend_test.exclude('test_scan_*')
 backend_test.exclude('test_split_to_sequence_*')
 backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121
 
+# TODO: not yet implemented
+backend_test.exclude('test_tensorscatter_*')
+backend_test.exclude('test_l1normalization_*')
+backend_test.exclude('test_l2normalization_*')
+backend_test.exclude('test_lpnormalization_*')
+backend_test.exclude('test_einsum_scalar_cpu')
+backend_test.exclude('test_mod_mixed_sign_float16_cpu')
+backend_test.exclude('test_qlinearmatmul_2D_uint8_float16_cpu')
+backend_test.exclude('test_qlinearmatmul_3D_uint8_float16_cpu')
+backend_test.exclude('test_attention_3d_*')
+backend_test.exclude('test_attention_4d_*')
+
+
 # rest of the failing tests
 backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # tf_crop_and_resize not implemented
 backend_test.exclude('test_resize_tf_crop_and_resize_axes_2_3_cpu') # tf_crop_and_resize not implemented
@@ -184,9 +197,6 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
 backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
 backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
 
-backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
-backend_test.exclude('test_scatternd_max_cpu') # max not yet supported
-
 # regression from removing StrEnum in Domain
 backend_test.exclude('test_adam_cpu')
 backend_test.exclude('test_gradient_of_add_and_mul_cpu')

test/external/external_test_optim.py

@@ -2,7 +2,7 @@
 import unittest, math
 import numpy as np
 import tensorflow as tf
-import tensorflow_addons as tfa
+from tensorflow.keras.optimizers import Lamb
 from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup
 
@@ -88,6 +88,8 @@ def create_tiny_lars(params, lr, skip_list=False):
   if skip_list: return OptimizerGroup(LARS([params[0]], lr), SGD([params[1]], lr, classic=True, weight_decay=0., momentum=.9))
   return LARS(params, lr)
 def create_tf_lars(lr, skip_list=False): return LARSOptimizer(lr, skip_list=["W"] if skip_list else None)
+def create_tf_lamb(lr=0.001, b1=0.9, b2=0.999, eps=1e-7, weight_decay=0.0):
+  return Lamb(learning_rate=float(lr), beta_1=b1, beta_2=b2, epsilon=eps, weight_decay=weight_decay)
 
 def create_tiny_polylr(optim, initial_lr, end_lr, train_steps, warmup, power=2, skip_list=False):
   assert power == 2
@@ -112,7 +114,7 @@ class ExternalTestOptim(unittest.TestCase):
                    step_tf(tensorflow_optim, steps=steps, kwargs=opts, scheduler=tf_sched, schedopts=schedopts, do_optim=do_optim)):
       np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
 
-  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, tfa.optimizers.LAMB, steps, opts, atol, rtol)
+  def _test_lamb(self, steps, opts, atol, rtol): self._test_optim(LAMB, create_tf_lamb, steps, opts, atol, rtol)
   def _test_lars(self, steps, opts, atol, rtol): self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol)
   def _test_lars_polylr(self, steps, opts, schedopts, atol, rtol, do_optim=True):
     self._test_optim(create_tiny_lars, create_tf_lars, steps, opts, atol, rtol,

test/external/external_test_process_count.py

@@ -0,0 +1,57 @@
+import os, sys, time, multiprocessing
+
+N = int(os.environ.get("NPROC", str(os.cpu_count())))
+DEVICE = os.environ.get("DEV", "AMD")
+
+# this tests the total number of processes that can be running tinygrad at a time
+def proc(i, device, stop_evt):
+  from tinygrad import Tensor
+
+  try:
+    a = Tensor.ones(2, device=device).contiguous()
+    b = Tensor.ones(2, device=device).contiguous()
+    c = (a + b).realize()
+    assert c.tolist() == [2, 2]
+  except Exception as e:
+    # fail if it fails
+    print(f"[child {i:2d}] tinygrad op failed: {e}", file=sys.stderr)
+    # non-zero exit code propagated back to parent
+    sys.exit(1)
+
+  # TODO: wait here for global exit if success. fail if it fails
+  # -> We wait on a global Event shared from the parent.
+  print(f"[child {i:2d}] success")
+  stop_evt.wait()
+  # Normal successful exit
+  sys.exit(0)
+
+if __name__ == "__main__":
+  print(f"testing {N} concurrent tinygrad processes")
+
+  # global exit event, shared by all children
+  stop_evt = multiprocessing.Event()
+  procs = []
+
+  # launch n proc of proc 1 per 200 ms
+  for i in range(N):
+    p = multiprocessing.Process(target=proc, args=(i, DEVICE, stop_evt), name=f"tinygrad-proc-{i}")
+    p.start()
+    procs.append(p)
+    time.sleep(0.1)  # 100 ms between launches
+
+  # signal global exit
+  time.sleep(0.5)
+  stop_evt.set()
+
+  # join all children
+  for p in procs: p.join()
+
+  # check for failures
+  failed = [p for p in procs if p.exitcode != 0]
+  if failed:
+    print(f"{len(failed)} / {len(procs)} processes failed "
+          f"with exit codes: {[p.exitcode for p in failed]}", file=sys.stderr)
+    sys.exit(1)
+
+  print(f"All {len(procs)} tinygrad processes ran successfully")
+  sys.exit(0)

test/external/external_uop_gc.py

@@ -1,5 +1,6 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
+from tinygrad.engine.schedule import schedule_cache
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general
@@ -68,6 +69,7 @@ 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()

test/external/mlperf_resnet/lars_optimizer.py

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

test/external/process_replay/process_replay.py

@@ -14,7 +14,6 @@ try:
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen.opt import Opt
   from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
-  from tinygrad.device import Device
 except ImportError as e:
   print(repr(e))
   exit(int(ASSERT_DIFF))
@@ -52,12 +51,10 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
     return "\n".join([f"{len(asts)} kernels", *asts])
   return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
   sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
   input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
-  # if no renderer was provided, open the device to get it
-  if renderer is None: renderer = Device[p.device].renderer
   p2 = get_program(input_ast, renderer=renderer)
   def to_str(ret:ProgramSpec) -> str:
     # PYTHON renderer pickles UOps, first unpickle and decode here

test/helpers.py

@@ -1,4 +1,4 @@
-import time, struct
+import time, struct, functools
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
@@ -60,5 +60,14 @@ def not_support_multi_device():
   # CL and CUDA don't support multi device if in CI
   return CI and REAL_DEV in ("CL", "CUDA")
 
+def needs_second_gpu(fn):
+  @functools.wraps(fn)
+  def wrapper(self, *args, **kwargs):
+    # check if there's a second GPU, if not, skip multi tests
+    try: Tensor.zeros(10, device=f"{Device.DEFAULT}:1").contiguous().realize()
+    except Exception as e: self.skipTest(f"second device not available: {e}")
+    return fn(self, *args, **kwargs)
+  return wrapper
+
 # NOTE: This will open REMOTE if it's the default device
 REAL_DEV = (Device.DEFAULT if Device.DEFAULT != "REMOTE" else Device['REMOTE'].properties.real_device)

test/models/test_end2end.py

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

test/models/test_real_world.py

@@ -28,13 +28,16 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
       model(*early_gen)
       Device[Device.DEFAULT].synchronize()
       tms.append(time.perf_counter_ns() - st)
-    mem_used = GlobalCounters.mem_used - global_mem_used
+    mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
 
     # TODO: jit should expose this correctly with graph
     kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
     print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
-    assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB - {mem_used/1e9:.2} GB used"
-    assert not kernels_used or kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
+    assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
+    assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
+    if kernels_used:
+      assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
+      assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used"
     if all_jitted:
       assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501
 
@@ -61,7 +64,7 @@ class TestRealWorld(unittest.TestCase):
     derandomize_model(model)
     @TinyJit
     def test(t, t2): return model(t, Tensor([801]), t2).realize()
-    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32),Tensor.randn(1, 77, params["ctx_dim"])), test, 18.0, 515)
+    helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
 
   def test_unet_resblock(self):
     model = [ResBlock(16, 24, 16) for _ in range(4)]
@@ -70,7 +73,7 @@ class TestRealWorld(unittest.TestCase):
     def test(t, t2):
       for l in model: t = l(t, t2)
       return t.realize()
-    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.01, 37)
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.0002, 37)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_llama(self):
@@ -82,7 +85,7 @@ class TestRealWorld(unittest.TestCase):
     @TinyJit
     def test(t): return model(t, 0).realize()
     # TODO: test first token vs rest properly
-    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.27, 168, all_jitted=True)
+    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_gpt2(self):
@@ -112,7 +115,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 103)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_forward_cifar(self):
@@ -122,7 +125,7 @@ class TestRealWorld(unittest.TestCase):
       model = SpeedyResNet(Tensor.ones((12,3,2,2)))
       @TinyJit
       def run(X): return model(X)
-      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, (1.0/48)*BS, 126)
+      helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CPU", "CL"}, "slow")
   def test_train_cifar(self):
@@ -139,7 +142,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, (1.0/48)*BS, 126)
+      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_train_cifar_hyp(self):
@@ -176,7 +179,7 @@ class TestRealWorld(unittest.TestCase):
       for v in data.values(): v.to_(Device.DEFAULT)
 
       helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
-          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 427)
+          data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 400)
 
 if __name__ == '__main__':
   unittest.main()

test/opt/test_gen_float4.py

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

test/test_arange.py

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

test/test_assign.py

@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import temp
+from tinygrad.helpers import temp, CI, CPU_LVP
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -180,6 +180,7 @@ class TestAssign(unittest.TestCase):
     new = a + times_a
     np.testing.assert_allclose(new.numpy(), 8)
 
+  @unittest.skipIf(CI and CPU_LVP, "flaky in CI")
   def test_double_assign(self):
     a = Tensor.ones(4).contiguous().realize()
     a += 1

test/test_custom_kernel.py

@@ -155,12 +155,17 @@ class TestCustomKernel(unittest.TestCase):
     self.assertTrue((b_p1 == 3).all().item())
 
   def test_sum(self):
-    # TODO: this only works for float, and silently fails with int
     a = Tensor([1.0, 2, 3, 4, 5])
     tst = Tensor.empty(1)
     b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
     self.assertEqual(b.item(), 15)
 
+  def test_sum_int(self):
+    a = Tensor([1, 2, 3, 4, 5])
+    tst = Tensor.empty(1, dtype=a.dtype)
+    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
+    self.assertEqual(b.item(), 15)
+
   def test_slice_sum(self):
     A = Tensor.randn(16, 16).contiguous()
     B = Tensor.empty(16)

test/test_dtype.py

@@ -17,8 +17,6 @@ pytestmark = pytest.mark.filterwarnings("ignore")
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
 
-if Device.DEFAULT == "CPU": core_dtypes.remove(dtypes.bfloat16)  # NOTE: this is for teenygrad, don't remove
-
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
   if not is_dtype_supported(dtype): return []
   # dont cast internal dtypes
@@ -435,6 +433,8 @@ class TestOpsBFloat16(unittest.TestCase):
     data = [60000.0, 70000.0, 80000.0]
     np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
 
+  # some CPUs there is no native bfloat16 sqrt
+  @unittest.skipIf(Device.DEFAULT == "CPU", "no approximation")
   def test_no_approximation(self):
     data = [326.0, 339.0, 10603200512.0]
     expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
@@ -442,3 +442,4 @@ class TestOpsBFloat16(unittest.TestCase):
 
 if __name__ == '__main__':
   unittest.main()
+

test/test_graph.py

@@ -8,6 +8,8 @@ from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import ExecItem, BufferXfer, get_runner, CompiledRunner
 
+from test.helpers import needs_second_gpu
+
 np.random.seed(1337)
 Tensor.manual_seed(1337)
 BUF_SIZE = 4096 if CI else 4096 * 128
@@ -154,6 +156,7 @@ class TestGraph(unittest.TestCase):
 
     helper_test_graphs(Device[d0].graph, graphs)
 
+  @needs_second_gpu
   def test_copies_2_devs(self):
     self.skip_if_not_multigraph()
 
@@ -167,6 +170,7 @@ class TestGraph(unittest.TestCase):
 
     helper_test_graphs(Device[d0].graph, graphs)
 
+  @needs_second_gpu
   def test_copies_after_graph_global(self):
     self.skip_if_not_multigraph()
 
@@ -215,6 +219,7 @@ class TestGraph(unittest.TestCase):
 
     helper_test_graphs(Device[d0].graph, graphs)
 
+  @needs_second_gpu
   def test_graph_after_copies_devs(self):
     self.skip_if_not_multigraph()
 

test/test_jit.py

@@ -3,7 +3,7 @@ import unittest, functools
 import numpy as np
 
 from hypothesis import given, settings, strategies as strat
-from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV
+from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV, needs_second_gpu
 from tinygrad.tensor import Tensor
 from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
@@ -439,6 +439,7 @@ class TestJit(unittest.TestCase):
       ja = jf(a)
       np.testing.assert_allclose(a.numpy(), ja.numpy(), atol=1e-4, rtol=1e-5)
 
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_jitted_transfers(self):
     d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@@ -472,6 +473,7 @@ class TestJit(unittest.TestCase):
       np.testing.assert_allclose((a.numpy()+b.numpy()), zc.numpy(), atol=1e-4, rtol=1e-5)
       np.testing.assert_allclose((a.numpy()*b.numpy()), wc.numpy(), atol=1e-4, rtol=1e-5)
 
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_jitted_view(self):
     d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@@ -499,6 +501,18 @@ class TestJit(unittest.TestCase):
     b = f(Tensor([2.0]))
     assert abs((a - b).item()) > 0.5
 
+  def test_jit_init_with_empty_different_size(self):
+    @TinyJit
+    def f(x:Tensor) -> Tensor: return (x + 1).realize()
+
+    f(Tensor.empty(1))
+    f(Tensor.empty(1))
+    # TODO: this should fail since input has a different size
+    f(Tensor(2.0)).item()
+    # TODO: this should not fail, and should return 3
+    with self.assertRaises(AssertionError):
+      f(Tensor([2.0])).item()
+
 @unittest.skip("Pending multioutput implementation #3607")
 class TestMultioutputJit(unittest.TestCase):
   def _test(self, f):

test/test_linearizer.py

@@ -45,7 +45,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
 
   @unittest.expectedFailure
@@ -53,7 +53,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -63,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.empty(16)
     out = img.conv2d(w, b)
     ast = helper_linearizer_opt(out)
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     # slice at the last loop end
     uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
     # only valid test if outermost range is the reduce
@@ -84,7 +84,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -92,16 +92,15 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
-  @unittest.expectedFailure # TODO: investigate
   def test_two_nested_range_alt_indexing(self):
     a = Tensor([2, 2]).realize()
     out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
     ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
     assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -114,7 +113,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # LOAD -> RANGE -> LOAD -> STORE
     assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -124,7 +123,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -135,7 +134,7 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -147,7 +146,7 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 
@@ -156,7 +155,8 @@ class TestLinearizer(unittest.TestCase):
     x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
     r = Tensor.conv2d(x,w,padding=1).relu()
 
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer,
+                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
     accs = [u for u in uops if u.op is Ops.DEFINE_REG]
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
@@ -179,7 +179,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(r.schedule()[-1].ast, opts=opts_to_apply)
+    program = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
 
     stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
@@ -194,7 +194,7 @@ class TestLinearizer(unittest.TestCase):
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
-    uops = get_program(r.schedule()[-1].ast, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
@@ -204,14 +204,14 @@ class TestLinearizer(unittest.TestCase):
       if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
         a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
         realized_ast = a.schedule()[-1].ast
-        program = get_program(realized_ast, opts=[])
+        program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
         local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
         assert local[0].dtype.base == acc_dtype
 
   def test_arg_acc_dtype(self):
     def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
       realized_ast = c.schedule()[-1].ast
-      program = get_program(realized_ast, opts=[])
+      program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
       local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
       self.assertEqual(local[0].dtype.base, expected_dtype)
 
@@ -239,7 +239,7 @@ class TestLinearizer(unittest.TestCase):
     opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
     ast = helper_linearizer_opt(r, [opt])
     # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = get_program(ast, opts=opt).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
     for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -353,7 +353,7 @@ class TestLinearizer(unittest.TestCase):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
     ast = helper_linearizer_opt(t+1)
-    uops = get_program(ast, opts=[]).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
     assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -386,13 +386,13 @@ class TestLinearizer(unittest.TestCase):
     sched_copy = sched[:]
     run_schedule(sched)
     np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
-    program = get_program(sched_copy[-1].ast, opts=())
+    program = get_program(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=())
     assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"
 
   def test_phi_simplification(self):
     def helper(t, max_ops=0):
       ast = helper_linearizer_opt(t)
-      uops = get_program(ast).uops
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
       # ignore kernel optimized IF statements for now
       if if_op:=next((u for u in uops if u.op is Ops.IF), None):
         uops = uops[:uops.index(if_op)]
@@ -425,7 +425,7 @@ class TestLinearizer(unittest.TestCase):
     out = x.matmul(y)
     with Context(TC=0):
       ast = helper_linearizer_opt(out)
-      uops = get_program(ast).uops
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
     # check that the float4 cast collapses
     store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
     for val in store_vals:
@@ -436,7 +436,7 @@ class TestLinearizer(unittest.TestCase):
     x = Tensor.randn((4,3,6,6)).realize()
     out = x.flip((0,1)).contiguous()
     ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in get_program(ast).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops if u.op is Ops.STORE][0]
     assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -449,7 +449,7 @@ class TestLinearizer(unittest.TestCase):
             Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
     ast = helper_linearizer_opt(out, opts=[opt])
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = get_program(ast, opts=opt).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
     local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
     global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_GLOBAL for x in get_recursive(u.src[0]))]
     barrier = [u for u in uops if u.op is Ops.BARRIER]
@@ -470,7 +470,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     ast = helper_linearizer_opt(r)
-    uops = get_program(ast).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
     stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
     # the float4 value stores directly in lds and we skip upcast
@@ -517,7 +517,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, opts=opts), device=device))
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/test_linearizer_failures.py

@@ -3,6 +3,7 @@ import unittest
 from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
+from tinygrad.device import Device
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
@@ -18,7 +19,7 @@ class TestLinearizerFailures(unittest.TestCase):
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
     ast = c10.sink()
-    get_program(ast)
+    get_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/test_multitensor.py

@@ -7,7 +7,7 @@ from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import REAL_DEV, not_support_multi_device
+from test.helpers import REAL_DEV, not_support_multi_device, needs_second_gpu
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -35,6 +35,9 @@ def _test_allreduce(t:Tensor):
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestMultiTensor(unittest.TestCase):
+  @needs_second_gpu
+  def setUp(self): pass
+
   def test_to(self):
     X = Tensor.ones(256).contiguous().realize()
     X.to_(devices_2)
@@ -54,6 +57,12 @@ class TestMultiTensor(unittest.TestCase):
       assert lb.shape == (128,)
     (X + X).realize()
 
+  def test_shard_empty(self):
+    GlobalCounters.reset()
+    X = Tensor.empty(256).shard(devices_2, 0).realize()
+    assert GlobalCounters.kernel_count == 0
+    (X + X).realize()
+
   def _test_shard_op(self, op, out, n=4):
     t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
     r = op(t).realize()
@@ -827,6 +836,7 @@ class TestMultiTensor(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestHandleData(unittest.TestCase):
+  @needs_second_gpu
   def test_copied_to_device(self):
     device = (d0, d1, d2, d3)
     t = Tensor([1, 2, 3, 4]).shard(device).realize()
@@ -851,6 +861,9 @@ class TestHandleData(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
+  @needs_second_gpu
+  def setUp(self): pass
+
   # shrink a multitensor on sharded axis
   def test_shrink_bad_args(self):
     t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
@@ -972,6 +985,9 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestBatchNorm(unittest.TestCase):
+  @needs_second_gpu
+  def setUp(self): pass
+
   def test_unsynced_backprop_conv_bn(self):
     with Tensor.train():
       from extra.lr_scheduler import OneCycleLR
@@ -1126,9 +1142,11 @@ def helper_test_shard_op(shps, fxn, atol=1e-6, rtol=1e-3):
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
 class TestTensorOps(unittest.TestCase):
+  @needs_second_gpu
   def test_interpolate(self):
     helper_test_shard_op([(4,16,16),(4,24,24)], lambda x: Tensor.interpolate(x, (19,19)))
 
+  @needs_second_gpu
   def test_bitcast(self):
     helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))
 
@@ -1171,6 +1189,7 @@ class TestMultiRamUsage(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
+  @needs_second_gpu
   def test_from_npy(self):
     t = Tensor(np.arange(100, dtype=np.uint32))
     ll = t.shard((d0, d1), axis=0) + 1
@@ -1180,6 +1199,9 @@ class TestMultiFromUnrenderable(unittest.TestCase):
 class TestMultiAssign(unittest.TestCase):
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
 
+  @needs_second_gpu
+  def setUp(self): pass
+
   def test_multi_assign_realized(self):
     out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
     ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
@@ -1242,6 +1264,7 @@ class TestMultiAssign(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
+  @needs_second_gpu
   def test_transformer(self):
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
 

test/test_nn.py

@@ -9,7 +9,7 @@ from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
 from tinygrad.engine.realize import run_schedule
-from test.helpers import not_support_multi_device
+from test.helpers import not_support_multi_device, needs_second_gpu
 
 @unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow")
 class TestNN(unittest.TestCase):
@@ -481,6 +481,21 @@ class TestNN(unittest.TestCase):
     np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
     np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
 
+  #https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/torch/nn/modules/module.py#L2425
+  def test_load_conv_num_batches_tracked(self):
+    layer = BatchNorm(sz=1, track_running_stats=False)
+    state_dict = {
+      'weight': Tensor.ones(1),
+      'bias': Tensor.ones(1),
+      'num_batches_tracked': Tensor.ones(1),
+    }
+    load_state_dict(layer, state_dict)
+    state_dict['num_batches_tracked'] = Tensor.empty()
+    load_state_dict(layer, state_dict)
+    layer.num_batches_tracked = Tensor.ones(1)
+    load_state_dict(layer, state_dict)
+
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_load_state_dict_sharded_model(self):
     devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
@@ -519,6 +534,7 @@ class TestNN(unittest.TestCase):
     np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
     np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
 
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_load_state_dict_sharded_model_dict_same_axis(self):
     devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")

test/test_ops.py

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM, EMULATE
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -16,6 +16,7 @@ if CI:
 
 FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
 PRINT_TENSORS = getenv("PRINT_TENSORS", 0)
+COMPILE_ONLY = Device.DEFAULT == "NULL" and not EMULATE
 
 def slow_test(test_func):
   return unittest.skipIf(getenv("SKIP_SLOW_TEST"), "Skipping slow test")(test_func)
@@ -38,6 +39,7 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
   tinygrad_fp = time.monotonic() - st
 
   def compare(s, tinygrad_output, torch_output, atol, rtol):
+    if COMPILE_ONLY: return
     if PRINT_TENSORS: print(s, tinygrad_output, torch_output)
     try:
       assert tinygrad_output.shape == torch_output.shape, f"shape mismatch: tinygrad={tinygrad_output.shape} | torch={torch_output.shape}"
@@ -421,8 +423,9 @@ class TestOps(unittest.TestCase):
   def test_isinf(self):
     val = [float('-inf'), 0., float('inf'), float('nan'), 1.1]
     helper_test_op(None, torch.isinf, Tensor.isinf, vals=[val], forward_only=True)
-    np.testing.assert_equal(Tensor(val).isinf(detect_positive=True, detect_negative=False).numpy(), [False, False, True, False, False])
-    np.testing.assert_equal(Tensor(val).isinf(detect_positive=False, detect_negative=True).numpy(), [True, False, False, False, False])
+    if not COMPILE_ONLY:
+      np.testing.assert_equal(Tensor(val).isinf(detect_positive=True, detect_negative=False).numpy(), [False, False, True, False, False])
+      np.testing.assert_equal(Tensor(val).isinf(detect_positive=False, detect_negative=True).numpy(), [True, False, False, False, False])
 
   def test_isnan(self):
     helper_test_op(None, torch.isnan, Tensor.isnan, vals=[[float('-inf'), 0., float('inf'), float('nan'), 1.1]], forward_only=True)
@@ -594,7 +597,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x//2, forward_only=True, vals=[[3, 4, 5]])
     helper_test_op(None, functools.partial(torch.div, rounding_mode="trunc"), Tensor.idiv, forward_only=True,
                    vals=[[-4, 7, 5, 4, -7, 8], [2, -3, 8, -2, 3, 5]])
-    if is_dtype_supported(dtypes.uint64):
+    if is_dtype_supported(dtypes.uint64) and not COMPILE_ONLY:
       x = Tensor(2**64 - 1, dtype=dtypes.uint64).idiv(1)
       np.testing.assert_equal(x.numpy(), 2**64 - 1)
 
@@ -679,6 +682,7 @@ class TestOps(unittest.TestCase):
     # float to power of int
     helper_test_op(None, lambda x: 0.7**x, vals=[[-2,-1,0,1,2,3]], forward_only=True)
 
+  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
   def test_pow_const_direct(self):
     # x ** c
     def get_tiny_gradient(x, c):
@@ -1070,8 +1074,7 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_cummax(self):
     helper_test_op([()], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
-    # TODO: torch allows this?
-    # self.helper_test_exception([()], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
+    self.helper_test_exception([()], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
     helper_test_op([(20,)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0))
     self.helper_test_exception([(20,)], lambda x: torch.cummax(x, dim=1).values, lambda x: Tensor.cummax(x, axis=1), expected=IndexError)
     self.helper_test_exception([(20,)], lambda x: torch.cummax(x, dim=-2).values, lambda x: Tensor.cummax(x, axis=-2), expected=IndexError)
@@ -1088,8 +1091,9 @@ class TestOps(unittest.TestCase):
     # check if it returns the first index for multiple occurences
     helper_test_op(None, lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[2, 2]])
     helper_test_op(None, lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[1, 2, 2]])
-    np.testing.assert_equal(Tensor([2,2]).argmax().numpy(), 0)
-    np.testing.assert_equal(Tensor([1,2,2]).argmax().numpy(), 1)
+    if not COMPILE_ONLY:
+      np.testing.assert_equal(Tensor([2,2]).argmax().numpy(), 0)
+      np.testing.assert_equal(Tensor([1,2,2]).argmax().numpy(), 1)
     helper_test_op([(10,20)], lambda x: x.argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True)
     helper_test_op([(10,20)], lambda x: x.argmax(0, False).type(torch.int32), lambda x: x.argmax(0, False), forward_only=True)
     helper_test_op([(10,20)], lambda x: x.argmax(1, False).type(torch.int32), lambda x: x.argmax(1, False), forward_only=True)
@@ -1107,8 +1111,9 @@ class TestOps(unittest.TestCase):
     # check if it returns the first index for multiple occurences
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[3, 2, 2]])
-    np.testing.assert_equal(Tensor([2,2]).argmin().numpy(), 0)
-    np.testing.assert_equal(Tensor([3,2,2]).argmin().numpy(), 1)
+    if not COMPILE_ONLY:
+      np.testing.assert_equal(Tensor([2,2]).argmin().numpy(), 0)
+      np.testing.assert_equal(Tensor([3,2,2]).argmin().numpy(), 1)
     helper_test_op([(10,20)], lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True)
     helper_test_op([(10,20)], lambda x: x.argmin(0, False).type(torch.int32), lambda x: x.argmin(0, False), forward_only=True)
     helper_test_op([(10,20)], lambda x: x.argmin(1, False).type(torch.int32), lambda x: x.argmin(1, False), forward_only=True)
@@ -1156,12 +1161,13 @@ class TestOps(unittest.TestCase):
                           lambda x: x.topk(4, dim, largest, sorted_).indices.type(torch.int32),
                           lambda x: x.topk(4, dim, largest, sorted_)[1], forward_only=True)
     # repeated values
-    value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
-    np.testing.assert_equal(value.numpy(), [1, 1, 1])
-    np.testing.assert_equal(indices.numpy(), [0, 1, 3])
-    value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3, largest=False)
-    np.testing.assert_equal(value.numpy(), [0, 0, 0])
-    np.testing.assert_equal(indices.numpy(), [2, 4, 6])
+    if not COMPILE_ONLY:
+      value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
+      np.testing.assert_equal(value.numpy(), [1, 1, 1])
+      np.testing.assert_equal(indices.numpy(), [0, 1, 3])
+      value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3, largest=False)
+      np.testing.assert_equal(value.numpy(), [0, 0, 0])
+      np.testing.assert_equal(indices.numpy(), [2, 4, 6])
     self.helper_test_exception([(4)], lambda x: x.topk(5), expected=(RuntimeError, ValueError))
 
   @slow_test
@@ -1313,6 +1319,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
   @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and CPU_LLVM) or IMAGE
   or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "not precise enough")
   def test_gemm_fp16(self):
     helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
   def test_gemm(self):
@@ -1723,6 +1730,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(7,5,10)], lambda x: x[1:5:2, 3, ::4])
     helper_test_op([(7,5,10)], lambda x: x[1:5:2, None, None, 3, None, ::4])
 
+  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
   def test_slice_negative_strides(self):
     # Torch doesn't support slicing with negative steps
     a = np.random.randn(10, 10, 10).astype(np.float32)
@@ -2729,6 +2737,9 @@ class TestOps(unittest.TestCase):
 
   def test_clip(self):
     helper_test_op([(45,65)], lambda x: x.clip(-2.3, 1.2))
+    # NOTE: torch set backward to 1 at the boundaries
+    # https://github.com/pytorch/pytorch/blob/7a41b66367c38d0af3e8a90f7be48d6b281e7bca/tools/autograd/derivatives.yaml#L421
+    helper_test_op(None, lambda x: x.clip(-2.5, 1.5), vals=[[-3.0, -2.5, 0, 1.5, 2]])
     helper_test_op([(45,65)], lambda x: x.clip(0, 0))
     helper_test_op([(45,65)], lambda x: x.clip(10, 100))
     helper_test_op([(45,65)], lambda x: x.clip(0, 0.1))
@@ -2749,6 +2760,7 @@ class TestOps(unittest.TestCase):
     n = Tensor([1, float("nan")]).max().numpy()
     assert math.isnan(n.item()), f"{n.item()} is not nan"
 
+  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
   def test_inf_where(self):
     x = Tensor.full((3, 3), float("inf"))
     n = (x < 0).where(x, 1).numpy()
@@ -3165,6 +3177,7 @@ class TestOps(unittest.TestCase):
 
   @unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "QCOM fails with: Resource deadlock avoided")
   def test_masked_select(self):
     helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
     helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)

test/test_opt_gemm.py

@@ -1,6 +1,6 @@
 import numpy as np
 import unittest
-from tinygrad import Tensor
+from tinygrad import Tensor, Device
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
@@ -17,7 +17,7 @@ class TestOptGemm(unittest.TestCase):
     t = self.a.T @ self.b.T
     # TODO: this should be a generic test helper
     si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(si.ast, opts=opts))
+    run = CompiledRunner(get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts))
     ExecItem(run, si.bufs).run()
     test = si.bufs[0].numpy().reshape(self.res.shape)
     np.testing.assert_allclose(self.res, test, atol=1e-4)

test/test_opts.py

@@ -13,7 +13,7 @@ class TestOpts(unittest.TestCase):
     s = out.schedule()
     self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
     if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
-      prg = get_program(s[-1].ast)
+      prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer)
       self.assertIn('float4', prg.src)
 
 if __name__ == '__main__':

test/test_quantize_onnx.py

@@ -38,7 +38,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
 
 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
   si = out.schedule()[-1]
-  prg = get_program(si.ast, opts=opts)
+  prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
   if replace_src is not None:
     old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
     prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))

test/test_randomness.py

@@ -7,7 +7,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import lower_schedule, CompiledRunner
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from test.helpers import not_support_multi_device
+from test.helpers import not_support_multi_device, needs_second_gpu
 
 import numpy as np
 import torch
@@ -141,6 +141,7 @@ class TestRandomness(unittest.TestCase):
     r = Tensor.rand(10).numpy()
     np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
 
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_threefry_tensors_cnt(self):
     Tensor.manual_seed(1337)
@@ -160,6 +161,7 @@ class TestRandomness(unittest.TestCase):
     assert len(Tensor._device_rng_counters) == 0
     assert len(Tensor._device_seeds) == 0
 
+  @needs_second_gpu
   @unittest.skipIf(not_support_multi_device(), "no multi")
   def test_threefry_same_kernels(self):
     Tensor.manual_seed(0)

test/test_schedule.py

@@ -12,8 +12,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
-from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.schedule.rangeify import Kernel
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 
 class KernelCountException(Exception): pass
@@ -24,13 +23,11 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
-    sink = UOp.sink(t) if t.op is not Ops.SINK else t
-    becomes_map = get_rangeify_map(sink)
-    sched, _ = create_schedule_with_vars(sink.substitute(becomes_map))
+    sched = Tensor(t).schedule()
   # test lowering all the ScheduleItems to ExecItems
   kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
   if kernel_cnt != allowed:
-    print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
     if DEBUG >= 3:
       for i,s in enumerate(sched):
         print("kernel", i+1)

test/test_subbuffer.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
 from tinygrad.helpers import Context
-from test.helpers import REAL_DEV
+from test.helpers import REAL_DEV, needs_second_gpu
 
 @unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
 class TestSubBuffer(unittest.TestCase):
@@ -41,6 +41,7 @@ class TestSubBuffer(unittest.TestCase):
     out = (vt + 100).tolist()
     assert out == [102, 103]
 
+  @needs_second_gpu
   @unittest.skipIf(REAL_DEV not in {"CUDA", "NV", "AMD"}, "only NV, AMD, CUDA")
   def test_subbuffer_transfer(self):
     t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()

test/test_tensor.py

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

@@ -1,14 +1,14 @@
 import unittest
 
-from tinygrad import Tensor, Device
+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"
-    hevc_tensor = Tensor.from_url(url, device="CPU")
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
 
-    dat = bytes(hevc_tensor.data())
     opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat, device=Device.DEFAULT)
 
     def _test_common(frame, bts):

test/testextra/test_tk.py

@@ -1,8 +1,9 @@
-import unittest, math
+import unittest, math, time
 
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import CI
 import numpy as np
 
@@ -12,11 +13,16 @@ from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout
 
 @unittest.skipIf(CI or Device.DEFAULT not in ["AMD"], "only amd")
 class TestTK(unittest.TestCase):
+  def setUp(self):
+    arch = Device["AMD"].arch
+    if not arch.startswith("gfx9"):
+      self.skipTest(f"arch {arch} not supported")
+
   @unittest.skipIf(CI, "no wmma in ci")
   def test_simple_matmul(self):
     N = 8192
     BLOCK_SIZE = 64
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel("simple_matmul", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -65,7 +71,7 @@ class TestTK(unittest.TestCase):
   def test_simple_matmul_transposed(self):
     N = 8192
     BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128
-    with Kernel((N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
+    with Kernel("simple_matmul_transposed", (N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       c = ker.gl((1, 1, N, N), dtypes.float32)
@@ -113,7 +119,7 @@ class TestTK(unittest.TestCase):
   def test_load_store(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+    with Kernel("load_store", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -150,7 +156,7 @@ class TestTK(unittest.TestCase):
   def test_load_store_group(self):
     N = 256
     BLOCK_SIZE = 64
-    with Kernel((N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
+    with Kernel("load_store_group", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS * 2) as ker:
       warp = ker.warp
       group = ker.group(2)
 
@@ -187,7 +193,7 @@ class TestTK(unittest.TestCase):
   def test_add(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("add", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -224,7 +230,7 @@ class TestTK(unittest.TestCase):
   def test_max(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("max", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -269,7 +275,7 @@ class TestTK(unittest.TestCase):
   def test_max_nonsquare(self):
     N, M = 32, 128
     BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("max_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -314,7 +320,7 @@ class TestTK(unittest.TestCase):
   def test_sum(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("sum", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, N), dtypes.float32)
@@ -359,7 +365,7 @@ class TestTK(unittest.TestCase):
   def test_sum_nonsquare(self):
     N, M = 32, 128
     BLOCK_N, BLOCK_M = 16, 64
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("sum_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, M), dtypes.float32)
@@ -404,7 +410,7 @@ class TestTK(unittest.TestCase):
   def test_softmax(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("softmax", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
@@ -422,28 +428,29 @@ class TestTK(unittest.TestCase):
       norm_vec = warp.zero(norm_vec)
 
       for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
+        a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg_ = warp.load(a_reg, a_smem_)
 
-        a_reg *= 1.0 / math.log(2)
+        a_reg_ *= 1.0 / math.log(2)
 
         max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
-        max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
-        a_reg = (a_reg - max_vec).exp2()
+        max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
+        a_reg_ = (a_reg_ - max_vec).exp2()
         max_vec_last = (max_vec_last - max_vec).exp2()
         norm_vec *= max_vec_last
-        norm_vec = warp.row_reduce(norm_vec, a_reg, lambda a, b: a + b)
+        norm_vec = warp.row_reduce(norm_vec, a_reg_, lambda a, b: a + b)
       norm_vec = ker.endrange()
+      max_vec = max_vec.after(norm_vec)
 
       for tile_col in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
-        a_reg = warp.load(a_reg.after(norm_vec), a_smem)
+        a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
+        a_reg_ = warp.load(a_reg, a_smem_)
 
-        a_reg *= 1.0 / math.log(2)
-        a_reg = (a_reg - max_vec).exp2()
-        a_reg /= norm_vec
+        a_reg_ *= 1.0 / math.log(2)
+        a_reg_ = (a_reg_ - max_vec).exp2()
+        a_reg_ /= norm_vec
 
-        b = warp.store(b, a_reg, (0, 0, 0, tile_col), (), axis=2)
+        b = warp.store(b, a_reg_, (0, 0, 0, tile_col), (), axis=2)
 
       sink = ker.finish()
 
@@ -463,7 +470,7 @@ class TestTK(unittest.TestCase):
   def test_softmax_col(self):
     N = 64
     BLOCK_SIZE = 32
-    with Kernel((1, 1, 1), WARP_THREADS) as ker:
+    with Kernel("softmax_col", (1, 1, 1), WARP_THREADS) as ker:
       warp = ker.warp
 
       b = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
@@ -481,28 +488,29 @@ class TestTK(unittest.TestCase):
       norm_vec = warp.zero(norm_vec)
 
       for tile_row in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
-        a_reg = warp.load(a_reg, a_smem)
+        a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
+        a_reg_ = warp.load(a_reg, a_smem_)
 
-        a_reg *= 1.0 / math.log(2)
+        a_reg_ *= 1.0 / math.log(2)
 
         max_vec_last = warp.copy(max_vec_last.after(tile_row), max_vec)
-        max_vec = warp.col_reduce(max_vec.after(max_vec_last), a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
-        a_reg = (a_reg - max_vec).exp2()
+        max_vec = warp.col_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
+        a_reg_ = (a_reg_ - max_vec).exp2()
         max_vec_last = (max_vec_last - max_vec).exp2()
         norm_vec *= max_vec_last
-        norm_vec = warp.col_reduce(norm_vec, a_reg, lambda a, b: a + b)
+        norm_vec = warp.col_reduce(norm_vec, a_reg_, lambda a, b: a + b)
       norm_vec = ker.endrange()
+      max_vec = max_vec.after(norm_vec)
 
       for tile_row in ker.range(N // BLOCK_SIZE):
-        a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
-        a_reg = warp.load(a_reg.after(norm_vec), a_smem)
+        a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
+        a_reg_ = warp.load(a_reg.after(norm_vec), a_smem_)
 
-        a_reg *= 1.0 / math.log(2)
-        a_reg = (a_reg - max_vec).exp2()
-        a_reg /= norm_vec
+        a_reg_ *= 1.0 / math.log(2)
+        a_reg_ = (a_reg_ - max_vec).exp2()
+        a_reg_ /= norm_vec
 
-        b = warp.store(b, a_reg, (0, 0, tile_row, 0), (), axis=2)
+        b = warp.store(b, a_reg_, (0, 0, tile_row, 0), (), axis=2)
 
       sink = ker.finish()
 
@@ -521,11 +529,11 @@ class TestTK(unittest.TestCase):
 
   def test_fa(self):
     NUM_WORKERS = 1
-    B, N, H, H_KV, D = 1, 8192, 32, 8, 128
+    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
     Q_BLOCK_SIZE = 16
     KV_BLOCK_SIZE = 16
     GROUP_SIZE = H // H_KV
-    with Kernel((H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
+    with Kernel("fa", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
       warp = ker.warp
 
       # kernel
@@ -605,6 +613,7 @@ class TestTK(unittest.TestCase):
         att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
         o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
       o_reg = ker.endrange()
+      norm_vec = norm_vec.after(o_reg)
 
       o_reg /= norm_vec
 
@@ -621,7 +630,12 @@ class TestTK(unittest.TestCase):
       Tensor.realize(q, k, v, out)
 
     ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (out, q, k, v)])
-    for _ in range(5): ei.run(wait=True)
+    for _ in range(5):
+      et = ei.run(wait=True)
+      attn_flops = 2 * B * H * N * N * D + \
+                   4 * B * H * N * N + \
+                   2 * B * H * N * N * D
+      print(f"{attn_flops/(et*1e9):2f} GFLOPS")
     out = out.float()
 
     q_permuted = q.permute(0, 2, 1, 3)
@@ -630,7 +644,36 @@ class TestTK(unittest.TestCase):
     ref = q_permuted.scaled_dot_product_attention(k_permuted, v_permuted, is_causal=True, enable_gqa=True).float()
     ref = ref.permute(0, 2, 1, 3)
 
-    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-2, rtol=1e-5)
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
+
+  def test_fast_fa(self):
+    from extra.thunder.tiny.fa import flash_attention
+
+    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
+
+    with Context(DEBUG=0):
+      q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
+      k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      Tensor.realize(q, k, v)
+
+    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+    fa_jitted = TinyJit(flash_attention)
+
+    for _ in range(10):
+      st = time.perf_counter()
+      out = fa_jitted(q, k, v, is_causal=True)
+      et = time.perf_counter() - st
+      attn_flops = 2 * B * H * N * N * D + \
+                   4 * B * H * N * N + \
+                   2 * B * H * N * N * D
+      print(f"{attn_flops/(et*1e9):2f} GFLOPS")
+    out = out.float().transpose(1, 2)
+
+    ref = q.scaled_dot_product_attention(k, v, is_causal=True, enable_gqa=True).float().transpose(1, 2)
+
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_autogen.py

@@ -44,6 +44,34 @@ class TestAutogen(unittest.TestCase):
       test.argtypes = [Baz]
       self.assertEqual(test(b), b.a + b.b + b.c + b.d)
 
+  # https://github.com/python/cpython/issues/90914
+  @unittest.skipIf(WIN, "doesn't compile on windows")
+  def test_bitfield_interop(self):
+    class Baz(Struct): pass
+    Baz._fields_ = [(chr(ord('a') + i), ctypes.c_bool, 1) for i in range(8)]
+    src = '''#include <stdbool.h>
+      struct baz {
+        bool a:1;
+        bool b:1;
+        bool c:1;
+        bool d:1;
+        bool e:1;
+        bool f:1;
+        bool g:1;
+        bool h:1;
+      };
+
+      int test(struct baz x) {
+        return x.c;
+      }
+    '''
+    args = ('-x', 'c', '-fPIC', '-shared')
+    with tempfile.NamedTemporaryFile(suffix=".so") as f:
+      subprocess.check_output(('clang',) + args + ('-', '-o', f.name), input=src.encode('utf-8'))
+      test = ctypes.CDLL(f.name).test
+      test.argtypes = [Baz]
+      for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
+
   @unittest.skipIf(WIN, "doesn't compile on windows")
   def test_packed_structs(self):
     NvU32 = ctypes.c_uint32
@@ -82,4 +110,53 @@ class TestAutogen(unittest.TestCase):
     assert frts_cmd.readVbiosDesc.__class__ is FWSECLIC_READ_VBIOS_DESC
     assert frts_cmd.frtsRegionDesc.__class__ is FWSECLIC_FRTS_REGION_DESC
 
+  def test_packed_fields(self):
+    uint8_t = ctypes.c_ubyte
+    uint16_t = ctypes.c_ushort
+    uint32_t = ctypes.c_uint32
+
+    class struct_die_info(Struct): pass
+    struct_die_info._packed_ = True
+    struct_die_info._fields_ = [
+      ('die_id', uint16_t),
+      ('die_offset', uint16_t),
+    ]
+    die_info = struct_die_info
+    class struct_ip_discovery_header(Struct): pass
+    class struct_ip_discovery_header_0(ctypes.Union): pass
+    class struct_ip_discovery_header_0_0(Struct): pass
+    uint8_t = ctypes.c_ubyte
+    struct_ip_discovery_header_0_0._fields_ = [
+      ('base_addr_64_bit', uint8_t,1),
+      ('reserved', uint8_t,7),
+      ('reserved2', uint8_t),
+    ]
+    struct_ip_discovery_header_0._anonymous_ = ['_0']
+    struct_ip_discovery_header_0._packed_ = True
+    struct_ip_discovery_header_0._fields_ = [
+      ('padding', (uint16_t * 1)),
+      ('_0', struct_ip_discovery_header_0_0),
+    ]
+    struct_ip_discovery_header._anonymous_ = ['_0']
+    struct_ip_discovery_header._packed_ = True
+    struct_ip_discovery_header._fields_ = [
+      ('signature', uint32_t),
+      ('version', uint16_t),
+      ('size', uint16_t),
+      ('id', uint32_t),
+      ('num_dies', uint16_t),
+      ('die_info', (die_info * 16)),
+      ('_0', struct_ip_discovery_header_0),
+    ]
+    ip_discovery_header = struct_ip_discovery_header
+
+    hdr = b'IPDS\x04\x00|\x1d\x80\x1a\xffd\x01\x00\x00\x00\x8c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00' # noqa: E501
+    ihdr = ip_discovery_header.from_buffer_copy(hdr)
+
+    assert ctypes.sizeof(ihdr) == 80
+    assert ihdr.signature == 0x53445049
+    assert ihdr.version == 0x0004
+    assert ihdr.num_dies == 1
+    assert ihdr.base_addr_64_bit == 1
+
 if __name__ == "__main__": unittest.main()

test/unit/test_device.py

@@ -42,12 +42,10 @@ class TestDevice(unittest.TestCase):
                         shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_LLVM": "1"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_LLVM": "0"})
-      subprocess.run([f'python3 -c "{imports}; {expect_failure}"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "0", "CPU_LLVM": "0"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "0"})
-      subprocess.run([f'python3 -c "{imports}; {expect_failure}"'],
-                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CLANGJIT": "1", "CPU_LLVM": "1"})
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CC": "LLVM"})
+      subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangJITCompiler)"'],
+                        shell=True, check=True, env={**os.environ, "DEV": "CPU", "CPU_CC": "CLANGJIT"})
     elif Device.DEFAULT == "AMD":
       from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
       try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
@@ -64,14 +62,20 @@ class TestDevice(unittest.TestCase):
                         shell=True, check=True, env={**os.environ, "DEV": "AMD", "AMD_HIP": "1", "AMD_LLVM": "1"})
     else: self.skipTest("only run on CPU/AMD")
 
-  def test_compiler_envvar(self):
-    d = Device[Device.DEFAULT]
-    dname = Device.DEFAULT.split(':')[0].upper()
-    assert d._get_compiler_envvar(type("Compiler", (), {})) == f"{dname}_COMPILER"
-    assert d._get_compiler_envvar(type("LLVMCompiler", (), {})) == f"{dname}_LLVM"
-    assert d._get_compiler_envvar(type("RandomCompiler", (), {})) == f"{dname}_RANDOM"
-    assert d._get_compiler_envvar(type(f"{dname}Compiler", (), {})) == f"{dname}_{dname}COMPILER" # do not repeat device name alone
-    assert d._get_compiler_envvar(type(f"{dname}LLVMCompiler", (), {})) == f"{dname}_LLVM" # do not repeat device name
+  @unittest.skipIf((WIN and CI) or (not Device.DEFAULT == "CPU"), "skipping windows test")
+  def test_env_online(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler
+    try: _, _ = CPULLVMCompiler(), ClangJITCompiler()
+    except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
+
+    with Context(CPU_LLVM=1):
+      inst = Device["CPU"].compiler
+      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
+    with Context(CPU_LLVM=0):
+      self.assertIsInstance(Device["CPU"].compiler, ClangJITCompiler)
+    with Context(CPU_LLVM=1):
+      self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
+      assert inst is Device["CPU"].compiler  # cached
 
 class MockCompiler(Compiler):
   def __init__(self, key): super().__init__(key)

test/unit/test_helpers.py

@@ -163,6 +163,14 @@ class TestFetch(unittest.TestCase):
     fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
           allow_caching=False)
 
+  def test_fetch_half_and_full_file(self):
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-10"}).read_bytes()
+    assert len(x) == 11, f"{len(x) != 11}"
+    x = fetch("https://csrc.nist.gov/CSRC/media/Projects/lightweight-cryptography/documents/finalist-round/updated-submissions/sparkle.zip",
+          headers={"Range": "bytes=0-100"}).read_bytes()
+    assert len(x) == 101, f"{len(x) != 101}"
+
 class TestFullyFlatten(unittest.TestCase):
   def test_fully_flatten(self):
     self.assertEqual(fully_flatten([[1, 3], [1, 2]]), [1, 3, 1, 2])

test/unit/test_indexing.py

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

test/unit/test_llm_server.py

@@ -0,0 +1,138 @@
+import unittest, threading, time
+from unittest.mock import Mock
+
+class TestLLMServer(unittest.TestCase):
+  """Integration tests using the real OpenAI client."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.role = Mock(return_value=[100, 101])
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="Hello")
+    cls.mock_tok.end_turn = Mock(return_value=[999])
+    cls.mock_tok.stop_tokens = Mock(return_value={999})
+
+    cls.mock_model = Mock()
+    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
+
+    cls.bos_id = 1
+    cls.eos_id = 999
+
+    import tinygrad.apps.llm as llm_module
+    llm_module.model = cls.mock_model
+    llm_module.tok = cls.mock_tok
+    llm_module.bos_id = cls.bos_id
+    llm_module.eos_id = cls.eos_id
+
+    from tinygrad.apps.llm import Handler
+    from tinygrad.helpers import TCPServerWithReuse
+
+    cls.port = 11435
+    cls.server = TCPServerWithReuse(('127.0.0.1', cls.port), Handler)
+    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def test_chat_completion_stream(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[0].choices[0].delta.role, "assistant")
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_openai_response_structure(self):
+    stream = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Test"}],
+      stream=True
+    )
+
+    for chunk in stream:
+      self.assertTrue(chunk.id.startswith("chatcmpl-"))
+      self.assertEqual(chunk.object, "chat.completion.chunk")
+      self.assertIsNotNone(chunk.choices)
+      self.assertIsNotNone(chunk.created)
+      self.assertIsInstance(chunk.created, int)
+      self.assertEqual(chunk.model, "test-model")
+
+  def test_stream_with_usage(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True,
+      stream_options={"include_usage": True}
+    )
+
+    chunks = list(stream)
+    last_chunk = chunks[-1]
+
+    self.assertIsNotNone(last_chunk.usage)
+    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
+    self.assertIsNotNone(last_chunk.usage.completion_tokens)
+    self.assertIsNotNone(last_chunk.usage.total_tokens)
+
+  def test_multi_turn_conversation(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[
+        {"role": "system", "content": "You are helpful."},
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi!"},
+        {"role": "user", "content": "How are you?"}
+      ],
+      stream=True
+    )
+
+    chunks = list(stream)
+    self.assertGreater(len(chunks), 0)
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
+
+  def test_content_is_streamed(self):
+    stream = self.client.chat.completions.create(
+      model="test",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=True
+    )
+
+    contents = []
+    for chunk in stream:
+      if chunk.choices and chunk.choices[0].delta.content:
+        contents.append(chunk.choices[0].delta.content)
+
+    self.assertGreater(len(contents), 0)
+
+  def test_non_streaming(self):
+    resp = self.client.chat.completions.create(
+      model="test-model",
+      messages=[{"role": "user", "content": "Hello"}],
+      stream=False
+    )
+
+    self.assertTrue(resp.id.startswith("chatcmpl-"))
+    self.assertEqual(resp.object, "chat.completion")
+    self.assertEqual(resp.model, "test-model")
+    self.assertIsNotNone(resp.created)
+    self.assertEqual(len(resp.choices), 1)
+    self.assertEqual(resp.choices[0].message.role, "assistant")
+    self.assertIsNotNone(resp.choices[0].message.content)
+    self.assertEqual(resp.choices[0].finish_reason, "stop")
+    self.assertIsNotNone(resp.usage)
+    self.assertIsNotNone(resp.usage.prompt_tokens)
+    self.assertIsNotNone(resp.usage.completion_tokens)
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_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_uop_symbolic.py

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

tinygrad/apps/llm.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers
+import sys, argparse, typing, re, unicodedata, json, uuid, time
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv
+from tinygrad.helpers import partition, TCPServerWithReuse, HTTPRequestHandler, tqdm, DEBUG
 
 class SimpleTokenizer:
-  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
+  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3"):
+    self._preset = preset
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
@@ -22,10 +24,11 @@ class SimpleTokenizer:
   @staticmethod
   def from_gguf_kv(kv:dict):
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
-    if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'")
+    preset = kv["tokenizer.ggml.pre"]
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
     vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
-    normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+    normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
+    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens), preset)
 
   def _encode_word(self, word:bytes) -> list[int]:
     if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
@@ -48,7 +51,16 @@ class SimpleTokenizer:
     return tokens + self._encode_sentence(text[pos:])
 
   def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode()
-  def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
+  def role(self, role:str):
+    if self._preset == "qwen2": return self.encode(f"<|im_start|>{role}\n")
+    return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
+  def end_turn(self):
+    if self._preset == "qwen2": return self.encode("<|im_end|>\n")
+    return self.encode("<|eot_id|>")
+  def stop_tokens(self) -> set[int]:
+    """Returns tokens that indicate end of generation (subset of end_turn)"""
+    if self._preset == "qwen2": return {self._special_tokens["<|im_end|>"]}
+    return {self._special_tokens["<|eot_id|>"]}
 
 def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
   B, H, T, Hd = x.shape
@@ -63,22 +75,28 @@ def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
                       x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)
 
 class TransformerBlock:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int=0):
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int=0, head_dim:int|None=None, rope_base:float=10000.0, qk_norm:bool=False):
     self.n_heads      = n_heads
     self.n_kv_heads   = n_kv_heads
-    self.head_dim     = dim // n_heads
+    self.head_dim     = head_dim if head_dim is not None else dim // n_heads
     self.max_context  = max_context
+    self.rope_base    = rope_base
 
     # --- attention projections (all linear, bias-free) ------------------
-    kv_proj_out      = self.head_dim * n_kv_heads    # Llama-3 uses the same dim for K/V
-    self.attn_q      = nn.Linear(dim, dim,         bias=False)
+    q_proj_out       = self.head_dim * n_heads
+    kv_proj_out      = self.head_dim * n_kv_heads
+    self.attn_q      = nn.Linear(dim, q_proj_out,  bias=False)
     self.attn_k      = nn.Linear(dim, kv_proj_out, bias=False)
     self.attn_v      = nn.Linear(dim, kv_proj_out, bias=False)
-    self.attn_output = nn.Linear(dim, dim,         bias=False)
+    self.attn_output = nn.Linear(q_proj_out, dim,  bias=False)
 
     # --- RMSNorms --------------------------------------------------------
     self.attn_norm   = nn.RMSNorm(dim, norm_eps)
     self.ffn_norm    = nn.RMSNorm(dim, norm_eps)
+    # QK normalization (Qwen3-style)
+    if qk_norm:
+      self.attn_q_norm = nn.RMSNorm(self.head_dim, norm_eps)
+      self.attn_k_norm = nn.RMSNorm(self.head_dim, norm_eps)
 
     # --- feed-forward ----------------------------------------------------
     self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
@@ -94,8 +112,12 @@ class TransformerBlock:
     k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
     v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
 
-    q = apply_rope(q, start_pos)
-    k = apply_rope(k, start_pos)
+    # QK normalization (applied before RoPE)
+    if hasattr(self, "attn_q_norm"): q = self.attn_q_norm(q)
+    if hasattr(self, "attn_k_norm"): k = self.attn_k_norm(k)
+
+    q = apply_rope(q, start_pos, self.rope_base)
+    k = apply_rope(k, start_pos, self.rope_base)
 
     # TODO: remove these kv cache realizes
     if not hasattr(self, "cache_kv"):
@@ -120,8 +142,8 @@ class TransformerBlock:
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:
-  def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, max_context):
-    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context) for _ in range(num_blocks)]
+  def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, max_context, head_dim=None, rope_base=10000.0, qk_norm=False):
+    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context, head_dim, rope_base, qk_norm) for _ in range(num_blocks)]
     self.token_embd  = nn.Embedding(vocab_size, dim)
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
@@ -139,7 +161,7 @@ class Transformer:
     return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)
 
   @staticmethod
-  def from_gguf(gguf:Tensor, max_context:int|None=None) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=True) -> tuple[Transformer, dict]:
     # TODO: remove the need for copy to default device
     kv, state_dict = nn.state.gguf_load(gguf.to(None))
 
@@ -151,12 +173,17 @@ class Transformer:
 
     arch = kv['general.architecture']
     max_context = min(max_context, kv[f'{arch}.context_length']) if max_context is not None else kv[f'{arch}.context_length']
+    head_dim = kv.get(f'{arch}.attention.key_length')  # None for models that use dim // n_heads
+    rope_base = kv.get(f'{arch}.rope.freq_base', 10000.0)
+    qk_norm = 'blk.0.attn_q_norm.weight' in state_dict  # detect QK normalization (Qwen3-style)
     model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'], hidden_dim=kv[f'{arch}.feed_forward_length'],
                         n_heads=kv[f'{arch}.attention.head_count'], n_kv_heads=kv[f'{arch}.attention.head_count_kv'],
-                        norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), max_context=max_context)
+                        norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), max_context=max_context,
+                        head_dim=head_dim, rope_base=rope_base, qk_norm=qk_norm)
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
-    for s in nn.state.get_parameters(model): s.replace(s.contiguous())
+    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
+    if realize: Tensor.realize(*params)
     return model, kv
 
   def generate(self, tokens:list[int], start_pos=0):
@@ -172,34 +199,90 @@ class Transformer:
       yield next_id
 
 models = {
-  "1B": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
-  "3B": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
-  "3B_f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
-  "8B": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
+  "llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
+  "llama3.2:3b": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
+  "llama3.2:3b-f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
+  "llama3.1:8b": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
+  "qwen3:0.6b": "https://huggingface.co/Qwen/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q8_0.gguf",
 }
 
+# *** simple OpenAI compatible server on 11434 to match ollama ***
+# OPENAI_BASE_URL=http://localhost:11434/v1 OPENAI_API_KEY=ollama uvx --from gpt-command-line gpt
+
+class Handler(HTTPRequestHandler):
+  def run_model(self, ids:list[int], model_name:str, include_usage=False):
+    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
+    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
+    out, stop_ids = [], tok.stop_tokens()
+    for next_id in tqdm(model.generate(ids), disable=not DEBUG>=1):
+      if next_id in stop_ids: break
+      out.append(next_id)
+      yield {"choices": [{"index":0, "delta":{"content":tok.decode([next_id])}, "finish_reason":None}], **tmpl}
+    yield {"choices": [{"index":0, "delta":{},"finish_reason":"stop"}], **tmpl}
+    if include_usage:
+      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+
+  def do_POST(self):
+    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
+    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
+    if DEBUG >= 1:
+      print(self.path)
+      print(json.dumps(body, indent=2))
+    if self.path == "/v1/chat/completions":
+      # extract tokens
+      ids = [bos_id]
+      for msg in body["messages"]:
+        ids += tok.role(msg["role"])
+        # content can be a str or a list
+        content = msg["content"]
+        if isinstance(content, str): ids += tok.encode(content)
+        elif isinstance(content, list):
+          for c in content:
+            if c["type"] == "text": ids += tok.encode(c["text"])
+            else: raise RuntimeError(f"unhandled type: {c['type']}")
+        else: raise RuntimeError(f"unknown content type: {type(content)}")
+        ids += tok.end_turn()
+      ids += tok.role("assistant")
+
+      # reply
+      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False))
+      if body.get("stream"): self.stream_json(chunks)
+      else:
+        out = []
+        for c in chunks: out.append(c["choices"][0]["delta"].get("content", "") if c["choices"] else "")
+        self.send_data(json.dumps({**c, "object":"chat.completion",
+          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":"stop"}]}).encode())
+    else:
+      raise RuntimeError(f"unhandled path {self.path}")
+
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--size", choices=list(models.keys()), default=list(models.keys())[0], help="Model size")
+  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
   parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
+  parser.add_argument("--serve", action="store_true", help="Run OpenAI compatible API")
   args = parser.parse_args()
 
   # load the model
-  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.size]), args.max_context)
+  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
+  if DEBUG >= 1: print(f"using model {args.model}")
 
   # extract some metadata
   tok = SimpleTokenizer.from_gguf_kv(kv)
   bos_id: int = kv['tokenizer.ggml.bos_token_id']
   eos_id: int = kv['tokenizer.ggml.eos_token_id']
 
+  # start server
+  if args.serve: TCPServerWithReuse(('', 11434), Handler).serve_forever()
+
+  stop_ids = tok.stop_tokens()
   ids: list[int] = [bos_id]
   while 1:
     start_pos = len(ids) - 1
     try:
-      ids += tok.role("user") + tok.encode(input('>>> ')) + [eos_id] + tok.role("assistant")
+      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn() + tok.role("assistant")
     except EOFError:
       break
     for next_id in model.generate(ids, start_pos):
-      sys.stdout.write(tok.decode([next_id]) if next_id != eos_id else "\n\n")
+      sys.stdout.write(tok.decode([next_id]) if next_id not in stop_ids else "\n\n")
       sys.stdout.flush()
-      if next_id == eos_id: break
+      if next_id in stop_ids: break

tinygrad/codegen/opt/postrange.py

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

tinygrad/codegen/opt/search.py

@@ -92,9 +92,9 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
 # *** external API ***
 
 # get dictionary of all possible actions
-def get_kernel_actions(s:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
+def get_kernel_actions(s:Scheduler, include_0=True) -> dict[int, Scheduler]:
   acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
-  kernel_actions = (actions if candidates is None else candidates).copy()
+  kernel_actions = actions.copy()
 
   for i,a in enumerate(kernel_actions):
     if a.axis is not None and a.op is not OptOps.TC:

tinygrad/codegen/opt/tc.py

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

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
 
@@ -277,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):
     """
@@ -325,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,
@@ -344,7 +366,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
     if device in ["CUDA", "NV"]: return not CI
     if device == "CPU" and CPU_LLVM: return OSX
     if device == "PYTHON": return sys.version_info >= (3, 12)
-  if dtype == dtypes.float64: return device != "METAL" and not (OSX and device == "CL")
+  if dtype == dtypes.float64: return device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
   return True
 
 if PROFILE:
@@ -367,20 +389,23 @@ def enumerate_devices_str() -> Generator[str, None, None]:
     compilers_results, any_works = [], False
     try:
       d = Device[device]
-      default_renderer, default_compiler = d.renderer, d.compiler
+      default_comp_pairs, default_compiler, cc_ctrl_var = d.comp_sets, d.compiler, d.comps_ctrl_var
       try:
-        for r,c in d.compilers:
+        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()
+            # 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)'
+            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.renderer, d.compiler = default_renderer, default_compiler
+        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/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

@@ -14,7 +14,7 @@ from tinygrad.codegen.opt import Opt
 # **************** Program Creation ****************
 
 @track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret), replay=True)
-def get_program(ast:UOp, renderer:Renderer|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.
 
@@ -30,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)))

tinygrad/engine/schedule.py

@@ -2,10 +2,11 @@ import time
 from typing import cast
 from dataclasses import dataclass, field, replace
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass
+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
+from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
 
 # **** ScheduleItem return type
 
@@ -113,23 +114,77 @@ 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()
 
-  # verify Tensors match the spec
-  if SPEC: type_verify(big_sink, tensor_spec)
+  # 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
 
-  # tensor map is what we return
-  tensor_map: dict[UOp, UOp] = {}
+  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)
 
-  if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
-    tensor_map |= get_multi_map(big_sink)
-    big_sink = big_sink.substitute(tensor_map, name="Apply Multi Map")
-    big_sink = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink.src]))
+    # hack to preserve metadata
+    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata")
 
-  tensor_map |= get_rangeify_map(big_sink)
-  big_sink = big_sink.substitute(tensor_map, name="Apply Kernelize Map")
+    # 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)
@@ -139,5 +194,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
   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)} kernels in {(time.perf_counter()-st)*1000:.2f} ms ({len(UOpMetaClass.ucache)} uops in cache)")
+    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

@@ -1,8 +1,9 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
-import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
+import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, socketserver, json
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
+from http.server import BaseHTTPRequestHandler
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -114,10 +115,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 +178,19 @@ 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, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0)
+NULL_IR3, NULL_NAK = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0)
+AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
+QCOM_CC = ContextVar("QCOM_CC", "")
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", VIZ.value)
@@ -193,6 +203,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:
@@ -372,14 +384,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)
@@ -391,6 +405,33 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
       if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
   return fp
 
+# NOTE: using HTTPServer forces a potentially slow socket.getfqdn
+class TCPServerWithReuse(socketserver.TCPServer):
+  allow_reuse_address = True
+  def __init__(self, server_address, RequestHandlerClass):
+    print(f"*** started server on http://127.0.0.1:{server_address[1]}")
+    super().__init__(server_address, RequestHandlerClass)
+
+class HTTPRequestHandler(BaseHTTPRequestHandler):
+  def send_data(self, data:bytes, content_type:str="application/json", status_code:int=200):
+    self.send_response(status_code)
+    self.send_header("Content-Type", content_type)
+    self.send_header("Content-Length", str(len(data)))
+    self.end_headers()
+    return self.wfile.write(data)
+  def stream_json(self, source:Generator):
+    try:
+      self.send_response(200)
+      self.send_header("Content-Type", "text/event-stream")
+      self.send_header("Cache-Control", "no-cache")
+      self.end_headers()
+      for r in source:
+        self.wfile.write(f"data: {json.dumps(r)}\n\n".encode("utf-8"))
+        self.wfile.flush()
+      self.wfile.write("data: [DONE]\n\n".encode("utf-8"))
+    # pass if client closed connection
+    except (BrokenPipeError, ConnectionResetError): return
+
 # *** Exec helpers
 
 def system(cmd:str, **kwargs) -> str:

tinygrad/nn/onnx.py

@@ -626,6 +626,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def ThresholdedRelu(X:Tensor, alpha:float=1.0): return (X > alpha).where(X, 0)
   def LogSoftmax(x: Tensor, axis:int=-1): return x.log_softmax(axis)
   def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
+  def Swish(x:Tensor, alpha:float=1.0): return x * (x * alpha).sigmoid()
 
   # ***** Unary Ops (broadcasted) *****
   def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
@@ -710,7 +711,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 +1159,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 +1168,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

@@ -224,8 +224,12 @@ class ClangRenderer(CStyleLanguage):
                  Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
                  Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
                  Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
-  # LLVM legalizes double => half cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
+
+  # LLVM legalizes double => half/bf16 cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
+  # there is also no native bfl16 <-> fp16 conversion on those CPUs
   extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
+                                 (UPat.var("x", dtypes.float64).cast(dtypes.bfloat16), lambda x: x.cast(dtypes.float32).cast(dtypes.bfloat16)),
+                                 (UPat.var("x", dtypes.bfloat16).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
     (UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu)]) + CStyleLanguage.extra_matcher
 
   if sys.platform == 'win32':
@@ -373,6 +377,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/renderer/nir.py

@@ -1,21 +1,18 @@
 from typing import Callable, cast, Any
-from tinygrad.dtype import AddrSpace, DType, PtrDType, dtypes
+from tinygrad.dtype import AddrSpace, DType, PtrDType, ImageDType, dtypes
 from tinygrad.helpers import DEBUG, OSX, unwrap, charptr
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa
-import base64, contextlib, ctypes, ctypes.util, struct, functools, inspect
+import base64, ctypes, ctypes.util, struct, functools, inspect, contextlib, itertools
 
 def g(s:str): return getattr(mesa, s)
 def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))
 
-# this is a ridiculous hack, but I can't find a better way to grab the glsl_type objects
-glsl_base = {**{d:g(f"GLSL_TYPE_{'U' if d in dtypes.uints else ''}INT{d.itemsize*8 if d.itemsize != 4 else ''}") for d in dtypes.ints},
-             **{getattr(dtypes,d):g(f"GLSL_TYPE_{d.upper()}") for d in ['double', 'float', 'float16']}, dtypes.bool: mesa.GLSL_TYPE_UINT8}
-def glsl_type(t:DType) -> mesa.struct_glsl_type:
-  if isinstance(t, PtrDType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents
-  return mesa.glsl_get_base_glsl_type(mesa.glsl_type(base_type=glsl_base[t])).contents
+def glsl_type(t:DType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents if isinstance(t, PtrDType) else {
+  **{getattr(dtypes,k):g(f"glsl_type_builtin_{v}") for k,v in [('double','double'),('float','float'),('float16','float16_t'),('bool','uint8_t')]},
+  **{d:g(f"glsl_type_builtin_{'u' * (d in dtypes.uints)}int{str(d.itemsize*8)+'_t' if d.itemsize != 4 else ''}") for d in dtypes.ints}}[t]
 
 # alu ops, aop[<dtype>][<op>]
 u_aop = { Ops.ADD: "iadd", Ops.MUL: "imul", Ops.IDIV: "udiv", Ops.MOD: "umod", Ops.CMPLT: "ult", Ops.CMPNE: "ine", Ops.CMPEQ: "ieq", Ops.OR: "ior",
@@ -52,7 +49,7 @@ def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=N
       if has_def: mesa.nir_def_init(instr.contents.instr, getattr(instr.contents, "def"), go(nc), go(bs))
       for k, v in go(intrins or {}).items():
         idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic.value].index_map[g(f"NIR_INTRINSIC_{k}")]
-        assert idx > 0
+        assert idx > 0, "invalid intrinsic. mesa version mismatch?"
         instr.contents.const_index[idx - 1] = go(v)
       for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src)
       for k,v in {k:vcomp for k,v in contents.items() if (vcomp:=go(v)) is not None}.items(): setattr(instr.contents, k, go(v))
@@ -70,11 +67,16 @@ def nchannel(b:mesa.nir_builder, src:mesa.nir_def, c:int):
   ctypes.cast(mov.contents.src, ctypes.POINTER(mesa.nir_alu_src))[0] = alu_src
   return mov
 
+def nimm_set(imm:mesa.nir_def, x, dtype:DType):
+  instr = ctypes.cast(imm.parent_instr, ctypes.POINTER(mesa.nir_load_const_instr))
+  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, x)
+
 @nir_instr(nc=1, bs=lambda dtype: 1 if dtype == dtypes.bool else dtype.itemsize * 8)
 def nimm(b:mesa.nir_builder, x, dtype:DType) -> mesa.nir_def:
-  instr = mesa.nir_load_const_instr_create(b.shader, 1, 1 if dtype == dtypes.bool else dtype.itemsize * 8)
-  struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, x)
+  nimm_set(getattr((instr:=mesa.nir_load_const_instr_create(b.shader, 1, 1 if dtype==dtypes.bool else dtype.itemsize * 8)).contents, "def"), x, dtype)
   return instr
+@nir_instr(nc=1, bs=lambda dtype: 1 if dtype == dtypes.bool else dtype.itemsize * 8)
+def nundef(b, dtype): return mesa.nir_undef_instr_create(b.shader, 1, 1 if dtype == dtypes.bool else dtype.itemsize * 8)
 
 deref_var = nir_instr(nc=1, bs=32, modes=lambda var:var.data.mode, type=lambda var:var.type, var=lambda var:ctypes.pointer(var))( # pylint: disable=W0108
   lambda b, var: mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_var))
@@ -90,6 +92,8 @@ nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.itemsize*8/
 
 ngid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_id))
 nlid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_local_invocation_id))
+ngsz = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_size))
+def nid(b): return nalu(b, "iadd", nalu(b, "imul", ngid(b), ngsz(b)), nlid(b))
 
 nbarrier = nir_instr(has_def=False, intrins={"EXECUTION_SCOPE":mesa.SCOPE_WORKGROUP})(
   lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_barrier))
@@ -126,16 +130,16 @@ class NIRRenderer(Renderer):
     (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
      lambda x: x.replace(src=x.src[0:1] + (x.src[1].cast(dtypes.uint8),) + x.src[2:])),
     # load/store use pointer arithmetic, and the cast does nothing
-    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"),
-     lambda x,buf,off: x.replace(src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op != Ops.CAST else None),
+    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
+      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.VECTORIZE) else None),
     (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
   ])
 
   def_rewrite = PatternMatcher([
     (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
-    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 8)),
-    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x.dtype, 4)),
-    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, ngid(ctx.b) if x.arg[0] == 'g' else nlid(ctx.b), int(x.arg[-1]))),
+    (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
+    (UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
+    (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val")), allow_any_len=True, name="x"),
      lambda ctx,x,buf,off,val: nstore(ctx.b, buf.ptrdtype.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
     (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"), UPat.var("gate"))), UPat.var("alt")), allow_any_len=True, name="x"),
@@ -157,13 +161,15 @@ class NIRRenderer(Renderer):
   def __init__(self): mesa.glsl_type_singleton_init_or_ref()
 
   def __del__(self):
-    with contextlib.suppress(AttributeError):mesa.glsl_type_singleton_decref()
+    with contextlib.suppress(AttributeError): mesa.glsl_type_singleton_decref()
 
   @property
   def nir_options(self): raise NotImplementedError("needs nir_options")
-  def param(self, b:mesa.nir_builder, dtype:DType, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
+  def param(self, b:mesa.nir_builder, x, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
   def prerender(self, uops:list[UOp]):
     self.b = mesa.nir_builder_init_simple_shader(mesa.MESA_SHADER_COMPUTE, mesa.nir_shader_compiler_options.from_buffer_copy(self.nir_options), None)
+    self.b.shader.contents.info.workgroup_size_variable = any([u.op == Ops.SPECIAL and u.arg[0] == 'i' for u in uops])
+  def postrender(self, uops:list[UOp]): pass
 
   def render(self, uops:list[UOp]):
     self.prerender(uops)
@@ -196,6 +202,7 @@ class NIRRenderer(Renderer):
       else:
         if (d:=self.def_rewrite.rewrite(u, ctx=self)) is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
         self.r[u] = cast(mesa.nir_def, d)
+    self.postrender(uops)
 
     mesa.nir_validate_shader(self.b.shader, b"after render")
     if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(ctypes.CDLL(ctypes.util.find_library('c')),
@@ -209,22 +216,23 @@ class NIRRenderer(Renderer):
 
     return ret
 
-class NAKRenderer(NIRRenderer):
-  device = "NV"
+class NIRRendererWithOpts(NIRRenderer):
   def __init__(self, dev=None, nir_options=None):
     self.dev, self._nir_options = dev, nir_options
     super().__init__()
 
-  def __reduce__(self): return NAKRenderer, (None, self.nir_options,)
+  def __reduce__(self): return self.__class__, (None, self.nir_options)
 
   @property
   def nir_options(self):
     if self._nir_options is None: self._nir_options = self.dev.compiler.nir_options
     return self._nir_options
 
+class NAKRenderer(NIRRendererWithOpts):
+  device = "NV"
   param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
     intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
-       lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
+       lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
 
 class LVPRenderer(NIRRenderer):
   device = "CPU"
@@ -235,9 +243,55 @@ class LVPRenderer(NIRRenderer):
 
   param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
     srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
-    setattr(self, "param_idx", self.param_idx+sz))(lambda self, b, dtype, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
+    setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
 
   def prerender(self, uops:list[UOp]):
     super().prerender(uops)
     self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
 
+# FIXME: this should be a rewrite rule
+def tovec(b, coord): return nalu(b, "vec4", nchannel(b, coord, 0), nchannel(b, coord, 1), nundef(b, dtypes.int), nundef(b, dtypes.int))
+def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
+nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=lambda val:val.num_components,
+  intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'SRC_TYPE':nfloat(dtype)},
+  srcs=lambda b,img,coord,val:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
+    lambda b,img,coord,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))
+
+_nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'DEST_TYPE':nfloat(dtype)},
+  nc=4, bs=32, num_components=4, srcs=lambda b,img,coord:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
+    lambda b,img,coord,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
+
+class IR3Renderer(NIRRendererWithOpts):
+  device = "QCOM"
+
+  def nload_img(ctx,img,coord):
+    ctx.texs.add(img)
+    return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype)
+
+  def_rewrite = PatternMatcher([
+    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")),
+          allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))),
+     lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img),
+  ]) + NIRRenderer.def_rewrite
+
+  _param = LVPRenderer.param
+  def _param_img(self, x):
+    self.img_idx += 1
+    return nimm(self.b, self.img_idx - 1, dtypes.int)
+
+  def param(self, b, x, sz): return self._param_img(x) if isinstance(x.dtype, ImageDType) else self._param(b, x, sz)
+
+  def prerender(self, uops:list[UOp]):
+    super().prerender(uops)
+    self.texs:set[UOp] = set()
+    self.uops, self.ibo_idx, self.img_idx = uops, 0, 0
+    self.param_sz = sum([8 if u.op == Ops.DEFINE_GLOBAL else u.dtype.itemsize for u in uops if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR)])
+
+  def postrender(self, uops:list[UOp]):
+    bufs, texs, imgs = [u for u in uops if u.op == Ops.DEFINE_GLOBAL], itertools.count().__next__, itertools.count().__next__
+    for b in filter(lambda b: isinstance(b.dtype, ImageDType), bufs): nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int)
+
+    self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)])
+    self.b.shader.contents.info.num_images = texs() + imgs()

tinygrad/runtime/autogen/__init__.py

@@ -5,6 +5,7 @@ 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):
@@ -39,7 +40,7 @@ def __getattr__(nm):
         *[root/"extra/nv_gpu_driver"/s for s in ["clc9b0.h", "clc6c0qmd.h","clcec0qmd.h", "nvdec_drv.h"]], "{}/kernel-open/common/inc/nvmisc.h",
         *[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0070", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
                                                                     "83de", "c6c0", "cdc0"]],
-        *[f"{{}}/kernel-open/nvidia-uvm/{s}.h" for s in ["clc6b5", "clc9b5", "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",
@@ -89,13 +90,15 @@ 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++"])
     case "kgsl": return load("kgsl", [], [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
-    case "adreno": return load("adreno", [], [root/"extra/qcom_gpu_driver/a6xx.xml.h"])
     case "qcom_dsp":
       return load("qcom_dsp", [], [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
     case "sqtt": return load("sqtt", [], [root/"extra/sqtt/sqtt.h"])
@@ -110,18 +113,20 @@ def __getattr__(nm):
         *[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
         *[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
                                                                      "_flow", "_const"]],
-        "{}/src/compiler/glsl_types.h", "{}/src/util/blob.h", "{}/src/util/ralloc.h"], args=lambda:[
+        *[f"{{}}/src/freedreno/{s}.h" for s in ["common/freedreno_dev_info", "ir3/ir3_compiler", "ir3/ir3_shader", "ir3/ir3_nir"]],
+        "{}/src/compiler/glsl_types.h", "{}/src/util/blob.h", "{}/src/util/ralloc.h", "{}/gen/ir3-isa.h", "{}/gen/builtin_types.h",
+        "{}/gen/a6xx.xml.h", "{}/gen/adreno_pm4.xml.h", "{}/gen/a6xx_enums.xml.h", "{}/gen/a6xx_descriptors.xml.h"], args=lambda:[
           "-DHAVE_ENDIAN_H", "-DHAVE_STRUCT_TIMESPEC", "-DHAVE_PTHREAD", "-DHAVE_FUNC_ATTRIBUTE_PACKED", "-I{}/src", "-I{}/include", "-I{}/gen",
-          "-I{}/src/compiler/nir", "-I{}/src/gallium/auxiliary", "-I{}/src/gallium/include", f"-I{system('llvm-config-20 --includedir')}"],
-        preprocess=lambda path: subprocess.run("""mkdir -p gen/util/format
-python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
-python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
-python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
-python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
-python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-python3 src/compiler/builtin_types_h.py gen/builtin_types.h""", cwd=path, shell=True, check=True),
-  tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.4/mesa-25.2.4.tar.gz",
+          "-I{}/src/compiler/nir", "-I{}/src/gallium/auxiliary", "-I{}/src/gallium/include", "-I{}/src/freedreno/common",
+          f"-I{system('llvm-config-20 --includedir')}"],
+        preprocess=lambda path: subprocess.run("\n".join(["mkdir -p gen/util/format", "python3 src/compiler/builtin_types_h.py gen/builtin_types.h",
+          "python3 src/compiler/isaspec/decode.py --xml src/freedreno/isa/ir3.xml --out-c /dev/null --out-h gen/ir3-isa.h",
+          "python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h",
+          *["python3 src/freedreno/registers/gen_header.py --rnn src/freedreno/registers/ --xml " +
+            f"src/freedreno/registers/adreno/{s}.xml c-defines > gen/{s}.xml.h" for s in ["a6xx", "adreno_pm4", "a6xx_enums", "a6xx_descriptors"]],
+          *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
+          *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
+  tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
   prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
       return load("libclang", ["os.getenv('LIBCLANG_PATH', find_library('clang-20'))"],
@@ -131,4 +136,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/am/__init__.py

@@ -11,13 +11,13 @@ def __getattr__(nm):
     case "pm4_soc15": return load("am/pm4_soc15", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], tarball=am_src)
     case "pm4_nv": return load("am/pm4_nv", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], tarball=am_src)
     case "sdma_4_0_0": return load("am/sdma_4_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
     case "sdma_5_0_0": return load("am/sdma_5_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
-    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}//amdgpu/sdma_v6_0_0_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
+    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
     case "smu_v13_0_0": return load("am/smu_v13_0_0",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
-                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], tarball=am_src),
+                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
     case "smu_v14_0_2": return load("am/smu_v14_0_2", [], [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
                                     "smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/autogen/am/pm4_nv.py

@@ -3,34 +3,15 @@ import ctypes
 from tinygrad.helpers import unwrap
 from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class union_PM4_MES_TYPE_3_HEADER(ctypes.Union): pass
-class _anonstruct0(Struct): pass
 enum_mes_set_resources_queue_type_enum = CEnum(ctypes.c_uint32)
 queue_type__mes_set_resources__kernel_interface_queue_kiq = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__kernel_interface_queue_kiq', 0)
 queue_type__mes_set_resources__hsa_interface_queue_hiq = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__hsa_interface_queue_hiq', 1)
 queue_type__mes_set_resources__hsa_debug_interface_queue = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__hsa_debug_interface_queue', 4)
 
 class struct_pm4_mes_set_resources(Struct): pass
-class _anonunion1(ctypes.Union): pass
-class _anonunion2(ctypes.Union): pass
-class _anonstruct3(Struct): pass
-class _anonunion4(ctypes.Union): pass
-class _anonstruct5(Struct): pass
-class _anonunion6(ctypes.Union): pass
-class _anonstruct7(Struct): pass
 class struct_pm4_mes_runlist(Struct): pass
-class _anonunion8(ctypes.Union): pass
-class _anonunion9(ctypes.Union): pass
-class _anonstruct10(Struct): pass
-class _anonunion11(ctypes.Union): pass
-class _anonstruct12(Struct): pass
 class struct_pm4_mes_map_process(Struct): pass
-class _anonunion13(ctypes.Union): pass
-class _anonunion14(ctypes.Union): pass
-class _anonstruct15(Struct): pass
-class _anonunion16(ctypes.Union): pass
-class _anonstruct17(Struct): pass
 class struct_PM4_MES_MAP_PROCESS_VM(Struct): pass
-class _anonunion18(ctypes.Union): pass
 enum_mes_map_queues_queue_sel_enum = CEnum(ctypes.c_uint32)
 queue_sel__mes_map_queues__map_to_specified_queue_slots_vi = enum_mes_map_queues_queue_sel_enum.define('queue_sel__mes_map_queues__map_to_specified_queue_slots_vi', 0)
 queue_sel__mes_map_queues__map_to_hws_determined_queue_slots_vi = enum_mes_map_queues_queue_sel_enum.define('queue_sel__mes_map_queues__map_to_hws_determined_queue_slots_vi', 1)
@@ -52,11 +33,6 @@ extended_engine_sel__mes_map_queues__sdma0_to_7_sel = enum_mes_map_queues_extend
 extended_engine_sel__mes_map_queues__sdma8_to_15_sel = enum_mes_map_queues_extended_engine_sel_enum.define('extended_engine_sel__mes_map_queues__sdma8_to_15_sel', 2)
 
 class struct_pm4_mes_map_queues(Struct): pass
-class _anonunion19(ctypes.Union): pass
-class _anonunion20(ctypes.Union): pass
-class _anonstruct21(Struct): pass
-class _anonunion22(ctypes.Union): pass
-class _anonstruct23(Struct): pass
 enum_mes_query_status_interrupt_sel_enum = CEnum(ctypes.c_uint32)
 interrupt_sel__mes_query_status__completion_status = enum_mes_query_status_interrupt_sel_enum.define('interrupt_sel__mes_query_status__completion_status', 0)
 interrupt_sel__mes_query_status__process_status = enum_mes_query_status_interrupt_sel_enum.define('interrupt_sel__mes_query_status__process_status', 1)
@@ -74,12 +50,6 @@ engine_sel__mes_query_status__sdma0_queue = enum_mes_query_status_engine_sel_enu
 engine_sel__mes_query_status__sdma1_queue = enum_mes_query_status_engine_sel_enum.define('engine_sel__mes_query_status__sdma1_queue', 3)
 
 class struct_pm4_mes_query_status(Struct): pass
-class _anonunion24(ctypes.Union): pass
-class _anonunion25(ctypes.Union): pass
-class _anonstruct26(Struct): pass
-class _anonunion27(ctypes.Union): pass
-class _anonstruct28(Struct): pass
-class _anonstruct29(Struct): pass
 enum_mes_unmap_queues_action_enum = CEnum(ctypes.c_uint32)
 action__mes_unmap_queues__preempt_queues = enum_mes_unmap_queues_action_enum.define('action__mes_unmap_queues__preempt_queues', 0)
 action__mes_unmap_queues__reset_queues = enum_mes_unmap_queues_action_enum.define('action__mes_unmap_queues__reset_queues', 1)
@@ -102,18 +72,6 @@ extended_engine_sel__mes_unmap_queues__legacy_engine_sel = enum_mes_unmap_queues
 extended_engine_sel__mes_unmap_queues__sdma0_to_7_sel = enum_mes_unmap_queues_extended_engine_sel_enum.define('extended_engine_sel__mes_unmap_queues__sdma0_to_7_sel', 1)
 
 class struct_pm4_mes_unmap_queues(Struct): pass
-class _anonunion30(ctypes.Union): pass
-class _anonunion31(ctypes.Union): pass
-class _anonstruct32(Struct): pass
-class _anonunion33(ctypes.Union): pass
-class _anonstruct34(Struct): pass
-class _anonstruct35(Struct): pass
-class _anonunion36(ctypes.Union): pass
-class _anonstruct37(Struct): pass
-class _anonunion38(ctypes.Union): pass
-class _anonstruct39(Struct): pass
-class _anonunion40(ctypes.Union): pass
-class _anonstruct41(Struct): pass
 enum_mec_release_mem_event_index_enum = CEnum(ctypes.c_uint32)
 event_index__mec_release_mem__end_of_pipe = enum_mec_release_mem_event_index_enum.define('event_index__mec_release_mem__end_of_pipe', 5)
 event_index__mec_release_mem__shader_done = enum_mec_release_mem_event_index_enum.define('event_index__mec_release_mem__shader_done', 6)
@@ -150,18 +108,6 @@ data_sel__mec_release_mem__send_cp_perfcounter_hi_lo = enum_mec_release_mem_data
 data_sel__mec_release_mem__store_gds_data_to_memory = enum_mec_release_mem_data_sel_enum.define('data_sel__mec_release_mem__store_gds_data_to_memory', 5)
 
 class struct_pm4_mec_release_mem(Struct): pass
-class _anonunion42(ctypes.Union): pass
-class _anonunion43(ctypes.Union): pass
-class _anonstruct44(Struct): pass
-class _anonunion45(ctypes.Union): pass
-class _anonstruct46(Struct): pass
-class _anonunion47(ctypes.Union): pass
-class _anonstruct48(Struct): pass
-class _anonstruct49(Struct): pass
-class _anonunion50(ctypes.Union): pass
-class _anonunion51(ctypes.Union): pass
-class _anonstruct52(Struct): pass
-class _anonunion53(ctypes.Union): pass
 enum_WRITE_DATA_dst_sel_enum = CEnum(ctypes.c_uint32)
 dst_sel___write_data__mem_mapped_register = enum_WRITE_DATA_dst_sel_enum.define('dst_sel___write_data__mem_mapped_register', 0)
 dst_sel___write_data__tc_l2 = enum_WRITE_DATA_dst_sel_enum.define('dst_sel___write_data__tc_l2', 2)
@@ -182,37 +128,8 @@ cache_policy___write_data__lru = enum_WRITE_DATA_cache_policy_enum.define('cache
 cache_policy___write_data__stream = enum_WRITE_DATA_cache_policy_enum.define('cache_policy___write_data__stream', 1)
 
 class struct_pm4_mec_write_data_mmio(Struct): pass
-class _anonunion54(ctypes.Union): pass
-class _anonunion55(ctypes.Union): pass
-class _anonunion55_bitfields2(Struct): pass
-_anonunion55_bitfields2._fields_ = [
-  ('reserved1', ctypes.c_uint32,8),
-  ('dst_sel', ctypes.c_uint32,4),
-  ('reserved2', ctypes.c_uint32,4),
-  ('addr_incr', ctypes.c_uint32,1),
-  ('reserved3', ctypes.c_uint32,2),
-  ('resume_vf', ctypes.c_uint32,1),
-  ('wr_confirm', ctypes.c_uint32,1),
-  ('reserved4', ctypes.c_uint32,4),
-  ('cache_policy', ctypes.c_uint32,2),
-  ('reserved5', ctypes.c_uint32,5),
-]
-_anonunion55._fields_ = [
-  ('bitfields2', _anonunion55_bitfields2),
-  ('ordinal2', ctypes.c_uint32),
-]
-class _anonunion56(ctypes.Union): pass
-class _anonunion56_bitfields3(Struct): pass
-_anonunion56_bitfields3._fields_ = [
-  ('dst_mmreg_addr', ctypes.c_uint32,18),
-  ('reserved6', ctypes.c_uint32,14),
-]
-_anonunion56._fields_ = [
-  ('bitfields3', _anonunion56_bitfields3),
-  ('ordinal3', ctypes.c_uint32),
-]
-_anonenum57 = CEnum(ctypes.c_uint32)
-CACHE_FLUSH_AND_INV_TS_EVENT = _anonenum57.define('CACHE_FLUSH_AND_INV_TS_EVENT', 20)
+_anonenum0 = CEnum(ctypes.c_uint32)
+CACHE_FLUSH_AND_INV_TS_EVENT = _anonenum0.define('CACHE_FLUSH_AND_INV_TS_EVENT', 20)
 
 PACKET_TYPE0 = 0
 PACKET_TYPE1 = 1

tinygrad/runtime/autogen/am/pm4_soc15.py

@@ -3,34 +3,15 @@ import ctypes
 from tinygrad.helpers import unwrap
 from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class union_PM4_MES_TYPE_3_HEADER(ctypes.Union): pass
-class _anonstruct0(Struct): pass
 enum_mes_set_resources_queue_type_enum = CEnum(ctypes.c_uint32)
 queue_type__mes_set_resources__kernel_interface_queue_kiq = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__kernel_interface_queue_kiq', 0)
 queue_type__mes_set_resources__hsa_interface_queue_hiq = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__hsa_interface_queue_hiq', 1)
 queue_type__mes_set_resources__hsa_debug_interface_queue = enum_mes_set_resources_queue_type_enum.define('queue_type__mes_set_resources__hsa_debug_interface_queue', 4)
 
 class struct_pm4_mes_set_resources(Struct): pass
-class _anonunion1(ctypes.Union): pass
-class _anonunion2(ctypes.Union): pass
-class _anonstruct3(Struct): pass
-class _anonunion4(ctypes.Union): pass
-class _anonstruct5(Struct): pass
-class _anonunion6(ctypes.Union): pass
-class _anonstruct7(Struct): pass
 class struct_pm4_mes_runlist(Struct): pass
-class _anonunion8(ctypes.Union): pass
-class _anonunion9(ctypes.Union): pass
-class _anonstruct10(Struct): pass
-class _anonunion11(ctypes.Union): pass
-class _anonstruct12(Struct): pass
 class struct_pm4_mes_map_process(Struct): pass
-class _anonunion13(ctypes.Union): pass
-class _anonunion14(ctypes.Union): pass
-class _anonstruct15(Struct): pass
-class _anonunion16(ctypes.Union): pass
-class _anonstruct17(Struct): pass
 class struct_PM4_MES_MAP_PROCESS_VM(Struct): pass
-class _anonunion18(ctypes.Union): pass
 enum_mes_map_queues_queue_sel_enum = CEnum(ctypes.c_uint32)
 queue_sel__mes_map_queues__map_to_specified_queue_slots_vi = enum_mes_map_queues_queue_sel_enum.define('queue_sel__mes_map_queues__map_to_specified_queue_slots_vi', 0)
 queue_sel__mes_map_queues__map_to_hws_determined_queue_slots_vi = enum_mes_map_queues_queue_sel_enum.define('queue_sel__mes_map_queues__map_to_hws_determined_queue_slots_vi', 1)
@@ -52,11 +33,6 @@ extended_engine_sel__mes_map_queues__sdma0_to_7_sel = enum_mes_map_queues_extend
 extended_engine_sel__mes_map_queues__sdma8_to_15_sel = enum_mes_map_queues_extended_engine_sel_enum.define('extended_engine_sel__mes_map_queues__sdma8_to_15_sel', 2)
 
 class struct_pm4_mes_map_queues(Struct): pass
-class _anonunion19(ctypes.Union): pass
-class _anonunion20(ctypes.Union): pass
-class _anonstruct21(Struct): pass
-class _anonunion22(ctypes.Union): pass
-class _anonstruct23(Struct): pass
 enum_mes_query_status_interrupt_sel_enum = CEnum(ctypes.c_uint32)
 interrupt_sel__mes_query_status__completion_status = enum_mes_query_status_interrupt_sel_enum.define('interrupt_sel__mes_query_status__completion_status', 0)
 interrupt_sel__mes_query_status__process_status = enum_mes_query_status_interrupt_sel_enum.define('interrupt_sel__mes_query_status__process_status', 1)
@@ -74,12 +50,6 @@ engine_sel__mes_query_status__sdma0_queue = enum_mes_query_status_engine_sel_enu
 engine_sel__mes_query_status__sdma1_queue = enum_mes_query_status_engine_sel_enum.define('engine_sel__mes_query_status__sdma1_queue', 3)
 
 class struct_pm4_mes_query_status(Struct): pass
-class _anonunion24(ctypes.Union): pass
-class _anonunion25(ctypes.Union): pass
-class _anonstruct26(Struct): pass
-class _anonunion27(ctypes.Union): pass
-class _anonstruct28(Struct): pass
-class _anonstruct29(Struct): pass
 enum_mes_unmap_queues_action_enum = CEnum(ctypes.c_uint32)
 action__mes_unmap_queues__preempt_queues = enum_mes_unmap_queues_action_enum.define('action__mes_unmap_queues__preempt_queues', 0)
 action__mes_unmap_queues__reset_queues = enum_mes_unmap_queues_action_enum.define('action__mes_unmap_queues__reset_queues', 1)
@@ -102,18 +72,6 @@ extended_engine_sel__mes_unmap_queues__legacy_engine_sel = enum_mes_unmap_queues
 extended_engine_sel__mes_unmap_queues__sdma0_to_7_sel = enum_mes_unmap_queues_extended_engine_sel_enum.define('extended_engine_sel__mes_unmap_queues__sdma0_to_7_sel', 1)
 
 class struct_pm4_mes_unmap_queues(Struct): pass
-class _anonunion30(ctypes.Union): pass
-class _anonunion31(ctypes.Union): pass
-class _anonstruct32(Struct): pass
-class _anonunion33(ctypes.Union): pass
-class _anonstruct34(Struct): pass
-class _anonstruct35(Struct): pass
-class _anonunion36(ctypes.Union): pass
-class _anonstruct37(Struct): pass
-class _anonunion38(ctypes.Union): pass
-class _anonstruct39(Struct): pass
-class _anonunion40(ctypes.Union): pass
-class _anonstruct41(Struct): pass
 enum_mec_release_mem_event_index_enum = CEnum(ctypes.c_uint32)
 event_index__mec_release_mem__end_of_pipe = enum_mec_release_mem_event_index_enum.define('event_index__mec_release_mem__end_of_pipe', 5)
 event_index__mec_release_mem__shader_done = enum_mec_release_mem_event_index_enum.define('event_index__mec_release_mem__shader_done', 6)
@@ -150,18 +108,6 @@ data_sel__mec_release_mem__send_cp_perfcounter_hi_lo = enum_mec_release_mem_data
 data_sel__mec_release_mem__store_gds_data_to_memory = enum_mec_release_mem_data_sel_enum.define('data_sel__mec_release_mem__store_gds_data_to_memory', 5)
 
 class struct_pm4_mec_release_mem(Struct): pass
-class _anonunion42(ctypes.Union): pass
-class _anonunion43(ctypes.Union): pass
-class _anonstruct44(Struct): pass
-class _anonunion45(ctypes.Union): pass
-class _anonstruct46(Struct): pass
-class _anonunion47(ctypes.Union): pass
-class _anonstruct48(Struct): pass
-class _anonstruct49(Struct): pass
-class _anonunion50(ctypes.Union): pass
-class _anonunion51(ctypes.Union): pass
-class _anonstruct52(Struct): pass
-class _anonunion53(ctypes.Union): pass
 enum_WRITE_DATA_dst_sel_enum = CEnum(ctypes.c_uint32)
 dst_sel___write_data__mem_mapped_register = enum_WRITE_DATA_dst_sel_enum.define('dst_sel___write_data__mem_mapped_register', 0)
 dst_sel___write_data__tc_l2 = enum_WRITE_DATA_dst_sel_enum.define('dst_sel___write_data__tc_l2', 2)
@@ -182,37 +128,8 @@ cache_policy___write_data__lru = enum_WRITE_DATA_cache_policy_enum.define('cache
 cache_policy___write_data__stream = enum_WRITE_DATA_cache_policy_enum.define('cache_policy___write_data__stream', 1)
 
 class struct_pm4_mec_write_data_mmio(Struct): pass
-class _anonunion54(ctypes.Union): pass
-class _anonunion55(ctypes.Union): pass
-class _anonunion55_bitfields2(Struct): pass
-_anonunion55_bitfields2._fields_ = [
-  ('reserved1', ctypes.c_uint32,8),
-  ('dst_sel', ctypes.c_uint32,4),
-  ('reserved2', ctypes.c_uint32,4),
-  ('addr_incr', ctypes.c_uint32,1),
-  ('reserved3', ctypes.c_uint32,2),
-  ('resume_vf', ctypes.c_uint32,1),
-  ('wr_confirm', ctypes.c_uint32,1),
-  ('reserved4', ctypes.c_uint32,4),
-  ('cache_policy', ctypes.c_uint32,2),
-  ('reserved5', ctypes.c_uint32,5),
-]
-_anonunion55._fields_ = [
-  ('bitfields2', _anonunion55_bitfields2),
-  ('ordinal2', ctypes.c_uint32),
-]
-class _anonunion56(ctypes.Union): pass
-class _anonunion56_bitfields3(Struct): pass
-_anonunion56_bitfields3._fields_ = [
-  ('dst_mmreg_addr', ctypes.c_uint32,18),
-  ('reserved6', ctypes.c_uint32,14),
-]
-_anonunion56._fields_ = [
-  ('bitfields3', _anonunion56_bitfields3),
-  ('ordinal3', ctypes.c_uint32),
-]
-_anonenum57 = CEnum(ctypes.c_uint32)
-CACHE_FLUSH_AND_INV_TS_EVENT = _anonenum57.define('CACHE_FLUSH_AND_INV_TS_EVENT', 20)
+_anonenum0 = CEnum(ctypes.c_uint32)
+CACHE_FLUSH_AND_INV_TS_EVENT = _anonenum0.define('CACHE_FLUSH_AND_INV_TS_EVENT', 20)
 
 GFX9_NUM_GFX_RINGS = 1
 GFX9_NUM_COMPUTE_RINGS = 8

tinygrad/runtime/autogen/am/smu_v13_0_0.py

@@ -77,8 +77,18 @@ I2C_CONTROLLER_PROTOCOL_INA3221 = I2cControllerProtocol_e.define('I2C_CONTROLLER
 I2C_CONTROLLER_PROTOCOL_TMP_MAX6604 = I2cControllerProtocol_e.define('I2C_CONTROLLER_PROTOCOL_TMP_MAX6604', 4)
 I2C_CONTROLLER_PROTOCOL_COUNT = I2cControllerProtocol_e.define('I2C_CONTROLLER_PROTOCOL_COUNT', 5)
 
-class _anonstruct0(Struct): pass
-I2cControllerConfig_t = _anonstruct0
+class I2cControllerConfig_t(Struct): pass
+uint8_t = ctypes.c_ubyte
+I2cControllerConfig_t._fields_ = [
+  ('Enabled', uint8_t),
+  ('Speed', uint8_t),
+  ('SlaveAddress', uint8_t),
+  ('ControllerPort', uint8_t),
+  ('ControllerName', uint8_t),
+  ('ThermalThrotter', uint8_t),
+  ('I2cProtocol', uint8_t),
+  ('PaddingConfig', uint8_t),
+]
 I2cPort_e = CEnum(ctypes.c_uint32)
 I2C_PORT_SVD_SCL = I2cPort_e.define('I2C_PORT_SVD_SCL', 0)
 I2C_PORT_GPIO = I2cPort_e.define('I2C_PORT_GPIO', 1)
@@ -97,16 +107,40 @@ I2C_CMD_READ = I2cCmdType_e.define('I2C_CMD_READ', 0)
 I2C_CMD_WRITE = I2cCmdType_e.define('I2C_CMD_WRITE', 1)
 I2C_CMD_COUNT = I2cCmdType_e.define('I2C_CMD_COUNT', 2)
 
-class _anonstruct1(Struct): pass
-SwI2cCmd_t = _anonstruct1
-class _anonstruct2(Struct): pass
-SwI2cRequest_t = _anonstruct2
-class _anonstruct3(Struct): pass
-SwI2cRequestExternal_t = _anonstruct3
-class _anonstruct4(Struct): pass
-EccInfo_t = _anonstruct4
-class _anonstruct5(Struct): pass
-EccInfoTable_t = _anonstruct5
+class SwI2cCmd_t(Struct): pass
+SwI2cCmd_t._fields_ = [
+  ('ReadWriteData', uint8_t),
+  ('CmdConfig', uint8_t),
+]
+class SwI2cRequest_t(Struct): pass
+SwI2cRequest_t._fields_ = [
+  ('I2CcontrollerPort', uint8_t),
+  ('I2CSpeed', uint8_t),
+  ('SlaveAddress', uint8_t),
+  ('NumCmds', uint8_t),
+  ('SwI2cCmds', (SwI2cCmd_t * 24)),
+]
+class SwI2cRequestExternal_t(Struct): pass
+uint32_t = ctypes.c_uint32
+SwI2cRequestExternal_t._fields_ = [
+  ('SwI2cRequest', SwI2cRequest_t),
+  ('Spare', (uint32_t * 8)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class EccInfo_t(Struct): pass
+uint64_t = ctypes.c_uint64
+uint16_t = ctypes.c_uint16
+EccInfo_t._fields_ = [
+  ('mca_umc_status', uint64_t),
+  ('mca_umc_addr', uint64_t),
+  ('ce_count_lo_chip', uint16_t),
+  ('ce_count_hi_chip', uint16_t),
+  ('eccPadding', uint32_t),
+]
+class EccInfoTable_t(Struct): pass
+EccInfoTable_t._fields_ = [
+  ('EccInfo', (EccInfo_t * 24)),
+]
 D3HOTSequence_e = CEnum(ctypes.c_uint32)
 BACO_SEQUENCE = D3HOTSequence_e.define('BACO_SEQUENCE', 0)
 MSR_SEQUENCE = D3HOTSequence_e.define('MSR_SEQUENCE', 1)
@@ -122,12 +156,23 @@ PowerGatingSettings_e = CEnum(ctypes.c_uint32)
 PG_POWER_DOWN = PowerGatingSettings_e.define('PG_POWER_DOWN', 0)
 PG_POWER_UP = PowerGatingSettings_e.define('PG_POWER_UP', 1)
 
-class _anonstruct6(Struct): pass
-QuadraticInt_t = _anonstruct6
-class _anonstruct7(Struct): pass
-LinearInt_t = _anonstruct7
-class _anonstruct8(Struct): pass
-DroopInt_t = _anonstruct8
+class QuadraticInt_t(Struct): pass
+QuadraticInt_t._fields_ = [
+  ('a', uint32_t),
+  ('b', uint32_t),
+  ('c', uint32_t),
+]
+class LinearInt_t(Struct): pass
+LinearInt_t._fields_ = [
+  ('m', uint32_t),
+  ('b', uint32_t),
+]
+class DroopInt_t(Struct): pass
+DroopInt_t._fields_ = [
+  ('a', uint32_t),
+  ('b', uint32_t),
+  ('c', uint32_t),
+]
 DCS_ARCH_e = CEnum(ctypes.c_uint32)
 DCS_ARCH_DISABLED = DCS_ARCH_e.define('DCS_ARCH_DISABLED', 0)
 DCS_ARCH_FADCS = DCS_ARCH_e.define('DCS_ARCH_FADCS', 1)
@@ -186,8 +231,19 @@ PWR_CONFIG_TGP = PwrConfig_e.define('PWR_CONFIG_TGP', 1)
 PWR_CONFIG_TCP_ESTIMATED = PwrConfig_e.define('PWR_CONFIG_TCP_ESTIMATED', 2)
 PWR_CONFIG_TCP_MEASURED = PwrConfig_e.define('PWR_CONFIG_TCP_MEASURED', 3)
 
-class _anonstruct9(Struct): pass
-DpmDescriptor_t = _anonstruct9
+class DpmDescriptor_t(Struct): pass
+DpmDescriptor_t._fields_ = [
+  ('Padding', uint8_t),
+  ('SnapToDiscrete', uint8_t),
+  ('NumDiscreteLevels', uint8_t),
+  ('CalculateFopt', uint8_t),
+  ('ConversionToAvfsClk', LinearInt_t),
+  ('Padding3', (uint32_t * 3)),
+  ('Padding4', uint16_t),
+  ('FoptimalDc', uint16_t),
+  ('FoptimalAc', uint16_t),
+  ('Padding2', uint16_t),
+]
 PPT_THROTTLER_e = CEnum(ctypes.c_uint32)
 PPT_THROTTLER_PPT0 = PPT_THROTTLER_e.define('PPT_THROTTLER_PPT0', 0)
 PPT_THROTTLER_PPT1 = PPT_THROTTLER_e.define('PPT_THROTTLER_PPT1', 1)
@@ -332,18 +388,70 @@ PP_GRTAVFS_FW_SEP_FUSE_FREQUENCY_TO_COUNT_SCALER_3 = PP_GRTAVFS_FW_SEP_FUSE_e.de
 PP_GRTAVFS_FW_SEP_FUSE_FREQUENCY_TO_COUNT_SCALER_4 = PP_GRTAVFS_FW_SEP_FUSE_e.define('PP_GRTAVFS_FW_SEP_FUSE_FREQUENCY_TO_COUNT_SCALER_4', 18)
 PP_GRTAVFS_FW_SEP_FUSE_COUNT = PP_GRTAVFS_FW_SEP_FUSE_e.define('PP_GRTAVFS_FW_SEP_FUSE_COUNT', 19)
 
-class _anonstruct10(Struct): pass
-SviTelemetryScale_t = _anonstruct10
+class SviTelemetryScale_t(Struct): pass
+int8_t = ctypes.c_byte
+SviTelemetryScale_t._fields_ = [
+  ('Offset', int8_t),
+  ('Padding', uint8_t),
+  ('MaxCurrent', uint16_t),
+]
 FanMode_e = CEnum(ctypes.c_uint32)
 FAN_MODE_AUTO = FanMode_e.define('FAN_MODE_AUTO', 0)
 FAN_MODE_MANUAL_LINEAR = FanMode_e.define('FAN_MODE_MANUAL_LINEAR', 1)
 
-class _anonstruct11(Struct): pass
-OverDriveTable_t = _anonstruct11
-class _anonstruct12(Struct): pass
-OverDriveTableExternal_t = _anonstruct12
-class _anonstruct13(Struct): pass
-OverDriveLimits_t = _anonstruct13
+class OverDriveTable_t(Struct): pass
+int16_t = ctypes.c_int16
+OverDriveTable_t._fields_ = [
+  ('FeatureCtrlMask', uint32_t),
+  ('VoltageOffsetPerZoneBoundary', (int16_t * 6)),
+  ('Reserved', uint32_t),
+  ('GfxclkFmin', int16_t),
+  ('GfxclkFmax', int16_t),
+  ('UclkFmin', uint16_t),
+  ('UclkFmax', uint16_t),
+  ('Ppt', int16_t),
+  ('Tdc', int16_t),
+  ('FanLinearPwmPoints', (uint8_t * 6)),
+  ('FanLinearTempPoints', (uint8_t * 6)),
+  ('FanMinimumPwm', uint16_t),
+  ('AcousticTargetRpmThreshold', uint16_t),
+  ('AcousticLimitRpmThreshold', uint16_t),
+  ('FanTargetTemperature', uint16_t),
+  ('FanZeroRpmEnable', uint8_t),
+  ('FanZeroRpmStopTemp', uint8_t),
+  ('FanMode', uint8_t),
+  ('MaxOpTemp', uint8_t),
+  ('Spare', (uint32_t * 13)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class OverDriveTableExternal_t(Struct): pass
+OverDriveTableExternal_t._fields_ = [
+  ('OverDriveTable', OverDriveTable_t),
+]
+class OverDriveLimits_t(Struct): pass
+OverDriveLimits_t._fields_ = [
+  ('FeatureCtrlMask', uint32_t),
+  ('VoltageOffsetPerZoneBoundary', int16_t),
+  ('Reserved1', uint16_t),
+  ('Reserved2', uint16_t),
+  ('GfxclkFmin', int16_t),
+  ('GfxclkFmax', int16_t),
+  ('UclkFmin', uint16_t),
+  ('UclkFmax', uint16_t),
+  ('Ppt', int16_t),
+  ('Tdc', int16_t),
+  ('FanLinearPwmPoints', uint8_t),
+  ('FanLinearTempPoints', uint8_t),
+  ('FanMinimumPwm', uint16_t),
+  ('AcousticTargetRpmThreshold', uint16_t),
+  ('AcousticLimitRpmThreshold', uint16_t),
+  ('FanTargetTemperature', uint16_t),
+  ('FanZeroRpmEnable', uint8_t),
+  ('FanZeroRpmStopTemp', uint8_t),
+  ('FanMode', uint8_t),
+  ('MaxOpTemp', uint8_t),
+  ('Spare', (uint32_t * 13)),
+]
 BOARD_GPIO_TYPE_e = CEnum(ctypes.c_uint32)
 BOARD_GPIO_SMUIO_0 = BOARD_GPIO_TYPE_e.define('BOARD_GPIO_SMUIO_0', 0)
 BOARD_GPIO_SMUIO_1 = BOARD_GPIO_TYPE_e.define('BOARD_GPIO_SMUIO_1', 1)
@@ -390,52 +498,522 @@ BOARD_GPIO_DC_GENLK_VSYNC = BOARD_GPIO_TYPE_e.define('BOARD_GPIO_DC_GENLK_VSYNC'
 BOARD_GPIO_DC_SWAPLOCK_A = BOARD_GPIO_TYPE_e.define('BOARD_GPIO_DC_SWAPLOCK_A', 42)
 BOARD_GPIO_DC_SWAPLOCK_B = BOARD_GPIO_TYPE_e.define('BOARD_GPIO_DC_SWAPLOCK_B', 43)
 
-class _anonstruct14(Struct): pass
-BootValues_t = _anonstruct14
-class _anonstruct15(Struct): pass
-MsgLimits_t = _anonstruct15
-class _anonstruct16(Struct): pass
-DriverReportedClocks_t = _anonstruct16
-class _anonstruct17(Struct): pass
-AvfsDcBtcParams_t = _anonstruct17
-class _anonstruct18(Struct): pass
-AvfsFuseOverride_t = _anonstruct18
-class _anonstruct19(Struct): pass
-SkuTable_t = _anonstruct19
-class _anonstruct20(Struct): pass
-BoardTable_t = _anonstruct20
-class _anonstruct21(Struct): pass
-PPTable_t = _anonstruct21
-class _anonstruct22(Struct): pass
-DriverSmuConfig_t = _anonstruct22
-class _anonstruct23(Struct): pass
-DriverSmuConfigExternal_t = _anonstruct23
-class _anonstruct24(Struct): pass
-DriverInfoTable_t = _anonstruct24
-class _anonstruct25(Struct): pass
-SmuMetrics_t = _anonstruct25
-class _anonstruct26(Struct): pass
-SmuMetricsExternal_t = _anonstruct26
-class _anonstruct27(Struct): pass
-WatermarkRowGeneric_t = _anonstruct27
+class BootValues_t(Struct): pass
+BootValues_t._fields_ = [
+  ('InitGfxclk_bypass', uint16_t),
+  ('InitSocclk', uint16_t),
+  ('InitMp0clk', uint16_t),
+  ('InitMpioclk', uint16_t),
+  ('InitSmnclk', uint16_t),
+  ('InitUcpclk', uint16_t),
+  ('InitCsrclk', uint16_t),
+  ('InitDprefclk', uint16_t),
+  ('InitDcfclk', uint16_t),
+  ('InitDtbclk', uint16_t),
+  ('InitDclk', uint16_t),
+  ('InitVclk', uint16_t),
+  ('InitUsbdfsclk', uint16_t),
+  ('InitMp1clk', uint16_t),
+  ('InitLclk', uint16_t),
+  ('InitBaco400clk_bypass', uint16_t),
+  ('InitBaco1200clk_bypass', uint16_t),
+  ('InitBaco700clk_bypass', uint16_t),
+  ('InitFclk', uint16_t),
+  ('InitGfxclk_clkb', uint16_t),
+  ('InitUclkDPMState', uint8_t),
+  ('Padding', (uint8_t * 3)),
+  ('InitVcoFreqPll0', uint32_t),
+  ('InitVcoFreqPll1', uint32_t),
+  ('InitVcoFreqPll2', uint32_t),
+  ('InitVcoFreqPll3', uint32_t),
+  ('InitVcoFreqPll4', uint32_t),
+  ('InitVcoFreqPll5', uint32_t),
+  ('InitVcoFreqPll6', uint32_t),
+  ('InitGfx', uint16_t),
+  ('InitSoc', uint16_t),
+  ('InitU', uint16_t),
+  ('Padding2', uint16_t),
+  ('Spare', (uint32_t * 8)),
+]
+class MsgLimits_t(Struct): pass
+MsgLimits_t._fields_ = [
+  ('Power', ((uint16_t * 2) * 4)),
+  ('Tdc', (uint16_t * 3)),
+  ('Temperature', (uint16_t * 13)),
+  ('PwmLimitMin', uint8_t),
+  ('PwmLimitMax', uint8_t),
+  ('FanTargetTemperature', uint8_t),
+  ('Spare1', (uint8_t * 1)),
+  ('AcousticTargetRpmThresholdMin', uint16_t),
+  ('AcousticTargetRpmThresholdMax', uint16_t),
+  ('AcousticLimitRpmThresholdMin', uint16_t),
+  ('AcousticLimitRpmThresholdMax', uint16_t),
+  ('PccLimitMin', uint16_t),
+  ('PccLimitMax', uint16_t),
+  ('FanStopTempMin', uint16_t),
+  ('FanStopTempMax', uint16_t),
+  ('FanStartTempMin', uint16_t),
+  ('FanStartTempMax', uint16_t),
+  ('PowerMinPpt0', (uint16_t * 2)),
+  ('Spare', (uint32_t * 11)),
+]
+class DriverReportedClocks_t(Struct): pass
+DriverReportedClocks_t._fields_ = [
+  ('BaseClockAc', uint16_t),
+  ('GameClockAc', uint16_t),
+  ('BoostClockAc', uint16_t),
+  ('BaseClockDc', uint16_t),
+  ('GameClockDc', uint16_t),
+  ('BoostClockDc', uint16_t),
+  ('Reserved', (uint32_t * 4)),
+]
+class AvfsDcBtcParams_t(Struct): pass
+AvfsDcBtcParams_t._fields_ = [
+  ('DcBtcEnabled', uint8_t),
+  ('Padding', (uint8_t * 3)),
+  ('DcTol', uint16_t),
+  ('DcBtcGb', uint16_t),
+  ('DcBtcMin', uint16_t),
+  ('DcBtcMax', uint16_t),
+  ('DcBtcGbScalar', LinearInt_t),
+]
+class AvfsFuseOverride_t(Struct): pass
+AvfsFuseOverride_t._fields_ = [
+  ('AvfsTemp', (uint16_t * 2)),
+  ('VftFMin', uint16_t),
+  ('VInversion', uint16_t),
+  ('qVft', (QuadraticInt_t * 2)),
+  ('qAvfsGb', QuadraticInt_t),
+  ('qAvfsGb2', QuadraticInt_t),
+]
+class SkuTable_t(Struct): pass
+int32_t = ctypes.c_int32
+SkuTable_t._fields_ = [
+  ('Version', uint32_t),
+  ('FeaturesToRun', (uint32_t * 2)),
+  ('TotalPowerConfig', uint8_t),
+  ('CustomerVariant', uint8_t),
+  ('MemoryTemperatureTypeMask', uint8_t),
+  ('SmartShiftVersion', uint8_t),
+  ('SocketPowerLimitAc', (uint16_t * 4)),
+  ('SocketPowerLimitDc', (uint16_t * 4)),
+  ('SocketPowerLimitSmartShift2', uint16_t),
+  ('EnableLegacyPptLimit', uint8_t),
+  ('UseInputTelemetry', uint8_t),
+  ('SmartShiftMinReportedPptinDcs', uint8_t),
+  ('PaddingPpt', (uint8_t * 1)),
+  ('VrTdcLimit', (uint16_t * 3)),
+  ('PlatformTdcLimit', (uint16_t * 3)),
+  ('TemperatureLimit', (uint16_t * 13)),
+  ('HwCtfTempLimit', uint16_t),
+  ('PaddingInfra', uint16_t),
+  ('FitControllerFailureRateLimit', uint32_t),
+  ('FitControllerGfxDutyCycle', uint32_t),
+  ('FitControllerSocDutyCycle', uint32_t),
+  ('FitControllerSocOffset', uint32_t),
+  ('GfxApccPlusResidencyLimit', uint32_t),
+  ('ThrottlerControlMask', uint32_t),
+  ('FwDStateMask', uint32_t),
+  ('UlvVoltageOffset', (uint16_t * 2)),
+  ('UlvVoltageOffsetU', uint16_t),
+  ('DeepUlvVoltageOffsetSoc', uint16_t),
+  ('DefaultMaxVoltage', (uint16_t * 2)),
+  ('BoostMaxVoltage', (uint16_t * 2)),
+  ('VminTempHystersis', (int16_t * 2)),
+  ('VminTempThreshold', (int16_t * 2)),
+  ('Vmin_Hot_T0', (uint16_t * 2)),
+  ('Vmin_Cold_T0', (uint16_t * 2)),
+  ('Vmin_Hot_Eol', (uint16_t * 2)),
+  ('Vmin_Cold_Eol', (uint16_t * 2)),
+  ('Vmin_Aging_Offset', (uint16_t * 2)),
+  ('Spare_Vmin_Plat_Offset_Hot', (uint16_t * 2)),
+  ('Spare_Vmin_Plat_Offset_Cold', (uint16_t * 2)),
+  ('VcBtcFixedVminAgingOffset', (uint16_t * 2)),
+  ('VcBtcVmin2PsmDegrationGb', (uint16_t * 2)),
+  ('VcBtcPsmA', (uint32_t * 2)),
+  ('VcBtcPsmB', (uint32_t * 2)),
+  ('VcBtcVminA', (uint32_t * 2)),
+  ('VcBtcVminB', (uint32_t * 2)),
+  ('PerPartVminEnabled', (uint8_t * 2)),
+  ('VcBtcEnabled', (uint8_t * 2)),
+  ('SocketPowerLimitAcTau', (uint16_t * 4)),
+  ('SocketPowerLimitDcTau', (uint16_t * 4)),
+  ('Vmin_droop', QuadraticInt_t),
+  ('SpareVmin', (uint32_t * 9)),
+  ('DpmDescriptor', (DpmDescriptor_t * 13)),
+  ('FreqTableGfx', (uint16_t * 16)),
+  ('FreqTableVclk', (uint16_t * 8)),
+  ('FreqTableDclk', (uint16_t * 8)),
+  ('FreqTableSocclk', (uint16_t * 8)),
+  ('FreqTableUclk', (uint16_t * 4)),
+  ('FreqTableDispclk', (uint16_t * 8)),
+  ('FreqTableDppClk', (uint16_t * 8)),
+  ('FreqTableDprefclk', (uint16_t * 8)),
+  ('FreqTableDcfclk', (uint16_t * 8)),
+  ('FreqTableDtbclk', (uint16_t * 8)),
+  ('FreqTableFclk', (uint16_t * 8)),
+  ('DcModeMaxFreq', (uint32_t * 13)),
+  ('Mp0clkFreq', (uint16_t * 2)),
+  ('Mp0DpmVoltage', (uint16_t * 2)),
+  ('GfxclkSpare', (uint8_t * 2)),
+  ('GfxclkFreqCap', uint16_t),
+  ('GfxclkFgfxoffEntry', uint16_t),
+  ('GfxclkFgfxoffExitImu', uint16_t),
+  ('GfxclkFgfxoffExitRlc', uint16_t),
+  ('GfxclkThrottleClock', uint16_t),
+  ('EnableGfxPowerStagesGpio', uint8_t),
+  ('GfxIdlePadding', uint8_t),
+  ('SmsRepairWRCKClkDivEn', uint8_t),
+  ('SmsRepairWRCKClkDivVal', uint8_t),
+  ('GfxOffEntryEarlyMGCGEn', uint8_t),
+  ('GfxOffEntryForceCGCGEn', uint8_t),
+  ('GfxOffEntryForceCGCGDelayEn', uint8_t),
+  ('GfxOffEntryForceCGCGDelayVal', uint8_t),
+  ('GfxclkFreqGfxUlv', uint16_t),
+  ('GfxIdlePadding2', (uint8_t * 2)),
+  ('GfxOffEntryHysteresis', uint32_t),
+  ('GfxoffSpare', (uint32_t * 15)),
+  ('DfllBtcMasterScalerM', uint32_t),
+  ('DfllBtcMasterScalerB', int32_t),
+  ('DfllBtcSlaveScalerM', uint32_t),
+  ('DfllBtcSlaveScalerB', int32_t),
+  ('DfllPccAsWaitCtrl', uint32_t),
+  ('DfllPccAsStepCtrl', uint32_t),
+  ('DfllL2FrequencyBoostM', uint32_t),
+  ('DfllL2FrequencyBoostB', uint32_t),
+  ('GfxGpoSpare', (uint32_t * 8)),
+  ('DcsGfxOffVoltage', uint16_t),
+  ('PaddingDcs', uint16_t),
+  ('DcsMinGfxOffTime', uint16_t),
+  ('DcsMaxGfxOffTime', uint16_t),
+  ('DcsMinCreditAccum', uint32_t),
+  ('DcsExitHysteresis', uint16_t),
+  ('DcsTimeout', uint16_t),
+  ('FoptEnabled', uint8_t),
+  ('DcsSpare2', (uint8_t * 3)),
+  ('DcsFoptM', uint32_t),
+  ('DcsFoptB', uint32_t),
+  ('DcsSpare', (uint32_t * 11)),
+  ('ShadowFreqTableUclk', (uint16_t * 4)),
+  ('UseStrobeModeOptimizations', uint8_t),
+  ('PaddingMem', (uint8_t * 3)),
+  ('UclkDpmPstates', (uint8_t * 4)),
+  ('FreqTableUclkDiv', (uint8_t * 4)),
+  ('MemVmempVoltage', (uint16_t * 4)),
+  ('MemVddioVoltage', (uint16_t * 4)),
+  ('FclkDpmUPstates', (uint8_t * 8)),
+  ('FclkDpmVddU', (uint16_t * 8)),
+  ('FclkDpmUSpeed', (uint16_t * 8)),
+  ('FclkDpmDisallowPstateFreq', uint16_t),
+  ('PaddingFclk', uint16_t),
+  ('PcieGenSpeed', (uint8_t * 3)),
+  ('PcieLaneCount', (uint8_t * 3)),
+  ('LclkFreq', (uint16_t * 3)),
+  ('FanStopTemp', (uint16_t * 13)),
+  ('FanStartTemp', (uint16_t * 13)),
+  ('FanGain', (uint16_t * 13)),
+  ('FanGainPadding', uint16_t),
+  ('FanPwmMin', uint16_t),
+  ('AcousticTargetRpmThreshold', uint16_t),
+  ('AcousticLimitRpmThreshold', uint16_t),
+  ('FanMaximumRpm', uint16_t),
+  ('MGpuAcousticLimitRpmThreshold', uint16_t),
+  ('FanTargetGfxclk', uint16_t),
+  ('TempInputSelectMask', uint32_t),
+  ('FanZeroRpmEnable', uint8_t),
+  ('FanTachEdgePerRev', uint8_t),
+  ('FanTargetTemperature', (uint16_t * 13)),
+  ('FuzzyFan_ErrorSetDelta', int16_t),
+  ('FuzzyFan_ErrorRateSetDelta', int16_t),
+  ('FuzzyFan_PwmSetDelta', int16_t),
+  ('FuzzyFan_Reserved', uint16_t),
+  ('FwCtfLimit', (uint16_t * 13)),
+  ('IntakeTempEnableRPM', uint16_t),
+  ('IntakeTempOffsetTemp', int16_t),
+  ('IntakeTempReleaseTemp', uint16_t),
+  ('IntakeTempHighIntakeAcousticLimit', uint16_t),
+  ('IntakeTempAcouticLimitReleaseRate', uint16_t),
+  ('FanAbnormalTempLimitOffset', int16_t),
+  ('FanStalledTriggerRpm', uint16_t),
+  ('FanAbnormalTriggerRpmCoeff', uint16_t),
+  ('FanAbnormalDetectionEnable', uint16_t),
+  ('FanIntakeSensorSupport', uint8_t),
+  ('FanIntakePadding', (uint8_t * 3)),
+  ('FanSpare', (uint32_t * 13)),
+  ('OverrideGfxAvfsFuses', uint8_t),
+  ('GfxAvfsPadding', (uint8_t * 3)),
+  ('L2HwRtAvfsFuses', (uint32_t * 32)),
+  ('SeHwRtAvfsFuses', (uint32_t * 32)),
+  ('CommonRtAvfs', (uint32_t * 13)),
+  ('L2FwRtAvfsFuses', (uint32_t * 19)),
+  ('SeFwRtAvfsFuses', (uint32_t * 19)),
+  ('Droop_PWL_F', (uint32_t * 5)),
+  ('Droop_PWL_a', (uint32_t * 5)),
+  ('Droop_PWL_b', (uint32_t * 5)),
+  ('Droop_PWL_c', (uint32_t * 5)),
+  ('Static_PWL_Offset', (uint32_t * 5)),
+  ('dGbV_dT_vmin', uint32_t),
+  ('dGbV_dT_vmax', uint32_t),
+  ('V2F_vmin_range_low', uint32_t),
+  ('V2F_vmin_range_high', uint32_t),
+  ('V2F_vmax_range_low', uint32_t),
+  ('V2F_vmax_range_high', uint32_t),
+  ('DcBtcGfxParams', AvfsDcBtcParams_t),
+  ('GfxAvfsSpare', (uint32_t * 32)),
+  ('OverrideSocAvfsFuses', uint8_t),
+  ('MinSocAvfsRevision', uint8_t),
+  ('SocAvfsPadding', (uint8_t * 2)),
+  ('SocAvfsFuseOverride', (AvfsFuseOverride_t * 3)),
+  ('dBtcGbSoc', (DroopInt_t * 3)),
+  ('qAgingGb', (LinearInt_t * 3)),
+  ('qStaticVoltageOffset', (QuadraticInt_t * 3)),
+  ('DcBtcSocParams', (AvfsDcBtcParams_t * 3)),
+  ('SocAvfsSpare', (uint32_t * 32)),
+  ('BootValues', BootValues_t),
+  ('DriverReportedClocks', DriverReportedClocks_t),
+  ('MsgLimits', MsgLimits_t),
+  ('OverDriveLimitsMin', OverDriveLimits_t),
+  ('OverDriveLimitsBasicMax', OverDriveLimits_t),
+  ('reserved', (uint32_t * 22)),
+  ('DebugOverrides', uint32_t),
+  ('TotalBoardPowerSupport', uint8_t),
+  ('TotalBoardPowerPadding', (uint8_t * 3)),
+  ('TotalIdleBoardPowerM', int16_t),
+  ('TotalIdleBoardPowerB', int16_t),
+  ('TotalBoardPowerM', int16_t),
+  ('TotalBoardPowerB', int16_t),
+  ('qFeffCoeffGameClock', (QuadraticInt_t * 2)),
+  ('qFeffCoeffBaseClock', (QuadraticInt_t * 2)),
+  ('qFeffCoeffBoostClock', (QuadraticInt_t * 2)),
+  ('TemperatureLimit_Hynix', uint16_t),
+  ('TemperatureLimit_Micron', uint16_t),
+  ('TemperatureFwCtfLimit_Hynix', uint16_t),
+  ('TemperatureFwCtfLimit_Micron', uint16_t),
+  ('Spare', (uint32_t * 41)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class BoardTable_t(Struct): pass
+BoardTable_t._fields_ = [
+  ('Version', uint32_t),
+  ('I2cControllers', (I2cControllerConfig_t * 8)),
+  ('VddGfxVrMapping', uint8_t),
+  ('VddSocVrMapping', uint8_t),
+  ('VddMem0VrMapping', uint8_t),
+  ('VddMem1VrMapping', uint8_t),
+  ('GfxUlvPhaseSheddingMask', uint8_t),
+  ('SocUlvPhaseSheddingMask', uint8_t),
+  ('VmempUlvPhaseSheddingMask', uint8_t),
+  ('VddioUlvPhaseSheddingMask', uint8_t),
+  ('SlaveAddrMapping', (uint8_t * 5)),
+  ('VrPsiSupport', (uint8_t * 5)),
+  ('PaddingPsi', (uint8_t * 5)),
+  ('EnablePsi6', (uint8_t * 5)),
+  ('SviTelemetryScale', (SviTelemetryScale_t * 5)),
+  ('VoltageTelemetryRatio', (uint32_t * 5)),
+  ('DownSlewRateVr', (uint8_t * 5)),
+  ('LedOffGpio', uint8_t),
+  ('FanOffGpio', uint8_t),
+  ('GfxVrPowerStageOffGpio', uint8_t),
+  ('AcDcGpio', uint8_t),
+  ('AcDcPolarity', uint8_t),
+  ('VR0HotGpio', uint8_t),
+  ('VR0HotPolarity', uint8_t),
+  ('GthrGpio', uint8_t),
+  ('GthrPolarity', uint8_t),
+  ('LedPin0', uint8_t),
+  ('LedPin1', uint8_t),
+  ('LedPin2', uint8_t),
+  ('LedEnableMask', uint8_t),
+  ('LedPcie', uint8_t),
+  ('LedError', uint8_t),
+  ('UclkTrainingModeSpreadPercent', uint8_t),
+  ('UclkSpreadPadding', uint8_t),
+  ('UclkSpreadFreq', uint16_t),
+  ('UclkSpreadPercent', (uint8_t * 16)),
+  ('GfxclkSpreadEnable', uint8_t),
+  ('FclkSpreadPercent', uint8_t),
+  ('FclkSpreadFreq', uint16_t),
+  ('DramWidth', uint8_t),
+  ('PaddingMem1', (uint8_t * 7)),
+  ('HsrEnabled', uint8_t),
+  ('VddqOffEnabled', uint8_t),
+  ('PaddingUmcFlags', (uint8_t * 2)),
+  ('PostVoltageSetBacoDelay', uint32_t),
+  ('BacoEntryDelay', uint32_t),
+  ('FuseWritePowerMuxPresent', uint8_t),
+  ('FuseWritePadding', (uint8_t * 3)),
+  ('BoardSpare', (uint32_t * 63)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class PPTable_t(Struct): pass
+PPTable_t._packed_ = True
+PPTable_t._fields_ = [
+  ('SkuTable', SkuTable_t),
+  ('BoardTable', BoardTable_t),
+]
+class DriverSmuConfig_t(Struct): pass
+DriverSmuConfig_t._fields_ = [
+  ('GfxclkAverageLpfTau', uint16_t),
+  ('FclkAverageLpfTau', uint16_t),
+  ('UclkAverageLpfTau', uint16_t),
+  ('GfxActivityLpfTau', uint16_t),
+  ('UclkActivityLpfTau', uint16_t),
+  ('SocketPowerLpfTau', uint16_t),
+  ('VcnClkAverageLpfTau', uint16_t),
+  ('VcnUsageAverageLpfTau', uint16_t),
+]
+class DriverSmuConfigExternal_t(Struct): pass
+DriverSmuConfigExternal_t._fields_ = [
+  ('DriverSmuConfig', DriverSmuConfig_t),
+  ('Spare', (uint32_t * 8)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class DriverInfoTable_t(Struct): pass
+DriverInfoTable_t._fields_ = [
+  ('FreqTableGfx', (uint16_t * 16)),
+  ('FreqTableVclk', (uint16_t * 8)),
+  ('FreqTableDclk', (uint16_t * 8)),
+  ('FreqTableSocclk', (uint16_t * 8)),
+  ('FreqTableUclk', (uint16_t * 4)),
+  ('FreqTableDispclk', (uint16_t * 8)),
+  ('FreqTableDppClk', (uint16_t * 8)),
+  ('FreqTableDprefclk', (uint16_t * 8)),
+  ('FreqTableDcfclk', (uint16_t * 8)),
+  ('FreqTableDtbclk', (uint16_t * 8)),
+  ('FreqTableFclk', (uint16_t * 8)),
+  ('DcModeMaxFreq', (uint16_t * 13)),
+  ('Padding', uint16_t),
+  ('Spare', (uint32_t * 32)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class SmuMetrics_t(Struct): pass
+SmuMetrics_t._fields_ = [
+  ('CurrClock', (uint32_t * 13)),
+  ('AverageGfxclkFrequencyTarget', uint16_t),
+  ('AverageGfxclkFrequencyPreDs', uint16_t),
+  ('AverageGfxclkFrequencyPostDs', uint16_t),
+  ('AverageFclkFrequencyPreDs', uint16_t),
+  ('AverageFclkFrequencyPostDs', uint16_t),
+  ('AverageMemclkFrequencyPreDs', uint16_t),
+  ('AverageMemclkFrequencyPostDs', uint16_t),
+  ('AverageVclk0Frequency', uint16_t),
+  ('AverageDclk0Frequency', uint16_t),
+  ('AverageVclk1Frequency', uint16_t),
+  ('AverageDclk1Frequency', uint16_t),
+  ('PCIeBusy', uint16_t),
+  ('dGPU_W_MAX', uint16_t),
+  ('padding', uint16_t),
+  ('MetricsCounter', uint32_t),
+  ('AvgVoltage', (uint16_t * 5)),
+  ('AvgCurrent', (uint16_t * 5)),
+  ('AverageGfxActivity', uint16_t),
+  ('AverageUclkActivity', uint16_t),
+  ('Vcn0ActivityPercentage', uint16_t),
+  ('Vcn1ActivityPercentage', uint16_t),
+  ('EnergyAccumulator', uint32_t),
+  ('AverageSocketPower', uint16_t),
+  ('AverageTotalBoardPower', uint16_t),
+  ('AvgTemperature', (uint16_t * 13)),
+  ('AvgTemperatureFanIntake', uint16_t),
+  ('PcieRate', uint8_t),
+  ('PcieWidth', uint8_t),
+  ('AvgFanPwm', uint8_t),
+  ('Padding', (uint8_t * 1)),
+  ('AvgFanRpm', uint16_t),
+  ('ThrottlingPercentage', (uint8_t * 22)),
+  ('VmaxThrottlingPercentage', uint8_t),
+  ('Padding1', (uint8_t * 3)),
+  ('D3HotEntryCountPerMode', (uint32_t * 4)),
+  ('D3HotExitCountPerMode', (uint32_t * 4)),
+  ('ArmMsgReceivedCountPerMode', (uint32_t * 4)),
+  ('ApuSTAPMSmartShiftLimit', uint16_t),
+  ('ApuSTAPMLimit', uint16_t),
+  ('AvgApuSocketPower', uint16_t),
+  ('AverageUclkActivity_MAX', uint16_t),
+  ('PublicSerialNumberLower', uint32_t),
+  ('PublicSerialNumberUpper', uint32_t),
+]
+class SmuMetricsExternal_t(Struct): pass
+SmuMetricsExternal_t._fields_ = [
+  ('SmuMetrics', SmuMetrics_t),
+  ('Spare', (uint32_t * 29)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class WatermarkRowGeneric_t(Struct): pass
+WatermarkRowGeneric_t._fields_ = [
+  ('WmSetting', uint8_t),
+  ('Flags', uint8_t),
+  ('Padding', (uint8_t * 2)),
+]
 WATERMARKS_FLAGS_e = CEnum(ctypes.c_uint32)
 WATERMARKS_CLOCK_RANGE = WATERMARKS_FLAGS_e.define('WATERMARKS_CLOCK_RANGE', 0)
 WATERMARKS_DUMMY_PSTATE = WATERMARKS_FLAGS_e.define('WATERMARKS_DUMMY_PSTATE', 1)
 WATERMARKS_MALL = WATERMARKS_FLAGS_e.define('WATERMARKS_MALL', 2)
 WATERMARKS_COUNT = WATERMARKS_FLAGS_e.define('WATERMARKS_COUNT', 3)
 
-class _anonstruct28(Struct): pass
-Watermarks_t = _anonstruct28
-class _anonstruct29(Struct): pass
-WatermarksExternal_t = _anonstruct29
-class _anonstruct30(Struct): pass
-AvfsDebugTable_t = _anonstruct30
-class _anonstruct31(Struct): pass
-AvfsDebugTableExternal_t = _anonstruct31
-class _anonstruct32(Struct): pass
-DpmActivityMonitorCoeffInt_t = _anonstruct32
-class _anonstruct33(Struct): pass
-DpmActivityMonitorCoeffIntExternal_t = _anonstruct33
+class Watermarks_t(Struct): pass
+Watermarks_t._fields_ = [
+  ('WatermarkRow', (WatermarkRowGeneric_t * 4)),
+]
+class WatermarksExternal_t(Struct): pass
+WatermarksExternal_t._fields_ = [
+  ('Watermarks', Watermarks_t),
+  ('Spare', (uint32_t * 16)),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class AvfsDebugTable_t(Struct): pass
+AvfsDebugTable_t._fields_ = [
+  ('avgPsmCount', (uint16_t * 214)),
+  ('minPsmCount', (uint16_t * 214)),
+  ('avgPsmVoltage', (ctypes.c_float * 214)),
+  ('minPsmVoltage', (ctypes.c_float * 214)),
+]
+class AvfsDebugTableExternal_t(Struct): pass
+AvfsDebugTableExternal_t._fields_ = [
+  ('AvfsDebugTable', AvfsDebugTable_t),
+  ('MmHubPadding', (uint32_t * 8)),
+]
+class DpmActivityMonitorCoeffInt_t(Struct): pass
+DpmActivityMonitorCoeffInt_t._fields_ = [
+  ('Gfx_ActiveHystLimit', uint8_t),
+  ('Gfx_IdleHystLimit', uint8_t),
+  ('Gfx_FPS', uint8_t),
+  ('Gfx_MinActiveFreqType', uint8_t),
+  ('Gfx_BoosterFreqType', uint8_t),
+  ('PaddingGfx', uint8_t),
+  ('Gfx_MinActiveFreq', uint16_t),
+  ('Gfx_BoosterFreq', uint16_t),
+  ('Gfx_PD_Data_time_constant', uint16_t),
+  ('Gfx_PD_Data_limit_a', uint32_t),
+  ('Gfx_PD_Data_limit_b', uint32_t),
+  ('Gfx_PD_Data_limit_c', uint32_t),
+  ('Gfx_PD_Data_error_coeff', uint32_t),
+  ('Gfx_PD_Data_error_rate_coeff', uint32_t),
+  ('Fclk_ActiveHystLimit', uint8_t),
+  ('Fclk_IdleHystLimit', uint8_t),
+  ('Fclk_FPS', uint8_t),
+  ('Fclk_MinActiveFreqType', uint8_t),
+  ('Fclk_BoosterFreqType', uint8_t),
+  ('PaddingFclk', uint8_t),
+  ('Fclk_MinActiveFreq', uint16_t),
+  ('Fclk_BoosterFreq', uint16_t),
+  ('Fclk_PD_Data_time_constant', uint16_t),
+  ('Fclk_PD_Data_limit_a', uint32_t),
+  ('Fclk_PD_Data_limit_b', uint32_t),
+  ('Fclk_PD_Data_limit_c', uint32_t),
+  ('Fclk_PD_Data_error_coeff', uint32_t),
+  ('Fclk_PD_Data_error_rate_coeff', uint32_t),
+  ('Mem_UpThreshold_Limit', (uint32_t * 4)),
+  ('Mem_UpHystLimit', (uint8_t * 4)),
+  ('Mem_DownHystLimit', (uint8_t * 4)),
+  ('Mem_Fps', uint16_t),
+  ('padding', (uint8_t * 2)),
+]
+class DpmActivityMonitorCoeffIntExternal_t(Struct): pass
+DpmActivityMonitorCoeffIntExternal_t._fields_ = [
+  ('DpmActivityMonitorCoeffInt', DpmActivityMonitorCoeffInt_t),
+  ('MmHubPadding', (uint32_t * 8)),
+]
 class struct_smu_hw_power_state(Struct): pass
 struct_smu_hw_power_state._fields_ = [
   ('magic', ctypes.c_uint32),

tinygrad/runtime/autogen/hip.py

@@ -1809,8 +1809,7 @@ except AttributeError: pass
 try: (hipMemcpy2DToArrayAsync:=dll.hipMemcpy2DToArrayAsync).restype, hipMemcpy2DToArrayAsync.argtypes = hipError_t, [hipArray_t, size_t, size_t, ctypes.c_void_p, size_t, size_t, size_t, hipMemcpyKind, hipStream_t]
 except AttributeError: pass
 
-class const_hipArray(Struct): pass
-hipArray_const_t = ctypes.POINTER(const_hipArray)
+hipArray_const_t = ctypes.POINTER(hipArray)
 try: (hipMemcpy2DArrayToArray:=dll.hipMemcpy2DArrayToArray).restype, hipMemcpy2DArrayToArray.argtypes = hipError_t, [hipArray_t, size_t, size_t, hipArray_const_t, size_t, size_t, size_t, size_t, hipMemcpyKind]
 except AttributeError: pass
 
@@ -2358,21 +2357,7 @@ except AttributeError: pass
 try: (hipFreeMipmappedArray:=dll.hipFreeMipmappedArray).restype, hipFreeMipmappedArray.argtypes = hipError_t, [hipMipmappedArray_t]
 except AttributeError: pass
 
-class const_hipMipmappedArray(Struct): pass
-const_hipMipmappedArray._fields_ = [
-  ('data', ctypes.c_void_p),
-  ('desc', hipChannelFormatDesc),
-  ('type', ctypes.c_uint32),
-  ('width', ctypes.c_uint32),
-  ('height', ctypes.c_uint32),
-  ('depth', ctypes.c_uint32),
-  ('min_mipmap_level', ctypes.c_uint32),
-  ('max_mipmap_level', ctypes.c_uint32),
-  ('flags', ctypes.c_uint32),
-  ('format', hipArray_Format),
-  ('num_channels', ctypes.c_uint32),
-]
-hipMipmappedArray_const_t = ctypes.POINTER(const_hipMipmappedArray)
+hipMipmappedArray_const_t = ctypes.POINTER(hipMipmappedArray)
 try: (hipGetMipmappedArrayLevel:=dll.hipGetMipmappedArrayLevel).restype, hipGetMipmappedArrayLevel.argtypes = hipError_t, [ctypes.POINTER(hipArray_t), hipMipmappedArray_const_t, ctypes.c_uint32]
 except AttributeError: pass
 

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,67 @@ 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
+struct_hsa_amd_aie_ert_packet_s._fields_ = [
+  ('header', hsa_amd_vendor_packet_header_t),
+  ('state', uint32_t,4),
+  ('custom', uint32_t,8),
+  ('count', uint32_t,11),
+  ('opcode', uint32_t,5),
+  ('type', uint32_t,4),
+  ('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 +1318,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 +1367,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 +1413,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 +1465,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 +1565,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 +1590,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 +1669,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 +1688,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 +1721,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 +1741,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 +1770,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_ = [
@@ -1422,13 +1778,7 @@ struct_hsa_amd_event_s._fields_ = [
   ('_0', struct_hsa_amd_event_s_0),
 ]
 hsa_amd_event_t = struct_hsa_amd_event_s
-class const_struct_hsa_amd_event_s(Struct): pass
-const_struct_hsa_amd_event_s._anonymous_ = ['_0']
-const_struct_hsa_amd_event_s._fields_ = [
-  ('event_type', hsa_amd_event_type_t),
-  ('_0', struct_hsa_amd_event_s_0),
-]
-hsa_amd_system_event_callback_t = ctypes.CFUNCTYPE(hsa_status_t, ctypes.POINTER(const_struct_hsa_amd_event_s), ctypes.c_void_p)
+hsa_amd_system_event_callback_t = ctypes.CFUNCTYPE(hsa_status_t, ctypes.POINTER(struct_hsa_amd_event_s), ctypes.c_void_p)
 try: (hsa_amd_register_system_event_handler:=dll.hsa_amd_register_system_event_handler).restype, hsa_amd_register_system_event_handler.argtypes = hsa_status_t, [hsa_amd_system_event_callback_t, ctypes.c_void_p]
 except AttributeError: pass
 
@@ -1441,6 +1791,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 +1851,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 +1925,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 +1963,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 +1997,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 +2049,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 +2254,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 +2591,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 +2636,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 +2831,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/autogen/io_uring.py

@@ -181,94 +181,14 @@ struct_statx._fields_ = [
   ('__spare3', (ctypes.c_uint64 * 12)),
 ]
 class struct_epoll_event(Struct): pass
-class _anonunion0(ctypes.Union): pass
-class _anonunion0_0(Struct): pass
-_anonunion0_0._fields_ = [
-  ('cmd_op', ctypes.c_uint32),
-  ('__pad1', ctypes.c_uint32),
-]
-_anonunion0._anonymous_ = ['_0']
-_anonunion0._fields_ = [
-  ('off', ctypes.c_uint64),
-  ('addr2', ctypes.c_uint64),
-  ('_0', _anonunion0_0),
-]
-class _anonunion1(ctypes.Union): pass
-class _anonunion1_0(Struct): pass
-_anonunion1_0._fields_ = [
-  ('level', ctypes.c_uint32),
-  ('optname', ctypes.c_uint32),
-]
-_anonunion1._anonymous_ = ['_0']
-_anonunion1._fields_ = [
-  ('addr', ctypes.c_uint64),
-  ('splice_off_in', ctypes.c_uint64),
-  ('_0', _anonunion1_0),
-]
-class _anonunion2(ctypes.Union): pass
-_anonunion2._fields_ = [
-  ('rw_flags', ctypes.c_int32),
-  ('fsync_flags', ctypes.c_uint32),
-  ('poll_events', ctypes.c_uint16),
-  ('poll32_events', ctypes.c_uint32),
-  ('sync_range_flags', ctypes.c_uint32),
-  ('msg_flags', ctypes.c_uint32),
-  ('timeout_flags', ctypes.c_uint32),
-  ('accept_flags', ctypes.c_uint32),
-  ('cancel_flags', ctypes.c_uint32),
-  ('open_flags', ctypes.c_uint32),
-  ('statx_flags', ctypes.c_uint32),
-  ('fadvise_advice', ctypes.c_uint32),
-  ('splice_flags', ctypes.c_uint32),
-  ('rename_flags', ctypes.c_uint32),
-  ('unlink_flags', ctypes.c_uint32),
-  ('hardlink_flags', ctypes.c_uint32),
-  ('xattr_flags', ctypes.c_uint32),
-  ('msg_ring_flags', ctypes.c_uint32),
-  ('uring_cmd_flags', ctypes.c_uint32),
-  ('waitid_flags', ctypes.c_uint32),
-  ('futex_flags', ctypes.c_uint32),
-  ('install_fd_flags', ctypes.c_uint32),
-]
-class _anonunion3(ctypes.Union): pass
-_anonunion3._packed_ = True
-_anonunion3._fields_ = [
-  ('buf_index', ctypes.c_uint16),
-  ('buf_group', ctypes.c_uint16),
-]
-class _anonunion4(ctypes.Union): pass
-class _anonunion4_0(Struct): pass
-_anonunion4_0._fields_ = [
-  ('addr_len', ctypes.c_uint16),
-  ('__pad3', (ctypes.c_uint16 * 1)),
-]
-_anonunion4._anonymous_ = ['_0']
-_anonunion4._fields_ = [
-  ('splice_fd_in', ctypes.c_int32),
-  ('file_index', ctypes.c_uint32),
-  ('optlen', ctypes.c_uint32),
-  ('_0', _anonunion4_0),
-]
-class _anonunion5(ctypes.Union): pass
-class _anonunion5_0(Struct): pass
-_anonunion5_0._fields_ = [
-  ('addr3', ctypes.c_uint64),
-  ('__pad2', (ctypes.c_uint64 * 1)),
-]
-_anonunion5._anonymous_ = ['_0']
-_anonunion5._fields_ = [
-  ('_0', _anonunion5_0),
-  ('optval', ctypes.c_uint64),
-  ('cmd', (ctypes.c_ubyte * 0)),
-]
-_anonenum6 = CEnum(ctypes.c_uint32)
-IOSQE_FIXED_FILE_BIT = _anonenum6.define('IOSQE_FIXED_FILE_BIT', 0)
-IOSQE_IO_DRAIN_BIT = _anonenum6.define('IOSQE_IO_DRAIN_BIT', 1)
-IOSQE_IO_LINK_BIT = _anonenum6.define('IOSQE_IO_LINK_BIT', 2)
-IOSQE_IO_HARDLINK_BIT = _anonenum6.define('IOSQE_IO_HARDLINK_BIT', 3)
-IOSQE_ASYNC_BIT = _anonenum6.define('IOSQE_ASYNC_BIT', 4)
-IOSQE_BUFFER_SELECT_BIT = _anonenum6.define('IOSQE_BUFFER_SELECT_BIT', 5)
-IOSQE_CQE_SKIP_SUCCESS_BIT = _anonenum6.define('IOSQE_CQE_SKIP_SUCCESS_BIT', 6)
+_anonenum0 = CEnum(ctypes.c_uint32)
+IOSQE_FIXED_FILE_BIT = _anonenum0.define('IOSQE_FIXED_FILE_BIT', 0)
+IOSQE_IO_DRAIN_BIT = _anonenum0.define('IOSQE_IO_DRAIN_BIT', 1)
+IOSQE_IO_LINK_BIT = _anonenum0.define('IOSQE_IO_LINK_BIT', 2)
+IOSQE_IO_HARDLINK_BIT = _anonenum0.define('IOSQE_IO_HARDLINK_BIT', 3)
+IOSQE_ASYNC_BIT = _anonenum0.define('IOSQE_ASYNC_BIT', 4)
+IOSQE_BUFFER_SELECT_BIT = _anonenum0.define('IOSQE_BUFFER_SELECT_BIT', 5)
+IOSQE_CQE_SKIP_SUCCESS_BIT = _anonenum0.define('IOSQE_CQE_SKIP_SUCCESS_BIT', 6)
 
 enum_io_uring_op = CEnum(ctypes.c_uint32)
 IORING_OP_NOP = enum_io_uring_op.define('IORING_OP_NOP', 0)
@@ -328,12 +248,12 @@ IORING_OP_FUTEX_WAITV = enum_io_uring_op.define('IORING_OP_FUTEX_WAITV', 53)
 IORING_OP_FIXED_FD_INSTALL = enum_io_uring_op.define('IORING_OP_FIXED_FD_INSTALL', 54)
 IORING_OP_LAST = enum_io_uring_op.define('IORING_OP_LAST', 55)
 
-_anonenum7 = CEnum(ctypes.c_uint32)
-IORING_MSG_DATA = _anonenum7.define('IORING_MSG_DATA', 0)
-IORING_MSG_SEND_FD = _anonenum7.define('IORING_MSG_SEND_FD', 1)
+_anonenum1 = CEnum(ctypes.c_uint32)
+IORING_MSG_DATA = _anonenum1.define('IORING_MSG_DATA', 0)
+IORING_MSG_SEND_FD = _anonenum1.define('IORING_MSG_SEND_FD', 1)
 
-_anonenum8 = CEnum(ctypes.c_uint32)
-IORING_CQE_BUFFER_SHIFT = _anonenum8.define('IORING_CQE_BUFFER_SHIFT', 16)
+_anonenum2 = CEnum(ctypes.c_uint32)
+IORING_CQE_BUFFER_SHIFT = _anonenum2.define('IORING_CQE_BUFFER_SHIFT', 16)
 
 class struct_io_sqring_offsets(Struct): pass
 struct_io_sqring_offsets._fields_ = [
@@ -372,40 +292,40 @@ struct_io_uring_params._fields_ = [
   ('sq_off', struct_io_sqring_offsets),
   ('cq_off', struct_io_cqring_offsets),
 ]
-_anonenum9 = CEnum(ctypes.c_uint32)
-IORING_REGISTER_BUFFERS = _anonenum9.define('IORING_REGISTER_BUFFERS', 0)
-IORING_UNREGISTER_BUFFERS = _anonenum9.define('IORING_UNREGISTER_BUFFERS', 1)
-IORING_REGISTER_FILES = _anonenum9.define('IORING_REGISTER_FILES', 2)
-IORING_UNREGISTER_FILES = _anonenum9.define('IORING_UNREGISTER_FILES', 3)
-IORING_REGISTER_EVENTFD = _anonenum9.define('IORING_REGISTER_EVENTFD', 4)
-IORING_UNREGISTER_EVENTFD = _anonenum9.define('IORING_UNREGISTER_EVENTFD', 5)
-IORING_REGISTER_FILES_UPDATE = _anonenum9.define('IORING_REGISTER_FILES_UPDATE', 6)
-IORING_REGISTER_EVENTFD_ASYNC = _anonenum9.define('IORING_REGISTER_EVENTFD_ASYNC', 7)
-IORING_REGISTER_PROBE = _anonenum9.define('IORING_REGISTER_PROBE', 8)
-IORING_REGISTER_PERSONALITY = _anonenum9.define('IORING_REGISTER_PERSONALITY', 9)
-IORING_UNREGISTER_PERSONALITY = _anonenum9.define('IORING_UNREGISTER_PERSONALITY', 10)
-IORING_REGISTER_RESTRICTIONS = _anonenum9.define('IORING_REGISTER_RESTRICTIONS', 11)
-IORING_REGISTER_ENABLE_RINGS = _anonenum9.define('IORING_REGISTER_ENABLE_RINGS', 12)
-IORING_REGISTER_FILES2 = _anonenum9.define('IORING_REGISTER_FILES2', 13)
-IORING_REGISTER_FILES_UPDATE2 = _anonenum9.define('IORING_REGISTER_FILES_UPDATE2', 14)
-IORING_REGISTER_BUFFERS2 = _anonenum9.define('IORING_REGISTER_BUFFERS2', 15)
-IORING_REGISTER_BUFFERS_UPDATE = _anonenum9.define('IORING_REGISTER_BUFFERS_UPDATE', 16)
-IORING_REGISTER_IOWQ_AFF = _anonenum9.define('IORING_REGISTER_IOWQ_AFF', 17)
-IORING_UNREGISTER_IOWQ_AFF = _anonenum9.define('IORING_UNREGISTER_IOWQ_AFF', 18)
-IORING_REGISTER_IOWQ_MAX_WORKERS = _anonenum9.define('IORING_REGISTER_IOWQ_MAX_WORKERS', 19)
-IORING_REGISTER_RING_FDS = _anonenum9.define('IORING_REGISTER_RING_FDS', 20)
-IORING_UNREGISTER_RING_FDS = _anonenum9.define('IORING_UNREGISTER_RING_FDS', 21)
-IORING_REGISTER_PBUF_RING = _anonenum9.define('IORING_REGISTER_PBUF_RING', 22)
-IORING_UNREGISTER_PBUF_RING = _anonenum9.define('IORING_UNREGISTER_PBUF_RING', 23)
-IORING_REGISTER_SYNC_CANCEL = _anonenum9.define('IORING_REGISTER_SYNC_CANCEL', 24)
-IORING_REGISTER_FILE_ALLOC_RANGE = _anonenum9.define('IORING_REGISTER_FILE_ALLOC_RANGE', 25)
-IORING_REGISTER_PBUF_STATUS = _anonenum9.define('IORING_REGISTER_PBUF_STATUS', 26)
-IORING_REGISTER_LAST = _anonenum9.define('IORING_REGISTER_LAST', 27)
-IORING_REGISTER_USE_REGISTERED_RING = _anonenum9.define('IORING_REGISTER_USE_REGISTERED_RING', 2147483648)
+_anonenum3 = CEnum(ctypes.c_uint32)
+IORING_REGISTER_BUFFERS = _anonenum3.define('IORING_REGISTER_BUFFERS', 0)
+IORING_UNREGISTER_BUFFERS = _anonenum3.define('IORING_UNREGISTER_BUFFERS', 1)
+IORING_REGISTER_FILES = _anonenum3.define('IORING_REGISTER_FILES', 2)
+IORING_UNREGISTER_FILES = _anonenum3.define('IORING_UNREGISTER_FILES', 3)
+IORING_REGISTER_EVENTFD = _anonenum3.define('IORING_REGISTER_EVENTFD', 4)
+IORING_UNREGISTER_EVENTFD = _anonenum3.define('IORING_UNREGISTER_EVENTFD', 5)
+IORING_REGISTER_FILES_UPDATE = _anonenum3.define('IORING_REGISTER_FILES_UPDATE', 6)
+IORING_REGISTER_EVENTFD_ASYNC = _anonenum3.define('IORING_REGISTER_EVENTFD_ASYNC', 7)
+IORING_REGISTER_PROBE = _anonenum3.define('IORING_REGISTER_PROBE', 8)
+IORING_REGISTER_PERSONALITY = _anonenum3.define('IORING_REGISTER_PERSONALITY', 9)
+IORING_UNREGISTER_PERSONALITY = _anonenum3.define('IORING_UNREGISTER_PERSONALITY', 10)
+IORING_REGISTER_RESTRICTIONS = _anonenum3.define('IORING_REGISTER_RESTRICTIONS', 11)
+IORING_REGISTER_ENABLE_RINGS = _anonenum3.define('IORING_REGISTER_ENABLE_RINGS', 12)
+IORING_REGISTER_FILES2 = _anonenum3.define('IORING_REGISTER_FILES2', 13)
+IORING_REGISTER_FILES_UPDATE2 = _anonenum3.define('IORING_REGISTER_FILES_UPDATE2', 14)
+IORING_REGISTER_BUFFERS2 = _anonenum3.define('IORING_REGISTER_BUFFERS2', 15)
+IORING_REGISTER_BUFFERS_UPDATE = _anonenum3.define('IORING_REGISTER_BUFFERS_UPDATE', 16)
+IORING_REGISTER_IOWQ_AFF = _anonenum3.define('IORING_REGISTER_IOWQ_AFF', 17)
+IORING_UNREGISTER_IOWQ_AFF = _anonenum3.define('IORING_UNREGISTER_IOWQ_AFF', 18)
+IORING_REGISTER_IOWQ_MAX_WORKERS = _anonenum3.define('IORING_REGISTER_IOWQ_MAX_WORKERS', 19)
+IORING_REGISTER_RING_FDS = _anonenum3.define('IORING_REGISTER_RING_FDS', 20)
+IORING_UNREGISTER_RING_FDS = _anonenum3.define('IORING_UNREGISTER_RING_FDS', 21)
+IORING_REGISTER_PBUF_RING = _anonenum3.define('IORING_REGISTER_PBUF_RING', 22)
+IORING_UNREGISTER_PBUF_RING = _anonenum3.define('IORING_UNREGISTER_PBUF_RING', 23)
+IORING_REGISTER_SYNC_CANCEL = _anonenum3.define('IORING_REGISTER_SYNC_CANCEL', 24)
+IORING_REGISTER_FILE_ALLOC_RANGE = _anonenum3.define('IORING_REGISTER_FILE_ALLOC_RANGE', 25)
+IORING_REGISTER_PBUF_STATUS = _anonenum3.define('IORING_REGISTER_PBUF_STATUS', 26)
+IORING_REGISTER_LAST = _anonenum3.define('IORING_REGISTER_LAST', 27)
+IORING_REGISTER_USE_REGISTERED_RING = _anonenum3.define('IORING_REGISTER_USE_REGISTERED_RING', 2147483648)
 
-_anonenum10 = CEnum(ctypes.c_uint32)
-IO_WQ_BOUND = _anonenum10.define('IO_WQ_BOUND', 0)
-IO_WQ_UNBOUND = _anonenum10.define('IO_WQ_UNBOUND', 1)
+_anonenum4 = CEnum(ctypes.c_uint32)
+IO_WQ_BOUND = _anonenum4.define('IO_WQ_BOUND', 0)
+IO_WQ_UNBOUND = _anonenum4.define('IO_WQ_UNBOUND', 1)
 
 class struct_io_uring_files_update(Struct): pass
 struct_io_uring_files_update._fields_ = [
@@ -496,8 +416,8 @@ struct_io_uring_buf_ring._anonymous_ = ['_0']
 struct_io_uring_buf_ring._fields_ = [
   ('_0', struct_io_uring_buf_ring_0),
 ]
-_anonenum11 = CEnum(ctypes.c_uint32)
-IOU_PBUF_RING_MMAP = _anonenum11.define('IOU_PBUF_RING_MMAP', 1)
+_anonenum5 = CEnum(ctypes.c_uint32)
+IOU_PBUF_RING_MMAP = _anonenum5.define('IOU_PBUF_RING_MMAP', 1)
 
 class struct_io_uring_buf_reg(Struct): pass
 struct_io_uring_buf_reg._fields_ = [
@@ -513,12 +433,12 @@ struct_io_uring_buf_status._fields_ = [
   ('head', ctypes.c_uint32),
   ('resv', (ctypes.c_uint32 * 8)),
 ]
-_anonenum12 = CEnum(ctypes.c_uint32)
-IORING_RESTRICTION_REGISTER_OP = _anonenum12.define('IORING_RESTRICTION_REGISTER_OP', 0)
-IORING_RESTRICTION_SQE_OP = _anonenum12.define('IORING_RESTRICTION_SQE_OP', 1)
-IORING_RESTRICTION_SQE_FLAGS_ALLOWED = _anonenum12.define('IORING_RESTRICTION_SQE_FLAGS_ALLOWED', 2)
-IORING_RESTRICTION_SQE_FLAGS_REQUIRED = _anonenum12.define('IORING_RESTRICTION_SQE_FLAGS_REQUIRED', 3)
-IORING_RESTRICTION_LAST = _anonenum12.define('IORING_RESTRICTION_LAST', 4)
+_anonenum6 = CEnum(ctypes.c_uint32)
+IORING_RESTRICTION_REGISTER_OP = _anonenum6.define('IORING_RESTRICTION_REGISTER_OP', 0)
+IORING_RESTRICTION_SQE_OP = _anonenum6.define('IORING_RESTRICTION_SQE_OP', 1)
+IORING_RESTRICTION_SQE_FLAGS_ALLOWED = _anonenum6.define('IORING_RESTRICTION_SQE_FLAGS_ALLOWED', 2)
+IORING_RESTRICTION_SQE_FLAGS_REQUIRED = _anonenum6.define('IORING_RESTRICTION_SQE_FLAGS_REQUIRED', 3)
+IORING_RESTRICTION_LAST = _anonenum6.define('IORING_RESTRICTION_LAST', 4)
 
 class struct_io_uring_getevents_arg(Struct): pass
 struct_io_uring_getevents_arg._fields_ = [
@@ -556,11 +476,11 @@ struct_io_uring_recvmsg_out._fields_ = [
   ('payloadlen', ctypes.c_uint32),
   ('flags', ctypes.c_uint32),
 ]
-_anonenum13 = CEnum(ctypes.c_uint32)
-SOCKET_URING_OP_SIOCINQ = _anonenum13.define('SOCKET_URING_OP_SIOCINQ', 0)
-SOCKET_URING_OP_SIOCOUTQ = _anonenum13.define('SOCKET_URING_OP_SIOCOUTQ', 1)
-SOCKET_URING_OP_GETSOCKOPT = _anonenum13.define('SOCKET_URING_OP_GETSOCKOPT', 2)
-SOCKET_URING_OP_SETSOCKOPT = _anonenum13.define('SOCKET_URING_OP_SETSOCKOPT', 3)
+_anonenum7 = CEnum(ctypes.c_uint32)
+SOCKET_URING_OP_SIOCINQ = _anonenum7.define('SOCKET_URING_OP_SIOCINQ', 0)
+SOCKET_URING_OP_SIOCOUTQ = _anonenum7.define('SOCKET_URING_OP_SIOCOUTQ', 1)
+SOCKET_URING_OP_GETSOCKOPT = _anonenum7.define('SOCKET_URING_OP_GETSOCKOPT', 2)
+SOCKET_URING_OP_SETSOCKOPT = _anonenum7.define('SOCKET_URING_OP_SETSOCKOPT', 3)
 
 _XOPEN_SOURCE = 500
 uring_unlikely = lambda cond: __builtin_expect( not  not (cond), 0)

tinygrad/runtime/autogen/libusb.py

@@ -9,13 +9,6 @@ def dll():
   return None
 dll = dll()
 
-class _anonunion0(ctypes.Union): pass
-uint8_t = ctypes.c_ubyte
-uint16_t = ctypes.c_uint16
-_anonunion0._fields_ = [
-  ('b8', (uint8_t * 2)),
-  ('b16', uint16_t),
-]
 enum_libusb_class_code = CEnum(ctypes.c_uint32)
 LIBUSB_CLASS_PER_INTERFACE = enum_libusb_class_code.define('LIBUSB_CLASS_PER_INTERFACE', 0)
 LIBUSB_CLASS_AUDIO = enum_libusb_class_code.define('LIBUSB_CLASS_AUDIO', 1)
@@ -122,6 +115,8 @@ LIBUSB_BT_CONTAINER_ID = enum_libusb_bos_type.define('LIBUSB_BT_CONTAINER_ID', 4
 LIBUSB_BT_PLATFORM_DESCRIPTOR = enum_libusb_bos_type.define('LIBUSB_BT_PLATFORM_DESCRIPTOR', 5)
 
 class struct_libusb_device_descriptor(Struct): pass
+uint8_t = ctypes.c_ubyte
+uint16_t = ctypes.c_uint16
 struct_libusb_device_descriptor._fields_ = [
   ('bLength', uint8_t),
   ('bDescriptorType', uint8_t),

tinygrad/runtime/autogen/nv_570.py

@@ -1323,17 +1323,6 @@ nvdec_crc_s._fields_ = [
   ('reserved1', ctypes.c_uint32,16),
   ('reserved2', (ctypes.c_ubyte * 56)),
 ]
-class _anonunion12(ctypes.Union): pass
-NvUPtr = ctypes.c_uint64
-_anonunion12._fields_ = [
-  ('v', NvUPtr),
-  ('p', ctypes.c_void_p),
-]
-class _anonunion13(ctypes.Union): pass
-_anonunion13._fields_ = [
-  ('v', NvUPtr),
-  ('p', ctypes.c_void_p),
-]
 class struct_NV0000_ALLOC_PARAMETERS(Struct): pass
 NvHandle = ctypes.c_uint32
 NvU32 = ctypes.c_uint32
@@ -6660,6 +6649,7 @@ struct_NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER._fields_ = [
 ]
 NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER = struct_NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER
 class struct_NV2080_CTRL_FLCN_USTREAMER_QUEUE_INFO_PARAMS(Struct): pass
+NvUPtr = ctypes.c_uint64
 struct_NV2080_CTRL_FLCN_USTREAMER_QUEUE_INFO_PARAMS._fields_ = [
   ('engine', NvU32),
   ('pageSize', NvU32),

tinygrad/runtime/autogen/nv_580.py

@@ -1323,17 +1323,6 @@ nvdec_crc_s._fields_ = [
   ('reserved1', ctypes.c_uint32,16),
   ('reserved2', (ctypes.c_ubyte * 56)),
 ]
-class _anonunion12(ctypes.Union): pass
-NvUPtr = ctypes.c_uint64
-_anonunion12._fields_ = [
-  ('v', NvUPtr),
-  ('p', ctypes.c_void_p),
-]
-class _anonunion13(ctypes.Union): pass
-_anonunion13._fields_ = [
-  ('v', NvUPtr),
-  ('p', ctypes.c_void_p),
-]
 class struct_NV0000_ALLOC_PARAMETERS(Struct): pass
 NvHandle = ctypes.c_uint32
 NvU32 = ctypes.c_uint32
@@ -7219,6 +7208,7 @@ struct_NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER._fields_ = [
 ]
 NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER = struct_NV2080_CTRL_FLCN_USTREAMER_EVENT_FILTER
 class struct_NV2080_CTRL_FLCN_USTREAMER_QUEUE_INFO_PARAMS(Struct): pass
+NvUPtr = ctypes.c_uint64
 struct_NV2080_CTRL_FLCN_USTREAMER_QUEUE_INFO_PARAMS._fields_ = [
   ('engine', NvU32),
   ('pageSize', NvU32),

tinygrad/runtime/autogen/webgpu.py

@@ -293,7 +293,7 @@ WGPUCompilationInfoRequestStatus_DeviceLost = enum_WGPUCompilationInfoRequestSta
 WGPUCompilationInfoRequestStatus_Unknown = enum_WGPUCompilationInfoRequestStatus.define('WGPUCompilationInfoRequestStatus_Unknown', 5)
 WGPUCompilationInfoRequestStatus_Force32 = enum_WGPUCompilationInfoRequestStatus.define('WGPUCompilationInfoRequestStatus_Force32', 2147483647)
 
-class const_struct_WGPUCompilationInfo(Struct): pass
+class struct_WGPUCompilationInfo(Struct): pass
 size_t = ctypes.c_uint64
 class struct_WGPUCompilationMessage(Struct): pass
 WGPUCompilationMessage = struct_WGPUCompilationMessage
@@ -322,12 +322,12 @@ struct_WGPUCompilationMessage._fields_ = [
   ('utf16Offset', uint64_t),
   ('utf16Length', uint64_t),
 ]
-const_struct_WGPUCompilationInfo._fields_ = [
+struct_WGPUCompilationInfo._fields_ = [
   ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
   ('messageCount', size_t),
   ('messages', ctypes.POINTER(WGPUCompilationMessage)),
 ]
-WGPUCompilationInfoCallback = ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(const_struct_WGPUCompilationInfo), ctypes.c_void_p)
+WGPUCompilationInfoCallback = ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(struct_WGPUCompilationInfo), ctypes.c_void_p)
 struct_WGPUCompilationInfoCallbackInfo._fields_ = [
   ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
   ('mode', WGPUCallbackMode),
@@ -1874,12 +1874,6 @@ struct_WGPUColorTargetState._fields_ = [
   ('blend', ctypes.POINTER(WGPUBlendState)),
   ('writeMask', WGPUColorWriteMask),
 ]
-class struct_WGPUCompilationInfo(Struct): pass
-struct_WGPUCompilationInfo._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('messageCount', size_t),
-  ('messages', ctypes.POINTER(WGPUCompilationMessage)),
-]
 class struct_WGPUComputeState(Struct): pass
 WGPUConstantEntry = struct_WGPUConstantEntry
 struct_WGPUComputeState._fields_ = [
@@ -2062,7 +2056,7 @@ WGPUDeviceLostCallback = ctypes.CFUNCTYPE(None, enum_WGPUDeviceLostReason, struc
 WGPULoggingCallback = ctypes.CFUNCTYPE(None, enum_WGPULoggingType, struct_WGPUStringView, ctypes.c_void_p)
 WGPUProc = ctypes.CFUNCTYPE(None, )
 WGPUBufferMapCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUMapAsyncStatus, struct_WGPUStringView, ctypes.c_void_p, ctypes.c_void_p)
-WGPUCompilationInfoCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(const_struct_WGPUCompilationInfo), ctypes.c_void_p, ctypes.c_void_p)
+WGPUCompilationInfoCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(struct_WGPUCompilationInfo), ctypes.c_void_p, ctypes.c_void_p)
 WGPUCreateComputePipelineAsyncCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPUComputePipelineImpl), struct_WGPUStringView, ctypes.c_void_p, ctypes.c_void_p)
 WGPUCreateRenderPipelineAsyncCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPURenderPipelineImpl), struct_WGPUStringView, ctypes.c_void_p, ctypes.c_void_p)
 WGPUPopErrorScopeCallback2 = ctypes.CFUNCTYPE(None, enum_WGPUPopErrorScopeStatus, enum_WGPUErrorType, struct_WGPUStringView, ctypes.c_void_p, ctypes.c_void_p)
@@ -2272,12 +2266,7 @@ WGPUSurfaceDescriptorFromXcbWindow = struct_WGPUSurfaceSourceXCBWindow
 WGPUSurfaceDescriptorFromXlibWindow = struct_WGPUSurfaceSourceXlibWindow
 WGPUProcAdapterInfoFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUAdapterInfo)
 WGPUProcAdapterPropertiesMemoryHeapsFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUAdapterPropertiesMemoryHeaps)
-class const_struct_WGPUInstanceDescriptor(Struct): pass
-const_struct_WGPUInstanceDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('features', WGPUInstanceFeatures),
-]
-WGPUProcCreateInstance = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(const_struct_WGPUInstanceDescriptor))
+WGPUProcCreateInstance = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPUInstanceDescriptor))
 WGPUProcDrmFormatCapabilitiesFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUDrmFormatCapabilities)
 WGPUProcGetInstanceFeatures = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUInstanceFeatures))
 WGPUProcGetProcAddress = ctypes.CFUNCTYPE(ctypes.CFUNCTYPE(None, ), struct_WGPUStringView)
@@ -2285,27 +2274,16 @@ WGPUProcSharedBufferMemoryEndAccessStateFreeMembers = ctypes.CFUNCTYPE(None, str
 WGPUProcSharedTextureMemoryEndAccessStateFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUSharedTextureMemoryEndAccessState)
 WGPUProcSupportedFeaturesFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUSupportedFeatures)
 WGPUProcSurfaceCapabilitiesFreeMembers = ctypes.CFUNCTYPE(None, struct_WGPUSurfaceCapabilities)
-class const_struct_WGPUDeviceDescriptor(Struct): pass
-const_struct_WGPUDeviceDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('requiredFeatureCount', size_t),
-  ('requiredFeatures', ctypes.POINTER(WGPUFeatureName)),
-  ('requiredLimits', ctypes.POINTER(WGPURequiredLimits)),
-  ('defaultQueue', WGPUQueueDescriptor),
-  ('deviceLostCallbackInfo2', WGPUDeviceLostCallbackInfo2),
-  ('uncapturedErrorCallbackInfo2', WGPUUncapturedErrorCallbackInfo2),
-]
-WGPUProcAdapterCreateDevice = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(const_struct_WGPUDeviceDescriptor))
+WGPUProcAdapterCreateDevice = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUDeviceDescriptor))
 WGPUProcAdapterGetFeatures = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUSupportedFeatures))
 WGPUProcAdapterGetFormatCapabilities = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUAdapterImpl), enum_WGPUTextureFormat, ctypes.POINTER(struct_WGPUFormatCapabilities))
 WGPUProcAdapterGetInfo = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUAdapterInfo))
 WGPUProcAdapterGetInstance = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPUAdapterImpl))
 WGPUProcAdapterGetLimits = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUSupportedLimits))
 WGPUProcAdapterHasFeature = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUAdapterImpl), enum_WGPUFeatureName)
-WGPUProcAdapterRequestDevice = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(const_struct_WGPUDeviceDescriptor), ctypes.CFUNCTYPE(None, enum_WGPURequestDeviceStatus, ctypes.POINTER(struct_WGPUDeviceImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
-WGPUProcAdapterRequestDevice2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(const_struct_WGPUDeviceDescriptor), struct_WGPURequestDeviceCallbackInfo2)
-WGPUProcAdapterRequestDeviceF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(const_struct_WGPUDeviceDescriptor), struct_WGPURequestDeviceCallbackInfo)
+WGPUProcAdapterRequestDevice = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUDeviceDescriptor), ctypes.CFUNCTYPE(None, enum_WGPURequestDeviceStatus, ctypes.POINTER(struct_WGPUDeviceImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
+WGPUProcAdapterRequestDevice2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUDeviceDescriptor), struct_WGPURequestDeviceCallbackInfo2)
+WGPUProcAdapterRequestDeviceF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUDeviceDescriptor), struct_WGPURequestDeviceCallbackInfo)
 WGPUProcAdapterAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUAdapterImpl))
 WGPUProcAdapterRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUAdapterImpl))
 WGPUProcBindGroupSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUBindGroupImpl), struct_WGPUStringView)
@@ -2330,53 +2308,14 @@ WGPUProcBufferRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUBufferI
 WGPUProcCommandBufferSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandBufferImpl), struct_WGPUStringView)
 WGPUProcCommandBufferAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandBufferImpl))
 WGPUProcCommandBufferRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandBufferImpl))
-class const_struct_WGPUComputePassDescriptor(Struct): pass
-const_struct_WGPUComputePassDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('timestampWrites', ctypes.POINTER(WGPUComputePassTimestampWrites)),
-]
-WGPUProcCommandEncoderBeginComputePass = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUComputePassEncoderImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPUComputePassDescriptor))
-class const_struct_WGPURenderPassDescriptor(Struct): pass
-const_struct_WGPURenderPassDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('colorAttachmentCount', size_t),
-  ('colorAttachments', ctypes.POINTER(WGPURenderPassColorAttachment)),
-  ('depthStencilAttachment', ctypes.POINTER(WGPURenderPassDepthStencilAttachment)),
-  ('occlusionQuerySet', WGPUQuerySet),
-  ('timestampWrites', ctypes.POINTER(WGPURenderPassTimestampWrites)),
-]
-WGPUProcCommandEncoderBeginRenderPass = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPURenderPassDescriptor))
+WGPUProcCommandEncoderBeginComputePass = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUComputePassEncoderImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUComputePassDescriptor))
+WGPUProcCommandEncoderBeginRenderPass = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPURenderPassDescriptor))
 WGPUProcCommandEncoderClearBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64, ctypes.c_uint64)
 WGPUProcCommandEncoderCopyBufferToBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64, ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64, ctypes.c_uint64)
-class const_struct_WGPUImageCopyBuffer(Struct): pass
-const_struct_WGPUImageCopyBuffer._fields_ = [
-  ('layout', WGPUTextureDataLayout),
-  ('buffer', WGPUBuffer),
-]
-class const_struct_WGPUImageCopyTexture(Struct): pass
-const_struct_WGPUImageCopyTexture._fields_ = [
-  ('texture', WGPUTexture),
-  ('mipLevel', uint32_t),
-  ('origin', WGPUOrigin3D),
-  ('aspect', WGPUTextureAspect),
-]
-class const_struct_WGPUExtent3D(Struct): pass
-const_struct_WGPUExtent3D._fields_ = [
-  ('width', uint32_t),
-  ('height', uint32_t),
-  ('depthOrArrayLayers', uint32_t),
-]
-WGPUProcCommandEncoderCopyBufferToTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPUImageCopyBuffer), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUExtent3D))
-WGPUProcCommandEncoderCopyTextureToBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUImageCopyBuffer), ctypes.POINTER(const_struct_WGPUExtent3D))
-WGPUProcCommandEncoderCopyTextureToTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUExtent3D))
-class const_struct_WGPUCommandBufferDescriptor(Struct): pass
-const_struct_WGPUCommandBufferDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcCommandEncoderFinish = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUCommandBufferImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(const_struct_WGPUCommandBufferDescriptor))
+WGPUProcCommandEncoderCopyBufferToTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUImageCopyBuffer), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUExtent3D))
+WGPUProcCommandEncoderCopyTextureToBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUImageCopyBuffer), ctypes.POINTER(struct_WGPUExtent3D))
+WGPUProcCommandEncoderCopyTextureToTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUExtent3D))
+WGPUProcCommandEncoderFinish = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUCommandBufferImpl), ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUCommandBufferDescriptor))
 WGPUProcCommandEncoderInjectValidationError = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), struct_WGPUStringView)
 WGPUProcCommandEncoderInsertDebugMarker = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl), struct_WGPUStringView)
 WGPUProcCommandEncoderPopDebugGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUCommandEncoderImpl))
@@ -2403,151 +2342,29 @@ WGPUProcComputePipelineGetBindGroupLayout = ctypes.CFUNCTYPE(ctypes.POINTER(stru
 WGPUProcComputePipelineSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUComputePipelineImpl), struct_WGPUStringView)
 WGPUProcComputePipelineAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUComputePipelineImpl))
 WGPUProcComputePipelineRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUComputePipelineImpl))
-class const_struct_WGPUBindGroupDescriptor(Struct): pass
-const_struct_WGPUBindGroupDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('layout', WGPUBindGroupLayout),
-  ('entryCount', size_t),
-  ('entries', ctypes.POINTER(WGPUBindGroupEntry)),
-]
-WGPUProcDeviceCreateBindGroup = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBindGroupImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUBindGroupDescriptor))
-class const_struct_WGPUBindGroupLayoutDescriptor(Struct): pass
-const_struct_WGPUBindGroupLayoutDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('entryCount', size_t),
-  ('entries', ctypes.POINTER(WGPUBindGroupLayoutEntry)),
-]
-WGPUProcDeviceCreateBindGroupLayout = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBindGroupLayoutImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUBindGroupLayoutDescriptor))
-class const_struct_WGPUBufferDescriptor(Struct): pass
-const_struct_WGPUBufferDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('usage', WGPUBufferUsage),
-  ('size', uint64_t),
-  ('mappedAtCreation', WGPUBool),
-]
-WGPUProcDeviceCreateBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUBufferDescriptor))
-class const_struct_WGPUCommandEncoderDescriptor(Struct): pass
-const_struct_WGPUCommandEncoderDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcDeviceCreateCommandEncoder = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUCommandEncoderDescriptor))
-class const_struct_WGPUComputePipelineDescriptor(Struct): pass
-const_struct_WGPUComputePipelineDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('layout', WGPUPipelineLayout),
-  ('compute', WGPUComputeState),
-]
-WGPUProcDeviceCreateComputePipeline = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUComputePipelineImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUComputePipelineDescriptor))
-WGPUProcDeviceCreateComputePipelineAsync = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUComputePipelineDescriptor), ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPUComputePipelineImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
-WGPUProcDeviceCreateComputePipelineAsync2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUComputePipelineDescriptor), struct_WGPUCreateComputePipelineAsyncCallbackInfo2)
-WGPUProcDeviceCreateComputePipelineAsyncF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUComputePipelineDescriptor), struct_WGPUCreateComputePipelineAsyncCallbackInfo)
-WGPUProcDeviceCreateErrorBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUBufferDescriptor))
+WGPUProcDeviceCreateBindGroup = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBindGroupImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUBindGroupDescriptor))
+WGPUProcDeviceCreateBindGroupLayout = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBindGroupLayoutImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUBindGroupLayoutDescriptor))
+WGPUProcDeviceCreateBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUBufferDescriptor))
+WGPUProcDeviceCreateCommandEncoder = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUCommandEncoderImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUCommandEncoderDescriptor))
+WGPUProcDeviceCreateComputePipeline = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUComputePipelineImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUComputePipelineDescriptor))
+WGPUProcDeviceCreateComputePipelineAsync = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUComputePipelineDescriptor), ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPUComputePipelineImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
+WGPUProcDeviceCreateComputePipelineAsync2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUComputePipelineDescriptor), struct_WGPUCreateComputePipelineAsyncCallbackInfo2)
+WGPUProcDeviceCreateComputePipelineAsyncF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUComputePipelineDescriptor), struct_WGPUCreateComputePipelineAsyncCallbackInfo)
+WGPUProcDeviceCreateErrorBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUBufferDescriptor))
 WGPUProcDeviceCreateErrorExternalTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUExternalTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl))
-class const_struct_WGPUShaderModuleDescriptor(Struct): pass
-const_struct_WGPUShaderModuleDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcDeviceCreateErrorShaderModule = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUShaderModuleDescriptor), struct_WGPUStringView)
-class const_struct_WGPUTextureDescriptor(Struct): pass
-const_struct_WGPUTextureDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('usage', WGPUTextureUsage),
-  ('dimension', WGPUTextureDimension),
-  ('size', WGPUExtent3D),
-  ('format', WGPUTextureFormat),
-  ('mipLevelCount', uint32_t),
-  ('sampleCount', uint32_t),
-  ('viewFormatCount', size_t),
-  ('viewFormats', ctypes.POINTER(WGPUTextureFormat)),
-]
-WGPUProcDeviceCreateErrorTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUTextureDescriptor))
-class const_struct_WGPUExternalTextureDescriptor(Struct): pass
-const_struct_WGPUExternalTextureDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('plane0', WGPUTextureView),
-  ('plane1', WGPUTextureView),
-  ('cropOrigin', WGPUOrigin2D),
-  ('cropSize', WGPUExtent2D),
-  ('apparentSize', WGPUExtent2D),
-  ('doYuvToRgbConversionOnly', WGPUBool),
-  ('yuvToRgbConversionMatrix', ctypes.POINTER(ctypes.c_float)),
-  ('srcTransferFunctionParameters', ctypes.POINTER(ctypes.c_float)),
-  ('dstTransferFunctionParameters', ctypes.POINTER(ctypes.c_float)),
-  ('gamutConversionMatrix', ctypes.POINTER(ctypes.c_float)),
-  ('mirrored', WGPUBool),
-  ('rotation', WGPUExternalTextureRotation),
-]
-WGPUProcDeviceCreateExternalTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUExternalTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUExternalTextureDescriptor))
-class const_struct_WGPUPipelineLayoutDescriptor(Struct): pass
-const_struct_WGPUPipelineLayoutDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('bindGroupLayoutCount', size_t),
-  ('bindGroupLayouts', ctypes.POINTER(WGPUBindGroupLayout)),
-  ('immediateDataRangeByteSize', uint32_t),
-]
-WGPUProcDeviceCreatePipelineLayout = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUPipelineLayoutImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUPipelineLayoutDescriptor))
-class const_struct_WGPUQuerySetDescriptor(Struct): pass
-const_struct_WGPUQuerySetDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('type', WGPUQueryType),
-  ('count', uint32_t),
-]
-WGPUProcDeviceCreateQuerySet = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUQuerySetImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUQuerySetDescriptor))
-class const_struct_WGPURenderBundleEncoderDescriptor(Struct): pass
-const_struct_WGPURenderBundleEncoderDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('colorFormatCount', size_t),
-  ('colorFormats', ctypes.POINTER(WGPUTextureFormat)),
-  ('depthStencilFormat', WGPUTextureFormat),
-  ('sampleCount', uint32_t),
-  ('depthReadOnly', WGPUBool),
-  ('stencilReadOnly', WGPUBool),
-]
-WGPUProcDeviceCreateRenderBundleEncoder = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPURenderBundleEncoderDescriptor))
-class const_struct_WGPURenderPipelineDescriptor(Struct): pass
-const_struct_WGPURenderPipelineDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('layout', WGPUPipelineLayout),
-  ('vertex', WGPUVertexState),
-  ('primitive', WGPUPrimitiveState),
-  ('depthStencil', ctypes.POINTER(WGPUDepthStencilState)),
-  ('multisample', WGPUMultisampleState),
-  ('fragment', ctypes.POINTER(WGPUFragmentState)),
-]
-WGPUProcDeviceCreateRenderPipeline = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderPipelineImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPURenderPipelineDescriptor))
-WGPUProcDeviceCreateRenderPipelineAsync = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPURenderPipelineDescriptor), ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPURenderPipelineImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
-WGPUProcDeviceCreateRenderPipelineAsync2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPURenderPipelineDescriptor), struct_WGPUCreateRenderPipelineAsyncCallbackInfo2)
-WGPUProcDeviceCreateRenderPipelineAsyncF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPURenderPipelineDescriptor), struct_WGPUCreateRenderPipelineAsyncCallbackInfo)
-class const_struct_WGPUSamplerDescriptor(Struct): pass
-const_struct_WGPUSamplerDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('addressModeU', WGPUAddressMode),
-  ('addressModeV', WGPUAddressMode),
-  ('addressModeW', WGPUAddressMode),
-  ('magFilter', WGPUFilterMode),
-  ('minFilter', WGPUFilterMode),
-  ('mipmapFilter', WGPUMipmapFilterMode),
-  ('lodMinClamp', ctypes.c_float),
-  ('lodMaxClamp', ctypes.c_float),
-  ('compare', WGPUCompareFunction),
-  ('maxAnisotropy', uint16_t),
-]
-WGPUProcDeviceCreateSampler = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSamplerImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUSamplerDescriptor))
-WGPUProcDeviceCreateShaderModule = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUShaderModuleDescriptor))
-WGPUProcDeviceCreateTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUTextureDescriptor))
+WGPUProcDeviceCreateErrorShaderModule = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUShaderModuleDescriptor), struct_WGPUStringView)
+WGPUProcDeviceCreateErrorTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUTextureDescriptor))
+WGPUProcDeviceCreateExternalTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUExternalTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUExternalTextureDescriptor))
+WGPUProcDeviceCreatePipelineLayout = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUPipelineLayoutImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUPipelineLayoutDescriptor))
+WGPUProcDeviceCreateQuerySet = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUQuerySetImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUQuerySetDescriptor))
+WGPUProcDeviceCreateRenderBundleEncoder = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPURenderBundleEncoderDescriptor))
+WGPUProcDeviceCreateRenderPipeline = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderPipelineImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPURenderPipelineDescriptor))
+WGPUProcDeviceCreateRenderPipelineAsync = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPURenderPipelineDescriptor), ctypes.CFUNCTYPE(None, enum_WGPUCreatePipelineAsyncStatus, ctypes.POINTER(struct_WGPURenderPipelineImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
+WGPUProcDeviceCreateRenderPipelineAsync2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPURenderPipelineDescriptor), struct_WGPUCreateRenderPipelineAsyncCallbackInfo2)
+WGPUProcDeviceCreateRenderPipelineAsyncF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPURenderPipelineDescriptor), struct_WGPUCreateRenderPipelineAsyncCallbackInfo)
+WGPUProcDeviceCreateSampler = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSamplerImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUSamplerDescriptor))
+WGPUProcDeviceCreateShaderModule = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUShaderModuleDescriptor))
+WGPUProcDeviceCreateTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUTextureDescriptor))
 WGPUProcDeviceDestroy = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl))
 WGPUProcDeviceForceLoss = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), enum_WGPUDeviceLostReason, struct_WGPUStringView)
 WGPUProcDeviceGetAHardwareBufferProperties = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.c_void_p, ctypes.POINTER(struct_WGPUAHardwareBufferProperties))
@@ -2558,24 +2375,9 @@ WGPUProcDeviceGetLimits = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struc
 WGPUProcDeviceGetLostFuture = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl))
 WGPUProcDeviceGetQueue = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(struct_WGPUDeviceImpl))
 WGPUProcDeviceHasFeature = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUDeviceImpl), enum_WGPUFeatureName)
-class const_struct_WGPUSharedBufferMemoryDescriptor(Struct): pass
-const_struct_WGPUSharedBufferMemoryDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcDeviceImportSharedBufferMemory = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUSharedBufferMemoryDescriptor))
-class const_struct_WGPUSharedFenceDescriptor(Struct): pass
-const_struct_WGPUSharedFenceDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcDeviceImportSharedFence = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedFenceImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUSharedFenceDescriptor))
-class const_struct_WGPUSharedTextureMemoryDescriptor(Struct): pass
-const_struct_WGPUSharedTextureMemoryDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcDeviceImportSharedTextureMemory = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUSharedTextureMemoryDescriptor))
+WGPUProcDeviceImportSharedBufferMemory = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryDescriptor))
+WGPUProcDeviceImportSharedFence = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedFenceImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUSharedFenceDescriptor))
+WGPUProcDeviceImportSharedTextureMemory = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryDescriptor))
 WGPUProcDeviceInjectError = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), enum_WGPUErrorType, struct_WGPUStringView)
 WGPUProcDevicePopErrorScope = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.CFUNCTYPE(None, enum_WGPUErrorType, struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
 WGPUProcDevicePopErrorScope2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUDeviceImpl), struct_WGPUPopErrorScopeCallbackInfo2)
@@ -2584,7 +2386,7 @@ WGPUProcDevicePushErrorScope = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPU
 WGPUProcDeviceSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), struct_WGPUStringView)
 WGPUProcDeviceSetLoggingCallback = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.CFUNCTYPE(None, enum_WGPULoggingType, struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
 WGPUProcDeviceTick = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl))
-WGPUProcDeviceValidateTextureDescriptor = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(const_struct_WGPUTextureDescriptor))
+WGPUProcDeviceValidateTextureDescriptor = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl), ctypes.POINTER(struct_WGPUTextureDescriptor))
 WGPUProcDeviceAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl))
 WGPUProcDeviceRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUDeviceImpl))
 WGPUProcExternalTextureDestroy = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUExternalTextureImpl))
@@ -2593,28 +2395,13 @@ WGPUProcExternalTextureRefresh = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WG
 WGPUProcExternalTextureSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUExternalTextureImpl), struct_WGPUStringView)
 WGPUProcExternalTextureAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUExternalTextureImpl))
 WGPUProcExternalTextureRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUExternalTextureImpl))
-class const_struct_WGPUSurfaceDescriptor(Struct): pass
-const_struct_WGPUSurfaceDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcInstanceCreateSurface = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(const_struct_WGPUSurfaceDescriptor))
+WGPUProcInstanceCreateSurface = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPUSurfaceDescriptor))
 WGPUProcInstanceEnumerateWGSLLanguageFeatures = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(enum_WGPUWGSLFeatureName))
 WGPUProcInstanceHasWGSLLanguageFeature = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUInstanceImpl), enum_WGPUWGSLFeatureName)
 WGPUProcInstanceProcessEvents = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUInstanceImpl))
-class const_struct_WGPURequestAdapterOptions(Struct): pass
-const_struct_WGPURequestAdapterOptions._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('compatibleSurface', WGPUSurface),
-  ('featureLevel', WGPUFeatureLevel),
-  ('powerPreference', WGPUPowerPreference),
-  ('backendType', WGPUBackendType),
-  ('forceFallbackAdapter', WGPUBool),
-  ('compatibilityMode', WGPUBool),
-]
-WGPUProcInstanceRequestAdapter = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(const_struct_WGPURequestAdapterOptions), ctypes.CFUNCTYPE(None, enum_WGPURequestAdapterStatus, ctypes.POINTER(struct_WGPUAdapterImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
-WGPUProcInstanceRequestAdapter2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(const_struct_WGPURequestAdapterOptions), struct_WGPURequestAdapterCallbackInfo2)
-WGPUProcInstanceRequestAdapterF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(const_struct_WGPURequestAdapterOptions), struct_WGPURequestAdapterCallbackInfo)
+WGPUProcInstanceRequestAdapter = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPURequestAdapterOptions), ctypes.CFUNCTYPE(None, enum_WGPURequestAdapterStatus, ctypes.POINTER(struct_WGPUAdapterImpl), struct_WGPUStringView, ctypes.c_void_p), ctypes.c_void_p)
+WGPUProcInstanceRequestAdapter2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPURequestAdapterOptions), struct_WGPURequestAdapterCallbackInfo2)
+WGPUProcInstanceRequestAdapterF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.POINTER(struct_WGPURequestAdapterOptions), struct_WGPURequestAdapterCallbackInfo)
 WGPUProcInstanceWaitAny = ctypes.CFUNCTYPE(enum_WGPUWaitStatus, ctypes.POINTER(struct_WGPUInstanceImpl), ctypes.c_uint64, ctypes.POINTER(struct_WGPUFutureWaitInfo), ctypes.c_uint64)
 WGPUProcInstanceAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUInstanceImpl))
 WGPUProcInstanceRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUInstanceImpl))
@@ -2627,41 +2414,15 @@ WGPUProcQuerySetGetType = ctypes.CFUNCTYPE(enum_WGPUQueryType, ctypes.POINTER(st
 WGPUProcQuerySetSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQuerySetImpl), struct_WGPUStringView)
 WGPUProcQuerySetAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQuerySetImpl))
 WGPUProcQuerySetRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQuerySetImpl))
-class const_struct_WGPUImageCopyExternalTexture(Struct): pass
-const_struct_WGPUImageCopyExternalTexture._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('externalTexture', WGPUExternalTexture),
-  ('origin', WGPUOrigin3D),
-  ('naturalSize', WGPUExtent2D),
-]
-class const_struct_WGPUCopyTextureForBrowserOptions(Struct): pass
-const_struct_WGPUCopyTextureForBrowserOptions._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('flipY', WGPUBool),
-  ('needsColorSpaceConversion', WGPUBool),
-  ('srcAlphaMode', WGPUAlphaMode),
-  ('srcTransferFunctionParameters', ctypes.POINTER(ctypes.c_float)),
-  ('conversionMatrix', ctypes.POINTER(ctypes.c_float)),
-  ('dstTransferFunctionParameters', ctypes.POINTER(ctypes.c_float)),
-  ('dstAlphaMode', WGPUAlphaMode),
-  ('internalUsage', WGPUBool),
-]
-WGPUProcQueueCopyExternalTextureForBrowser = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(const_struct_WGPUImageCopyExternalTexture), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUExtent3D), ctypes.POINTER(const_struct_WGPUCopyTextureForBrowserOptions))
-WGPUProcQueueCopyTextureForBrowser = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.POINTER(const_struct_WGPUExtent3D), ctypes.POINTER(const_struct_WGPUCopyTextureForBrowserOptions))
+WGPUProcQueueCopyExternalTextureForBrowser = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(struct_WGPUImageCopyExternalTexture), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUExtent3D), ctypes.POINTER(struct_WGPUCopyTextureForBrowserOptions))
+WGPUProcQueueCopyTextureForBrowser = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.POINTER(struct_WGPUExtent3D), ctypes.POINTER(struct_WGPUCopyTextureForBrowserOptions))
 WGPUProcQueueOnSubmittedWorkDone = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.CFUNCTYPE(None, enum_WGPUQueueWorkDoneStatus, ctypes.c_void_p), ctypes.c_void_p)
 WGPUProcQueueOnSubmittedWorkDone2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUQueueImpl), struct_WGPUQueueWorkDoneCallbackInfo2)
 WGPUProcQueueOnSubmittedWorkDoneF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUQueueImpl), struct_WGPUQueueWorkDoneCallbackInfo)
 WGPUProcQueueSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), struct_WGPUStringView)
 WGPUProcQueueSubmit = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(struct_WGPUCommandBufferImpl)))
 WGPUProcQueueWriteBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64, ctypes.c_void_p, ctypes.c_uint64)
-class const_struct_WGPUTextureDataLayout(Struct): pass
-const_struct_WGPUTextureDataLayout._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('offset', uint64_t),
-  ('bytesPerRow', uint32_t),
-  ('rowsPerImage', uint32_t),
-]
-WGPUProcQueueWriteTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(const_struct_WGPUImageCopyTexture), ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(const_struct_WGPUTextureDataLayout), ctypes.POINTER(const_struct_WGPUExtent3D))
+WGPUProcQueueWriteTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl), ctypes.POINTER(struct_WGPUImageCopyTexture), ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(struct_WGPUTextureDataLayout), ctypes.POINTER(struct_WGPUExtent3D))
 WGPUProcQueueAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl))
 WGPUProcQueueRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUQueueImpl))
 WGPUProcRenderBundleSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleImpl), struct_WGPUStringView)
@@ -2671,12 +2432,7 @@ WGPUProcRenderBundleEncoderDraw = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_W
 WGPUProcRenderBundleEncoderDrawIndexed = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_int32, ctypes.c_uint32)
 WGPUProcRenderBundleEncoderDrawIndexedIndirect = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64)
 WGPUProcRenderBundleEncoderDrawIndirect = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.c_uint64)
-class const_struct_WGPURenderBundleDescriptor(Struct): pass
-const_struct_WGPURenderBundleDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-]
-WGPUProcRenderBundleEncoderFinish = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderBundleImpl), ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(const_struct_WGPURenderBundleDescriptor))
+WGPUProcRenderBundleEncoderFinish = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPURenderBundleImpl), ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), ctypes.POINTER(struct_WGPURenderBundleDescriptor))
 WGPUProcRenderBundleEncoderInsertDebugMarker = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), struct_WGPUStringView)
 WGPUProcRenderBundleEncoderPopDebugGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl))
 WGPUProcRenderBundleEncoderPushDebugGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderBundleEncoderImpl), struct_WGPUStringView)
@@ -2702,14 +2458,7 @@ WGPUProcRenderPassEncoderPixelLocalStorageBarrier = ctypes.CFUNCTYPE(None, ctype
 WGPUProcRenderPassEncoderPopDebugGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl))
 WGPUProcRenderPassEncoderPushDebugGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), struct_WGPUStringView)
 WGPUProcRenderPassEncoderSetBindGroup = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.c_uint32, ctypes.POINTER(struct_WGPUBindGroupImpl), ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint32))
-class const_struct_WGPUColor(Struct): pass
-const_struct_WGPUColor._fields_ = [
-  ('r', ctypes.c_double),
-  ('g', ctypes.c_double),
-  ('b', ctypes.c_double),
-  ('a', ctypes.c_double),
-]
-WGPUProcRenderPassEncoderSetBlendConstant = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(const_struct_WGPUColor))
+WGPUProcRenderPassEncoderSetBlendConstant = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(struct_WGPUColor))
 WGPUProcRenderPassEncoderSetIndexBuffer = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(struct_WGPUBufferImpl), enum_WGPUIndexFormat, ctypes.c_uint64, ctypes.c_uint64)
 WGPUProcRenderPassEncoderSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), struct_WGPUStringView)
 WGPUProcRenderPassEncoderSetPipeline = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPURenderPassEncoderImpl), ctypes.POINTER(struct_WGPURenderPipelineImpl))
@@ -2727,22 +2476,14 @@ WGPUProcRenderPipelineRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGP
 WGPUProcSamplerSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSamplerImpl), struct_WGPUStringView)
 WGPUProcSamplerAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSamplerImpl))
 WGPUProcSamplerRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSamplerImpl))
-WGPUProcShaderModuleGetCompilationInfo = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(const_struct_WGPUCompilationInfo), ctypes.c_void_p), ctypes.c_void_p)
+WGPUProcShaderModuleGetCompilationInfo = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUShaderModuleImpl), ctypes.CFUNCTYPE(None, enum_WGPUCompilationInfoRequestStatus, ctypes.POINTER(struct_WGPUCompilationInfo), ctypes.c_void_p), ctypes.c_void_p)
 WGPUProcShaderModuleGetCompilationInfo2 = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUShaderModuleImpl), struct_WGPUCompilationInfoCallbackInfo2)
 WGPUProcShaderModuleGetCompilationInfoF = ctypes.CFUNCTYPE(struct_WGPUFuture, ctypes.POINTER(struct_WGPUShaderModuleImpl), struct_WGPUCompilationInfoCallbackInfo)
 WGPUProcShaderModuleSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUShaderModuleImpl), struct_WGPUStringView)
 WGPUProcShaderModuleAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUShaderModuleImpl))
 WGPUProcShaderModuleRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUShaderModuleImpl))
-class const_struct_WGPUSharedBufferMemoryBeginAccessDescriptor(Struct): pass
-const_struct_WGPUSharedBufferMemoryBeginAccessDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('initialized', WGPUBool),
-  ('fenceCount', size_t),
-  ('fences', ctypes.POINTER(WGPUSharedFence)),
-  ('signaledValues', ctypes.POINTER(uint64_t)),
-]
-WGPUProcSharedBufferMemoryBeginAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(const_struct_WGPUSharedBufferMemoryBeginAccessDescriptor))
-WGPUProcSharedBufferMemoryCreateBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(const_struct_WGPUBufferDescriptor))
+WGPUProcSharedBufferMemoryBeginAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryBeginAccessDescriptor))
+WGPUProcSharedBufferMemoryCreateBuffer = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUBufferDescriptor))
 WGPUProcSharedBufferMemoryEndAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUBufferImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryEndAccessState))
 WGPUProcSharedBufferMemoryGetProperties = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl), ctypes.POINTER(struct_WGPUSharedBufferMemoryProperties))
 WGPUProcSharedBufferMemoryIsDeviceLost = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUSharedBufferMemoryImpl))
@@ -2752,37 +2493,15 @@ WGPUProcSharedBufferMemoryRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct
 WGPUProcSharedFenceExportInfo = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedFenceImpl), ctypes.POINTER(struct_WGPUSharedFenceExportInfo))
 WGPUProcSharedFenceAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedFenceImpl))
 WGPUProcSharedFenceRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedFenceImpl))
-class const_struct_WGPUSharedTextureMemoryBeginAccessDescriptor(Struct): pass
-const_struct_WGPUSharedTextureMemoryBeginAccessDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('concurrentRead', WGPUBool),
-  ('initialized', WGPUBool),
-  ('fenceCount', size_t),
-  ('fences', ctypes.POINTER(WGPUSharedFence)),
-  ('signaledValues', ctypes.POINTER(uint64_t)),
-]
-WGPUProcSharedTextureMemoryBeginAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(const_struct_WGPUSharedTextureMemoryBeginAccessDescriptor))
-WGPUProcSharedTextureMemoryCreateTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(const_struct_WGPUTextureDescriptor))
+WGPUProcSharedTextureMemoryBeginAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryBeginAccessDescriptor))
+WGPUProcSharedTextureMemoryCreateTexture = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUTextureDescriptor))
 WGPUProcSharedTextureMemoryEndAccess = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryEndAccessState))
 WGPUProcSharedTextureMemoryGetProperties = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), ctypes.POINTER(struct_WGPUSharedTextureMemoryProperties))
 WGPUProcSharedTextureMemoryIsDeviceLost = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl))
 WGPUProcSharedTextureMemorySetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl), struct_WGPUStringView)
 WGPUProcSharedTextureMemoryAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl))
 WGPUProcSharedTextureMemoryRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSharedTextureMemoryImpl))
-class const_struct_WGPUSurfaceConfiguration(Struct): pass
-const_struct_WGPUSurfaceConfiguration._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('device', WGPUDevice),
-  ('format', WGPUTextureFormat),
-  ('usage', WGPUTextureUsage),
-  ('viewFormatCount', size_t),
-  ('viewFormats', ctypes.POINTER(WGPUTextureFormat)),
-  ('alphaMode', WGPUCompositeAlphaMode),
-  ('width', uint32_t),
-  ('height', uint32_t),
-  ('presentMode', WGPUPresentMode),
-]
-WGPUProcSurfaceConfigure = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(const_struct_WGPUSurfaceConfiguration))
+WGPUProcSurfaceConfigure = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(struct_WGPUSurfaceConfiguration))
 WGPUProcSurfaceGetCapabilities = ctypes.CFUNCTYPE(enum_WGPUStatus, ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(struct_WGPUAdapterImpl), ctypes.POINTER(struct_WGPUSurfaceCapabilities))
 WGPUProcSurfaceGetCurrentTexture = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl), ctypes.POINTER(struct_WGPUSurfaceTexture))
 WGPUProcSurfacePresent = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl))
@@ -2790,21 +2509,8 @@ WGPUProcSurfaceSetLabel = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfa
 WGPUProcSurfaceUnconfigure = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl))
 WGPUProcSurfaceAddRef = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl))
 WGPUProcSurfaceRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUSurfaceImpl))
-class const_struct_WGPUTextureViewDescriptor(Struct): pass
-const_struct_WGPUTextureViewDescriptor._fields_ = [
-  ('nextInChain', ctypes.POINTER(WGPUChainedStruct)),
-  ('label', WGPUStringView),
-  ('format', WGPUTextureFormat),
-  ('dimension', WGPUTextureViewDimension),
-  ('baseMipLevel', uint32_t),
-  ('mipLevelCount', uint32_t),
-  ('baseArrayLayer', uint32_t),
-  ('arrayLayerCount', uint32_t),
-  ('aspect', WGPUTextureAspect),
-  ('usage', WGPUTextureUsage),
-]
-WGPUProcTextureCreateErrorView = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureViewImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(const_struct_WGPUTextureViewDescriptor))
-WGPUProcTextureCreateView = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureViewImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(const_struct_WGPUTextureViewDescriptor))
+WGPUProcTextureCreateErrorView = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureViewImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUTextureViewDescriptor))
+WGPUProcTextureCreateView = ctypes.CFUNCTYPE(ctypes.POINTER(struct_WGPUTextureViewImpl), ctypes.POINTER(struct_WGPUTextureImpl), ctypes.POINTER(struct_WGPUTextureViewDescriptor))
 WGPUProcTextureDestroy = ctypes.CFUNCTYPE(None, ctypes.POINTER(struct_WGPUTextureImpl))
 WGPUProcTextureGetDepthOrArrayLayers = ctypes.CFUNCTYPE(ctypes.c_uint32, ctypes.POINTER(struct_WGPUTextureImpl))
 WGPUProcTextureGetDimension = ctypes.CFUNCTYPE(enum_WGPUTextureDimension, ctypes.POINTER(struct_WGPUTextureImpl))