more complex

also added Tensor.shard_like

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

@@ -70,13 +70,13 @@ runs:
       uses: actions/cache@v4
       with:
         path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
     - name: Cache downloads (macOS)
       if: inputs.key != '' && runner.os == 'macOS'
       uses: actions/cache@v4
       with:
         path: ~/Library/Caches/tinygrad/downloads/
-        key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
 
     # **** Python deps ****
 
@@ -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

@@ -14,12 +14,6 @@ on:
       - update_benchmark
       - update_benchmark_staging
   workflow_dispatch:
-    inputs:
-      run_process_replay:
-        description: "Run process replay tests"
-        required: false
-        default: false
-        type: boolean
 
 jobs:
   testmacbenchmark:
@@ -124,14 +118,6 @@ jobs:
     # TODO: too slow
     # - name: Run 10 CIFAR training steps w winograd
     #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
-    - name: UsbGPU copy speeds
-      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
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -165,6 +151,37 @@ jobs:
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py
 
+  testusbgpu:
+    name: UsbGPU Benchmark
+    env:
+      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 10
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: UsbGPU boot time
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU tiny tests
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+    - name: UsbGPU copy speeds
+      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
+
   testnvidiabenchmark:
     name: tinybox green Benchmark
     runs-on: [self-hosted, Linux, tinyboxgreen]
@@ -342,7 +359,7 @@ jobs:
       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)
@@ -537,8 +554,6 @@ jobs:
       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)
@@ -550,7 +565,6 @@ jobs:
           train_cifar_wino.txt
           train_cifar_one_gpu.txt
           train_cifar_six_gpu.txt
-          train_cifar_six_gpu_remote.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -594,7 +608,7 @@ jobs:
       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)
@@ -629,6 +643,8 @@ jobs:
       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

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '14'
+  CACHE_VERSION: '15'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -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:
@@ -310,7 +310,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -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
 
@@ -721,71 +721,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  amdremote:
-    name: Linux (remote)
-    runs-on: ubuntu-22.04
-    timeout-minutes: 20
-    env:
-      REMOTE: 1
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: linux-remote
-          deps: testing_minimal
-          amd: 'true'
-          llvm: 'true'
-          opencl: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "CL" 7667
-          start_server "remote-server-cpu" "CPU" 8667
-      - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
-      - name: Run REMOTE=1 Test (CL)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
-
 # ****** OSX Tests ******
 
   testmetal:
@@ -883,30 +818,6 @@ jobs:
     - name: Test ONNX Runner (WEBGPU)
       run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py
 
-  osxremote:
-   name: MacOS (remote metal)
-   runs-on: macos-15
-   timeout-minutes: 10
-   env:
-     REMOTE: 1
-     REMOTEDEV: METAL
-   steps:
-     - name: Checkout Code
-       uses: actions/checkout@v4
-     - name: Setup Environment
-       uses: ./.github/actions/setup-tinygrad
-       with:
-         key: macos-remote
-         deps: testing_minimal
-     - name: Check Device.DEFAULT and print some source
-       run: |
-         python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-         python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
-         DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-     - name: Run REMOTE=1 Test
-       run: |
-         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
-
   osxtests:
     strategy:
       fail-fast: false
@@ -972,3 +883,33 @@ jobs:
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
+
+# ****** Compile-only Tests ******
+
+  compiletests:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [ir3, nak]
+    name: Compile-only (${{ matrix.backend }})
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-${{ matrix.backend }}
+          deps: testing_minimal
+          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
+          python-version: '3.14'
+      - name: Set env
+        shell: bash
+        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/test_ops.py --durations=20

CLAUDE.md

@@ -0,0 +1,227 @@
+# Claude Code Guide for tinygrad
+
+## Architecture Overview
+
+tinygrad compiles tensor operations into optimized kernels. The pipeline:
+
+1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
+2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
+3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
+4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
+5. **Runtime** (`runtime/`) - Device-specific execution
+
+## Key Concepts
+
+### UOp (Universal Operation)
+Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
+- `op`: The operation type (Ops enum)
+- `dtype`: Data type
+- `src`: Tuple of source UOps
+- `arg`: Operation-specific argument
+- `tag`: Optional tag for graph transformations
+
+UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
+
+### PatternMatcher
+Used extensively for graph transformations:
+```python
+pm = PatternMatcher([
+  (UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
+])
+result = graph_rewrite(uop, pm)
+```
+
+### Schedule Cache
+Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
+
+## Directory Structure
+
+```
+tinygrad/
+├── tensor.py          # Tensor class, user API
+├── device.py          # Buffer, device management
+├── dtype.py           # Data types
+├── helpers.py         # Utilities, environment vars
+├── uop/
+│   ├── ops.py         # UOp class, Ops enum, PatternMatcher
+│   ├── spec.py        # UOp type verification
+│   └── symbolic.py    # Symbolic math simplification
+├── engine/
+│   ├── schedule.py    # Schedule creation, caching
+│   ├── realize.py     # Tensor realization
+│   ├── jit.py         # JIT compilation
+│   └── memory.py      # Memory planning
+├── schedule/
+│   ├── rangeify.py    # Convert movements to ranges
+│   └── indexing.py    # Index calculations
+├── codegen/
+│   ├── kernel.py      # Kernel optimization
+│   └── uopgraph.py    # UOp graph transformations
+├── renderer/          # Code generation (CUDA, Metal, etc.)
+└── runtime/           # Device backends
+```
+
+## Testing
+
+```bash
+# Run specific test
+python -m pytest test/unit/test_schedule_cache.py -xvs
+
+# Run with timeout
+python -m pytest test/test_symbolic_ops.py -x --timeout=60
+
+# Debug with print
+DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
+
+# Visualize UOp graphs
+VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
+```
+
+## Common Environment Variables
+
+- `DEBUG=1-4` - Increasing verbosity
+- `VIZ=1` - Enable graph visualization
+- `SPEC=1` - Enable UOp spec verification
+- `NOOPT=1` - Disable optimizations
+- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
+
+## Debugging Tips
+
+1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
+2. **Check schedule**: `tensor.schedule()` returns list of ScheduleItems
+3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
+4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
+
+## Workflow Rules
+
+- **NEVER commit without explicit user approval** - always show the diff and wait for approval
+- **NEVER amend commits** - always create a new commit instead
+- Run `pre-commit run --all-files` before committing to catch linting/type errors
+- Run tests before proposing commits
+- Test with `SPEC=2` when modifying UOp-related code
+
+## Style Notes
+
+- 2-space indentation, 150 char line limit
+- PatternMatchers should be defined at module level (slow to construct)
+- Prefer `graph_rewrite` over manual graph traversal
+- UOp methods like `.replace()` preserve tags unless explicitly changed
+- Use `.rtag(value)` to add tags to UOps
+
+## Lessons Learned
+
+### UOp ucache Behavior
+UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
+- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
+- Two UOps with same structure are identical (`is` comparison works)
+
+### Spec Validation
+When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
+```bash
+SPEC=2 python3 test/unit/test_something.py
+```
+Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
+
+### Schedule Cache Key Normalization
+The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
+- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
+- `pm_post_sched_cache`: restores original BIND from context
+- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
+- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
+
+### Avoiding Extra Work
+- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
+- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
+- PatternMatchers are slow to construct - define at module level, not in functions
+
+### Readability Over Speed
+Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
+```python
+# BAD: "optimized" with extra complexity
+if has_afters:  # skip toposort if no AFTERs
+  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
+
+# GOOD: simple, always works
+after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
+```
+The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
+
+### Testing LLM Changes
+```bash
+# Quick smoke test
+echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
+
+# Check cache hits (should see "cache hit" after warmup)
+echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
+
+# Test with beam search
+echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
+```
+
+## Common Patterns
+
+### Graph Transformation
+```python
+def my_transform(ctx, x):
+  # Return new UOp or None to skip
+  return x.replace(arg=new_arg)
+
+pm = PatternMatcher([
+  (UPat(Ops.SOMETHING, name="x"), my_transform),
+])
+result = graph_rewrite(input_uop, pm, ctx={})
+```
+
+### Finding Variables
+```python
+# Get all variables in a UOp graph
+variables = uop.variables()
+
+# Get bound variable values
+var, val = bind_uop.unbind()
+```
+
+### Shape Handling
+```python
+# Shapes can be symbolic (contain UOps)
+shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
+```
+
+## Performance Optimization
+
+When optimizing tinygrad internals:
+
+1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
+
+2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
+
+3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
+
+4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
+
+5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
+
+6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
+
+7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
+
+8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
+
+9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
+
+## Pattern Matching Profiling
+
+Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
+
+```bash
+TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
+```
+
+Output format: `matches / attempts -- match_time / total_time ms -- location`
+
+Key patterns to watch (from ResNet50 benchmark):
+- `split_load_store`: ~146ms, 31% match rate - does real work
+- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
+- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
+
+Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.

examples/mlperf/helpers.py

@@ -223,13 +223,13 @@ def get_mlperf_bert_model():
 
 def get_fake_data_bert(BS:int):
   return {
-    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
-    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
+    "input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
+    "input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
+    "segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
+    "masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
+    "masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
+    "masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
+    "next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
   }
 
 def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:

examples/mlperf/initializers.py

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

examples/mlperf/model_train.py

@@ -918,40 +918,6 @@ def train_rnnt():
   # TODO: RNN-T
   pass
 
-@TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
-  optimizer.zero_grad()
-
-  for i in range(grad_acc):
-    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
-    # NOTE: these two have different names
-    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
-    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
-
-    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-      else: t.to_(GPUS[0])
-
-    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-    (loss * loss_scaler).backward()
-    # TODO: OOM without this realize with large grad_acc
-    Tensor.realize(*[p.grad for p in optimizer.params])
-
-  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
-  for p in optimizer.params:
-    p.grad = p.grad / loss_scaler
-    global_norm += p.grad.float().square().sum()
-  global_norm = global_norm.sqrt().contiguous()
-  for p in optimizer.params:
-    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
-
-  optimizer.step()
-  scheduler.step()
-  # TODO: no to("CPU") here because it blocks and messes the python time
-  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
-  return loss, global_norm, optimizer.optimizers[0].lr
-
 @TinyJit
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                    masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
@@ -1014,7 +980,8 @@ def train_bert():
   # ** hyperparameters **
   BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  # TODO: mlperf logging
+  # TODO: implement grad accumulation + mlperf logging
+  assert grad_acc == 1
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
   max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
@@ -1073,8 +1040,8 @@ def train_bert():
 
   # ** Optimizer **
   parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
-  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  parameters_wd = [x for x in parameters if x not in set(parameters_no_wd)]
+  optimizer_wd = LAMB(parameters_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
   optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
   optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)
 
@@ -1131,12 +1098,38 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
+  i, train_data = start_step, next(train_it)
 
   if RUNMLPERF:
     if MLLOGGER:
       MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
 
+  @TinyJit
+  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
+                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+    optimizer_group.zero_grad()
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+
+    global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer_group[0].device)
+    for p in optimizer_group.params:
+      p.grad = p.grad / loss_scaler
+      global_norm += p.grad.float().square().sum()
+    global_norm = global_norm.sqrt().contiguous()
+    for p in optimizer_group.params:
+      p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
+
+    optimizer_group.step()
+    scheduler_group.step()
+    # TODO: no to("CPU") here because it blocks and messes the python time
+    Tensor.realize(loss, global_norm, optimizer_group.optimizers[0].lr)
+    return loss, global_norm, optimizer_group.optimizers[0].lr
+
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
       Tensor.training = True
@@ -1144,16 +1137,12 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
-        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)
+        loss, global_norm, lr = train_step_bert(
+          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
+          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
 
         pt = time.perf_counter()
-
-        try:
-          next_data = [next(train_it) for _ in range(grad_acc)]
-        except StopIteration:
-          next_data = None
-
+        next_data = next(train_it)
         dt = time.perf_counter()
 
         device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
@@ -1188,8 +1177,8 @@ def train_bert():
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
       if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+      elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
+        # TODO: this hangs on tiny green after 90 minutes of training
         train_step_bert.captured.free_intermediates()
       eval_lm_losses = []
       eval_clsf_losses = []
@@ -1224,7 +1213,7 @@ def train_bert():
           return
 
       if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE") and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
 
       del eval_data
       avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1300,6 +1289,7 @@ def train_llama3():
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  assert grad_acc == 1, f"{grad_acc=} is not supported"
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1324,12 +1314,14 @@ def train_llama3():
   opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
   opt_end_learning_rate = getenv("END_LR", 8e-7)
 
-  # TODO: confirm weights are in bf16
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  if not SMALL: model_params |= {"vocab_size": 32000}
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  params = get_parameters(model)
+  # weights are all bfloat16 for now
+  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
 
   if getenv("FAKEDATA"):
     for v in get_parameters(model):
@@ -1374,20 +1366,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
@@ -1422,18 +1411,18 @@ def train_llama3():
   # ** data iters **
   def fake_data(bs, samples):
     for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
 
   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 +1440,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

extra/amdpci/hive_reset.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+
+from tinygrad.helpers import Context
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.am.amdev import AMDev
+
+if __name__ == "__main__":
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1])])
+  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  amdevs = []
+  with Context(DEBUG=2):
+    for pcidev in pcidevs:
+      amdevs.append(AMDev(pcidev, reset_mode=True))
+    for amdev in amdevs: amdev.smu.mode1_reset()

extra/amdpci/proclogs.py

@@ -1,48 +1,65 @@
 import re, ctypes, sys, importlib
+from tinygrad.helpers import getenv
 
 from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+
+class GFXFake:
+  def __init__(self): self.xccs = 8
+
 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()
 
+    self.gfx = GFXFake()
+
 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 +68,4 @@ if __name__ == '__main__':
   if len(sys.argv) != 3:
     print("Usage: <input_file_path> <output_file_path>")
     sys.exit(1)
-
-  main()
+  main()

extra/models/bert.py

@@ -242,7 +242,8 @@ class BertIntermediate:
   def __call__(self, hidden_states):
     x = self.dense(hidden_states)
     # tinygrad gelu is openai gelu but we need the original bert gelu
-    return gelu(x)
+    # NOTE: contiguous for speed
+    return gelu(x).contiguous()
 
 class BertAttention:
   def __init__(self, hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob):

extra/sqtt/roc.py

@@ -158,14 +158,13 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   t.join()
   return ROCParseCtx
 
-def print_pmc(ev:ProfilePMCEvent) -> None:
-  ptr = 0
-  view = memoryview(ev.blob).cast('Q')
-  for s in ev.sched:
-    print(f"\t{s.name}")
-    for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-      print(f"\t\tXCC {xcc} Inst {inst:<2} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
-      ptr += 1
+def print_pmc(events:list[ProfilePMCEvent]) -> None:
+  from tinygrad.viz.serve import unpack_pmc
+  from tabulate import tabulate
+  for e in events:
+    print("**", e.kern)
+    data = unpack_pmc(e)
+    print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
@@ -176,7 +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}")
-    print_pmc(ev)
+  print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])

extra/sqtt/test_pmc.py

@@ -6,8 +6,11 @@ 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.dtype import dtypes, AddrSpace
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.runtime.ops_amd import ProfilePMCEvent
+from tinygrad.engine.realize import get_runner
+from tinygrad.viz.serve import unpack_pmc
 from extra.sqtt.roc import print_pmc
 
 def copy_kernel(B, A, stride=1):
@@ -19,6 +22,16 @@ def copy_kernel(B, A, stride=1):
   index = (i * stride) % A.size
   return B[index].store(A[index]).sink(arg=KernelInfo(name=f"copy_{A.size}_stride_{stride}", opts_to_apply=()))
 
+def lds_kernel(offset:UOp, size:int, inst:str) -> UOp:
+  tid = UOp.range(offset.size, 0, AxisType.LOCAL)
+  dst = UOp.placeholder((size,), dtypes.float32, 1, AddrSpace.REG)
+  #lds = UOp.placeholder((1024,), dtypes.float32, 2, AddrSpace.LOCAL)
+  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_waitcnt(0);')
+  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_barrier();', src=(u,))
+  u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_sched_barrier(0);', src=(u,))
+  u = UOp(Ops.CUSTOM, arg=f'asm volatile("{inst} '+'%0, %1" : "=v"({0}) : "v"({1}));', src=(dst, offset[tid], u))
+  return UOp.sink(u, arg=KernelInfo(name="test_lds", opts_to_apply=()))
+
 dev = Device[Device.DEFAULT]
 
 @contextlib.contextmanager
@@ -40,10 +53,35 @@ class TestPMC(unittest.TestCase):
     b = Tensor.custom_kernel(b, a, fxn=functools.partial(copy_kernel, stride=stride))[0]
     with save_pmc() as pmc:
       b.realize()
-    print_pmc(pmc[0])
+    print_pmc(pmc)
     np.testing.assert_equal(a.numpy(), b.numpy())
 
   def test_copy_uncoalesced(self): return self.test_copy(stride=17)
 
+  # test with two threads issuing ds_reads at different offsets
+  def test_ds_read(self, size=1, inst='ds_read_b32'):
+    test_banks = 256
+    offsets = [Tensor([0, b*4]) for b in range(1, test_banks)]
+    with Context(DEBUG=0): Tensor.realize(*offsets)
+    k = Tensor.custom_kernel(offsets[0], fxn=functools.partial(lds_kernel, size=size, inst=inst))[0]
+    # sample all kernels
+    with save_pmc() as pmc_events:
+      runner = get_runner(Device.DEFAULT, k.schedule()[0].ast)
+      # TODO: llvm eliminates lds definition from the ELF, is there another way to pin lds size?
+      runner._prg.group_segment_size = 1024
+      for offset in offsets: runner([offset.uop.buffer])
+    # find read offsets that created bank conflicts from the pmc counters
+    found:list[Tensor] = []
+    for i,e in enumerate(pmc_events):
+      pmc = unpack_pmc(e)["rows"]
+      # SQ on gfx9, renamed to SQC after gfx10
+      val = next(total for name,total,_all_instances in pmc if name in {"SQ_LDS_BANK_CONFLICT", "SQC_LDS_BANK_CONFLICT"})
+      if val > 0: found.append(offsets[i])
+    print("Found bank conflicts at offsets:", [s.numpy() for s in found])
+
+  def test_ds_read_b64(self): self.test_ds_read(2, 'ds_read_b64')
+
+  def test_ds_read_b128(self): self.test_ds_read(4, 'ds_read_b128')
+
 if __name__ == "__main__":
   unittest.main()

extra/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], track=False):
+        src_i += outer * reductions + (laneid % reductions)
+
+        src_load = srcf[src_i]
+        if src.dtype.base != dst.dtype.base:
+          src_load = src_load.cast(dst.dtype.base)
+        dst_store = dst[outer, 0].store(src_load).end(outer)
     else:
-      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+      raise NotImplementedError(f"load from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(dst)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
@@ -380,8 +403,30 @@ class Group:
   def store(self, dst:ALL_TILES, src:ALL_TILES, idxs:tuple[UOp|int,...]=(), src_idxs:tuple[UOp|int,...]=(), axis:int=0):
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
-    dst_dtype, src_dtype = cast(PtrDType, dst.dtype), cast(PtrDType, src.dtype)
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL:
+    dst_dtype, src_dtype = dst.dtype, src.dtype
+    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.LOCAL:
+      laneid = self.ker.laneid
+      st, rt = cast(ST, dst), cast(RT, src)
+      elements_per_thread = rt.base_shape.elements_per_thread
+
+      for height in self.ker.range(src.shape[-3], track=False):
+        for width in self.ker.range(src.shape[-2], track=False):
+          for inner in self.ker.range(elements_per_thread, track=False):
+            if rt.layout != st.layout:
+              row = rt.base_shape.stride * (laneid // rt.base_shape.cols) + inner
+              col = laneid % rt.base_shape.cols
+            else:
+              row = laneid % rt.base_shape.rows
+              col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
+
+            srow, scol = cast(ST, dst).swizzle(row, col)
+
+            src_load = src[*src_idxs, height, width, inner]
+            if src.dtype.base != dst.dtype.base:
+              src_load = src_load.cast(dst.dtype.base)
+            dst_store = dst[*idxs[:-2], height, width, srow, scol].store(src_load)
+            dst_store = dst_store.end(height, width, inner)
+    elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
       dstf = dst.flatten()
       row_stride = prod(dst.shape[axis+1:])
 
@@ -414,8 +459,28 @@ class Group:
             if src.dtype.base != dst.dtype.base:
               src_load = src_load.cast(dst.dtype.base)
             dst_store = dstf[dst_i].store(src_load).end(height, width, inner)
+    elif src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RV):
+      dstf = dst.flatten()
+      row_stride = prod(dst.shape[axis+1:])
+
+      laneid = self.ker.laneid
+      rv = cast(RV, src)
+      reductions = rv.base_shape.rows
+
+      assert rv.layout == VecLayout.ORTHO, "only ortho layout supported"
+
+      idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
+      dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
+
+      for outer in self.ker.range(src.shape[-2]):
+        dst_i += outer * reductions + (laneid % reductions)
+
+        src_load = src[outer, 0]
+        if src.dtype.base != dst.dtype.base:
+          src_load = src_load.cast(dst.dtype.base)
+        dst_store = dstf[dst_i].store(src_load).end(outer)
     else:
-      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented")
+      raise NotImplementedError(f"store from {src_dtype.addrspace} to {dst_dtype.addrspace} not implemented for {type(src)=}")
 
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)

extra/thunder/tiny/tk/kernel.py

@@ -25,14 +25,14 @@ class Kernel(AbstractContextManager):
     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.range_id = 0
-    self.allocs = {}
+    self.allocs: dict[tuple[str, tuple], UOp] = {}
 
   @property
   def warpid(self): return self.threadIdx_x // WARP_THREADS
@@ -82,18 +82,27 @@ class Kernel(AbstractContextManager):
 
   def push_store(self, store:UOp, uop:UOp): self.store_stack.append((store, uop))
 
-  def finish(self):
+  def finish(self, stores:int=1):
     # end all ranges
     rngs = []
     while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
 
-    last_store = self.store_stack.pop()[0]
-    if hasattr(last_store, '_uop'): uop = last_store._uop
-    else: uop = last_store
+    # end stores stores
+    store_uops = []
+    for _ in range(stores):
+      store = self.store_stack.pop()[0]
+      if hasattr(store, '_uop'): store_uops.append(store._uop)
+      else: store_uops.append(store)
+    uop = UOp.group(*store_uops)
 
     return uop.end(*rngs).sink(arg=KernelInfo(name=self.name, opts_to_apply=())).simplify()
 
-  def endrange(self):
+  def endrange(self, ranges:int=1):
     last_store = self.store_stack.pop()
-    last_range = self.range_stack.pop()
-    return last_store[1].after(last_store[0].end(last_range._rng)).reshape(last_store[1].shape)
+
+    rngs = []
+    for _ in range(ranges):
+      last_range = self.range_stack.pop()
+      rngs.append(last_range._rng)
+
+    return last_store[1].after(last_store[0].end(*rngs)).reshape(last_store[1].shape)

extra/thunder/tiny/tk/tiles.py

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

extra/usbgpu/patch.py

@@ -3,6 +3,13 @@ import sys, os, zlib, struct, hashlib
 from tinygrad.helpers import DEBUG, getenv, fetch
 from tinygrad.runtime.support.usb import USB3
 
+SUPPORTED_CONTROLLERS = [
+  (0x174C, 0x2464),
+  (0x174C, 0x2463),
+  (0xADD1, 0x0001),
+]
+if getenv("USBDEV", ""): SUPPORTED_CONTROLLERS.insert(0, (int(x, 16) for x in getenv("USBDEV", "").split(":")))
+
 def patch(input_filepath, file_hash, patches):
   with open(input_filepath, 'rb') as infile: data = bytearray(infile.read())
 
@@ -40,10 +47,14 @@ if not os.path.exists(file_path):
 patches = [(0x2a0d + 1 + 4, b'\x0a', b'\x05')]
 patched_fw = patch(file_path, file_hash, patches)
 
-vendor, device = [int(x, base=16) for x in getenv("USBDEV", "174C:2464").split(":")]
-try: dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
-except RuntimeError as e:
-  raise RuntimeError(f'{e}. You can set USBDEV environment variable to your device\'s vendor and device ID (e.g., USBDEV="174C:2464")') from e
+dev = None
+for vendor, device in SUPPORTED_CONTROLLERS:
+  try:
+    dev = USB3(vendor, device, 0x81, 0x83, 0x02, 0x04)
+    break
+  except RuntimeError: pass
+if dev is None:
+  raise RuntimeError('Could not open controller. You can set USBDEV environment variable to your device\'s vendor and device ID (e.g., USBDEV="174C:2464")')
 
 config1 = bytes([
   0xFF, 0xFF, 0xFF, 0xFF, 0x41, 0x41, 0x41, 0x41, 0x42, 0x42, 0x42, 0x42, 0x30, 0x30, 0x36, 0x30,

extra/weekly_commits_table.py

@@ -1,7 +1,7 @@
 # extra/weekly_commits_table.py
 import os, subprocess, datetime as dt
 
-NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot"]
+NAMES = ["chenyu","George Hotz","nimlgen","qazal","wozeparrot","Christopher Milan"]
 REPO  = os.environ.get("REPO_PATH",".")
 today = dt.date.today()
 days  = [(today - dt.timedelta(i)).strftime("%Y-%m-%d") for i in range(6,-1,-1)]

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

@@ -8,7 +8,7 @@ def multidevice_test(fxn):
   def ret(self):
     for device in Device._devices:
       # broken on OSX USB AMD, why?
-      if device in ["REMOTE", "DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
+      if device in ["DISK", "NPY", "FAKE", "DSP", "NULL"] or (OSX and device in ["AMD"]): continue
       if not CI: print(device)
       if device in exclude_devices:
         if not CI: print(f"WARNING: {device} test is excluded")

test/external/external_test_onnx_backend.py

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

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

@@ -69,5 +69,4 @@ def needs_second_gpu(fn):
     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)
+REAL_DEV = Device.DEFAULT

test/models/test_onnx.py

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

test/models/test_real_world.py

@@ -20,7 +20,7 @@ global_mem_used = 0
 def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jitted=False):
   with Context(JIT=2):
     tms = []
-    for _ in range(4):
+    for _ in range(2 if CI else 4):
       early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
       GlobalCounters.reset()
       Device[Device.DEFAULT].synchronize()

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

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp, Context
+from tinygrad import Tensor, UOp
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import KernelInfo, AxisType
 
@@ -117,6 +117,17 @@ class TestCustomKernel(unittest.TestCase):
     out = c.flatten().tolist()
     assert all(x == 2 for x in out), "all 2"
 
+  def test_simple_sharded(self):
+    devs = ("CPU:0", "CPU:1")
+
+    a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
+    b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
+    # ugly construction to get a sharded empty tensor
+    c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
+    c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
+    out = c.flatten().tolist()
+    assert all(x == 2 for x in out), "all 2"
+
   def test_multioutput(self):
     a = Tensor.full((16, 16), 3.).contiguous()
     b = Tensor.full((16, 16), 3.).contiguous()
@@ -184,7 +195,6 @@ class TestCustomKernel(unittest.TestCase):
 
   def test_gemm_backward_custom(self): self.test_gemm_backward(True)
   # NOTE: grad_fxn doesn't work with pyrender
-  @Context(SPEC=1)
   def test_gemm_backward(self, custom_backward_gemm=False):
     N = 4
     a_rand = Tensor.randn(N, 8)

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

@@ -501,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_jit_footguns.py

@@ -0,0 +1,270 @@
+#!/usr/bin/env python
+"""
+JIT Footguns: Documenting unexpected behavior changes when using @TinyJit
+
+Each test shows behavior that works without JIT but changes with JIT.
+Comments marked "should be X!" indicate the intuitively expected value.
+
+SILENT MISMATCHES (highest priority - wrong results, no error):
+  tensors_in_containers_ignored      EASY   only checks t.__class__ is Tensor, could scan lists/dicts
+  non_tensor_outputs_frozen          EASY   could warn/error if return contains non-Tensor values
+  class_method_shared_across_instances EASY could check if first arg is self and warn
+  output_buffer_reuse                MED    performance tradeoff, could add option or better docs
+  python_constants_frozen            HARD   inherent to tracing JITs
+  conditional_branches_frozen        HARD   inherent to tracing JITs
+
+ERRORS RAISED (lower priority - at least users know):
+  positional_kwargs_cannot_mix       EASY   normalize positional args to kwargs using function signature
+  duplicate_inputs_fail              MED    would need to handle aliasing in input_replace
+  nested_jit_fails_on_second_call    MED    could fail on first call instead of second
+"""
+import unittest
+import numpy as np
+from tinygrad import Tensor, TinyJit
+
+class TestJitFootguns(unittest.TestCase):
+
+  def test_output_buffer_reuse(self):
+    """Output tensors share buffer after capture - old references get overwritten."""
+    @TinyJit
+    def f(x): return x.sum().realize()
+
+    r1 = f(Tensor([1, 1]))  # warmup
+    r2 = f(Tensor([2, 2]))  # capture
+    r3 = f(Tensor([3, 3]))  # jit exec
+
+    self.assertEqual(r1.item(), 2)  # warmup result independent
+    self.assertEqual(r3.item(), 6)  # latest is correct
+    self.assertEqual(r2.item(), 6)  # should be 4! (overwritten by r3)
+
+  def test_output_buffer_workaround(self):
+    """Use .clone().realize() to get independent copies."""
+    @TinyJit
+    def f(x): return x.sum().realize()
+
+    r1 = f(Tensor([1, 1])).clone().realize()
+    r2 = f(Tensor([2, 2])).clone().realize()
+    r3 = f(Tensor([3, 3])).clone().realize()
+
+    self.assertEqual([r1.item(), r2.item(), r3.item()], [2, 4, 6])
+
+  def test_non_tensor_outputs_frozen(self):
+    """Non-tensor return values are frozen at capture time."""
+    @TinyJit
+    def f(x, mult): return (x * 2).realize(), mult * 10
+
+    # collect results, copying tensor values immediately (buffer reuse!)
+    results = []
+    for i in range(5):
+      t, s = f(Tensor([i]), i)
+      results.append((t.item(), s))
+
+    # tensor outputs work correctly
+    self.assertEqual([r[0] for r in results[2:]], [4, 6, 8])
+    # scalar outputs frozen at capture (i=1) - should be 20, 30, 40!
+    self.assertEqual([r[1] for r in results[2:]], [10, 10, 10])
+
+  def test_duplicate_inputs_fail(self):
+    """JIT cannot handle the same tensor passed as multiple arguments."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    x = Tensor([1, 2, 3])
+    with self.assertRaises(AssertionError):
+      f(x, x)
+
+  def test_tensors_in_containers_ignored(self):
+    """Tensors inside lists/dicts are not tracked as inputs."""
+    @TinyJit
+    def f(a, arr): return (a + arr[0]).realize()
+
+    results = []
+    for i in range(4):
+      a, b = Tensor([1, 1, 1]).realize(), Tensor([i, i, i]).realize()
+      results.append(f(a, [b]).numpy().copy())
+
+    np.testing.assert_array_equal(results[0], [1, 1, 1])  # warmup
+    np.testing.assert_array_equal(results[1], [2, 2, 2])  # capture
+    np.testing.assert_array_equal(results[2], [2, 2, 2])  # should be [3,3,3]!
+    np.testing.assert_array_equal(results[3], [2, 2, 2])  # should be [4,4,4]!
+
+  def test_nested_jit_fails_on_second_call(self):
+    """Nested JIT works on first call but fails on second."""
+    @TinyJit
+    def inner(t): return t + 1
+    @TinyJit
+    def outer(t): return inner(t) * 3
+
+    self.assertEqual(outer(Tensor([1])).realize().item(), 6)  # works!
+    with self.assertRaises(RuntimeError):
+      outer(Tensor([2])).realize()  # fails
+
+  def test_implicit_inputs_need_realize(self):
+    """Closure tensors must be realized before JIT call."""
+    x = Tensor([0])
+
+    @TinyJit
+    def f(): return (x * 2).realize()
+
+    for i in range(5):
+      x.assign(Tensor([i])).realize()  # must realize!
+      self.assertEqual(f().item(), i * 2)
+
+  def test_views_with_different_offsets_fail(self):
+    """JIT requires consistent tensor views across calls."""
+    @TinyJit
+    def f(a): return (a + 1).realize()
+
+    base = Tensor.randn(10, 10).realize()
+    with self.assertRaises(AssertionError):
+      for i in range(1, 5):
+        f(base[:, i:i+2])  # different offset each time
+
+  def test_shape_change_after_capture_fails(self):
+    """Shapes are locked at capture time."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # warmup
+    f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # capture
+
+    with self.assertRaises(AssertionError):
+      f(Tensor.randn(20, 20), Tensor.randn(20, 20))
+
+  def test_python_constants_frozen(self):
+    """Python variables inside JIT use capture-time values."""
+    mult = 1
+
+    @TinyJit
+    def f(x): return (x * mult).realize()
+
+    results = []
+    for i in range(5):
+      mult = i + 1
+      results.append(f(Tensor([10])).item())
+
+    self.assertEqual(results[0], 10)   # warmup, mult=1
+    self.assertEqual(results[1], 20)   # capture, mult=2
+    self.assertEqual(results[2], 20)   # should be 30!
+    self.assertEqual(results[3], 20)   # should be 40!
+
+  def test_conditional_branches_frozen(self):
+    """Only the branch taken during capture runs thereafter."""
+    @TinyJit
+    def f(x, use_square):
+      if use_square:
+        return (x * x).realize()
+      return (x * 2).realize()
+
+    f(Tensor([3]), True)   # warmup
+    f(Tensor([3]), False)  # capture (False branch)
+
+    result = f(Tensor([3]), True)  # passing True but False branch runs
+    self.assertEqual(result.item(), 6)  # should be 9!
+
+  def test_positional_kwargs_cannot_mix(self):
+    """Must use same calling convention after capture."""
+    @TinyJit
+    def f(a, b): return (a + b).realize()
+
+    f(Tensor([1]), Tensor([2]))  # warmup with positional
+    f(Tensor([1]), Tensor([2]))  # capture with positional
+
+    with self.assertRaises(AssertionError):
+      f(a=Tensor([3]), b=Tensor([4]))  # kwargs fail
+
+  def test_class_method_shared_across_instances(self):
+    """JIT on instance methods is shared at class level."""
+    class Model:
+      def __init__(self, scale):
+        self.scale = Tensor([scale])
+      @TinyJit
+      def forward(self, x):
+        return (x * self.scale).realize()
+
+    m1, m2 = Model(2), Model(3)
+
+    m1.forward(Tensor([5]))  # warmup
+    m1.forward(Tensor([5]))  # capture with m1.scale=2
+
+    self.assertEqual(m1.forward(Tensor([5])).item(), 10)
+    self.assertEqual(m2.forward(Tensor([5])).item(), 10)  # should be 15!
+
+  def test_side_effects_only_during_capture(self):
+    """Function body not executed during JIT replay."""
+    call_count = [0]
+
+    @TinyJit
+    def f(x):
+      call_count[0] += 1
+      return (x * 2).realize()
+
+    f(Tensor([1]))  # warmup
+    f(Tensor([2]))  # capture
+    self.assertEqual(call_count[0], 2)
+
+    f(Tensor([3]))
+    f(Tensor([4]))
+    f(Tensor([5]))
+    self.assertEqual(call_count[0], 2)  # still 2, not 5!
+
+  def test_nothing_realized_fails(self):
+    """Must JIT at least one kernel."""
+    @TinyJit
+    def f(a, b): return None
+
+    with self.assertRaises(AssertionError):
+      for _ in range(3):
+        f(Tensor([1]), Tensor([2]))
+
+
+class TestJitCorrectBehavior(unittest.TestCase):
+  """Behaviors that work correctly - documented for clarity."""
+
+  def test_random_regenerates(self):
+    """Random tensors regenerate each call."""
+    @TinyJit
+    def f(x):
+      return (x + Tensor.rand(3)).realize()
+
+    f(Tensor([0, 0, 0]))  # warmup
+    f(Tensor([0, 0, 0]))  # capture
+
+    results = {tuple(f(Tensor([0, 0, 0])).numpy().tolist()) for _ in range(5)}
+    self.assertEqual(len(results), 5)
+
+  def test_unrealized_return_auto_realized(self):
+    """Unrealized return tensors are auto-realized."""
+    @TinyJit
+    def f(a, b): return a + b  # no explicit realize
+
+    for _ in range(5):
+      a, b = Tensor.randn(10), Tensor.randn(10)
+      np.testing.assert_allclose(f(a, b).numpy(), a.numpy() + b.numpy(), atol=1e-5)
+
+  def test_kwargs_order_doesnt_matter(self):
+    """Kwargs are sorted by name, so order doesn't matter."""
+    @TinyJit
+    def f(first, second): return (first / second).realize()
+
+    for _ in range(3):
+      a, b = Tensor.randn(10), Tensor.randn(10) + 1
+      np.testing.assert_allclose(f(second=b, first=a).numpy(), a.numpy() / b.numpy(), atol=1e-4)
+      np.testing.assert_allclose(f(first=a, second=b).numpy(), a.numpy() / b.numpy(), atol=1e-4)
+
+  def test_input_mutation_consistent(self):
+    """Input mutation via assign works consistently."""
+    @TinyJit
+    def f(x):
+      x += 1
+      x.realize()
+      return x
+
+    a = Tensor([0]).contiguous().realize()
+    for _ in range(5):
+      f(a)
+    self.assertEqual(a.item(), 5)
+
+
+if __name__ == '__main__':
+  unittest.main()

test/test_multitensor.py

@@ -57,6 +57,34 @@ 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_like(self):
+    X = Tensor.ones(256).shard(devices_2, 0)
+    Y = Tensor.zeros(256).shard_like(X)
+    self.assertEqual(Y.device, X.device)
+    self.assertEqual(Y.uop.axis, 0)
+    # also test with axis=None
+    X2 = Tensor.ones(256).shard(devices_2, axis=None)
+    Y2 = Tensor.zeros(256).shard_like(X2)
+    self.assertEqual(Y2.device, X2.device)
+    self.assertEqual(Y2.uop.axis, None)
+    # test with single device
+    X3 = Tensor.ones(256)
+    Y3 = Tensor.zeros(256).shard_like(X3)
+    self.assertEqual(Y3.device, X3.device)
+    # cannot shard_like multi unless it's a no-op
+    X4 = Tensor.ones(256).shard(devices_2, 0)
+    Y4 = Tensor.ones(256).shard(devices_2, 0).shard_like(X4)
+    self.assertEqual(Y4.device, X4.device)
+    self.assertEqual(Y4.uop.axis, 0)
+    with self.assertRaises(RuntimeError):
+      Tensor.ones(256).shard(devices_2, None).shard_like(X4)
+
   def _test_shard_op(self, op, out, n=4):
     t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
     r = op(t).realize()

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)
@@ -2752,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()
@@ -3168,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_optim.py

@@ -2,9 +2,10 @@ import numpy as np
 import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
+from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
+from test.helpers import needs_second_gpu
 
 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -163,5 +164,31 @@ class TestOptim(unittest.TestCase):
     optimizer.step()
     Tensor.training = old_state
 
+  def test_lamb_cpu_offload(self):
+    # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
+    t = Tensor(x_init.copy(), requires_grad=True)
+    opt = LAMB([t])
+    # move optimizer state to CPU
+    for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
+    # run a step
+    t.sum().backward()
+    opt.step()
+    self.assertEqual(t.device, Device.DEFAULT)
+    self.assertEqual(opt.m[0].device, "CPU")
+
+  @needs_second_gpu
+  def test_lamb_cpu_offload_multi(self):
+    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    t = Tensor(x_init.copy(), requires_grad=True).shard(ds, axis=1)
+    ds = t.device
+    opt = LAMB([t])
+    # move optimizer state to CPU
+    for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
+    # run a step
+    t.sum().backward()
+    opt.step()
+    self.assertEqual(t.device, ds)
+    self.assertEqual(opt.m[0].device, "CPU")
+
 if __name__ == '__main__':
   unittest.main()

test/test_rangeify.py

@@ -339,7 +339,7 @@ class TestRangeify(unittest.TestCase):
   def test_transformer_ffn(self):
     from tinygrad.apps.llm import TransformerBlock
     from tinygrad import nn
-    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5)
+    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5, head_dim=1024, rope_theta=10000.0)
     for p in nn.state.get_parameters(blk): p.replace(Tensor.empty(p.shape))
 
     x = Tensor.empty(128, 1024)

test/test_remote.py

@@ -1,101 +0,0 @@
-import numpy as np, unittest, string
-from hypothesis import given, strategies as st
-from tinygrad import Device, Tensor, TinyJit, dtypes
-from tinygrad.runtime.ops_remote import RemoteDevice, parse_hosts
-from tinygrad.runtime.graph.remote import RemoteGraph
-from tinygrad.helpers import LazySeq, all_same, Context
-
-def multihost_env(devices):
-  def same_hosts(devices): return all_same([h for h,_ in devices])
-  return isinstance(devices, list) and len(devices) >= 12 and not same_hosts(devices[0:12]) and same_hosts(devices[0:6]) and same_hosts(devices[6:12])
-
-@unittest.skipUnless(Device.DEFAULT == "REMOTE" and multihost_env(RemoteDevice.devices), "Requires special environment")
-class TestRemoteMultiHost(unittest.TestCase):
-  def test_mutlihost_transfer(self):
-    a = Tensor.arange(0, 16, device='REMOTE:0').contiguous().realize()
-    b = a.to('REMOTE:6').contiguous().realize()
-    np.testing.assert_equal(b.numpy(), np.arange(0, 16))
-
-  @Context(JIT_BATCH_SIZE=2**32)
-  @unittest.skip("kernel must all be multibuffer")
-  def test_multihost_matmul_jit_graph(self):
-    @TinyJit
-    def do(a:Tensor, b:Tensor): return (a @ b).contiguous().realize()
-
-    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
-    for _ in range(3):
-      na, nb = np.random.rand(128, 128).astype(np.float32), np.random.rand(128, 128).astype(np.float32)
-      a, b = Tensor(na).shard(ds, 0).contiguous().realize(), Tensor(nb).shard(ds, 0).contiguous().realize()
-      nc = na @ nb
-      c = do(a, b)
-      np.testing.assert_allclose(nc, c.numpy(), rtol=3e-2, atol=1e-4) # tolerances from extra/gemm/simple_matmul.py
-
-    # Verify that everything is in one big cross-host graph
-    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
-
-  @Context(JIT_BATCH_SIZE=2**32)
-  @unittest.skip("assign target and input devices mismatch")
-  def test_multihost_aware_schedule(self):
-    @TinyJit
-    def do(*ts:Tensor):
-      acc = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
-      for t in ts: acc += t.sum()
-      return acc.realize()
-
-    def do_np(*ts:np.ndarray):
-      acc = np.zeros(1, np.float32)
-      for t in ts: acc += t.sum()
-      return acc
-
-    ds = ('REMOTE:0', 'REMOTE:1', 'REMOTE:6', 'REMOTE:7')
-    TS = 64
-    for _ in range(3):
-      inp_np = [np.random.rand(256).astype(np.float32) for _ in range(TS)]
-      inp = [Tensor(inp).shard(ds, 0).contiguous().realize() for inp in inp_np]
-      out_np = do_np(*inp_np)
-      out = do(*inp)
-      np.testing.assert_allclose(out_np, out.numpy(), rtol=3e-2, atol=1e-4)
-
-    # Verify that everything is in one big cross-host graph and that the scheduling is reasonable
-    assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured)
-    # At the time of writing this: 2050 graph breaks without multihost aware scheduling, 14 with it. I've set fail threshold to 28 to not fail on
-    # unrelated scheduling changes. Maybe 2x is a bit too pessimistic, but remote should perform just fine as long as this is not like a half hundred
-    # or more here.
-    self.assertLess(len(do.captured._jit_cache[0].prg.template), 28, "Very bad scheduling! Many unnecesary graph breaks!")
-
-class TestParseHosts(unittest.TestCase):
-  def assert_seq(self, result:LazySeq, host:str):
-    self.assertIsInstance(result, LazySeq)
-    for i in [0, 1, 5, 10]: self.assertEqual(result[i], (host, i))
-
-  @given(st.sampled_from(["", "localhost", "192.168.1.1:8080", "host"]))
-  def test_single_host_no_count(self, host:str):
-    self.assert_seq(parse_hosts(host), host)
-
-  @given(host=st.sampled_from(["localhost", "host", "192.168.1.1:8080"]), count=st.integers(0, 10))
-  def test_single_host_with_count(self, host:str, count:int):
-    self.assertEqual(parse_hosts(f"{host}*{count}"), [(host, i) for i in range(count)])
-
-  def test_multiple_hosts_with_counts_simple(self):
-    self.assertEqual(parse_hosts("host1*2,host2*3"), [("host1", i) for i in range(2)] + [("host2", i) for i in range(3)])
-
-  @given(st.lists(st.tuples(st.text(alphabet=string.ascii_letters + string.digits + ".-:"), st.integers(1, 16)), min_size=1))
-  def test_multiple_hosts_with_counts_sampled(self, host_count_pairs):
-    hosts_str = ",".join(f"{host}*{count}" for host, count in host_count_pairs)
-    expected = [(host, i) for host, count in host_count_pairs for i in range(count)]
-    self.assertEqual(parse_hosts(hosts_str), expected)
-
-  @given(st.sampled_from(["host1*2,host2", "a*1,b", "x*3,y*2,z"]))
-  def test_mixed_hosts_fails(self, hosts):
-    with self.assertRaises(AssertionError): parse_hosts(hosts)
-
-  @given(st.sampled_from(["host*abc", "test*xyz", "a*1.5"]))
-  def test_invalid_count_fails(self, hosts):
-    with self.assertRaises(ValueError): parse_hosts(hosts)
-
-  @given(st.sampled_from(["host*2*3", "a*1*2*3", "test*x*y"]))
-  def test_multiple_asterisks_fails(self, hosts):
-    with self.assertRaises(ValueError): parse_hosts(hosts)
-
-if __name__ == '__main__':
-  unittest.main()

test/test_tensor.py

@@ -70,15 +70,15 @@ class TestTinygrad(unittest.TestCase):
     out = out.log_softmax()
     out = out.mul(m).add(m).sum()
     out.backward()
-    xgrad,wgrad = x.grad, W.grad
+    xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
     out.backward()
-    xgrad2,wgrad2 = x.grad, W.grad
+    xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
     out.backward() # no need to retain again since we will not re-run backward
-    xgrad3,wgrad3 = x.grad, W.grad
-    np.testing.assert_allclose(xgrad3.numpy(), xgrad.numpy() * 3., atol=1e-6)
-    np.testing.assert_allclose(wgrad3.numpy(), wgrad.numpy() * 3., atol=1e-6)
-    np.testing.assert_allclose(xgrad2.numpy(), xgrad.numpy() * 2., atol=1e-6)
-    np.testing.assert_allclose(wgrad2.numpy(), wgrad.numpy() * 2., atol=1e-6)
+    xgrad3, wgrad3 = x.grad.numpy(), W.grad.numpy()
+    np.testing.assert_allclose(xgrad3, xgrad * 3., atol=1e-6)
+    np.testing.assert_allclose(wgrad3, wgrad * 3., atol=1e-6)
+    np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6)
+    np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6)
 
   def test_second_order_backward_pass(self):
     def test_pytorch():
@@ -829,6 +829,7 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(len(si.metadata), 3)
     self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
 
+  @unittest.skip("metadata is no longer promised to be exact with schedulecache")
   def test_complex_backward(self):
     x = Tensor.rand(3, requires_grad=True).realize()
     y = Tensor.rand(3, requires_grad=True).realize()

test/test_tensor_variable.py

@@ -95,6 +95,37 @@ class TestTensorVariable(unittest.TestCase):
     assert t.uop.base.buffer.size == 30
     assert t.uop.shape == (3, vb)
 
+  def test_symbolic_chunk(self):
+    # chunk should work when split dimension is concrete, even if other dims are symbolic
+    vv = Variable("a", 1, 10).bind(4)
+    t = Tensor.ones(10, 8).contiguous()[:vv, :]  # shape (vv, 8)
+    chunks = t.chunk(2, dim=-1)  # split along concrete dim 8
+    assert len(chunks) == 2
+    assert chunks[0].shape[1] == 4
+    assert chunks[1].shape[1] == 4
+    # verify the values by shrinking to concrete shape first
+    np.testing.assert_equal(chunks[0].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
+    np.testing.assert_equal(chunks[1].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
+
+  def test_symbolic_split(self):
+    # split should work when split dimension is concrete, even if other dims are symbolic
+    vv = Variable("a", 1, 10).bind(3)
+    t = Tensor.arange(30).reshape(10, 3).contiguous()[:, :vv]  # shape (10, vv)
+    splits = t.split(5, dim=0)  # split along concrete dim 10
+    assert len(splits) == 2
+    assert splits[0].shape[0] == 5
+    assert splits[1].shape[0] == 5
+    # verify the values by shrinking to concrete shape first
+    np.testing.assert_equal(splits[0].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[:5, :3])
+    np.testing.assert_equal(splits[1].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[5:, :3])
+
+  def test_symbolic_chunk_error_on_symbolic_dim(self):
+    # chunk should fail when trying to split along a symbolic dimension
+    vv = Variable("a", 1, 10).bind(4)
+    t = Tensor.ones(10, 8).contiguous()[:vv, :]  # shape (vv, 8)
+    with self.assertRaises(AssertionError):
+      t.chunk(2, dim=0)  # can't split along symbolic dim
+
 
 if __name__ == '__main__':
   unittest.main()

test/testextra/test_tk.py

@@ -1,8 +1,9 @@
-import unittest, math
+import unittest, math, time
 
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import CI
 import numpy as np
 
@@ -30,14 +31,16 @@ class TestTK(unittest.TestCase):
 
       a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
       b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
+      c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
 
       a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
       b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
-      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      c_reg_col = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
+      c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
 
       col, row = ker.blockIdx_x, ker.blockIdx_y
 
-      c_reg = warp.zero(c_reg)
+      c_reg_col = warp.zero(c_reg_col)
       for tile in ker.range(N // BLOCK_SIZE):
         a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
         b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
@@ -45,8 +48,11 @@ class TestTK(unittest.TestCase):
         a_reg = warp.load(a_reg, a_smem)
         b_reg = warp.load(b_reg, b_smem)
 
-        c_reg = warp.mma_AB(c_reg, a_reg, b_reg)
-      c_reg = ker.endrange()
+        c_reg_col = warp.mma_AB(c_reg_col, a_reg, b_reg)
+      c_reg_col = ker.endrange()
+
+      c_smem = warp.store(c_smem, c_reg_col)
+      c_reg = warp.load(c_reg, c_smem)
 
       c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
 
@@ -151,6 +157,89 @@ class TestTK(unittest.TestCase):
 
     np.testing.assert_allclose(b.numpy(), ref.numpy())
 
+  def test_load_store_local_hop(self):
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("load_store_local_hop", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
+      a_reg = warp.load(a_reg, a_smem)
+      b_reg = warp.copy(b_reg, a_reg)
+      b_smem = warp.store(b_smem, b_reg)
+      b_reg = warp.load(b_reg, b_smem)
+      b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish()
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+
+    ref = a.float()
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+
+  def test_load_store_multioutput(self):
+    N = 64
+    BLOCK_SIZE = 32
+    with Kernel("load_store_multioutput", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
+      warp = ker.warp
+
+      b = ker.gl((1, 1, N, N), dtypes.float32)
+      c = ker.gl((1, 1, N, N), dtypes.float32)
+      a = ker.gl((1, 1, N, N), dtypes.float32)
+
+      a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+      b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
+
+      col, row = ker.blockIdx_x, ker.blockIdx_y
+
+      a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
+      a_reg = warp.load(a_reg, a_smem)
+      b_reg = warp.copy(b_reg, a_reg)
+      b_smem = warp.store(b_smem, b_reg)
+      b_reg = warp.load(b_reg, b_smem)
+      b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
+      c = warp.store(c, b_reg, (0, 0, row, col), (), axis=2)
+
+      sink = ker.finish(2)
+
+    with Context(DEBUG=0):
+      a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
+      b = Tensor.empty(1, 1, N, N, dtype="float32")
+      c = Tensor.empty(1, 1, N, N, dtype="float32")
+      Tensor.realize(a, b, c)
+
+    ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in (b, c, a)])
+    for _ in range(5): ei.run(wait=True)
+    b = b.float()
+    c = c.float()
+
+    ref = a.float()
+
+    np.testing.assert_allclose(b.numpy(), ref.numpy())
+    np.testing.assert_allclose(c.numpy(), ref.numpy())
+
   @unittest.skip("TODO")
   def test_load_store_group(self):
     N = 256
@@ -629,7 +718,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)
@@ -640,5 +734,34 @@ class TestTK(unittest.TestCase):
 
     np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
 
+  def test_fast_fa(self):
+    from extra.thunder.tiny.fa import flash_attention
+
+    B, N, H, H_KV, D = 2, 8192, 32, 8, 128
+
+    with Context(DEBUG=0):
+      q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
+      k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
+      Tensor.realize(q, k, v)
+
+    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+    fa_jitted = TinyJit(flash_attention)
+
+    for _ in range(10):
+      st = time.perf_counter()
+      out = fa_jitted(q, k, v, is_causal=True)
+      et = time.perf_counter() - st
+      attn_flops = 2 * B * H * N * N * D + \
+                   4 * B * H * N * N + \
+                   2 * B * H * N * N * D
+      print(f"{attn_flops/(et*1e9):2f} GFLOPS")
+    out = out.float().transpose(1, 2)
+
+    ref = q.scaled_dot_product_attention(k, v, is_causal=True, enable_gqa=True).float().transpose(1, 2)
+
+    np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=2e-2, rtol=2e-2)
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,8 +1,14 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.apps.llm import apply_rope
+from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
 #from tinygrad.engine.realize import run_schedule
 
+def apply_rope(x:Tensor, start_pos:int):
+  B, H, T, Hd = x.shape
+  precompute_freqs_cis.cache_clear()
+  freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
+  return apply_rope_new(x, freqs_cis)
+
 # TODO: test_scheduler, but just in uint
 class TestAttention(unittest.TestCase):
   def test_half_qkv_buffers(self):
@@ -39,7 +45,7 @@ class TestAttention(unittest.TestCase):
     prune_size = len(rope_prune.captured.jit_cache)
 
     self.assertGreater(noprune_size, prune_size)
-    self.assertGreaterEqual(noprune_size, 3)
+    self.assertGreaterEqual(noprune_size, 2)
     self.assertEqual(prune_size, 1)
 
 if __name__ == '__main__':

test/unit/test_autogen.py

@@ -110,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_gguf.py

@@ -58,6 +58,7 @@ class TestGGUF(unittest.TestCase):
   def test_dequantization_q4_0(self): self._test_dequantization(ggml.GGML_TYPE_Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(ggml.GGML_TYPE_Q4_1)
   def test_dequantization_q8_0(self): self._test_dequantization(ggml.GGML_TYPE_Q8_0)
+  def test_dequantization_q4_k(self): self._test_dequantization(ggml.GGML_TYPE_Q4_K)
   def test_dequantization_q6_k(self): self._test_dequantization(ggml.GGML_TYPE_Q6_K)
   def test_dequantization_mxfp4(self):
     MXFP4 = 39

test/unit/test_gradient.py

@@ -110,6 +110,18 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
+  def test_multiple_backward(self):
+    x = Tensor([3.], requires_grad=True)
+    (x*2)[0].backward()
+    np.testing.assert_allclose(x.grad.numpy(), [2.0])
+    old_grad = x.grad
+    (x*3)[0].backward()
+    np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0])
+    self.assertIs(x.grad, old_grad)
+    (x*x)[0].backward()
+    np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
+    self.assertIs(x.grad, old_grad)
+
 class TestRealizeMeansRealize(unittest.TestCase):
   def test_randn_realizes(self):
     x = Tensor.randn(2, 3, 64, 64, requires_grad=True).realize()

test/unit/test_hashing.py

@@ -57,14 +57,6 @@ class TestKeccak(unittest.TestCase):
     # https://www.di-mgt.com.au/sha_testvectors.html
     self.assertEqual(bytes(Tensor(b"abc").keccak().tolist()),
                      bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-    self.assertEqual(bytes(Tensor(b"").keccak().tolist()),
-                     bytearray.fromhex("a7ffc6f8bf1ed766 51c14756a061d662 f580ff4de43b49fa 82d80a4b80f8434a"))
-    t = Tensor(b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu").keccak()
-    self.assertEqual(bytes(t.tolist()),
-                     bytearray.fromhex("916f6061fe879741 ca6469b43971dfdb 28b1a32dc36cb325 4e812be27aad1d18"))
-    # TODO: this does not run or very slow
-    # self.assertEqual(bytes(Tensor(b"a" * 1000000).keccak().tolist()),
-    #                  bytearray.fromhex("5c8875ae474a3634 ba4fd55ec85bffd6 61f32aca75c6d699 d0cdcb6c115891c1"))
 
   @unittest.skipIf(CI, "times out in ci")
   def test_long(self):
@@ -75,24 +67,14 @@ class TestKeccak(unittest.TestCase):
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
   def test_variable_bs(self):
-    data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
-
-    bs = UOp.variable("bs", 1, 4096).bind(1)
-    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(1, 32)
-    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-
-    bs = UOp.variable("bs", 1, 4096).bind(2)
-    out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(2, 32)
-    self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-    self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-
-    bs = UOp.variable("bs", 1, 4096).bind(3)
     data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
+    bs = UOp.variable("bs", 1, 4096).bind(3)
     out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32)
     self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
     self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
 
+  @unittest.skipIf(CI, "redundant with test_variable_bs")
   def test_variable_bs_jit(self):
     def f(data):
       return data.keccak()

test/unit/test_llm_server.py

@@ -0,0 +1,137 @@
+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=[998])
+
+    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,47 @@
+import unittest
+from tinygrad import Tensor, Variable
+from tinygrad.engine.schedule import schedule_cache
+
+class TestScheduleCache(unittest.TestCase):
+  def test_bound_variable_reuses_cache(self):
+    schedule_cache.clear()
+    v = Variable('v', 1, 100)
+    x = Tensor.ones(10).contiguous().realize()
+
+    # first run with v=5
+    t1 = (x + Tensor(v.bind(5))).sum()
+    self.assertEqual(t1.item(), 60.0)
+    cache_size_after_first = len(schedule_cache)
+
+    # second run with v=10 should reuse cache
+    t2 = (x + Tensor(v.bind(10))).sum()
+    self.assertEqual(t2.item(), 110.0)
+    self.assertEqual(len(schedule_cache), cache_size_after_first)
+
+  def test_bound_variable_var_vals(self):
+    v = Variable('pos', 1, 100)
+    x = Tensor.ones(10).contiguous().realize()
+
+    t = x + Tensor(v.bind(42))
+    _, var_vals = t.schedule_with_vars()
+    self.assertEqual(var_vals, {'pos': 42})
+
+  def test_simple(self):
+    a = Tensor.ones(10).contiguous()
+    b = Tensor.ones(10).contiguous()
+    Tensor.realize(a, b)
+
+    # warm up
+    for _ in range(2):
+      num = (a.sum().contiguous()+b.sum().contiguous()).item()
+      print(num)
+
+    # confirm schedule cache doesn't grow
+    start_len_schedule_cache = len(schedule_cache)
+    for _ in range(3):
+      num = (a.sum().contiguous()+b.sum().contiguous()).item()
+      print(num)
+    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_viz.py

@@ -167,7 +167,8 @@ class TestViz(BaseTestViz):
       (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
       (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
     ])
-    with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
+    # use smaller stack limit for faster test (default is 250000)
+    with Context(REWRITE_STACK_LIMIT=100): self.assertRaises(RuntimeError, exec_rewrite, a, [pm])
     graphs = flatten(x["graph"].values() for x in get_viz_details(0, 0))
     self.assertEqual(graphs[0], uop_to_json(a)[id(a)])
     self.assertEqual(graphs[1], uop_to_json(b)[id(b)])
@@ -232,7 +233,7 @@ import gc
 
 def bufs_allocated() -> int:
   gc.collect()
-  return sum([isinstance(x, Buffer) for x in gc.get_objects()])
+  return sum([type(x).__name__ == "Buffer" and type(x).__module__ == "tinygrad.device" for x in gc.get_objects()])
 
 class TestVizGC(BaseTestViz):
   def test_gc(self):

tinygrad/apps/llm.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers
+import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv
+from tinygrad.helpers import partition, TCPServerWithReuse, HTTPRequestHandler, DEBUG, Timing, GlobalCounters, stderr_log, colored
 
 class SimpleTokenizer:
-  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
+  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3"):
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2"): raise ValueError(f"Invalid tokenizer 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)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
+    # TODO: ucat_range is slow
     def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
@@ -18,14 +21,14 @@ class SimpleTokenizer:
     self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
     self._special_tokens = special_tokens
     self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
+    self.preset = preset
 
   @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']}'")
     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), kv["tokenizer.ggml.pre"])
 
   def _encode_word(self, word:bytes) -> list[int]:
     if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
@@ -47,38 +50,45 @@ class SimpleTokenizer:
       pos = match.end(0)
     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 decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
+  def role(self, role:str):
+    if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
+    return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
+  def end_turn(self, eos_id:int): return [eos_id] + self.encode("\n") if self.preset == 'qwen2' else [eos_id]
 
-def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
-  B, H, T, Hd = x.shape
-  assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
-  half = Hd // 2
-  t_start_pos = start_pos if isinstance(start_pos, int) else Tensor(start_pos)
-  angles = (Tensor.arange(T, dtype="float32") + t_start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
-  # contiguous here allows RoPE to be pruned in the JIT
-  cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous()
-  x_pairs = x.reshape(B, H, T, half, 2)
-  return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin,
-                      x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)
+@functools.cache
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
+
+def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
+  assert x.shape[-1] % 2 == 0
+  cos, sin = freqs_cis.reshape(1, 1, x.shape[2], -1).chunk(2, dim=-1)
+  x1, x2 = x.chunk(2, dim=-1)
+  return (x1 * cos - x2 * sin).cat(x2 * cos + x1 * sin, dim=-1)
 
 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, head_dim:int, rope_theta:float,
+               max_context:int=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
     self.max_context  = max_context
+    self.rope_theta   = rope_theta
 
     # --- 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)
+    if qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(self.head_dim, norm_eps), nn.RMSNorm(self.head_dim, norm_eps)
 
     # --- feed-forward ----------------------------------------------------
     self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
@@ -94,8 +104,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)
+    if hasattr(self, 'attn_q_norm'): q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+
+    # TODO: make UOp have SupportsIndex
+    freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)[start_pos:start_pos+T]  # type: ignore
+    q = apply_rope(q, freqs_cis)
+    k = apply_rope(k, freqs_cis)
 
     # TODO: remove these kv cache realizes
     if not hasattr(self, "cache_kv"):
@@ -113,15 +127,18 @@ class TransformerBlock:
 
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
-    gated  = self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)
+    # TODO: remove the need for this contiguous
+    gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
     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, head_dim:int, rope_theta:float,
+               max_context:int=0, qk_norm:bool=False):
+    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, 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 +156,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 +168,22 @@ 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']
+    n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
+
+    # permute Q/K weights from interleaved to half-split RoPE layout: [0,1,2,3,4,5...] -> [0,2,4,...,1,3,5,...]
+    if arch != 'qwen3':
+      for name in state_dict:
+        if 'attn_q.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
+        if 'attn_k.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
+
     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)
+                        n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'],
+                        vocab_size=len(kv['tokenizer.ggml.tokens']), head_dim=kv[f'{arch}.attention.key_length'],
+                        rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context, qk_norm='blk.0.attn_q_norm.weight' in state_dict)
     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,31 +199,138 @@ 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:1b-q4": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.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",
+  "qwen3:1.7b": "https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/main/Qwen3-1.7B-Q4_K_M.gguf",
+  "qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.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
+
+CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
+  * { margin: 0 }
+  body { background: #212121; color: #e3e3e3; font-family: system-ui;
+         height: 100vh; display: flex; flex-direction: column }
+  #chat { flex: 1; overflow-y: auto; padding: 20px }
+  .msg { padding: 10px 16px; margin: 8px 0; white-space: pre-wrap; border-radius: 18px }
+  .user { background: #2f2f2f; margin-left: auto; width: fit-content; max-width: 70% }
+  #input { max-width: 768px; width: 100%; margin: 20px auto; padding: 14px 20px;
+           background: #2f2f2f; color: inherit; font: inherit;
+           border: none; outline: none; resize: none; border-radius: 24px; field-sizing: content }
+</style></head><body><div id="chat"></div>
+<textarea id="input" rows="1" placeholder="Ask anything"></textarea>
+<script>
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send() } }
+  const msgs = [];
+  async function send() {
+    if (!input.value.trim()) return;
+    msgs.push({role: 'user', content: input.value.trim()});
+    chat.innerHTML += '<div class="msg user">' + input.value.trim().replace(/</g, '&lt;') + '</div>';
+    input.value = '';
+    const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
+    const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
+      body: JSON.stringify({model: 'llama', messages: msgs, stream: true})});
+    for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
+      const {done, value} = await rd.read();
+      if (done) break;
+      for (const ln of dec.decode(value).split('\\n'))
+        if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
+          try { d.textContent += JSON.parse(ln.slice(6)).choices[0]?.delta?.content || '' } catch {}
+      chat.scrollTop = chat.scrollHeight;
+    }
+    msgs.push({role: 'assistant', content: d.textContent});
+  }
+</script></body></html>'''
+
+class Handler(HTTPRequestHandler):
+  def log_request(self, code='-', size='-'): pass
+  def do_GET(self): self.send_data(CHAT_HTML, content_type="text/html")
+  def run_model(self, ids:list[int], model_name:str, include_usage=False):
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  in:{len(ids):5d}  {colored('--', 'BLACK')}  ")
+    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
+    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
+    out: list[int] = []
+    st = time.perf_counter()
+    for next_id in model.generate(ids):
+      if len(out) == 0: stderr_log(f"prefill:{len(ids)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if next_id == eos_id: break
+      out.append(next_id)
+      yield {"choices": [{"index":0, "delta":{"content":tok.decode([next_id])}, "finish_reason":None}], **tmpl}
+    yield {"choices": [{"index":0, "delta":{},"finish_reason":"stop"}], **tmpl}
+    if include_usage:
+      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+    stderr_log(f"out:{len(out):5d}  {colored('--', 'BLACK')}  gen: {len(out)/(time.perf_counter()-pt):4.0f} tok/s\n")
+
+  def do_POST(self):
+    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
+    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
+    if DEBUG >= 1: print(json.dumps(body, indent=2))
+    if self.path == "/v1/chat/completions":
+      # extract tokens
+      ids: list[int] = [bos_id] if bos_id is not None else []
+      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(eos_id)
+      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", nargs='?', type=int, const=11434, metavar="PORT", help="Run OpenAI compatible API (optional port, default 11434)")
+  parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
   args = parser.parse_args()
 
   # load the model
-  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.size]), args.max_context)
+  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
+  if DEBUG >= 1: print(f"using model {args.model}")
+
+  # do benchmark
+  if args.benchmark:
+    param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
+    gen = model.generate([0], 0)
+    for _ in range(args.benchmark):
+      GlobalCounters.reset()
+      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s, param {param_bytes/x:7.2f} GB/s"): next(gen)
+    exit(0)
 
   # extract some metadata
   tok = SimpleTokenizer.from_gguf_kv(kv)
-  bos_id: int = kv['tokenizer.ggml.bos_token_id']
+  bos_id: int|None = kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None
   eos_id: int = kv['tokenizer.ggml.eos_token_id']
 
-  ids: list[int] = [bos_id]
+  # start server
+  if args.serve: TCPServerWithReuse(('', args.serve), Handler).serve_forever()
+
+  ids: list[int] = [bos_id] if bos_id is not None else []
   while 1:
-    start_pos = len(ids) - 1
+    start_pos = max(len(ids) - 1, 0)
     try:
-      ids += tok.role("user") + tok.encode(input('>>> ')) + [eos_id] + tok.role("assistant")
+      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn(eos_id) + tok.role("assistant")
     except EOFError:
       break
     for next_id in model.generate(ids, start_pos):

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

@@ -366,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:

tinygrad/engine/schedule.py

@@ -1,8 +1,9 @@
 import time
 from typing import cast
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites
+from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
@@ -12,19 +13,17 @@ from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, fla
 @dataclass(frozen=True)
 class ScheduleItem:
   ast: UOp
-  bufs: tuple[Buffer, ...]
+  bufs: tuple[Buffer, ...] = ()
   metadata: tuple[Metadata, ...] = ()
   fixedvars: dict[str, int] = field(default_factory=dict)
-  bound_ranges: tuple[UOp, ...] = ()
 
 # **** schedule linearizer
 
-def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[str, int]]:
+def create_schedule(sched_sink:UOp) -> tuple[list[ScheduleItem], UOp]:
   with cpu_profile(TracingKey("toposort sched_sink")):
     # construct the KERNEL children graph based on assigns
     children: dict[UOp, list[UOp]] = {}
     in_degree: dict[UOp, int] = {}
-    var_vals: dict[str, int] = {}
     for u in sched_sink.toposort():
       if u.op is Ops.RANGE:
         in_degree.setdefault(u, 0)
@@ -43,44 +42,26 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
               assert ss.op is Ops.AFTER, f"ss.op is not AFTER, it's {ss.op}"
               children.setdefault(ss.src[1], []).append(k)
               in_degree[k] += 1
-        elif s.op is Ops.BUFFER:
-          pass  # a BUFFER is already realized, nothing to do here
-        elif s.op is Ops.BIND:
-          # for RANGE this is in fixedvars
-          if s.src[1].op is not Ops.RANGE:
-            var, val = s.unbind()
-            assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
-            var_vals[var.expr] = val
+        elif s.op in {Ops.BUFFER, Ops.BIND}:
+          pass  # a BUFFER is already realized, BINDs are handled in complete_create_schedule_with_vars
         else:
           raise RuntimeError(f"input to kernel must be AFTER or BUFFER, not {s.op}")
 
-  with cpu_profile(TracingKey("linearize to ScheduleItem")):
+  with cpu_profile(TracingKey("linearize schedule")):
     queue: deque[UOp] = deque()
     for k,v in in_degree.items():
       if v == 0: queue.append(k)
 
-    schedule: list[ScheduleItem|UOp] = []
+    schedule: list[tuple|UOp] = []
     while len(queue):
       k = rk = queue.popleft()
       if k.op is Ops.END: k = k.src[0]
       if k.op is Ops.RANGE: schedule.append(k)
       elif k.op is Ops.KERNEL:
         ast = k.arg.ast
-        # create subbuffers if needed
-        if ast.op is Ops.BUFFER_VIEW:
-          base = k.src[1].buf_uop.buffer
-          assert isinstance(base, Buffer), "base can't be MultiBuffer"
-          buffers[k.src[0]] = base.view(k.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-        ubufs = tuple(s.buf_uop.buffer for s in k.src if s.op is not Ops.BIND)
-        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and s.src[1].op is Ops.RANGE)
-        if any(isinstance(x, MultiBuffer) for x in ubufs):
-          assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
-          dnums = [x for x in ast.variables() if x.arg[0] == '_device_num']
-          for i,bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-            schedule.append(ScheduleItem(ast, bufs, k.arg.metadata, {dnums[0].expr:i} if len(dnums) else {}, bound_ranges=bound_ranges))
-        else:
-          # ONE -> ONE
-          schedule.append(ScheduleItem(ast, cast(tuple[Buffer, ...], ubufs), k.arg.metadata, bound_ranges=bound_ranges))
+        buf_uops = tuple(s.buf_uop for s in k.src if s.op is not Ops.BIND)
+        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
+        schedule.append((ast, buf_uops, k.arg.metadata, {}, bound_ranges))
         if rk.op is Ops.END: schedule.append(rk)
       else:
         raise RuntimeError(f"can't schedule {k.op}")
@@ -89,10 +70,11 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
         if in_degree[x] == 0: queue.append(x)
 
   with cpu_profile(TracingKey("expand ranges")):
-    real_schedule: list[ScheduleItem] = []
+    pre_schedule: list[ScheduleItem] = []
+    buf_uops_list: list[UOp] = []
     sched_ptr = 0
-    in_ranges = {}
-    range_ptrs = {}
+    in_ranges: dict[UOp, int] = {}
+    range_ptrs: dict[UOp, int] = {}
     while sched_ptr < len(schedule):
       si = schedule[sched_ptr]
       if isinstance(si, UOp):
@@ -105,40 +87,129 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
             sched_ptr = range_ptrs[si.src[1]]
             continue
       else:
-        real_schedule.append(replace(si, fixedvars=si.fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in si.bound_ranges}, bound_ranges=()))
+        ast, buf_uops, metadata, fixedvars, bound_ranges = si
+        fixedvars = fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
+        pre_schedule.append(ScheduleItem(ast, (), metadata, fixedvars))
+        buf_uops_list.append(UOp.sink(*buf_uops))
       sched_ptr += 1
-  return real_schedule, var_vals
+  return pre_schedule, UOp.sink(*buf_uops_list)
 
 from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.schedule.multi import get_multi_map
 
+def replace_input_buffer(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
+    if b.op is Ops.BUFFER:
+      ctx[b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=len(ctx)), b.src[1]))
+    else:
+      # TODO: flip args in CONST
+      assert b.op is Ops.CONST
+      ctx[b] = ret = b.replace(src=(b.src[0], UOp(Ops.LUNIQUE, arg=len(ctx))))
+  return ret
+
+pm_pre_sched_cache = PatternMatcher([
+  # replace input buffers
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
+  # remove unique consts
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="b"), replace_input_buffer),
+  # strip value from BIND for cache key normalization, so different values hit same cache
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), lambda ctx,b: ctx.setdefault(b, b.replace(src=(b.src[0],)))),
+])
+
+def replace_input_buffer_back(ctx:dict[UOp, UOp], b:UOp):
+  if (ret:=ctx.get(b, None)) is None:
+    assert b.op is Ops.BUFFER
+    # if it's not in the cache, create a new buffer
+    ctx[b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
+  return ret
+
+pm_post_sched_cache = PatternMatcher([
+  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer_back),
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE), UPat(Ops.LUNIQUE)), name="b"), replace_input_buffer_back),
+  # restore BIND value stripped in pm_pre_sched_cache
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)),
+])
+
+schedule_cache: dict[bytes, tuple[list[ScheduleItem], 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, strip BIND values for cache key
+  input_buffers: dict[UOp, UOp] = {}
+  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache")
+  sched_cache_key = big_sink_cache.key
 
-  # 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] = {}
 
-  # create the schedule
-  schedule, var_vals = create_schedule_with_vars(big_sink)
+    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")
+
+    pre_schedule, buf_uops_sink = create_schedule(big_sink)
+
+    # save in schedule cache (include AFTERs in tensor_map so we don't need big_sink)
+    after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
+    tensor_map_sink = UOp.sink(*flatten([(k,v) for k,v in tensor_map.items()]), *flatten(after_map))
+    combined_sink = UOp.sink(tensor_map_sink, buf_uops_sink)
+    schedule_cache[sched_cache_key] = (pre_schedule, combined_sink)
+  else:
+    # schedule cache hit
+    del big_sink_cache
+    pre_schedule, combined_sink = sc_ret
+
+  # replace all the LUNIQUEs with UNIQUEs (single graph_rewrite for everything)
+  input_buffers_reverse = {v:k for k,v in input_buffers.items()}
+  combined = graph_rewrite(combined_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite combined")
+  tensor_map_sink, buf_uops_sink = combined.src
+  tm_src = tensor_map_sink.src
+  tensor_map = {tm_src[i]:tm_src[i+1] for i in range(0, len(tm_src), 2)}
+
+  # add bufs to pre_schedule
+  schedule: list[ScheduleItem] = []
+  for i, si in enumerate(pre_schedule):
+    buf_uops = buf_uops_sink.src[i].src
+    # create subbuffers if needed
+    if si.ast.op is Ops.BUFFER_VIEW:
+      base = buf_uops[1].buffer
+      assert isinstance(base, Buffer), "base can't be MultiBuffer"
+      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, si.ast.dtype, si.ast.arg[1]*base.dtype.itemsize)
+    ubufs = tuple(b.buffer for b in buf_uops)
+    if any(isinstance(x, MultiBuffer) for x in ubufs):
+      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
+      dnums = [x for x in si.ast.variables() if x.arg[0] == '_device_num']
+      for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
+        schedule.append(ScheduleItem(si.ast, bufs, si.metadata, si.fixedvars | ({dnums[0].expr:j} if len(dnums) else {})))
+    else:
+      # ONE -> ONE
+      schedule.append(ScheduleItem(si.ast, cast(tuple[Buffer, ...], ubufs), si.metadata, si.fixedvars))
   with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
 
-  # remove all AFTERs, after scheduling, the tensors are just buffers
-  tensor_map |= {u:u.buf_uop for u in big_sink.toposort() if u.op is Ops.AFTER}
+  # extract var_vals from BINDs that were stripped (only if there are kernels)
+  var_vals: dict[str, int] = {}
+  if schedule:
+    for u in input_buffers:
+      if u.op is Ops.BIND:
+        var, val = u.unbind()
+        assert var.expr not in var_vals or var_vals[var.expr] == val, f"bind mismatch on {var}, {var_vals[var.expr]} != {val}"
+        var_vals[var.expr] = val
 
   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/gradient.py

@@ -42,7 +42,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   # NOTE: this is only correct when the KERNEL has a single output
   (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
+  (UPat(Ops.CUSTOM_KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])

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")
@@ -148,6 +149,10 @@ def getenv(key:str, default:Any=0): return type(default)(os.getenv(key, default)
 def temp(x:str, append_user:bool=False) -> str:
   return (pathlib.Path(tempfile.gettempdir()) / (f"{x}.{getpass.getuser()}" if append_user else x)).as_posix()
 
+def stderr_log(msg):
+  sys.stderr.write(msg)
+  sys.stderr.flush()
+
 class Context(contextlib.ContextDecorator):
   def __init__(self, **kwargs): self.kwargs = kwargs
   def __enter__(self):
@@ -186,8 +191,10 @@ EMULATE = ContextVar("EMULATE", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 # Compilers
 CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-NV_PTX, CUDA_PTX, NV_NAK = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0)
+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)
@@ -402,6 +409,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/mixin/math.py

@@ -1,3 +1,4 @@
+import math
 from typing import Self
 from tinygrad.uop import Ops
 from tinygrad.dtype import dtypes, ConstType
@@ -258,6 +259,13 @@ class MathMixin:
     return self.alu(Ops.RECIPROCAL)
 
   def trunc(self):
+    """
+    Truncates the tensor element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).trunc().numpy())
+    ```
+    """
     return self.alu(Ops.TRUNC)
 
   def sqrt(self):
@@ -277,3 +285,232 @@ class MathMixin:
 
   def __pow__(self, x: Self | ConstType):
     return self.pow(x)
+
+  def square(self):
+    """
+    Squares the tensor element-wise.
+    Equivalent to `self*self`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).square().numpy())
+    ```
+    """
+    return self * self
+
+  def clamp(self, min_=None, max_=None):
+    """
+    Clips (clamps) the values in the tensor between `min_` and `max_` element-wise.
+    If `min_` is `None`, there is no lower bound. If `max_` is None, there is no upper bound.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).clip(-1, 1).numpy())
+    ```
+    """
+    if min_ is None and max_ is None: raise RuntimeError("at least one of 'min_' or 'max_' must not be None")
+    ret = (self < min_).where(min_, self) if min_ is not None else self
+    return (ret > max_).where(max_, ret) if max_ is not None else ret
+
+  def clip(self, min_=None, max_=None):
+    """Alias for `Tensor.clamp`."""
+    return self.clamp(min_, max_)
+
+  def isnan(self):
+    """
+    Checks the tensor element-wise to return True where the element is NaN, otherwise returns False
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isnan().numpy())
+    ```
+    """
+    return self != self
+
+  def isinf(self, detect_positive: bool = True, detect_negative: bool = True):
+    """
+    Checks the tensor element-wise to return True where the element is infinity, otherwise returns False
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isinf().numpy())
+    ```
+    """
+    return self.eq(float("inf")) * detect_positive + self.eq(float("-inf")) * detect_negative
+
+  def isfinite(self):
+    """
+    Checks the tensor element-wise to return True where the element is finite, otherwise returns False
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isfinite().numpy())
+    ```
+    """
+    return (self.isinf() | self.isnan()).logical_not()
+
+  def ceil(self):
+    """
+    Rounds the tensor element-wise towards positive infinity.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).ceil().numpy())
+    ```
+    """
+    return (self > (b := self.trunc())).where(b+1, b)
+
+  def floor(self):
+    """
+    Rounds the tensor element-wise towards negative infinity.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).floor().numpy())
+    ```
+    """
+    return (self < (b := self.trunc())).where(b-1, b)
+
+  def relu(self):
+    """
+    Applies the Rectified Linear Unit (ReLU) function element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).relu().numpy())
+    ```
+    """
+    # NOTE: if you write this as self.maximum(0) the gradient is wrong, passing through half when self is 0
+    return (self > 0).where(self, 0)
+
+  def sigmoid(self):
+    """
+    Applies the Sigmoid function element-wise.
+
+    - Described: https://en.wikipedia.org/wiki/Sigmoid_function
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).sigmoid().numpy())
+    ```
+    """
+    return (1 + (self * (-1/math.log(2))).exp2()).reciprocal()
+
+  def relu6(self):
+    """
+    Applies the ReLU6 function element-wise.
+
+    - Paper: https://arxiv.org/abs/1704.04861v1
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-9., -6., -3., 0., 3., 6., 9.]).relu6().numpy())
+    ```
+    """
+    return self.relu() - (self-6).relu()
+
+  def hardswish(self):
+    """
+    Applies the Hardswish function element-wise.
+
+    - Paper: https://arxiv.org/abs/1905.02244v5
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).hardswish().numpy())
+    ```
+    """
+    return self * (self+3).relu6() * (1/6)
+
+  def hardsigmoid(self, alpha: float = 1/6, beta: float = 0.5):
+    """
+    Applies the Hardsigmoid function element-wise.
+    NOTE: default `alpha` and `beta` values are taken from torch
+
+    - See: https://pytorch.org/docs/stable/generated/torch.nn.functional.hardsigmoid.html
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).hardsigmoid().numpy())
+    ```
+    """
+    return (alpha * self + beta).relu() - (alpha * self + beta - 1).relu()
+
+  def hardtanh(self, min_val=-1, max_val=1):
+    """
+    Applies the Hardtanh function element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-1.5, -1.0, -0.5, 0., 0.5, 1.0, 1.5]).hardtanh().numpy())
+    ```
+    """
+    return self.clip(min_val, max_val)
+
+  def leaky_relu(self, neg_slope=0.01):
+    """
+    Applies the Leaky ReLU function element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).leaky_relu().numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).leaky_relu(neg_slope=0.42).numpy())
+    ```
+    """
+    return (self < 0).where(neg_slope*self, self)
+
+  def tanh(self):
+    """
+    Applies the Hyperbolic Tangent (tanh) function element-wise.
+
+    - Described: https://en.wikipedia.org/wiki/Hyperbolic_functions#Tanh
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).tanh().numpy())
+    ```
+    """
+    return 2.0 * ((2.0 * self).sigmoid()) - 1.0
+
+  def quick_gelu(self):
+    """
+    Applies the Sigmoid GELU approximation element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).quick_gelu().numpy())
+    ```
+    """
+    return self * (self * 1.702).sigmoid()
+
+  def gelu(self):
+    """
+    Applies the Gaussian Error Linear Unit (GELU) function element-wise.
+
+    - Paper: https://arxiv.org/abs/1606.08415v5
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).gelu().numpy())
+    ```
+    """
+    return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self ** 3)).tanh())
+
+  def swish(self):
+    """
+    See `.silu()`
+
+    - Paper: https://arxiv.org/abs/1710.05941v1
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).swish().numpy())
+    ```
+    """
+    return self * self.sigmoid()
+
+  def silu(self):
+    """
+    Applies the Sigmoid Linear Unit (SiLU) function element-wise.
+
+    - Paper: https://arxiv.org/abs/1606.08415
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).silu().numpy())
+    ```
+    """
+    return self.swish()  # The SiLU function is also known as the swish function.
+
+  def rsqrt(self):
+    """
+    Computes the reciprocal of the square root of the tensor element-wise.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1., 2., 3., 4.]).rsqrt().numpy())
+    ```
+    """
+    return self.sqrt().reciprocal()

tinygrad/nn/onnx.py

@@ -498,13 +498,15 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)
 
   # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
-  def _onnx_pads_to_tiny_pads(pads): return tuple(flatten(reversed(list(zip(pads, pads[len(pads)//2:])))))
+  def _onnx_pads_to_tiny_pads(pads):
+    n = len(pads) // 2
+    return tuple(x for i in range(n-1, -1, -1) for x in (pads[i], pads[i+n]))
 
   AUTO_PAD_OPTIONS = Literal["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"]
   # (padding_height, padding_width) -> (padding_top, padding_left, padding_bottom, padding_right)
   def _auto_pad(pads, auto_pad: AUTO_PAD_OPTIONS):
-    if auto_pad == "SAME_UPPER": return [pads[i]//2 for i in range(len(pads))] + [pads[i]-pads[i]//2 for i in range(len(pads))]
-    return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))]
+    first = [p//2 for p in pads] if auto_pad == "SAME_UPPER" else [p - p//2 for p in pads]
+    return first + [p - f for p, f in zip(pads, first)]
 
   def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS):
     if auto_pad == "VALID": return [0]*(len(k_)*2)
@@ -626,6 +628,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
@@ -646,7 +649,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def Mod(x:Tensor,y:Tensor,fmod=0): return x - x.div(y, rounding_mode="trunc") * y if fmod else x % y
 
   # ***** Casting Ops *****
-  # TODO: saturate
+  # TODO: saturate parameter is ignored in Cast and CastLike
   def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_fallback(OnnxDataType(to).to_dtype(), "Cast op"))
   def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
 
@@ -698,8 +701,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
   def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
     axes = axes or list(range(data.ndim))
-    steps = steps or [1]*data.ndim
-    slices = [slice(0,x,1) for x in data.shape]
+    steps = steps or [1] * data.ndim
+    slices = [slice(None)] * data.ndim
     for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
     return data[tuple(slices)]
 
@@ -809,7 +812,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
     input_shape = cast(tuple[int, ...], X.shape[2:])
     if scales is not None: assert all(sc==1 for sc in scales[:-len(input_shape)]), "resizing batch_size dim or channel dim not supported"
-    if sizes is not None: assert tuple(sizes[:-2]) == tuple(X.shape[X.ndim-len(sizes):-2]),  "resizing batch_size dim or channel dim not supported"
+    if sizes is not None: assert tuple(sizes[:-2]) == tuple(X.shape[X.ndim-len(sizes):-2]), "resizing batch_size dim or channel dim not supported"
 
     scales, sizes = (None if scales is None else scales[-len(input_shape):]), (None if sizes is None else sizes[-len(input_shape):])
     if sizes is not None:
@@ -933,11 +936,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     # https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
     assert (segment_ids is None) is (segment_embedding is None)
     assert mask is None and not mask_index_type, "functionality not supported yet"  # TODO
-    input_shape = input_ids.shape
-    seq_length = input_shape[1]
-    compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
+    input_shape, seq_length = input_ids.shape, input_ids.shape[1]
     vocab_size, max_position_embeddings = word_embedding.shape[0], position_embedding.shape[0]
-    type_vocab_size  = (segment_embedding.shape[0] if compute_seg_emb else None)
 
     def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
       return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
@@ -946,10 +946,9 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape)
     wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding)
     pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding)
-    seg_embedding_res = embedding(segment_ids, type_vocab_size, segment_embedding) if compute_seg_emb else None
 
     embedding_sum = wrd_embedding_res + pos_embedding_res
-    if seg_embedding_res is not None: embedding_sum = embedding_sum + seg_embedding_res
+    if segment_embedding is not None: embedding_sum = embedding_sum + embedding(segment_ids, segment_embedding.shape[0], segment_embedding)
     out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
     return out, None, embedding_sum
   def MeanVarianceNormalization(x:Tensor, axis:list[int]|None=None):
@@ -1003,7 +1002,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)
 
   def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
-                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
+                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None, do_rotary:int=0, mask_filter_value:float=-10000.0,
                         num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
                         rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
     assert not do_rotary and not attention_bias, "TODO"
@@ -1287,7 +1286,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     # Tensor ops
     **{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
     "Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
-    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh",  "Elu", "Celu", "Selu", "Round", "Erf")},
+    "Tanh", "Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Round", "Erf")},
     # Implemented ops
     **{name:obj for name,obj in locals().items() if isinstance(obj, types.FunctionType) and not name.startswith("_") and name[0].isupper()},
     # Version ops

tinygrad/nn/optim.py

@@ -163,11 +163,12 @@ class LAMB(Optimizer):
     self.b1_t *= self.b1
     self.b2_t *= self.b2
     for i, (t, g) in enumerate(zip(params, grads)):
+      if g.device != self.m[i].device: g = g.contiguous().to(self.m[i].device)
       self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
       self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
       m_hat = self.m[i] / (1.0 - self.b1_t)
       v_hat = self.v[i] / (1.0 - self.b2_t)
-      up = (m_hat / (v_hat.sqrt() + self.eps)) + self.wd * t.detach()
+      up = (m_hat / (v_hat.sqrt() + self.eps)).shard_like(t) + self.wd * t.detach()
       if not self.adam:
         r1 = t.detach().square().sum().sqrt()
         r2 = up.square().sum().sqrt()

tinygrad/nn/state.py

@@ -308,7 +308,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   Converts ggml tensor data to a tinygrad tensor.
 
   Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 16), int16 (id: 17), int32 (id: 18)
-  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q6_K (id: 14), MXFP4 (id: 39)
+  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q4_K (id: 12), Q6_K (id: 14), MXFP4 (id: 39)
   """
   # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
 
@@ -322,13 +322,20 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     return t.unsqueeze(-1).expand((*t.shape,8//b)).idiv(shift_tensor).bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   # map to (number of elements, number of bytes)
-  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 14: (256, 210), 8: (32, 34), 39: (32, 17) }.get(ggml_type)) is not None:
+  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 8: (32, 34), 12: (256, 144), 14: (256, 210), 39: (32, 17) }.get(ggml_type)) is not None:
     blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
       d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
       return q_to_uint8(blocks[:,4:], 4).bitcast(dtypes.int8) * d + m
     if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
+    if ggml_type == 12:  # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
+      d, dmin = (blocks[:,i:i+2].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1) for i in [0, 2])
+      s = blocks[:,4:16]  # 12 bytes: 6-bit scales[0-3], 6-bit mins[0-3], high bits[4-7]
+      sc = s[:,0:4].bitwise_and(63).cat(s[:,8:12].bitwise_and(0xF).bitwise_or(s[:,0:4].rshift(6).lshift(4)), dim=-1)
+      mn = s[:,4:8].bitwise_and(63).cat(s[:,8:12].rshift(4).bitwise_or(s[:,4:8].rshift(6).lshift(4)), dim=-1)
+      q = Tensor.stack((qs:=blocks[:,16:144].reshape(-1,4,32)).bitwise_and(0xF), qs.rshift(4), dim=2).reshape(-1,8,32).cast(dtypes.float32)
+      return (d * sc.unsqueeze(-1) * q - dmin * mn.unsqueeze(-1)).flatten(-2)
     if ggml_type == 14:
       xl, xh = q_to_uint8(blocks[:,:128].reshape((-1, 2, 64)), 4), q_to_uint8(blocks[:,128:192].reshape((-1, 2, 32)), 2).lshift(4)
       scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))

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':

tinygrad/renderer/llvmir.py

@@ -224,10 +224,6 @@ class AMDLLVMRenderer(LLVMRenderer):
       f"  {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x.src[0]]}_i32, i32 0)"),
   ]) + base_rewrite
   extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
-    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))),
-      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
-    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
-      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
     # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
     (UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2),
     (UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2),
@@ -269,7 +265,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
     if self.arch.split(":")[0] == "gfx1100":
       self.extra_matcher += PatternMatcher([
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half.vec(8)),
-          lambda x: UOp(Ops.WMMA, dtypes.half.vec(16), (x.src[0], x.src[1], x.src[2].cast(dtypes.half.vec(16))), (*x.arg,)).cast(dtypes.half.vec(8))),
+          lambda x: UOp(Ops.VECTORIZE, dtypes.half.vec(8), tuple(UOp(Ops.WMMA, dtypes.half.vec(16), (x.src[0], x.src[1],
+            UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(x.src[2].gep(i//2) if i%2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
+            (*x.arg,)).gep(i*2) for i in range(8)))),
         (UPat(Ops.WMMA, name="x"), lambda x: UOp(Ops.WMMA, x.dtype, (x.src[0].bitcast(dtypes.uint16.vec(16)), x.src[1].bitcast(dtypes.uint16.vec(16)),
           x.src[2]), x.arg) if x.src[0].dtype == dtypes.bfloat16.vec(16) else None),
       ])

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

@@ -8,6 +8,12 @@ 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"
 
+llvm_lib = (r"'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll' if WIN else '/opt/homebrew/opt/llvm@20/lib/libLLVM.dylib' if OSX else " +
+            repr(['LLVM'] + [f'LLVM-{i}' for i in reversed(range(14, 21+1))]))
+
+webgpu_lib = "os.path.join(sysconfig.get_paths()['purelib'], 'pydawn', 'lib', 'libwebgpu_dawn.dll') if WIN else 'webgpu_dawn'"
+nv_lib_path = "f'/usr/local/cuda/targets/{sysconfig.get_config_var(\"MULTIARCH\").rsplit(\"-\", 1)[0]}/lib'"
+
 def load(name, dll, files, **kwargs):
   if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
     files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
@@ -21,22 +27,22 @@ def load(name, dll, files, **kwargs):
       if (preprocess:=kwargs.pop('preprocess', None)): preprocess(base)
     files = flatten(sorted(glob.glob(p, recursive=True)) if isinstance(p, str) and '*' in p else [p] for p in files)
     kwargs['epilog'] = (epi(base) if tarball else epi()) if callable(epi:=kwargs.get('epilog', [])) else epi
-    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(dll, files, **kwargs))
+    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
   return importlib.import_module(f"{path}.{name.replace('/', '.')}")
 
 def __getattr__(nm):
   match nm:
-    case "libc": return load("libc", ["find_library('c')"], lambda: (
+    case "libc": return load("libc", "'c'", lambda: (
       [i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
-      ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), use_errno=True)
-    case "avcodec": return load("avcodec", [], ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], tarball=ffmpeg_src)
-    case "opencl": return load("opencl", ["find_library('OpenCL')"], ["/usr/include/CL/cl.h"])
-    case "cuda": return load("cuda", ["find_library('cuda')"], ["/usr/include/cuda.h"], args=["-D__CUDA_API_VERSION_INTERNAL"], parse_macros=False)
-    case "nvrtc": return load("nvrtc", ["find_library('nvrtc')"], ["/usr/include/nvrtc.h"])
-    case "nvjitlink": load("nvjitlink", ["find_library('nvJitLink')"], [root/"extra/nvJitLink.h"])
-    case "kfd": return load("kfd", [], ["/usr/include/linux/kfd_ioctl.h"])
+      ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), errno=True)
+    case "avcodec": return load("avcodec", None, ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], tarball=ffmpeg_src)
+    case "opencl": return load("opencl", "'OpenCL'", ["/usr/include/CL/cl.h"])
+    case "cuda": return load("cuda", "'cuda'", ["/usr/include/cuda.h"], args=["-D__CUDA_API_VERSION_INTERNAL"], parse_macros=False)
+    case "nvrtc": return load("nvrtc", "'nvrtc'", ["/usr/include/nvrtc.h"], paths=nv_lib_path, prolog=["import sysconfig"])
+    case "nvjitlink": load("nvjitlink", "'nvJitLink'", [root/"extra/nvJitLink.h"], paths=nv_lib_path, prolog=["import sysconfig"])
+    case "kfd": return load("kfd", None, ["/usr/include/linux/kfd_ioctl.h"])
     case "nv_570" | "nv_580":
-      return load(nm, [], [
+      return load(nm, None, [
         *[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"]],
@@ -51,7 +57,7 @@ def __getattr__(nm):
         "-include", "{}/src/common/sdk/nvidia/inc/nvtypes.h", "-I{}/src/common/inc", "-I{}/kernel-open/nvidia-uvm", "-I{}/kernel-open/common/inc",
         "-I{}/src/common/sdk/nvidia/inc", "-I{}/src/nvidia/arch/nvalloc/unix/include", "-I{}/src/common/sdk/nvidia/inc/ctrl"
       ], rules=[(r'MW\(([^:]+):(.+)\)',r'(\1, \2)')], tarball=nv_src[nm], anon_names={"{}/kernel-open/common/inc/nvstatus.h:37":"nv_status_codes"})
-    case "nv": return load("nv", [], [
+    case "nv": return load("nv", None, [
       *[f"{{}}/src/nvidia/inc/kernel/gpu/{s}.h" for s in ["fsp/kern_fsp_cot_payload", "gsp/gsp_init_args"]],
       *[f"{{}}/src/nvidia/arch/nvalloc/common/inc/{s}.h" for s in ["gsp/gspifpub", "gsp/gsp_fw_wpr_meta", "gsp/gsp_fw_sr_meta", "rmRiscvUcode",
                                                                    "fsp/fsp_nvdm_format"]],
@@ -69,70 +75,68 @@ def __getattr__(nm):
       "{}/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h:244": "rpc_events"
     })
     # this defines all syscall numbers. should probably unify linux autogen?
-    case "io_uring": return load("io_uring", [], ["/usr/include/liburing.h", "/usr/include/linux/io_uring.h", "/usr/include/asm-generic/unistd.h"],
+    case "io_uring": return load("io_uring", None, ["/usr/include/liburing.h", "/usr/include/linux/io_uring.h", "/usr/include/asm-generic/unistd.h"],
                                  rules=[('__NR', 'NR')])
-    case "ib": return load("ib", ["ibverbs"], ["/usr/include/infiniband/verbs.h", "/usr/include/infiniband/verbs_api.h",
-                                               "/usr/include/infiniband/ib_user_ioctl_verbs.h","/usr/include/rdma/ib_user_verbs.h"], use_errno=True)
-    case "llvm": return load("llvm", ["LLVM_PATH"], lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"],
-                             args=lambda: system("llvm-config-20 --cflags").split(), recsym=True,
-                             prolog=["from tinygrad.runtime.support.llvm import LLVM_PATH"])
-    case "pci": return load("pci", [], ["/usr/include/linux/pci_regs.h"])
-    case "vfio": return load("vfio", [], ["/usr/include/linux/vfio.h"])
+    case "ib": return load("ib", "'ibverbs'", ["/usr/include/infiniband/verbs.h", "/usr/include/infiniband/verbs_api.h",
+                                               "/usr/include/infiniband/ib_user_ioctl_verbs.h","/usr/include/rdma/ib_user_verbs.h"], errno=True)
+    case "llvm": return load("llvm", llvm_lib, lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"],
+                             args=lambda: system("llvm-config-20 --cflags").split(), recsym=True, prolog=["from tinygrad.helpers import WIN, OSX"])
+    case "pci": return load("pci", None, ["/usr/include/linux/pci_regs.h"])
+    case "vfio": return load("vfio", None, ["/usr/include/linux/vfio.h"])
     # could add rule: WGPU_COMMA -> ','
-    case "webgpu":
-      return load("webgpu", ["WEBGPU_PATH"], [root/"extra/webgpu/webgpu.h"], prolog=["from tinygrad.runtime.support.webgpu import WEBGPU_PATH"])
-    case "libusb": return load("libusb", ["os.getenv('LIBUSB_PATH', find_library('usb-1.0'))"], ["/usr/include/libusb-1.0/libusb.h"])
-    case "hip": return load("hip", ["os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'"], ["/opt/rocm/include/hip/hip_ext.h",
+    case "webgpu": return load("webgpu", webgpu_lib, [root/"extra/webgpu/webgpu.h"],
+                               prolog=["from tinygrad.helpers import WIN, OSX", "import sysconfig, os"])
+    case "libusb": return load("libusb", "'usb-1.0'", ["/usr/include/libusb-1.0/libusb.h"])
+    case "hip": return load("hip", "os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'", ["/opt/rocm/include/hip/hip_ext.h",
                             "/opt/rocm/include/hip/hiprtc.h", "/opt/rocm/include/hip/hip_runtime_api.h", "/opt/rocm/include/hip/driver_types.h"],
-                            args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"])
+                            args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
     case "comgr" | "comgr_3":
-      return load("comgr_3" if nm == "comgr_3" else "comgr", [
-        "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')"], [
+      return load("comgr_3" if nm == "comgr_3" else "comgr", "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
+                  ["/opt/rocm/include/amd_comgr/amd_comgr.h"], args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"],
+                  prolog=["import os"])
+    case "hsa": return load("hsa", "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so', 'hsa-runtime64']", [
       *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
       *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue",
                                                                               "amd_hsa_kernel_code", "hsa_ext_finalize",
                                                                               "hsa_ext_image", "hsa_ven_amd_aqlprofile"]]],
-      tarball=rocr_src, args=["-DLITTLEENDIAN_CPU"])
-    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"]],
+      tarball=rocr_src, args=["-DLITTLEENDIAN_CPU"], prolog=["import os"])
+    case "amd_gpu": return load("amd_gpu", None, [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 "kgsl": return load("kgsl", None, [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
     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"])
+      return load("qcom_dsp", None, [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
+    case "sqtt": return load("sqtt", None, [root/"extra/sqtt/sqtt.h"])
     case "rocprof":
-      return load("rocprof", ["find_library('rocprof-trace-decoder')", p:="'/usr/local/lib/rocprof-trace-decoder.so'", p.replace('so','dylib')],
+      return load("rocprof", "['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')]",
                   [f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
                   tarball="https://github.com/ROCm/rocprof-trace-decoder/archive/dd0485100971522cc4cd8ae136bdda431061a04d.tar.gz")
-    case "mesa": return load("mesa", ["find_library('tinymesa_cpu')",
-      "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if OSX else '.so')",
-      "f'{BASE}/libtinymesa{EXT}'", "'/opt/homebrew/lib/libtinymesa_cpu.dylib'", "'/opt/homebrew/lib/libtinymesa.dylib'"], [
+    case "mesa": return load("mesa", "['tinymesa_cpu', 'tinymesa']", [
         *[f"{{}}/src/compiler/nir/{s}.h" for s in ["nir", "nir_builder", "nir_shader_compiler_options", "nir_serialize"]], "{}/gen/nir_intrinsics.h",
         *[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",
-  prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
+          "-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"], 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'))"],
+      return load("libclang", "'clang-20'",
                   lambda: [f"{system('llvm-config-20 --includedir')}/clang-c/{s}.h" for s in ["Index", "CXString", "CXSourceLocation", "CXFile"]],
                   args=lambda: system("llvm-config-20 --cflags").split())
     case "metal":
-      return load("metal", ["find_library('Metal')"],[f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
+      return load("metal", "'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

@@ -6,18 +6,21 @@ inc = ["-include", "stdint.h"]
 
 def __getattr__(nm):
   match nm:
-    case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm", "discovery",
-  "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]]+[f"{AMD}/include/soc15_ih_clientid.h"], args=inc, tarball=am_src)
+    case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
+      "discovery", "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]] + \
+      [f"{AMD}/include/{s}.h" for s in ["v9_structs", "soc15_ih_clientid"]], args=inc, tarball=am_src)
     case "pm4_soc15": return load("am/pm4_soc15", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], tarball=am_src)
     case "pm4_nv": return load("am/pm4_nv", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], tarball=am_src)
     case "sdma_4_0_0": return load("am/sdma_4_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
     case "sdma_5_0_0": return load("am/sdma_5_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
-    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}//amdgpu/sdma_v6_0_0_pkt_open.h"],
-                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src),
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
+    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
+                                   args=["-I/opt/rocm/include", "-x", "c++"], tarball=am_src)
     case "smu_v13_0_0": return load("am/smu_v13_0_0",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
-                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], tarball=am_src),
+                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
+    case "smu_v13_0_6": return load("am/smu_v13_0_6",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu13_driver_if_v13_0_6"]]
+                                    +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, tarball=am_src)
     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/am.py

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

tinygrad/runtime/autogen/am/pm4_nv.py

@@ -1,36 +1,16 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, 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 +32,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 +49,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 +71,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 +107,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 +127,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

@@ -1,36 +1,16 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, 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 +32,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 +49,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 +71,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 +107,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 +127,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/sdma_4_0_0.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG(Struct): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION(ctypes.Union): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION_0(Struct): pass

tinygrad/runtime/autogen/am/sdma_5_0_0.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG(Struct): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION(ctypes.Union): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION_0(Struct): pass

tinygrad/runtime/autogen/am/sdma_6_0_0.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG(Struct): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION(ctypes.Union): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION_0(Struct): pass

tinygrad/runtime/autogen/am/smu_v13_0_0.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 FEATURE_PWR_DOMAIN_e = CEnum(ctypes.c_uint32)
 FEATURE_PWR_ALL = FEATURE_PWR_DOMAIN_e.define('FEATURE_PWR_ALL', 0)
 FEATURE_PWR_S5 = FEATURE_PWR_DOMAIN_e.define('FEATURE_PWR_S5', 1)
@@ -77,8 +76,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 +106,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 +155,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 +230,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 +387,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 +497,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/am/smu_v13_0_6.py

@@ -0,0 +1,531 @@
+# mypy: ignore-errors
+import ctypes
+from tinygrad.helpers import unwrap
+from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+PPSMC_Result = ctypes.c_uint32
+PPSMC_MSG = ctypes.c_uint32
+I2cControllerPort_e = CEnum(ctypes.c_uint32)
+I2C_CONTROLLER_PORT_0 = I2cControllerPort_e.define('I2C_CONTROLLER_PORT_0', 0)
+I2C_CONTROLLER_PORT_1 = I2cControllerPort_e.define('I2C_CONTROLLER_PORT_1', 1)
+I2C_CONTROLLER_PORT_COUNT = I2cControllerPort_e.define('I2C_CONTROLLER_PORT_COUNT', 2)
+
+I2cSpeed_e = CEnum(ctypes.c_uint32)
+UNSUPPORTED_1 = I2cSpeed_e.define('UNSUPPORTED_1', 0)
+I2C_SPEED_STANDARD_100K = I2cSpeed_e.define('I2C_SPEED_STANDARD_100K', 1)
+I2C_SPEED_FAST_400K = I2cSpeed_e.define('I2C_SPEED_FAST_400K', 2)
+I2C_SPEED_FAST_PLUS_1M = I2cSpeed_e.define('I2C_SPEED_FAST_PLUS_1M', 3)
+UNSUPPORTED_2 = I2cSpeed_e.define('UNSUPPORTED_2', 4)
+UNSUPPORTED_3 = I2cSpeed_e.define('UNSUPPORTED_3', 5)
+I2C_SPEED_COUNT = I2cSpeed_e.define('I2C_SPEED_COUNT', 6)
+
+I2cCmdType_e = CEnum(ctypes.c_uint32)
+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)
+
+ERR_CODE_e = CEnum(ctypes.c_uint32)
+CODE_DAGB0 = ERR_CODE_e.define('CODE_DAGB0', 0)
+CODE_EA0 = ERR_CODE_e.define('CODE_EA0', 5)
+CODE_UTCL2_ROUTER = ERR_CODE_e.define('CODE_UTCL2_ROUTER', 10)
+CODE_VML2 = ERR_CODE_e.define('CODE_VML2', 11)
+CODE_VML2_WALKER = ERR_CODE_e.define('CODE_VML2_WALKER', 12)
+CODE_MMCANE = ERR_CODE_e.define('CODE_MMCANE', 13)
+CODE_VIDD = ERR_CODE_e.define('CODE_VIDD', 14)
+CODE_VIDV = ERR_CODE_e.define('CODE_VIDV', 15)
+CODE_JPEG0S = ERR_CODE_e.define('CODE_JPEG0S', 16)
+CODE_JPEG0D = ERR_CODE_e.define('CODE_JPEG0D', 17)
+CODE_JPEG1S = ERR_CODE_e.define('CODE_JPEG1S', 18)
+CODE_JPEG1D = ERR_CODE_e.define('CODE_JPEG1D', 19)
+CODE_JPEG2S = ERR_CODE_e.define('CODE_JPEG2S', 20)
+CODE_JPEG2D = ERR_CODE_e.define('CODE_JPEG2D', 21)
+CODE_JPEG3S = ERR_CODE_e.define('CODE_JPEG3S', 22)
+CODE_JPEG3D = ERR_CODE_e.define('CODE_JPEG3D', 23)
+CODE_JPEG4S = ERR_CODE_e.define('CODE_JPEG4S', 24)
+CODE_JPEG4D = ERR_CODE_e.define('CODE_JPEG4D', 25)
+CODE_JPEG5S = ERR_CODE_e.define('CODE_JPEG5S', 26)
+CODE_JPEG5D = ERR_CODE_e.define('CODE_JPEG5D', 27)
+CODE_JPEG6S = ERR_CODE_e.define('CODE_JPEG6S', 28)
+CODE_JPEG6D = ERR_CODE_e.define('CODE_JPEG6D', 29)
+CODE_JPEG7S = ERR_CODE_e.define('CODE_JPEG7S', 30)
+CODE_JPEG7D = ERR_CODE_e.define('CODE_JPEG7D', 31)
+CODE_MMSCHD = ERR_CODE_e.define('CODE_MMSCHD', 32)
+CODE_SDMA0 = ERR_CODE_e.define('CODE_SDMA0', 33)
+CODE_SDMA1 = ERR_CODE_e.define('CODE_SDMA1', 34)
+CODE_SDMA2 = ERR_CODE_e.define('CODE_SDMA2', 35)
+CODE_SDMA3 = ERR_CODE_e.define('CODE_SDMA3', 36)
+CODE_HDP = ERR_CODE_e.define('CODE_HDP', 37)
+CODE_ATHUB = ERR_CODE_e.define('CODE_ATHUB', 38)
+CODE_IH = ERR_CODE_e.define('CODE_IH', 39)
+CODE_XHUB_POISON = ERR_CODE_e.define('CODE_XHUB_POISON', 40)
+CODE_SMN_SLVERR = ERR_CODE_e.define('CODE_SMN_SLVERR', 40)
+CODE_WDT = ERR_CODE_e.define('CODE_WDT', 41)
+CODE_UNKNOWN = ERR_CODE_e.define('CODE_UNKNOWN', 42)
+CODE_COUNT = ERR_CODE_e.define('CODE_COUNT', 43)
+
+GC_ERROR_CODE_e = CEnum(ctypes.c_uint32)
+SH_FED_CODE = GC_ERROR_CODE_e.define('SH_FED_CODE', 0)
+GCEA_CODE = GC_ERROR_CODE_e.define('GCEA_CODE', 1)
+SQ_CODE = GC_ERROR_CODE_e.define('SQ_CODE', 2)
+LDS_CODE = GC_ERROR_CODE_e.define('LDS_CODE', 3)
+GDS_CODE = GC_ERROR_CODE_e.define('GDS_CODE', 4)
+SP0_CODE = GC_ERROR_CODE_e.define('SP0_CODE', 5)
+SP1_CODE = GC_ERROR_CODE_e.define('SP1_CODE', 6)
+TCC_CODE = GC_ERROR_CODE_e.define('TCC_CODE', 7)
+TCA_CODE = GC_ERROR_CODE_e.define('TCA_CODE', 8)
+TCX_CODE = GC_ERROR_CODE_e.define('TCX_CODE', 9)
+CPC_CODE = GC_ERROR_CODE_e.define('CPC_CODE', 10)
+CPF_CODE = GC_ERROR_CODE_e.define('CPF_CODE', 11)
+CPG_CODE = GC_ERROR_CODE_e.define('CPG_CODE', 12)
+SPI_CODE = GC_ERROR_CODE_e.define('SPI_CODE', 13)
+RLC_CODE = GC_ERROR_CODE_e.define('RLC_CODE', 14)
+SQC_CODE = GC_ERROR_CODE_e.define('SQC_CODE', 15)
+TA_CODE = GC_ERROR_CODE_e.define('TA_CODE', 16)
+TD_CODE = GC_ERROR_CODE_e.define('TD_CODE', 17)
+TCP_CODE = GC_ERROR_CODE_e.define('TCP_CODE', 18)
+TCI_CODE = GC_ERROR_CODE_e.define('TCI_CODE', 19)
+GC_ROUTER_CODE = GC_ERROR_CODE_e.define('GC_ROUTER_CODE', 20)
+VML2_CODE = GC_ERROR_CODE_e.define('VML2_CODE', 21)
+VML2_WALKER_CODE = GC_ERROR_CODE_e.define('VML2_WALKER_CODE', 22)
+ATCL2_CODE = GC_ERROR_CODE_e.define('ATCL2_CODE', 23)
+GC_CANE_CODE = GC_ERROR_CODE_e.define('GC_CANE_CODE', 24)
+MP5_CODE_SMN_SLVERR = GC_ERROR_CODE_e.define('MP5_CODE_SMN_SLVERR', 40)
+MP5_CODE_UNKNOWN = GC_ERROR_CODE_e.define('MP5_CODE_UNKNOWN', 42)
+
+class SwI2cCmd_t(Struct): pass
+uint8_t = ctypes.c_ubyte
+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)),
+]
+PPCLK_e = CEnum(ctypes.c_uint32)
+PPCLK_VCLK = PPCLK_e.define('PPCLK_VCLK', 0)
+PPCLK_DCLK = PPCLK_e.define('PPCLK_DCLK', 1)
+PPCLK_SOCCLK = PPCLK_e.define('PPCLK_SOCCLK', 2)
+PPCLK_UCLK = PPCLK_e.define('PPCLK_UCLK', 3)
+PPCLK_FCLK = PPCLK_e.define('PPCLK_FCLK', 4)
+PPCLK_LCLK = PPCLK_e.define('PPCLK_LCLK', 5)
+PPCLK_COUNT = PPCLK_e.define('PPCLK_COUNT', 6)
+
+GpioIntPolarity_e = CEnum(ctypes.c_uint32)
+GPIO_INT_POLARITY_ACTIVE_LOW = GpioIntPolarity_e.define('GPIO_INT_POLARITY_ACTIVE_LOW', 0)
+GPIO_INT_POLARITY_ACTIVE_HIGH = GpioIntPolarity_e.define('GPIO_INT_POLARITY_ACTIVE_HIGH', 1)
+
+UCLK_DPM_MODE_e = CEnum(ctypes.c_uint32)
+UCLK_DPM_MODE_BANDWIDTH = UCLK_DPM_MODE_e.define('UCLK_DPM_MODE_BANDWIDTH', 0)
+UCLK_DPM_MODE_LATENCY = UCLK_DPM_MODE_e.define('UCLK_DPM_MODE_LATENCY', 1)
+
+class AvfsDebugTableAid_t(Struct): pass
+uint16_t = ctypes.c_uint16
+AvfsDebugTableAid_t._fields_ = [
+  ('avgPsmCount', (uint16_t * 30)),
+  ('minPsmCount', (uint16_t * 30)),
+  ('avgPsmVoltage', (ctypes.c_float * 30)),
+  ('minPsmVoltage', (ctypes.c_float * 30)),
+]
+class AvfsDebugTableXcd_t(Struct): pass
+AvfsDebugTableXcd_t._fields_ = [
+  ('avgPsmCount', (uint16_t * 30)),
+  ('minPsmCount', (uint16_t * 30)),
+  ('avgPsmVoltage', (ctypes.c_float * 30)),
+  ('minPsmVoltage', (ctypes.c_float * 30)),
+]
+class struct_smu_hw_power_state(Struct): pass
+struct_smu_hw_power_state._fields_ = [
+  ('magic', ctypes.c_uint32),
+]
+class struct_smu_power_state(Struct): pass
+enum_smu_state_ui_label = CEnum(ctypes.c_uint32)
+SMU_STATE_UI_LABEL_NONE = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_NONE', 0)
+SMU_STATE_UI_LABEL_BATTERY = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_BATTERY', 1)
+SMU_STATE_UI_TABEL_MIDDLE_LOW = enum_smu_state_ui_label.define('SMU_STATE_UI_TABEL_MIDDLE_LOW', 2)
+SMU_STATE_UI_LABEL_BALLANCED = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_BALLANCED', 3)
+SMU_STATE_UI_LABEL_MIDDLE_HIGHT = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_MIDDLE_HIGHT', 4)
+SMU_STATE_UI_LABEL_PERFORMANCE = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_PERFORMANCE', 5)
+SMU_STATE_UI_LABEL_BACO = enum_smu_state_ui_label.define('SMU_STATE_UI_LABEL_BACO', 6)
+
+enum_smu_state_classification_flag = CEnum(ctypes.c_uint32)
+SMU_STATE_CLASSIFICATION_FLAG_BOOT = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_BOOT', 1)
+SMU_STATE_CLASSIFICATION_FLAG_THERMAL = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_THERMAL', 2)
+SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE', 4)
+SMU_STATE_CLASSIFICATION_FLAG_RESET = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_RESET', 8)
+SMU_STATE_CLASSIFICATION_FLAG_FORCED = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_FORCED', 16)
+SMU_STATE_CLASSIFICATION_FLAG_USER_3D_PERFORMANCE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_USER_3D_PERFORMANCE', 32)
+SMU_STATE_CLASSIFICATION_FLAG_USER_2D_PERFORMANCE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_USER_2D_PERFORMANCE', 64)
+SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE', 128)
+SMU_STATE_CLASSIFICATION_FLAG_AC_OVERDIRVER_TEMPLATE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_AC_OVERDIRVER_TEMPLATE', 256)
+SMU_STATE_CLASSIFICATION_FLAG_UVD = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_UVD', 512)
+SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE_LOW = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE_LOW', 1024)
+SMU_STATE_CLASSIFICATION_FLAG_ACPI = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_ACPI', 2048)
+SMU_STATE_CLASSIFICATION_FLAG_HD2 = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_HD2', 4096)
+SMU_STATE_CLASSIFICATION_FLAG_UVD_HD = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_UVD_HD', 8192)
+SMU_STATE_CLASSIFICATION_FLAG_UVD_SD = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_UVD_SD', 16384)
+SMU_STATE_CLASSIFICATION_FLAG_USER_DC_PERFORMANCE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_USER_DC_PERFORMANCE', 32768)
+SMU_STATE_CLASSIFICATION_FLAG_DC_OVERDIRVER_TEMPLATE = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_DC_OVERDIRVER_TEMPLATE', 65536)
+SMU_STATE_CLASSIFICATION_FLAG_BACO = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_BACO', 131072)
+SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE2 = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE2', 262144)
+SMU_STATE_CLASSIFICATION_FLAG_ULV = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_ULV', 524288)
+SMU_STATE_CLASSIFICATION_FLAG_UVD_MVC = enum_smu_state_classification_flag.define('SMU_STATE_CLASSIFICATION_FLAG_UVD_MVC', 1048576)
+
+class struct_smu_state_classification_block(Struct): pass
+struct_smu_state_classification_block._fields_ = [
+  ('ui_label', enum_smu_state_ui_label),
+  ('flags', enum_smu_state_classification_flag),
+  ('bios_index', ctypes.c_int32),
+  ('temporary_state', ctypes.c_bool),
+  ('to_be_deleted', ctypes.c_bool),
+]
+class struct_smu_state_pcie_block(Struct): pass
+struct_smu_state_pcie_block._fields_ = [
+  ('lanes', ctypes.c_uint32),
+]
+enum_smu_refreshrate_source = CEnum(ctypes.c_uint32)
+SMU_REFRESHRATE_SOURCE_EDID = enum_smu_refreshrate_source.define('SMU_REFRESHRATE_SOURCE_EDID', 0)
+SMU_REFRESHRATE_SOURCE_EXPLICIT = enum_smu_refreshrate_source.define('SMU_REFRESHRATE_SOURCE_EXPLICIT', 1)
+
+class struct_smu_state_display_block(Struct): pass
+struct_smu_state_display_block._fields_ = [
+  ('disable_frame_modulation', ctypes.c_bool),
+  ('limit_refreshrate', ctypes.c_bool),
+  ('refreshrate_source', enum_smu_refreshrate_source),
+  ('explicit_refreshrate', ctypes.c_int32),
+  ('edid_refreshrate_index', ctypes.c_int32),
+  ('enable_vari_bright', ctypes.c_bool),
+]
+class struct_smu_state_memory_block(Struct): pass
+struct_smu_state_memory_block._fields_ = [
+  ('dll_off', ctypes.c_bool),
+  ('m3arb', ctypes.c_ubyte),
+  ('unused', (ctypes.c_ubyte * 3)),
+]
+class struct_smu_state_software_algorithm_block(Struct): pass
+struct_smu_state_software_algorithm_block._fields_ = [
+  ('disable_load_balancing', ctypes.c_bool),
+  ('enable_sleep_for_timestamps', ctypes.c_bool),
+]
+class struct_smu_temperature_range(Struct): pass
+struct_smu_temperature_range._fields_ = [
+  ('min', ctypes.c_int32),
+  ('max', ctypes.c_int32),
+  ('edge_emergency_max', ctypes.c_int32),
+  ('hotspot_min', ctypes.c_int32),
+  ('hotspot_crit_max', ctypes.c_int32),
+  ('hotspot_emergency_max', ctypes.c_int32),
+  ('mem_min', ctypes.c_int32),
+  ('mem_crit_max', ctypes.c_int32),
+  ('mem_emergency_max', ctypes.c_int32),
+  ('software_shutdown_temp', ctypes.c_int32),
+  ('software_shutdown_temp_offset', ctypes.c_int32),
+]
+class struct_smu_state_validation_block(Struct): pass
+struct_smu_state_validation_block._fields_ = [
+  ('single_display_only', ctypes.c_bool),
+  ('disallow_on_dc', ctypes.c_bool),
+  ('supported_power_levels', ctypes.c_ubyte),
+]
+class struct_smu_uvd_clocks(Struct): pass
+struct_smu_uvd_clocks._fields_ = [
+  ('vclk', ctypes.c_uint32),
+  ('dclk', ctypes.c_uint32),
+]
+enum_smu_power_src_type = CEnum(ctypes.c_uint32)
+SMU_POWER_SOURCE_AC = enum_smu_power_src_type.define('SMU_POWER_SOURCE_AC', 0)
+SMU_POWER_SOURCE_DC = enum_smu_power_src_type.define('SMU_POWER_SOURCE_DC', 1)
+SMU_POWER_SOURCE_COUNT = enum_smu_power_src_type.define('SMU_POWER_SOURCE_COUNT', 2)
+
+enum_smu_ppt_limit_type = CEnum(ctypes.c_uint32)
+SMU_DEFAULT_PPT_LIMIT = enum_smu_ppt_limit_type.define('SMU_DEFAULT_PPT_LIMIT', 0)
+SMU_FAST_PPT_LIMIT = enum_smu_ppt_limit_type.define('SMU_FAST_PPT_LIMIT', 1)
+
+enum_smu_ppt_limit_level = CEnum(ctypes.c_int32)
+SMU_PPT_LIMIT_MIN = enum_smu_ppt_limit_level.define('SMU_PPT_LIMIT_MIN', -1)
+SMU_PPT_LIMIT_CURRENT = enum_smu_ppt_limit_level.define('SMU_PPT_LIMIT_CURRENT', 0)
+SMU_PPT_LIMIT_DEFAULT = enum_smu_ppt_limit_level.define('SMU_PPT_LIMIT_DEFAULT', 1)
+SMU_PPT_LIMIT_MAX = enum_smu_ppt_limit_level.define('SMU_PPT_LIMIT_MAX', 2)
+
+enum_smu_memory_pool_size = CEnum(ctypes.c_uint32)
+SMU_MEMORY_POOL_SIZE_ZERO = enum_smu_memory_pool_size.define('SMU_MEMORY_POOL_SIZE_ZERO', 0)
+SMU_MEMORY_POOL_SIZE_256_MB = enum_smu_memory_pool_size.define('SMU_MEMORY_POOL_SIZE_256_MB', 268435456)
+SMU_MEMORY_POOL_SIZE_512_MB = enum_smu_memory_pool_size.define('SMU_MEMORY_POOL_SIZE_512_MB', 536870912)
+SMU_MEMORY_POOL_SIZE_1_GB = enum_smu_memory_pool_size.define('SMU_MEMORY_POOL_SIZE_1_GB', 1073741824)
+SMU_MEMORY_POOL_SIZE_2_GB = enum_smu_memory_pool_size.define('SMU_MEMORY_POOL_SIZE_2_GB', 2147483648)
+
+enum_smu_clk_type = CEnum(ctypes.c_uint32)
+SMU_GFXCLK = enum_smu_clk_type.define('SMU_GFXCLK', 0)
+SMU_VCLK = enum_smu_clk_type.define('SMU_VCLK', 1)
+SMU_DCLK = enum_smu_clk_type.define('SMU_DCLK', 2)
+SMU_VCLK1 = enum_smu_clk_type.define('SMU_VCLK1', 3)
+SMU_DCLK1 = enum_smu_clk_type.define('SMU_DCLK1', 4)
+SMU_ECLK = enum_smu_clk_type.define('SMU_ECLK', 5)
+SMU_SOCCLK = enum_smu_clk_type.define('SMU_SOCCLK', 6)
+SMU_UCLK = enum_smu_clk_type.define('SMU_UCLK', 7)
+SMU_DCEFCLK = enum_smu_clk_type.define('SMU_DCEFCLK', 8)
+SMU_DISPCLK = enum_smu_clk_type.define('SMU_DISPCLK', 9)
+SMU_PIXCLK = enum_smu_clk_type.define('SMU_PIXCLK', 10)
+SMU_PHYCLK = enum_smu_clk_type.define('SMU_PHYCLK', 11)
+SMU_FCLK = enum_smu_clk_type.define('SMU_FCLK', 12)
+SMU_SCLK = enum_smu_clk_type.define('SMU_SCLK', 13)
+SMU_MCLK = enum_smu_clk_type.define('SMU_MCLK', 14)
+SMU_PCIE = enum_smu_clk_type.define('SMU_PCIE', 15)
+SMU_LCLK = enum_smu_clk_type.define('SMU_LCLK', 16)
+SMU_OD_CCLK = enum_smu_clk_type.define('SMU_OD_CCLK', 17)
+SMU_OD_SCLK = enum_smu_clk_type.define('SMU_OD_SCLK', 18)
+SMU_OD_MCLK = enum_smu_clk_type.define('SMU_OD_MCLK', 19)
+SMU_OD_VDDC_CURVE = enum_smu_clk_type.define('SMU_OD_VDDC_CURVE', 20)
+SMU_OD_RANGE = enum_smu_clk_type.define('SMU_OD_RANGE', 21)
+SMU_OD_VDDGFX_OFFSET = enum_smu_clk_type.define('SMU_OD_VDDGFX_OFFSET', 22)
+SMU_OD_FAN_CURVE = enum_smu_clk_type.define('SMU_OD_FAN_CURVE', 23)
+SMU_OD_ACOUSTIC_LIMIT = enum_smu_clk_type.define('SMU_OD_ACOUSTIC_LIMIT', 24)
+SMU_OD_ACOUSTIC_TARGET = enum_smu_clk_type.define('SMU_OD_ACOUSTIC_TARGET', 25)
+SMU_OD_FAN_TARGET_TEMPERATURE = enum_smu_clk_type.define('SMU_OD_FAN_TARGET_TEMPERATURE', 26)
+SMU_OD_FAN_MINIMUM_PWM = enum_smu_clk_type.define('SMU_OD_FAN_MINIMUM_PWM', 27)
+SMU_CLK_COUNT = enum_smu_clk_type.define('SMU_CLK_COUNT', 28)
+
+class struct_smu_user_dpm_profile(Struct): pass
+struct_smu_user_dpm_profile._fields_ = [
+  ('fan_mode', ctypes.c_uint32),
+  ('power_limit', ctypes.c_uint32),
+  ('fan_speed_pwm', ctypes.c_uint32),
+  ('fan_speed_rpm', ctypes.c_uint32),
+  ('flags', ctypes.c_uint32),
+  ('user_od', ctypes.c_uint32),
+  ('clk_mask', (ctypes.c_uint32 * 28)),
+  ('clk_dependency', ctypes.c_uint32),
+]
+class struct_smu_table(Struct): pass
+class struct_amdgpu_bo(Struct): pass
+struct_smu_table._fields_ = [
+  ('size', ctypes.c_uint64),
+  ('align', ctypes.c_uint32),
+  ('domain', ctypes.c_ubyte),
+  ('mc_address', ctypes.c_uint64),
+  ('cpu_addr', ctypes.c_void_p),
+  ('bo', ctypes.POINTER(struct_amdgpu_bo)),
+  ('version', ctypes.c_uint32),
+]
+enum_smu_perf_level_designation = CEnum(ctypes.c_uint32)
+PERF_LEVEL_ACTIVITY = enum_smu_perf_level_designation.define('PERF_LEVEL_ACTIVITY', 0)
+PERF_LEVEL_POWER_CONTAINMENT = enum_smu_perf_level_designation.define('PERF_LEVEL_POWER_CONTAINMENT', 1)
+
+class struct_smu_performance_level(Struct): pass
+struct_smu_performance_level._fields_ = [
+  ('core_clock', ctypes.c_uint32),
+  ('memory_clock', ctypes.c_uint32),
+  ('vddc', ctypes.c_uint32),
+  ('vddci', ctypes.c_uint32),
+  ('non_local_mem_freq', ctypes.c_uint32),
+  ('non_local_mem_width', ctypes.c_uint32),
+]
+class struct_smu_clock_info(Struct): pass
+struct_smu_clock_info._fields_ = [
+  ('min_mem_clk', ctypes.c_uint32),
+  ('max_mem_clk', ctypes.c_uint32),
+  ('min_eng_clk', ctypes.c_uint32),
+  ('max_eng_clk', ctypes.c_uint32),
+  ('min_bus_bandwidth', ctypes.c_uint32),
+  ('max_bus_bandwidth', ctypes.c_uint32),
+]
+class struct_smu_bios_boot_up_values(Struct): pass
+struct_smu_bios_boot_up_values._fields_ = [
+  ('revision', ctypes.c_uint32),
+  ('gfxclk', ctypes.c_uint32),
+  ('uclk', ctypes.c_uint32),
+  ('socclk', ctypes.c_uint32),
+  ('dcefclk', ctypes.c_uint32),
+  ('eclk', ctypes.c_uint32),
+  ('vclk', ctypes.c_uint32),
+  ('dclk', ctypes.c_uint32),
+  ('vddc', ctypes.c_uint16),
+  ('vddci', ctypes.c_uint16),
+  ('mvddc', ctypes.c_uint16),
+  ('vdd_gfx', ctypes.c_uint16),
+  ('cooling_id', ctypes.c_ubyte),
+  ('pp_table_id', ctypes.c_uint32),
+  ('format_revision', ctypes.c_uint32),
+  ('content_revision', ctypes.c_uint32),
+  ('fclk', ctypes.c_uint32),
+  ('lclk', ctypes.c_uint32),
+  ('firmware_caps', ctypes.c_uint32),
+]
+enum_smu_table_id = CEnum(ctypes.c_uint32)
+SMU_TABLE_PPTABLE = enum_smu_table_id.define('SMU_TABLE_PPTABLE', 0)
+SMU_TABLE_WATERMARKS = enum_smu_table_id.define('SMU_TABLE_WATERMARKS', 1)
+SMU_TABLE_CUSTOM_DPM = enum_smu_table_id.define('SMU_TABLE_CUSTOM_DPM', 2)
+SMU_TABLE_DPMCLOCKS = enum_smu_table_id.define('SMU_TABLE_DPMCLOCKS', 3)
+SMU_TABLE_AVFS = enum_smu_table_id.define('SMU_TABLE_AVFS', 4)
+SMU_TABLE_AVFS_PSM_DEBUG = enum_smu_table_id.define('SMU_TABLE_AVFS_PSM_DEBUG', 5)
+SMU_TABLE_AVFS_FUSE_OVERRIDE = enum_smu_table_id.define('SMU_TABLE_AVFS_FUSE_OVERRIDE', 6)
+SMU_TABLE_PMSTATUSLOG = enum_smu_table_id.define('SMU_TABLE_PMSTATUSLOG', 7)
+SMU_TABLE_SMU_METRICS = enum_smu_table_id.define('SMU_TABLE_SMU_METRICS', 8)
+SMU_TABLE_DRIVER_SMU_CONFIG = enum_smu_table_id.define('SMU_TABLE_DRIVER_SMU_CONFIG', 9)
+SMU_TABLE_ACTIVITY_MONITOR_COEFF = enum_smu_table_id.define('SMU_TABLE_ACTIVITY_MONITOR_COEFF', 10)
+SMU_TABLE_OVERDRIVE = enum_smu_table_id.define('SMU_TABLE_OVERDRIVE', 11)
+SMU_TABLE_I2C_COMMANDS = enum_smu_table_id.define('SMU_TABLE_I2C_COMMANDS', 12)
+SMU_TABLE_PACE = enum_smu_table_id.define('SMU_TABLE_PACE', 13)
+SMU_TABLE_ECCINFO = enum_smu_table_id.define('SMU_TABLE_ECCINFO', 14)
+SMU_TABLE_COMBO_PPTABLE = enum_smu_table_id.define('SMU_TABLE_COMBO_PPTABLE', 15)
+SMU_TABLE_WIFIBAND = enum_smu_table_id.define('SMU_TABLE_WIFIBAND', 16)
+SMU_TABLE_COUNT = enum_smu_table_id.define('SMU_TABLE_COUNT', 17)
+
+PPSMC_Result_OK = 0x1
+PPSMC_Result_Failed = 0xFF
+PPSMC_Result_UnknownCmd = 0xFE
+PPSMC_Result_CmdRejectedPrereq = 0xFD
+PPSMC_Result_CmdRejectedBusy = 0xFC
+PPSMC_MSG_TestMessage = 0x1
+PPSMC_MSG_GetSmuVersion = 0x2
+PPSMC_MSG_GfxDriverReset = 0x3
+PPSMC_MSG_GetDriverIfVersion = 0x4
+PPSMC_MSG_EnableAllSmuFeatures = 0x5
+PPSMC_MSG_DisableAllSmuFeatures = 0x6
+PPSMC_MSG_RequestI2cTransaction = 0x7
+PPSMC_MSG_GetMetricsVersion = 0x8
+PPSMC_MSG_GetMetricsTable = 0x9
+PPSMC_MSG_GetEccInfoTable = 0xA
+PPSMC_MSG_GetEnabledSmuFeaturesLow = 0xB
+PPSMC_MSG_GetEnabledSmuFeaturesHigh = 0xC
+PPSMC_MSG_SetDriverDramAddrHigh = 0xD
+PPSMC_MSG_SetDriverDramAddrLow = 0xE
+PPSMC_MSG_SetToolsDramAddrHigh = 0xF
+PPSMC_MSG_SetToolsDramAddrLow = 0x10
+PPSMC_MSG_SetSystemVirtualDramAddrHigh = 0x11
+PPSMC_MSG_SetSystemVirtualDramAddrLow = 0x12
+PPSMC_MSG_SetSoftMinByFreq = 0x13
+PPSMC_MSG_SetSoftMaxByFreq = 0x14
+PPSMC_MSG_GetMinDpmFreq = 0x15
+PPSMC_MSG_GetMaxDpmFreq = 0x16
+PPSMC_MSG_GetDpmFreqByIndex = 0x17
+PPSMC_MSG_SetPptLimit = 0x18
+PPSMC_MSG_GetPptLimit = 0x19
+PPSMC_MSG_DramLogSetDramAddrHigh = 0x1A
+PPSMC_MSG_DramLogSetDramAddrLow = 0x1B
+PPSMC_MSG_DramLogSetDramSize = 0x1C
+PPSMC_MSG_GetDebugData = 0x1D
+PPSMC_MSG_HeavySBR = 0x1E
+PPSMC_MSG_SetNumBadHbmPagesRetired = 0x1F
+PPSMC_MSG_DFCstateControl = 0x20
+PPSMC_MSG_GetGmiPwrDnHyst = 0x21
+PPSMC_MSG_SetGmiPwrDnHyst = 0x22
+PPSMC_MSG_GmiPwrDnControl = 0x23
+PPSMC_MSG_EnterGfxoff = 0x24
+PPSMC_MSG_ExitGfxoff = 0x25
+PPSMC_MSG_EnableDeterminism = 0x26
+PPSMC_MSG_DisableDeterminism = 0x27
+PPSMC_MSG_DumpSTBtoDram = 0x28
+PPSMC_MSG_STBtoDramLogSetDramAddrHigh = 0x29
+PPSMC_MSG_STBtoDramLogSetDramAddrLow = 0x2A
+PPSMC_MSG_STBtoDramLogSetDramSize = 0x2B
+PPSMC_MSG_SetSystemVirtualSTBtoDramAddrHigh = 0x2C
+PPSMC_MSG_SetSystemVirtualSTBtoDramAddrLow = 0x2D
+PPSMC_MSG_GfxDriverResetRecovery = 0x2E
+PPSMC_MSG_TriggerVFFLR = 0x2F
+PPSMC_MSG_SetSoftMinGfxClk = 0x30
+PPSMC_MSG_SetSoftMaxGfxClk = 0x31
+PPSMC_MSG_GetMinGfxDpmFreq = 0x32
+PPSMC_MSG_GetMaxGfxDpmFreq = 0x33
+PPSMC_MSG_PrepareForDriverUnload = 0x34
+PPSMC_MSG_ReadThrottlerLimit = 0x35
+PPSMC_MSG_QueryValidMcaCount = 0x36
+PPSMC_MSG_McaBankDumpDW = 0x37
+PPSMC_MSG_GetCTFLimit = 0x38
+PPSMC_MSG_ClearMcaOnRead = 0x39
+PPSMC_MSG_QueryValidMcaCeCount = 0x3A
+PPSMC_MSG_McaBankCeDumpDW = 0x3B
+PPSMC_MSG_SelectPLPDMode = 0x40
+PPSMC_MSG_RmaDueToBadPageThreshold = 0x43
+PPSMC_MSG_SelectPstatePolicy = 0x44
+PPSMC_MSG_SetPhsDetWRbwThreshold = 0x45
+PPSMC_MSG_SetPhsDetWRbwFreqHigh = 0x46
+PPSMC_MSG_SetPhsDetWRbwFreqLow = 0x47
+PPSMC_MSG_SetPhsDetWRbwHystDown = 0x48
+PPSMC_MSG_SetPhsDetWRbwAlpha = 0x49
+PPSMC_MSG_SetPhsDetOnOff = 0x4A
+PPSMC_MSG_GetPhsDetResidency = 0x4B
+PPSMC_Message_Count = 0x4C
+PPSMC_RESET_TYPE_DRIVER_MODE_1_RESET = 0x1
+PPSMC_RESET_TYPE_DRIVER_MODE_2_RESET = 0x2
+PPSMC_RESET_TYPE_DRIVER_MODE_3_RESET = 0x3
+PPSMC_THROTTLING_LIMIT_TYPE_SOCKET = 0x1
+PPSMC_THROTTLING_LIMIT_TYPE_HBM = 0x2
+PPSMC_AID_THM_TYPE = 0x1
+PPSMC_CCD_THM_TYPE = 0x2
+PPSMC_XCD_THM_TYPE = 0x3
+PPSMC_HBM_THM_TYPE = 0x4
+PPSMC_PLPD_MODE_DEFAULT = 0x1
+PPSMC_PLPD_MODE_OPTIMIZED = 0x2
+SMU13_0_6_DRIVER_IF_VERSION = 0x08042024
+NUM_I2C_CONTROLLERS = 8
+I2C_CONTROLLER_ENABLED = 1
+I2C_CONTROLLER_DISABLED = 0
+MAX_SW_I2C_COMMANDS = 24
+CMDCONFIG_STOP_BIT = 0
+CMDCONFIG_RESTART_BIT = 1
+CMDCONFIG_READWRITE_BIT = 2
+CMDCONFIG_STOP_MASK = (1 << CMDCONFIG_STOP_BIT)
+CMDCONFIG_RESTART_MASK = (1 << CMDCONFIG_RESTART_BIT)
+CMDCONFIG_READWRITE_MASK = (1 << CMDCONFIG_READWRITE_BIT)
+IH_INTERRUPT_ID_TO_DRIVER = 0xFE
+IH_INTERRUPT_CONTEXT_ID_THERMAL_THROTTLING = 0x7
+THROTTLER_PROCHOT_BIT = 0
+THROTTLER_PPT_BIT = 1
+THROTTLER_THERMAL_SOCKET_BIT = 2
+THROTTLER_THERMAL_VR_BIT = 3
+THROTTLER_THERMAL_HBM_BIT = 4
+ClearMcaOnRead_UE_FLAG_MASK = 0x1
+ClearMcaOnRead_CE_POLL_MASK = 0x2
+int32_t = int
+SMU_THERMAL_MINIMUM_ALERT_TEMP = 0
+SMU_THERMAL_MAXIMUM_ALERT_TEMP = 255
+SMU_TEMPERATURE_UNITS_PER_CENTIGRADES = 1000
+SMU_FW_NAME_LEN = 0x24
+SMU_DPM_USER_PROFILE_RESTORE = (1 << 0)
+SMU_CUSTOM_FAN_SPEED_RPM = (1 << 1)
+SMU_CUSTOM_FAN_SPEED_PWM = (1 << 2)
+SMU_THROTTLER_PPT0_BIT = 0
+SMU_THROTTLER_PPT1_BIT = 1
+SMU_THROTTLER_PPT2_BIT = 2
+SMU_THROTTLER_PPT3_BIT = 3
+SMU_THROTTLER_SPL_BIT = 4
+SMU_THROTTLER_FPPT_BIT = 5
+SMU_THROTTLER_SPPT_BIT = 6
+SMU_THROTTLER_SPPT_APU_BIT = 7
+SMU_THROTTLER_TDC_GFX_BIT = 16
+SMU_THROTTLER_TDC_SOC_BIT = 17
+SMU_THROTTLER_TDC_MEM_BIT = 18
+SMU_THROTTLER_TDC_VDD_BIT = 19
+SMU_THROTTLER_TDC_CVIP_BIT = 20
+SMU_THROTTLER_EDC_CPU_BIT = 21
+SMU_THROTTLER_EDC_GFX_BIT = 22
+SMU_THROTTLER_APCC_BIT = 23
+SMU_THROTTLER_TEMP_GPU_BIT = 32
+SMU_THROTTLER_TEMP_CORE_BIT = 33
+SMU_THROTTLER_TEMP_MEM_BIT = 34
+SMU_THROTTLER_TEMP_EDGE_BIT = 35
+SMU_THROTTLER_TEMP_HOTSPOT_BIT = 36
+SMU_THROTTLER_TEMP_SOC_BIT = 37
+SMU_THROTTLER_TEMP_VR_GFX_BIT = 38
+SMU_THROTTLER_TEMP_VR_SOC_BIT = 39
+SMU_THROTTLER_TEMP_VR_MEM0_BIT = 40
+SMU_THROTTLER_TEMP_VR_MEM1_BIT = 41
+SMU_THROTTLER_TEMP_LIQUID0_BIT = 42
+SMU_THROTTLER_TEMP_LIQUID1_BIT = 43
+SMU_THROTTLER_VRHOT0_BIT = 44
+SMU_THROTTLER_VRHOT1_BIT = 45
+SMU_THROTTLER_PROCHOT_CPU_BIT = 46
+SMU_THROTTLER_PROCHOT_GFX_BIT = 47
+SMU_THROTTLER_PPM_BIT = 56
+SMU_THROTTLER_FIT_BIT = 57

tinygrad/runtime/autogen/am/smu_v14_0_2.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class struct_SMU14_Firmware_Footer(Struct): pass
 uint32_t = ctypes.c_uint32
 struct_SMU14_Firmware_Footer._packed_ = True

tinygrad/runtime/autogen/amd_gpu.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG(Struct): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION(ctypes.Union): pass
 class rocr_AMD_SDMA_PKT_COPY_LINEAR_TAG_HEADER_UNION_0(Struct): pass

tinygrad/runtime/autogen/comgr.py

@@ -1,17 +1,8 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/usr/local/lib/libamd_comgr.dylib'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/opt/homebrew/lib/libamd_comgr.dylib'))
-  except: pass
-  return None
-dll = dll()
-
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import os
+dll = DLL('comgr', [os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr'])
 amd_comgr_status_s = CEnum(ctypes.c_uint32)
 AMD_COMGR_STATUS_SUCCESS = amd_comgr_status_s.define('AMD_COMGR_STATUS_SUCCESS', 0)
 AMD_COMGR_STATUS_ERROR = amd_comgr_status_s.define('AMD_COMGR_STATUS_ERROR', 1)

tinygrad/runtime/autogen/comgr_3.py

@@ -1,17 +1,8 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/usr/local/lib/libamd_comgr.dylib'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/opt/homebrew/lib/libamd_comgr.dylib'))
-  except: pass
-  return None
-dll = dll()
-
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import os
+dll = DLL('comgr_3', [os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr'])
 amd_comgr_status_s = CEnum(ctypes.c_uint32)
 AMD_COMGR_STATUS_SUCCESS = amd_comgr_status_s.define('AMD_COMGR_STATUS_SUCCESS', 0)
 AMD_COMGR_STATUS_ERROR = amd_comgr_status_s.define('AMD_COMGR_STATUS_ERROR', 1)

tinygrad/runtime/autogen/cuda.py

@@ -1,14 +1,7 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('cuda')))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('cuda', 'cuda')
 cuuint32_t = ctypes.c_uint32
 cuuint64_t = ctypes.c_uint64
 CUdeviceptr_v2 = ctypes.c_uint64

tinygrad/runtime/autogen/hip.py

@@ -1,13 +1,8 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'))
-  except: pass
-  return None
-dll = dll()
-
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import os
+dll = DLL('hip', os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so')
 hipError_t = CEnum(ctypes.c_uint32)
 hipSuccess = hipError_t.define('hipSuccess', 0)
 hipErrorInvalidValue = hipError_t.define('hipErrorInvalidValue', 1)
@@ -1809,8 +1804,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 +2352,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

@@ -1,16 +1,8 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so'))
-  except: pass
-  try: return ctypes.CDLL(unwrap(find_library('hsa-runtime64')))
-  except: pass
-  return None
-dll = dll()
-
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import os
+dll = DLL('hsa', [os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so', 'hsa-runtime64'])
 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)
@@ -1294,18 +1286,13 @@ struct_hsa_amd_aie_ert_start_kernel_data_s._fields_ = [
 ]
 hsa_amd_aie_ert_start_kernel_data_t = struct_hsa_amd_aie_ert_start_kernel_data_s
 class struct_hsa_amd_aie_ert_packet_s(Struct): pass
-class struct_hsa_amd_aie_ert_packet_s_0(Struct): pass
-struct_hsa_amd_aie_ert_packet_s_0._fields_ = [
+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),
-]
-struct_hsa_amd_aie_ert_packet_s._anonymous_ = ['_0']
-struct_hsa_amd_aie_ert_packet_s._fields_ = [
-  ('header', hsa_amd_vendor_packet_header_t),
-  ('_0', struct_hsa_amd_aie_ert_packet_s_0),
   ('reserved0', uint64_t),
   ('reserved1', uint64_t),
   ('reserved2', uint64_t),
@@ -1783,13 +1770,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
 

tinygrad/runtime/autogen/ib.py

@@ -1,13 +1,7 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-def dll():
-  try: return ctypes.CDLL(unwrap(ibverbs), use_errno=True)
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('ib', 'ibverbs', use_errno=True)
 class union_ibv_gid(ctypes.Union): pass
 uint8_t = ctypes.c_ubyte
 class union_ibv_gid_global(Struct): pass

tinygrad/runtime/autogen/io_uring.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class struct_io_uring_sq(Struct): pass
 class struct_io_uring_sqe(Struct): pass
 __u8 = ctypes.c_ubyte
@@ -181,94 +180,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 +247,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 +291,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 +415,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 +432,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 +475,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/kfd.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class struct_kfd_ioctl_get_version_args(Struct): pass
 __u32 = ctypes.c_uint32
 struct_kfd_ioctl_get_version_args._fields_ = [

tinygrad/runtime/autogen/kgsl.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 enum_kgsl_user_mem_type = CEnum(ctypes.c_uint32)
 KGSL_USER_MEM_TYPE_PMEM = enum_kgsl_user_mem_type.define('KGSL_USER_MEM_TYPE_PMEM', 0)
 KGSL_USER_MEM_TYPE_ASHMEM = enum_kgsl_user_mem_type.define('KGSL_USER_MEM_TYPE_ASHMEM', 1)

tinygrad/runtime/autogen/libc.py

@@ -1,14 +1,7 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('c')), use_errno=True)
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('libc', 'c', use_errno=True)
 off_t = ctypes.c_int64
 mode_t = ctypes.c_uint32
 size_t = ctypes.c_uint64

tinygrad/runtime/autogen/libclang.py

@@ -1,14 +1,7 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('LIBCLANG_PATH', find_library('clang-20'))))
-  except: pass
-  return None
-dll = dll()
-
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('libclang', 'clang-20')
 CXIndex = ctypes.c_void_p
 class struct_CXTargetInfoImpl(Struct): pass
 CXTargetInfo = ctypes.POINTER(struct_CXTargetInfoImpl)

tinygrad/runtime/autogen/libusb.py

@@ -1,21 +1,7 @@
 # mypy: ignore-errors
-import ctypes, os
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(os.getenv('LIBUSB_PATH', find_library('usb-1.0'))))
-  except: pass
-  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),
-]
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('libusb', 'usb-1.0')
 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 +108,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/llvm.py

@@ -1,14 +1,8 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from tinygrad.runtime.support.llvm import LLVM_PATH
-def dll():
-  try: return ctypes.CDLL(unwrap(LLVM_PATH))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.helpers import WIN, OSX
+dll = DLL('llvm', 'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll' if WIN else '/opt/homebrew/opt/llvm@20/lib/libLLVM.dylib' if OSX else ['LLVM', 'LLVM-21', 'LLVM-20', 'LLVM-19', 'LLVM-18', 'LLVM-17', 'LLVM-16', 'LLVM-15', 'LLVM-14'])
 intmax_t = ctypes.c_int64
 try: (imaxabs:=dll.imaxabs).restype, imaxabs.argtypes = intmax_t, [intmax_t]
 except AttributeError: pass

tinygrad/runtime/autogen/metal.py

@@ -1,15 +1,8 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import objc
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('Metal')))
-  except: pass
-  return None
-dll = dll()
-
+dll = DLL('metal', 'Metal')
 class MTLDispatchThreadgroupsIndirectArguments(Struct): pass
 uint32_t = ctypes.c_uint32
 MTLDispatchThreadgroupsIndirectArguments._fields_ = [

tinygrad/runtime/autogen/nv.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 class MCTP_HEADER(Struct): pass
 NvU32 = ctypes.c_uint32
 NvU8 = ctypes.c_ubyte

tinygrad/runtime/autogen/nv_570.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 _anonenum0 = CEnum(ctypes.c_uint32)
 AES128_NONE = _anonenum0.define('AES128_NONE', 0)
 AES128_CTR = _anonenum0.define('AES128_CTR', 1)
@@ -1323,17 +1322,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 +6648,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

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 _anonenum0 = CEnum(ctypes.c_uint32)
 AES128_NONE = _anonenum0.define('AES128_NONE', 0)
 AES128_CTR = _anonenum0.define('AES128_CTR', 1)
@@ -1323,17 +1322,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 +7207,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/nvjitlink.py

@@ -1,14 +1,8 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('nvJitLink')))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import sysconfig
+dll = DLL('nvjitlink', 'nvJitLink', f'/usr/local/cuda/targets/{sysconfig.get_config_var("MULTIARCH").rsplit("-", 1)[0]}/lib')
 nvJitLinkResult = CEnum(ctypes.c_uint32)
 NVJITLINK_SUCCESS = nvJitLinkResult.define('NVJITLINK_SUCCESS', 0)
 NVJITLINK_ERROR_UNRECOGNIZED_OPTION = nvJitLinkResult.define('NVJITLINK_ERROR_UNRECOGNIZED_OPTION', 1)

tinygrad/runtime/autogen/nvrtc.py

@@ -1,14 +1,8 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('nvrtc')))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+import sysconfig
+dll = DLL('nvrtc', 'nvrtc', f'/usr/local/cuda/targets/{sysconfig.get_config_var("MULTIARCH").rsplit("-", 1)[0]}/lib')
 nvrtcResult = CEnum(ctypes.c_uint32)
 NVRTC_SUCCESS = nvrtcResult.define('NVRTC_SUCCESS', 0)
 NVRTC_ERROR_OUT_OF_MEMORY = nvrtcResult.define('NVRTC_ERROR_OUT_OF_MEMORY', 1)

tinygrad/runtime/autogen/opencl.py

@@ -1,14 +1,7 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('OpenCL')))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('opencl', 'OpenCL')
 class struct__cl_platform_id(Struct): pass
 cl_platform_id = ctypes.POINTER(struct__cl_platform_id)
 class struct__cl_device_id(Struct): pass

tinygrad/runtime/autogen/pci.py

@@ -1,8 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 PCI_CFG_SPACE_SIZE = 256
 PCI_CFG_SPACE_EXP_SIZE = 4096
 PCI_STD_HEADER_SIZEOF = 64

tinygrad/runtime/autogen/qcom_dsp.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import ctypes
-from tinygrad.helpers import unwrap
-from tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 ion_user_handle_t = ctypes.c_int32
 enum_ion_heap_type = CEnum(ctypes.c_uint32)
 ION_HEAP_TYPE_SYSTEM = enum_ion_heap_type.define('ION_HEAP_TYPE_SYSTEM', 0)