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

.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,18 +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
-    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -169,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]
@@ -541,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)
@@ -554,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
 

.github/workflows/test.yml

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

CLAUDE.md

@@ -95,6 +95,8 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
 ## 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
 
@@ -132,6 +134,18 @@ The schedule cache strips values from BIND nodes so different bound values (e.g.
 - 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

examples/mlperf/model_train.py

@@ -1314,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):
@@ -1409,7 +1411,7 @@ 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):

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

@@ -2,6 +2,10 @@ 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, pci_dev, dma_regions=None):
     self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
@@ -9,6 +13,8 @@ class AMDFake(AMDev):
     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

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/thunder/tiny/tk/group.py

@@ -387,7 +387,7 @@ class Group:
       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]):
+      for outer in self.ker.range(dst.shape[-2], track=False):
         src_i += outer * reductions + (laneid % reductions)
 
         src_load = srcf[src_i]
@@ -404,7 +404,29 @@ class Group:
     dst, src = cast(UOp, dst), cast(UOp, src)
     assert isinstance(dst.dtype, PtrDType) and isinstance(src.dtype, PtrDType)
     dst_dtype, src_dtype = dst.dtype, src.dtype
-    if src_dtype.addrspace == AddrSpace.REG and dst_dtype.addrspace == AddrSpace.GLOBAL and isinstance(src, RT):
+    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:])
 

extra/thunder/tiny/tk/kernel.py

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

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

@@ -63,6 +63,28 @@ class TestMultiTensor(unittest.TestCase):
     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_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_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

@@ -31,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)
@@ -46,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)
 
@@ -152,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

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

@@ -10,6 +10,7 @@ class TestLLMServer(unittest.TestCase):
     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]))

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

@@ -4,7 +4,8 @@ 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)}
@@ -20,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 = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+    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]
@@ -49,41 +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]
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
   freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
   freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).contiguous()
+  return freqs.cos().cat(freqs.sin(), dim=-1).contiguous()
 
 def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
-  B, H, T, Hd = x.shape
-  assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
-  x_pairs     = x.reshape(B, H, T, Hd//2, 2)
-  cos = freqs_cis.reshape(1, 1, T, Hd//2, 2)[..., 0]
-  sin = freqs_cis.reshape(1, 1, T, Hd//2, 2)[..., 1]
-  return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin,
-                      x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd)
+  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)
@@ -99,8 +104,10 @@ class TransformerBlock:
     k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
     v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
 
+    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)[start_pos:start_pos+T]  # type: ignore
+    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)
 
@@ -128,8 +135,10 @@ class TransformerBlock:
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:
-  def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, max_context):
-    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context) for _ in range(num_blocks)]
+  def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, 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)
@@ -159,9 +168,18 @@ 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 (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
@@ -182,16 +200,56 @@ class Transformer:
 
 models = {
   "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}
@@ -214,7 +272,7 @@ class Handler(HTTPRequestHandler):
     if DEBUG >= 1: print(json.dumps(body, indent=2))
     if self.path == "/v1/chat/completions":
       # extract tokens
-      ids = [bos_id]
+      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
@@ -225,7 +283,8 @@ class Handler(HTTPRequestHandler):
             if c["type"] == "text": ids += tok.encode(c["text"])
             else: raise RuntimeError(f"unhandled type: {c['type']}")
         else: raise RuntimeError(f"unknown content type: {type(content)}")
-        ids += tok.role("assistant")
+        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))
@@ -242,8 +301,8 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
   parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
-  parser.add_argument("--serve", action="store_true", help="Run OpenAI compatible API")
-  parser.add_argument("--benchmark", action="store_true", help="Benchmark tok/s")
+  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
@@ -254,24 +313,24 @@ if __name__ == "__main__":
   if args.benchmark:
     param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
     gen = model.generate([0], 0)
-    for _ in range(20):
+    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']
 
   # start server
-  if args.serve: TCPServerWithReuse(('', 11434), Handler).serve_forever()
+  if args.serve: TCPServerWithReuse(('', args.serve), Handler).serve_forever()
 
-  ids: list[int] = [bos_id]
+  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/engine/schedule.py

@@ -1,6 +1,6 @@
 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
@@ -13,14 +13,13 @@ 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(sched_sink:UOp) -> list[ScheduleItem]:
+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]] = {}
@@ -48,33 +47,21 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
         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)
+        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)
-        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))
+        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}")
@@ -83,10 +70,11 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
         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):
@@ -99,9 +87,12 @@ def create_schedule(sched_sink:UOp) -> list[ScheduleItem]:
             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
+  return pre_schedule, UOp.sink(*buf_uops_list)
 
 from tinygrad.engine.memory import memory_planner
 from tinygrad.schedule.rangeify import get_rangeify_map
@@ -140,7 +131,7 @@ pm_post_sched_cache = PatternMatcher([
   (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)),
 ])
 
-schedule_cache: dict[bytes, tuple[UOp, UOp]] = {}
+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
@@ -169,22 +160,43 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
     tensor_map |= get_rangeify_map(big_sink_cache)
     big_sink = big_sink_cache.substitute(tensor_map, name="Apply Kernelize Map")
 
-    # save in schedule cache
-    tensor_map_sink = UOp.sink(*flatten([(k,v) for k,v in tensor_map.items()]))
-    schedule_cache[sched_cache_key] = (big_sink, tensor_map_sink)
+    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
-    big_sink, tensor_map_sink = sc_ret
+    pre_schedule, combined_sink = sc_ret
 
-  # replace all the LUNIQUEs with UNIQUEs
+  # replace all the LUNIQUEs with UNIQUEs (single graph_rewrite for everything)
   input_buffers_reverse = {v:k for k,v in input_buffers.items()}
-  big_sink = graph_rewrite(big_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for sched cache")
-  tm_src = graph_rewrite(tensor_map_sink, pm_post_sched_cache, ctx=input_buffers_reverse, name="unrewrite for tensor map").src
+  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)}
 
-  # create the schedule
-  schedule = create_schedule(big_sink)
+  # 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)
 
   # extract var_vals from BINDs that were stripped (only if there are kernels)
@@ -196,9 +208,6 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li
         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
 
-  # remove all AFTERs, after scheduling, the tensors are just buffers
-  tensor_map |= {u:u.buf_uop for u in big_sink.toposort() if u.op is Ops.AFTER}
-
   if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
     print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
           f" | {' cache hit' if sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\

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/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/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/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/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,46 +75,43 @@ 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 "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",
@@ -127,13 +130,13 @@ def __getattr__(nm):
           *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
-  prolog=["import gzip, base64", "from tinygrad.helpers import OSX"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
+  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}")

tinygrad/runtime/autogen/am/__init__.py

@@ -6,8 +6,9 @@ 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"],
@@ -18,6 +19,8 @@ def __getattr__(nm):
                                    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"], 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,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 union_PM4_MES_TYPE_3_HEADER(ctypes.Union): 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)

tinygrad/runtime/autogen/am/pm4_soc15.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 union_PM4_MES_TYPE_3_HEADER(ctypes.Union): 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)

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)

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)

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)

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

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,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('LIBUSB_PATH', find_library('usb-1.0'))))
-  except: pass
-  return None
-dll = dll()
-
+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)

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

@@ -1,24 +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
+import ctypes
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
 import gzip, base64
-from tinygrad.helpers import OSX
-from ctypes.util import find_library
-def dll():
-  try: return ctypes.CDLL(unwrap(find_library('tinymesa_cpu')))
-  except: pass
-  try: return ctypes.CDLL(unwrap((BASE:=os.getenv('MESA_PATH', f"/usr{'/local/' if OSX else '/'}lib"))+'/libtinymesa_cpu'+(EXT:='.dylib' if OSX else '.so')))
-  except: pass
-  try: return ctypes.CDLL(unwrap(f'{BASE}/libtinymesa{EXT}'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/opt/homebrew/lib/libtinymesa_cpu.dylib'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/opt/homebrew/lib/libtinymesa.dylib'))
-  except: pass
-  return None
-dll = dll()
-
+dll = DLL('mesa', ['tinymesa_cpu', 'tinymesa'])
 class struct_u_printf_info(Struct): pass
 u_printf_info = struct_u_printf_info
 uint32_t = ctypes.c_uint32

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)

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)

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)

tinygrad/runtime/autogen/rocprof.py

@@ -1,18 +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('rocprof-trace-decoder')))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/usr/local/lib/rocprof-trace-decoder.so'))
-  except: pass
-  try: return ctypes.CDLL(unwrap('/usr/local/lib/rocprof-trace-decoder.dylib'))
-  except: pass
-  return None
-dll = dll()
-
+from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
+dll = DLL('rocprof', ['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')])
 rocprofiler_thread_trace_decoder_status_t = CEnum(ctypes.c_uint32)
 ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS', 0)
 ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = rocprofiler_thread_trace_decoder_status_t.define('ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR', 1)

tinygrad/runtime/autogen/sqtt.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_sqtt_data_info(Struct): pass
 uint32_t = ctypes.c_uint32
 class struct_sqtt_data_info_0(ctypes.Union): pass

tinygrad/runtime/autogen/vfio.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_vfio_info_cap_header(Struct): pass
 __u16 = ctypes.c_uint16
 __u32 = ctypes.c_uint32

tinygrad/runtime/autogen/webgpu.py

@@ -1,14 +1,9 @@
 # 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.webgpu import WEBGPU_PATH
-def dll():
-  try: return ctypes.CDLL(unwrap(WEBGPU_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
+import sysconfig, os
+dll = DLL('webgpu', os.path.join(sysconfig.get_paths()['purelib'], 'pydawn', 'lib', 'libwebgpu_dawn.dll') if WIN else 'webgpu_dawn')
 WGPUFlags = ctypes.c_uint64
 WGPUBool = ctypes.c_uint32
 class struct_WGPUAdapterImpl(Struct): pass

tinygrad/runtime/graph/remote.py

@@ -1,113 +0,0 @@
-import time, itertools
-from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.engine.realize import CompiledRunner, BufferXfer, ExecItem
-from tinygrad.device import Device, Compiled, Buffer
-from tinygrad.runtime.ops_remote import RemoteDevice, RemoteConnection, RemoteRequest, GraphComputeItem, Transfer, GraphAlloc, GraphFree, GraphExec
-from tinygrad.runtime.ops_remote import BatchTransfer, Event, Wait
-from tinygrad.helpers import unwrap, flatten, dedup
-from enum import Enum, auto
-from dataclasses import replace
-from collections import defaultdict
-from typing import cast
-
-class StagingType(Enum): NONE = auto(); GRAPH = auto(); TRANSFER = auto() # noqa: E702
-
-def rd(dev:Compiled) -> RemoteDevice: return cast(RemoteDevice, dev)
-def dev_key(dev:RemoteDevice): return dev.conn if dev.properties.graph_supports_multi else dev
-def map_rawbuf(rawbuf:Buffer): return (cast(RemoteDevice, Device[rawbuf.device]).session, rawbuf._buf)
-
-class RemoteGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: list[ExecItem], rawbufs: list[Buffer], var_vals: dict[str, int]):
-    super().__init__(jit_cache, rawbufs, var_vals)
-    devices = dedup(flatten([[Device[unwrap(buf).device] for buf in ji.bufs] for ji in jit_cache]))
-    c2d = {device.conn: device for device in devices}
-    self.handle_indexes = {map_rawbuf(rawbufs[i]): i for i in sorted(dedup(self.input_replace.values()))}
-
-    self.template: list[RemoteRequest] = []
-
-    stagings: dict[RemoteDevice|RemoteConnection, list[GraphComputeItem|Transfer]] = defaultdict(list)
-    clobbered_buffers: set[Buffer] = set()
-    cur_staging_type: StagingType = StagingType.NONE
-
-    def _flush(new_staging_type:StagingType, force_break:bool=False):
-      nonlocal cur_staging_type
-      if cur_staging_type == new_staging_type and not force_break: return
-      # Pre-sync
-      if cur_staging_type == StagingType.TRANSFER:
-        for sdev,ddev in itertools.permutations(c2d.values(), 2):
-          self.template.append(Event(ddev.session, event:=next(ddev.event_num), session=sdev.session))
-          self.template.append(Wait(event, session=ddev.session))
-      # Flush
-      for dev in devices:
-        dk = dev_key(dev)
-        staging = stagings[dk]
-        if not staging: continue
-        match cur_staging_type:
-          case StagingType.GRAPH:
-            bufs = tuple(map_rawbuf(rawbufs[i]) for i in sorted(dedup(self.input_replace.values())) if dev_key(rd(Device[rawbufs[i].device])) == dk)
-            dev.q(GraphAlloc(graph_num:=next(dev.graph_num), tuple(staging), tuple(bufs), var_vals))
-            self.template.append(GraphExec(graph_num, bufs, var_vals, wait=False, session=dev.session))
-          case StagingType.TRANSFER:
-            st = cast(list[Transfer], staging)
-            for host in dedup(t.dsession.host for t in st):
-              sbuffer_nums = [(unwrap(t.session), t.buffer_num) for t in st if t.dsession.host == host]
-              dbuffer_nums = [(t.dsession, t.dbuffer_num) for t in st if t.dsession.host == host]
-              self.template.append(BatchTransfer(sbuffer_nums, dbuffer_nums, session=dev.session))
-        staging.clear()
-      # Post-sync
-      if cur_staging_type == StagingType.TRANSFER:
-        for sdev,ddev in itertools.permutations(c2d.values(), 2):
-          self.template.append(Event(ddev.session, event:=next(ddev.event_num), session=sdev.session))
-          self.template.append(Wait(event, session=ddev.session))
-      cur_staging_type = new_staging_type
-      clobbered_buffers.clear()
-
-    for ji in jit_cache:
-      match ji.prg:
-        case CompiledRunner():
-          _flush(StagingType.GRAPH)
-          gi = GraphComputeItem(ji.prg.dev.session, ji.prg._prg.name, ji.prg._prg.datahash, tuple(unwrap(buf)._buf for buf in ji.bufs),
-                                tuple(ji.prg.p.vars), ji.fixedvars, tuple(ji.prg.p.ins), tuple(ji.prg.p.outs),
-                                tuple(ji.prg.p.global_size) if ji.prg.p.global_size is not None else None,
-                                tuple(ji.prg.p.local_size) if ji.prg.p.local_size is not None else None)
-          stagings[dev_key(ji.prg.dev)].append(gi)
-        case BufferXfer():
-          dest, src = ji.bufs[0:2]
-          dest_dev, src_dev = cast(RemoteDevice, Device[unwrap(dest).device]), cast(RemoteDevice, Device[unwrap(src).device])
-          assert dest is not None and src is not None, ji
-          ti = Transfer(session=src_dev.session, buffer_num=src._buf, dsession=dest_dev.session, dbuffer_num=dest._buf)
-          if dev_key(dest_dev) == dev_key(src_dev):
-            _flush(StagingType.GRAPH)
-            stagings[dev_key(src_dev)].append(ti)
-          elif dest_dev.conn == src_dev.conn:
-            _flush(StagingType.NONE)
-            self.template.append(ti)
-          else:
-            _flush(StagingType.TRANSFER, force_break=src in clobbered_buffers)
-            clobbered_buffers.add(dest)
-            stagings[dev_key(src_dev)].append(ti)
-        case _: raise NotImplementedError(ji.prg)
-    _flush(StagingType.NONE)
-  def __del__(self):
-    for req in self.template:
-      match req:
-        case GraphExec(): RemoteConnection(unwrap(req.session).host).q(GraphFree(req.graph_num, session=req.session))
-  def __call__(self, rawbufs: list[Buffer], var_vals: dict[str, int], wait=False):
-    if wait: st = time.perf_counter()
-    rmap = {orig: map_rawbuf(rawbufs[replace_idx]) for orig,replace_idx in self.handle_indexes.items()}
-    for req in self.template:
-      match req:
-        case GraphExec():
-          req = replace(req, bufs=tuple(rmap[buf] for buf in req.bufs), var_vals=var_vals, wait=wait)
-        case Transfer():
-          if (req.session, req.buffer_num) in rmap: req = replace(req, buffer_num=rmap[(req.session, req.buffer_num)][1])
-          if (req.dsession, req.dbuffer_num) in rmap: req = replace(req, dbuffer_num=rmap[(req.dsession, req.dbuffer_num)][1])
-        case BatchTransfer():
-          req = replace(req, sbuffer_nums=[rmap.get(b, b) for b in req.sbuffer_nums], dbuffer_nums=[rmap.get(b, b) for b in req.dbuffer_nums])
-        case Event()|Wait():
-          pass # event number can be reused
-        case _: raise NotImplementedError(req)
-      RemoteConnection(unwrap(req.session).host).q(req)
-    if wait:
-      RemoteConnection(unwrap(req.session).host).batch_submit()
-      return time.perf_counter() - st

tinygrad/runtime/ops_amd.py

@@ -796,7 +796,7 @@ class PCIIface(PCIIfaceBase):
   gpus:ClassVar[list[str]] = []
 
   def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x744c, 0x7480, 0x7550, 0x7590])], bars=[0, 2, 5], vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x74a1, 0x744c, 0x7480, 0x7550, 0x7590, 0x75a0])], bars=[0, 2, 5], vram_bar=0,
       va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
     self._setup_adev(self.pci_dev)
     self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
@@ -808,26 +808,32 @@ class PCIIface(PCIIfaceBase):
     self.ip_versions = self.dev_impl.ip_ver
 
     gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
-    array_count = self.dev_impl.gc_info.gc_num_sa_per_se * self.dev_impl.gc_info.gc_num_se
-    self.props = {'cu_per_simd_array': (cu_per_sa:=2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)),
-      'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count, 'gfx_target_version': gfxver,
+    if self.dev_impl.gc_info.header.version_major == 2:
+      cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se
+    else:
+      cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se
+
+    array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs
+    self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count,
       'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
-      'simd_arrays_per_engine': self.dev_impl.gc_info.gc_num_sa_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size,
-      'num_xcc': self.dev_impl.gfx.xccs}
+      'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
+      'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
 
   def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
     assert cwsr_buffer is None, "no cwsr buffer for am"
 
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
+      pv = self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
                                     doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
     else:
-      self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
+      pv = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
         eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0,
         aql=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL))
 
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
-      read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')])
+      read_ptrs=[gart.cpu_view().view(offset=rptr, size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=wptr, size=8, fmt='Q')], put_value=pv)
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):

tinygrad/runtime/ops_cuda.py

@@ -11,7 +11,7 @@ if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore
 
 def check(status):
   if status != 0:
-    error = ctypes.string_at(init_c_var(ctypes.c_char_p(), lambda x: cuda.cuGetErrorString(status, x))).decode()
+    error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, x))).decode()
     raise RuntimeError(f"CUDA Error {status}, {error}")
 
 def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]:

tinygrad/runtime/ops_remote.py

@@ -1,491 +0,0 @@
-# the REMOTE=1 device is a process boundary between the frontend/runtime
-# normally tinygrad is    frontend <-> middleware <-> runtime <-> hardware
-# with REMOTE tinygrad is  frontend <-> middleware <-> RemoteDevice ///HTTP/// remote_server <-> runtime <-> hardware
-# this client and server can be on the same machine, same network, or just same internet
-# it should be a secure (example: no use of pickle) boundary. HTTP is used for RPC
-
-from __future__ import annotations
-from typing import Callable, Iterator, Any, cast
-from collections import defaultdict
-from dataclasses import dataclass, field, replace
-import multiprocessing, threading, functools, itertools, asyncio, http, http.client, hashlib, time, os, binascii, struct, ast, contextlib, weakref
-import traceback, builtins
-from tinygrad.renderer import Renderer, ProgramSpec
-from tinygrad.dtype import DTYPES_DICT, dtypes
-from tinygrad.uop.ops import UOp, Ops, Variable, sint
-from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, LazySeq, Timing
-from tinygrad.engine.jit import GraphRunner, MultiGraphRunner, ExecItem, graph_class
-from tinygrad.engine.realize import CompiledRunner, BufferXfer
-from tinygrad.device import Compiled, Buffer, Allocator, Compiler, Device, BufferSpec, CompilerSet, CompilerPair
-from tinygrad.runtime.support.ib import IBCtx, IBConn, SGE
-
-# ***** API *****
-
-@dataclass(frozen=True)
-class SessionKey: host: str; idx: int; nonce: str # noqa: E702
-
-@dataclass(frozen=True)
-class RemoteRequest: session: SessionKey|None = field(default=None, kw_only=True)
-
-@dataclass(frozen=True)
-class SessionFree(RemoteRequest): pass
-
-@dataclass(frozen=True)
-class RemoteProperties:
-  real_device: str
-  renderer: tuple[str, str, tuple[Any, ...]]
-  offset_supported: bool
-  graph_supported: bool
-  graph_supports_multi: bool
-  ib_gid: bytes|None
-
-@dataclass(frozen=True)
-class RemoteException:
-  exc: Exception
-  trace: str = ""
-
-@dataclass(frozen=True)
-class GetProperties(RemoteRequest): pass
-
-@dataclass(frozen=True)
-class Event(RemoteRequest): event_session: SessionKey; event: int # noqa: E702
-
-@dataclass(frozen=True)
-class Wait(RemoteRequest): event: int
-
-@dataclass(frozen=True)
-class IBConnect(RemoteRequest): host: str; gid: bytes; qp_num: int # noqa: E702
-
-@dataclass(frozen=True)
-class BufferAlloc(RemoteRequest): buffer_num: int; size: int; options: BufferSpec # noqa: E702
-
-@dataclass(frozen=True)
-class BufferOffset(RemoteRequest): buffer_num: int; size: int; offset: int; sbuffer_num: int # noqa: E702
-
-@dataclass(frozen=True)
-class BufferIOVAS(RemoteRequest): buffer_nums: list[tuple[SessionKey, int]] # noqa: E702
-
-@dataclass(frozen=True)
-class BufferFree(RemoteRequest): buffer_num: int # noqa: E702
-
-@dataclass(frozen=True)
-class CopyIn(RemoteRequest): buffer_num: int; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class CopyOut(RemoteRequest): buffer_num: int
-
-@dataclass(frozen=True)
-class Transfer(RemoteRequest): buffer_num: int; dsession: SessionKey; dbuffer_num: int # noqa: E702
-
-@dataclass(frozen=True)
-class BatchTransfer(RemoteRequest):
-  sbuffer_nums: list[tuple[SessionKey, int]]
-  dbuffer_nums: list[tuple[SessionKey, int]]
-
-@dataclass(frozen=True)
-class ProgramAlloc(RemoteRequest): name: str; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class ProgramFree(RemoteRequest): name: str; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class ProgramExec(RemoteRequest):
-  name: str; datahash: str; bufs: tuple[int, ...]; vals: tuple[int, ...] # noqa: E702
-  global_size: tuple[int, ...]|None; local_size: tuple[int, ...]|None; wait: bool # noqa: E702
-
-@dataclass(frozen=True)
-class GraphComputeItem:
-  session: SessionKey
-  name: str
-  datahash: str
-  bufs: tuple[int, ...]
-  vars: tuple[Variable, ...]
-  fixedvars: dict[str, int]
-  ins: tuple[int, ...]
-  outs: tuple[int, ...]
-  global_size: tuple[sint, ...]|None
-  local_size: tuple[sint, ...]|None
-
-@dataclass(frozen=True)
-class GraphAlloc(RemoteRequest):
-  graph_num: int
-  jit_cache: tuple[GraphComputeItem|Transfer, ...]
-  bufs: tuple[tuple[SessionKey, int], ...]
-  var_vals: dict[str, int]
-
-@dataclass(frozen=True)
-class GraphFree(RemoteRequest):
-  graph_num: int
-
-@dataclass(frozen=True)
-class GraphExec(RemoteRequest):
-  graph_num: int
-  bufs: tuple[tuple[SessionKey, int], ...]
-  var_vals: dict[str, int]
-  wait: bool
-
-# for safe deserialization
-eval_excs = [v for k,v in builtins.__dict__.items() if isinstance(v, type) and issubclass(v, Exception) and not k.endswith("Warning")]
-eval_globals = {x.__name__:x for x in [SessionKey, SessionFree, RemoteProperties, GetProperties, Event, Wait, BufferAlloc, BufferOffset, BufferIOVAS,
-                                       BufferFree, CopyIn, CopyOut, Transfer, BatchTransfer, IBConnect, ProgramAlloc, ProgramFree, ProgramExec,
-                                       GraphComputeItem, GraphAlloc, GraphFree, GraphExec, BufferSpec, UOp, Ops, dtypes, RemoteException] + eval_excs}
-attribute_whitelist: dict[Any, set[str]] = {dtypes: {*DTYPES_DICT.keys(), 'imagef', 'imageh'}, Ops: {x.name for x in Ops}}
-eval_fxns = {ast.Constant: lambda x: x.value, ast.Tuple: lambda x: tuple(map(safe_eval, x.elts)), ast.List: lambda x: list(map(safe_eval, x.elts)),
-  ast.Dict: lambda x: {safe_eval(k):safe_eval(v) for k,v in zip(x.keys, x.values)},
-  ast.Call: lambda x: safe_eval(x.func)(*[safe_eval(arg) for arg in x.args], **{kwarg.arg: safe_eval(kwarg.value) for kwarg in x.keywords}),
-  ast.Name: lambda x: eval_globals[x.id], ast.Attribute: lambda x: safe_getattr(safe_eval(x.value), x.attr)}
-def safe_getattr(value, attr):
-  assert attr in attribute_whitelist.get(value, set()), f'getattr({value}, {repr(attr)}) is not whitelisted'
-  return getattr(value, attr)
-def safe_eval(node): return eval_fxns[node.__class__](node)
-
-class BatchRequest:
-  def __init__(self):
-    self._q: list[RemoteRequest] = []
-    self._h: dict[str, bytes] = {}
-  def h(self, d:bytes|memoryview) -> str:
-    datahash = hashlib.sha256(d).hexdigest() # NOTE: this is very slow, should use blake3 on gpu instead
-    if datahash not in self._h:
-      self._h[datahash] = bytes.fromhex(datahash)+struct.pack("<Q", len(d))+bytes(d)
-    return datahash
-  def q(self, x:RemoteRequest): self._q.append(x)
-  def serialize(self) -> bytes:
-    self.h(repr(self._q).encode())
-    return b''.join(self._h.values())
-  def deserialize(self, dat:bytes) -> BatchRequest:
-    ptr = 0
-    while ptr < len(dat):
-      datahash, datalen = binascii.hexlify(dat[ptr:ptr+0x20]).decode(), struct.unpack("<Q", dat[ptr+0x20:ptr+0x28])[0]
-      self._h[datahash] = dat[ptr+0x28:ptr+0x28+datalen]
-      ptr += 0x28+datalen
-    self._q = safe_eval(ast.parse(self._h[datahash], mode="eval").body)
-    return self
-
-# ***** backend *****
-
-@dataclass
-class RemoteSession:
-  programs: dict[tuple[str, str], Any] = field(default_factory=dict)
-  graphs: dict[int, GraphRunner] = field(default_factory=dict)
-  buffers: dict[int, Buffer] = field(default_factory=dict)
-  events: defaultdict[int, asyncio.Event] = field(default_factory=functools.partial(defaultdict, asyncio.Event))
-
-class RemoteHandler:
-  def __init__(self, base_device: str):
-    self.base_device = base_device
-    self.sessions: defaultdict[SessionKey, RemoteSession] = defaultdict(RemoteSession)
-
-    try: self.ib_ctx: IBCtx|None = IBCtx(getenv("IB_DEV", 0))
-    except (RuntimeError, IndexError, AttributeError): self.ib_ctx = None
-    self.ib_lock = asyncio.Lock()
-    self.ib_conns: dict[str, IBConn|None] = {}
-    self.iova_cache: dict[tuple[SessionKey, int], tuple[int, int, int]] = {}
-
-  async def __call__(self, reader:asyncio.StreamReader, writer:asyncio.StreamWriter):
-    while (req_hdr:=(await reader.readline()).decode().strip()):
-      req_method, req_path, _ = req_hdr.split(' ')
-      req_headers = {}
-      while (hdr:=(await reader.readline()).decode().strip()):
-        key, value = hdr.split(':', 1)
-        req_headers[key.lower()] = value.strip()
-      req_body = await reader.readexactly(int(req_headers.get("content-length", "0")))
-      try: res_status, res_body = await self.handle(req_method, req_path, req_body)
-      except Exception as e:
-        res_status, res_body = http.HTTPStatus.INTERNAL_SERVER_ERROR, repr(RemoteException(e, traceback.format_exc())).encode()
-        print(f"{traceback.format_exc()}", flush=True)
-      writer.write(f"HTTP/1.1 {res_status.value} {res_status.phrase}\r\nContent-Length: {len(res_body)}\r\n\r\n".encode() + res_body)
-
-  async def ib_connect(self, ssession:SessionKey, dsession:SessionKey) -> IBConn|None:
-    if self.ib_ctx is None: return None
-    await self.ib_lock.acquire()
-    conn = RemoteConnection(dsession.host)
-    if dsession.host not in self.ib_conns:
-      props = safe_eval(ast.parse(conn.q(GetProperties(session=dsession), wait=True), mode="eval").body)
-      if props.ib_gid is not None:
-        self.ib_conns[dsession.host] = ib_conn = IBConn(self.ib_ctx)
-        ibxc_ret = conn.q(IBConnect(ssession.host, ib_conn.gid, ib_conn.qp_num, session=dsession), wait=True)
-        ib_conn.connect(*struct.unpack('<16sQ', ibxc_ret))
-      else:
-        self.ib_conns[dsession.host] = None
-    self.ib_lock.release()
-    return self.ib_conns[dsession.host]
-
-  async def get_iovas(self, bufs:list[tuple[SessionKey, int]]) -> list[tuple[int, int, int]]:
-    await self.ib_lock.acquire()
-    if (rbufs:=[buf for buf in bufs if buf not in self.iova_cache]):
-      conn = RemoteConnection(rbufs[0][0].host)
-      resp = await conn.aq(BufferIOVAS(rbufs, session=rbufs[0][0]), wait=True)
-      self.iova_cache.update({rbuf: struct.unpack('<QQQ', resp[i*24:(i+1)*24]) for i,rbuf in enumerate(rbufs)})
-    self.ib_lock.release()
-    return [self.iova_cache[buf] for buf in bufs]
-
-  async def handle(self, method:str, path:str, body:bytes) -> tuple[http.HTTPStatus, bytes]:
-    status, ret = http.HTTPStatus.OK, b""
-    if path == "/batch" and method == "POST":
-      # TODO: streaming deserialize?
-      req = BatchRequest().deserialize(body)
-      # the cmds are always last (currently in datahash)
-      for c in req._q:
-        if DEBUG >= 1: print(c)
-        session, dev = self.sessions[unwrap(c.session)], Device[f"{self.base_device}:{unwrap(c.session).idx}"]
-        match c:
-          case SessionFree(): del self.sessions[unwrap(c.session)]
-          case GetProperties():
-            cls, args = dev.renderer.__reduce__()
-            graph_cls = graph_class(Device[self.base_device])
-            rp = RemoteProperties(
-              real_device=dev.device, renderer=(cls.__module__, cls.__name__, args), offset_supported=hasattr(dev.allocator, '_offset'),
-              graph_supported=graph_cls is not None,
-              graph_supports_multi=graph_cls is not None and issubclass(graph_cls, MultiGraphRunner) and hasattr(dev.allocator, '_transfer'),
-              ib_gid=bytes(self.ib_ctx.gid_attr.raw) if self.ib_ctx is not None else None,
-            )
-            ret = repr(rp).encode()
-          case Event():
-            if c.session == c.event_session:
-              session.events[c.event].set()
-            else:
-              for d in Device._opened_devices: Device[d].synchronize() # wait for device*s* to finish executing previous stuff
-              # TODO: don't wait, just send
-              await RemoteConnection(c.event_session.host).aq(Event(c.event_session, c.event, session=c.event_session), wait=True)
-          case Wait():
-            assert await session.events[c.event].wait()
-            del session.events[c.event] # do not leak memory
-          case IBConnect():
-            self.ib_conns[c.host] = ibc = IBConn(unwrap(self.ib_ctx))
-            ibc.connect(c.gid, c.qp_num)
-            ret = struct.pack('<16sQ', ibc.gid, ibc.qp_num)
-          case BufferAlloc():
-            assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated"
-            session.buffers[c.buffer_num] = Buffer(dev.device, c.size, dtypes.uint8, options=c.options, preallocate=True)
-          case BufferIOVAS():
-            rets = []
-            for buffer_session,buffer_num in c.buffer_nums:
-              iova, mr = unwrap(self.ib_ctx).reg(buf:=self.sessions[buffer_session].buffers[buffer_num])
-              rets.append(struct.pack("<QQQ", iova, mr.contents.rkey, buf.nbytes))
-            ret = b"".join(rets)
-          case BufferOffset():
-            assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already exists"
-            session.buffers[c.buffer_num] = session.buffers[c.sbuffer_num].view(c.size, dtypes.uint8, c.offset).allocate()
-          case BufferFree(): del session.buffers[c.buffer_num]
-          case CopyIn(): session.buffers[c.buffer_num].copyin(memoryview(bytearray(req._h[c.datahash])))
-          case CopyOut(): session.buffers[c.buffer_num].copyout(memoryview(ret:=bytearray(session.buffers[c.buffer_num].nbytes)))
-          case Transfer():
-            if c.dsession.host == unwrap(c.session).host:
-              dsession, ddev = self.sessions[c.dsession], Device[f"{self.base_device}:{unwrap(c.dsession).idx}"]
-              dbuf, sbuf = dsession.buffers[c.dbuffer_num], session.buffers[c.buffer_num]
-              if hasattr(ddev.allocator, '_transfer'):
-                assert dbuf.nbytes == sbuf.nbytes, f"{dbuf.nbytes} != {sbuf.nbytes}"
-                ddev.allocator._transfer(dbuf._buf, sbuf._buf, dbuf.nbytes, dest_dev=ddev, src_dev=dev)
-              else:
-                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-                dbuf.copyin(data)
-            else:
-              conn, ib_conn = RemoteConnection(c.dsession.host), await self.ib_connect(unwrap(c.session), c.dsession)
-              sbuf = session.buffers[c.buffer_num]
-              if ib_conn is not None:
-                src_iova, src_mr = unwrap(self.ib_ctx).reg(sbuf)
-                dst_iova, dst_key, dst_size = (await self.get_iovas([(c.dsession, c.dbuffer_num)]))[0]
-                assert sbuf.nbytes == dst_size, f"{sbuf.nbytes} != {dst_size}"
-                for d in Device._opened_devices: Device[d].synchronize()
-                ib_conn.rdma_write([SGE(dst_iova, dst_key, src_iova, src_mr.contents.lkey, dst_size)])
-              else:
-                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-                await conn.aq(CopyIn(c.dbuffer_num, conn.req.h(data), session=c.dsession), wait=True)
-          case BatchTransfer():
-            conn, ib_conn = RemoteConnection(c.dbuffer_nums[0][0].host), await self.ib_connect(c.sbuffer_nums[0][0], c.dbuffer_nums[0][0])
-            if ib_conn is not None:
-              sbufs = [unwrap(self.ib_ctx).reg(self.sessions[s].buffers[bi]) for s,bi in c.sbuffer_nums]
-              dbufs = await self.get_iovas(c.dbuffer_nums)
-              for d in Device._opened_devices: Device[d].synchronize()
-              ib_conn.rdma_write([SGE(di, dk, si, sm.contents.lkey, ds) for (di,dk,ds),(si,sm) in zip(dbufs, sbufs)])
-            else:
-              for (sbuf_session,sbuf_num),(dbuf_session,dbuf_num) in zip(c.sbuffer_nums, c.dbuffer_nums):
-                sbuf = self.sessions[sbuf_session].buffers[sbuf_num]
-                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-                await conn.aq(CopyIn(dbuf_num, conn.req.h(data), session=dbuf_session), wait=True)
-          case ProgramAlloc():
-            lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
-            session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)
-          case ProgramFree():
-            key = (c.name, c.datahash)
-            # WORKAROUND: should be unconditional once the protocol supports proper exception handling
-            if key in session.programs: del session.programs[key]
-          case ProgramExec():
-            bufs = [session.buffers[x]._buf for x in c.bufs]
-            extra_args = {k:v for k,v in [("global_size", c.global_size), ("local_size", c.local_size)] if v is not None}
-            r = session.programs[(c.name, c.datahash)](*bufs, vals=c.vals, wait=c.wait, **extra_args)
-            if r is not None: ret = str(r).encode()
-          case GraphAlloc():
-            graph_fn: Callable = unwrap(dev.graph)
-            def _parse_ji(gi: GraphComputeItem|Transfer):
-              match gi:
-                case GraphComputeItem():
-                  prg = self.sessions[gi.session].programs[(gi.name, gi.datahash)]
-                  ps = ProgramSpec(gi.name, '', f"{self.base_device}:{gi.session.idx}", UOp(Ops.NOOP),
-                                   vars=list(gi.vars), ins=list(gi.ins), outs=list(gi.outs),
-                                   global_size=list(cast(tuple[int], gi.global_size)) if gi.global_size is not None else None,
-                                   local_size=list(cast(tuple[int], gi.local_size)) if gi.local_size is not None else None)
-                  return ExecItem(CompiledRunner(ps, precompiled=b'', prg=prg), [self.sessions[gi.session].buffers[buf] for buf in gi.bufs],
-                                  fixedvars=gi.fixedvars)
-                case Transfer():
-                  dbuf, sbuf = self.sessions[gi.dsession].buffers[gi.dbuffer_num], self.sessions[unwrap(gi.session)].buffers[gi.buffer_num]
-                  assert dbuf.nbytes == sbuf.nbytes, f"{dbuf.nbytes} != {sbuf.nbytes}"
-                  return ExecItem(BufferXfer(dbuf.nbytes, dbuf.device, sbuf.device), [dbuf, sbuf])
-            assert c.graph_num not in session.graphs, f"graph {c.graph_num} already allocated"
-            session.graphs[c.graph_num] = graph_fn(list(map(_parse_ji, c.jit_cache)), [self.sessions[s].buffers[i] for s,i in c.bufs], c.var_vals)
-          case GraphFree(): del session.graphs[c.graph_num]
-          case GraphExec():
-            r = session.graphs[c.graph_num]([self.sessions[s].buffers[i] for s,i in c.bufs], c.var_vals, wait=c.wait)
-            if r is not None: ret = str(r).encode()
-    else: status, ret = http.HTTPStatus.NOT_FOUND, b"Not Found"
-    return status, ret
-
-def remote_server(port:int):
-  device = getenv("REMOTEDEV", next(Device.get_available_devices()) if Device.DEFAULT == "REMOTE" else Device.DEFAULT)
-  async def _inner_async(port:int, device:str):
-    print(f"start remote server on {port} with device {device}")
-    await (await asyncio.start_server(RemoteHandler(device), host='', port=port)).serve_forever()
-  asyncio.run(_inner_async(port, device))
-
-# ***** frontend *****
-
-class RemoteAllocator(Allocator['RemoteDevice']):
-  def __init__(self, dev:RemoteDevice):
-    if dev.properties.offset_supported: self._offset = self._dyn_offset
-    super().__init__(dev)
-  # TODO: ideally we shouldn't have to deal with images here
-  def _alloc(self, size:int, options:BufferSpec) -> int:
-    self.dev.q(BufferAlloc(buffer_num:=next(self.dev.buffer_num), size, options))
-    return buffer_num
-  # TODO: options should not be here in any Allocator
-  def _free(self, opaque:int, options):
-    try: self.dev.q(BufferFree(opaque))
-    except (TypeError, AttributeError): pass
-  def _copyin(self, dest:int, src:memoryview): self.dev.q(CopyIn(dest, self.dev.conn.req.h(src)))
-  def _copyout(self, dest:memoryview, src:int):
-    resp = self.dev.q(CopyOut(src), wait=True)
-    assert len(resp) == len(dest), f"buffer length mismatch {len(resp)} != {len(dest)}"
-    dest[:] = resp
-  def _transfer(self, dest, src, sz, src_dev, dest_dev):
-    if dest_dev.conn != src_dev.conn:
-      dest_dev.q(Event(src_dev.session, start_event:=next(src_dev.event_num)))
-      src_dev.q(Wait(start_event))
-    src_dev.q(Transfer(src, dest_dev.session, dest))
-    if dest_dev.conn != src_dev.conn:
-      src_dev.q(Event(dest_dev.session, end_event:=next(dest_dev.event_num)))
-      dest_dev.q(Wait(end_event))
-    if DEBUG >= 2: dest_dev.conn.batch_submit()
-  def _dyn_offset(self, opaque:int, size:int, offset:int) -> int:
-    self.dev.q(BufferOffset(buffer_num:=next(self.dev.buffer_num), size, offset, opaque))
-    return buffer_num
-
-class RemoteProgram:
-  def __init__(self, dev:RemoteDevice, name:str, lib:bytes):
-    self.dev, self.name = dev, name
-    self.datahash = self.dev.conn.req.h(lib)
-    self.dev.q(ProgramAlloc(self.name, self.datahash))
-    super().__init__()
-    weakref.finalize(self, self._fini, self.dev, self.name, self.datahash)
-
-  @staticmethod
-  def _fini(dev:RemoteDevice, name:str, datahash:str): dev.q(ProgramFree(name, datahash))
-
-  def __call__(self, *bufs, global_size=None, local_size=None, vals:tuple[int, ...]=(), wait=False):
-    ret = self.dev.q(ProgramExec(self.name, self.datahash, bufs, vals, global_size, local_size, wait), wait=wait)
-    if wait: return float(ret)
-
-@functools.cache
-class RemoteConnection:
-  q_lock = threading.Lock()
-  all: dict[RemoteConnection, None] = {} # dict instead of set for deterministic ordering
-
-  def __init__(self, host:str):
-    if DEBUG >= 1: print(f"remote with host {host}")
-    while 1:
-      try:
-        self.conn = http.client.HTTPConnection(host, timeout=getenv("REMOTE_TIMEOUT", 300.0))
-        self.conn.connect()
-        break
-      except Exception as e:
-        print(e)
-        time.sleep(0.1)
-    self.req: BatchRequest = BatchRequest()
-    RemoteConnection.all[self] = None
-
-  def q(self, x:RemoteRequest, wait:bool=False):
-    with RemoteConnection.q_lock:
-      self.req.q(x)
-      if wait: return self.batch_submit(take_q=False)
-
-  async def aq(self, x:RemoteRequest, wait:bool=False): return await asyncio.to_thread(self.q, x, wait=wait)
-
-  def batch_submit(self, take_q:bool=True):
-    if take_q: RemoteConnection.q_lock.acquire()
-    conns = RemoteConnection.all.keys()
-    datas = {conn: conn.req.serialize() for conn in conns}
-    reqs, hashes, hash_datas = sum(len(c.req._q) for c in conns), sum(len(c.req._h) for c in conns), sum(len(data) for data in datas.values())
-    ret, resps = None, []
-    with Timing(f"*** send {reqs:-3d} requests {hashes:-3d} hashes with len {hash_datas/1024:.2f} kB in ", enabled=DEBUG>=3):
-      for conn,data in datas.items(): conn.conn.request("POST", "/batch", data)
-      for conn in datas.keys():
-        resp = conn.conn.getresponse()
-        body = resp.read()
-        resps.append((conn, resp, body))
-        conn.req = BatchRequest()
-    if take_q: RemoteConnection.q_lock.release()
-    for conn,resp,body in resps:
-      match resp.status:
-        case http.HTTPStatus.OK: pass
-        case http.HTTPStatus.INTERNAL_SERVER_ERROR:
-          exc_wrapper = safe_eval(ast.parse(body.decode(), mode="eval").body)
-          exc_wrapper.exc.add_note(exc_wrapper.trace)
-          raise exc_wrapper.exc
-        case code: raise RuntimeError(f"POST /batch failed with {code}: {body.decode()}")
-      if conn == self: ret = body
-    return ret
-
-def parse_hosts(hs:str) -> list[tuple[str, int]]|LazySeq[tuple[str, int]]:
-  hosts = [(unwrap(h), int(c) if c is not None else c) for h,c in ((h.split("*", maxsplit=1)+[None,])[:2] for h in hs.split(","))]
-  if len(hosts) == 1 and hosts[0][1] is None: return LazySeq(lambda idx: (hosts[0][0], idx))
-  return [(h, i) for h,c in hosts for i in range(unwrap(c))]
-
-class RemoteDevice(Compiled):
-  devices = parse_hosts(getenv("HOST", ""))
-
-  def __init__(self, device:str):
-    host, idx = RemoteDevice.devices[int(device.split(":")[1]) if ":" in device else 0]
-
-    # connection is shared between sessions on the same host
-    self.session: SessionKey = SessionKey(host or RemoteDevice.local_server(), idx, binascii.hexlify(os.urandom(0x10)).decode())
-    self.conn: RemoteConnection = RemoteConnection(self.session.host)
-
-    # state for the session
-    self.buffer_num: Iterator[int] = itertools.count(0)
-    self.graph_num: Iterator[int] = itertools.count(0)
-    self.event_num: Iterator[int] = itertools.count(0)
-
-    self.properties: RemoteProperties = safe_eval(ast.parse(self.q(GetProperties(), wait=True), mode="eval").body)
-    if DEBUG >= 1: print(f"remote has device {self.properties.real_device}")
-    # TODO: how to we have BEAM be cached on the backend? this should just send a specification of the compute. rethink what goes in Renderer
-    renderer = self.properties.renderer
-    if not renderer[0].startswith("tinygrad.") or not renderer[1].endswith("Renderer"): raise RuntimeError(f"bad renderer {renderer}")
-    renderer_class = fromimport(renderer[0], renderer[1])  # TODO: is this secure?
-    if not issubclass(renderer_class, Renderer): raise RuntimeError(f"renderer isn't a Renderer {renderer}")
-
-    graph = fromimport('tinygrad.runtime.graph.remote', "RemoteGraph") if self.properties.graph_supported else None
-    compilers = CompilerSet([CompilerPair(functools.partial(renderer_class, *renderer[2]), Compiler)])
-    super().__init__(device, RemoteAllocator(self), compilers, functools.partial(RemoteProgram, self), graph, id(self.conn))
-    self.renderer.device = device
-
-  def finalize(self):
-    with contextlib.suppress(ConnectionError, http.client.HTTPException): self.q(SessionFree(), wait=True)
-
-  def q(self, x:RemoteRequest, wait:bool=False): return self.conn.q(replace(x, session=self.session), wait=wait)
-
-  @functools.cache
-  @staticmethod
-  def local_server():
-    multiprocessing.Process(target=remote_server, args=(6667,), name="MainProcess", daemon=True).start()
-    return "127.0.0.1:6667"
-
-if __name__ == "__main__": remote_server(getenv("PORT", 6667))

tinygrad/runtime/support/am/amdev.py

@@ -41,40 +41,63 @@ class AMFirmware:
     self.ucode_start: dict[str, int] = {}
     self.descs: list[tuple[list[int], memoryview]] = []
 
-    blob, hdr = self.load_fw(f"smu_{fmt_ver(am.MP1_HWIP)}.bin", am.struct_smc_firmware_header_v1_0)
-    self.smu_psp_desc = self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.header.ucode_size_bytes, am.GFX_FW_TYPE_SMU)
+    # SMU firmware
+    if adev.ip_ver[am.MP1_HWIP] != (13,0,12):
+      blob, hdr = self.load_fw(f"smu_{fmt_ver(am.MP1_HWIP)}.bin", versioned_header="struct_smc_firmware_header")
+      if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0):
+        self.smu_psp_desc = self.desc(blob, hdr.v1_0.header.ucode_array_offset_bytes, hdr.v1_0.header.ucode_size_bytes, am.GFX_FW_TYPE_SMU)
+      else:
+        p2stables = (am.struct_smc_soft_pptable_entry * hdr.pptable_count).from_buffer(blob[hdr.pptable_entry_offset:])
+        for p2stable in p2stables:
+          if p2stable.id == (__P2S_TABLE_ID_X:=0x50325358):
+            self.descs += [self.desc(blob, p2stable.ppt_offset_bytes, p2stable.ppt_size_bytes, am.GFX_FW_TYPE_P2S_TABLE)]
 
     # SDMA firmware
-    blob, hdr = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", versioned_header='struct_sdma_firmware_header')
-    if hdr.header.header_version_major < 3:
+    blob, hdr = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", versioned_header="struct_sdma_firmware_header")
+    if hdr.header.header_version_major == 1:
+      self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.header.ucode_size_bytes, am.GFX_FW_TYPE_SDMA0)]
+    elif hdr.header.header_version_major == 2:
       self.descs += [self.desc(blob, hdr.ctl_ucode_offset, hdr.ctl_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH1)]
       self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ctx_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
     else: self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
 
     # PFP, ME, MEC firmware
     for (fw_name, fw_cnt) in ([('PFP', 1), ('ME', 1)] if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else []) + [('MEC', 1)]:
-      blob, hdr = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_{fw_name.lower()}.bin", am.struct_gfx_firmware_header_v2_0)
+      blob, hdr = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_{fw_name.lower()}.bin", versioned_header="struct_gfx_firmware_header")
 
-      # Code part
-      self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ucode_size_bytes, getattr(am, f'GFX_FW_TYPE_RS64_{fw_name}'))]
-
-      # Stack
-      stack_fws = [getattr(am, f'GFX_FW_TYPE_RS64_{fw_name}_P{fwnum}_STACK') for fwnum in range(fw_cnt)]
-      self.descs += [self.desc(blob, hdr.data_offset_bytes, hdr.data_size_bytes, *stack_fws)]
-      self.ucode_start[fw_name] = hdr.ucode_start_addr_lo | (hdr.ucode_start_addr_hi << 32)
+      ucode_off = hdr.header.ucode_array_offset_bytes
+      if hdr.header.header_version_major == 1:
+        # Code
+        self.descs += [self.desc(blob, ucode_off, hdr.header.ucode_size_bytes - hdr.jt_size * 4, getattr(am, f'GFX_FW_TYPE_CP_{fw_name}'))]
+        # JT
+        self.descs += [self.desc(blob, ucode_off + hdr.jt_offset * 4, hdr.jt_size * 4, getattr(am, f'GFX_FW_TYPE_CP_{fw_name}_ME1'))]
+      else:
+        # Code
+        self.descs += [self.desc(blob, ucode_off, hdr.ucode_size_bytes, getattr(am, f'GFX_FW_TYPE_RS64_{fw_name}'))]
+        # Stack
+        stack_fws = [getattr(am, f'GFX_FW_TYPE_RS64_{fw_name}_P{fwnum}_STACK') for fwnum in range(fw_cnt)]
+        self.descs += [self.desc(blob, hdr.data_offset_bytes, hdr.data_size_bytes, *stack_fws)]
+        self.ucode_start[fw_name] = hdr.ucode_start_addr_lo | (hdr.ucode_start_addr_hi << 32)
 
     # IMU firmware
-    blob, hdr = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_imu.bin", am.struct_imu_firmware_header_v1_0)
-    imu_i_off, imu_i_sz, imu_d_sz = hdr.header.ucode_array_offset_bytes, hdr.imu_iram_ucode_size_bytes, hdr.imu_dram_ucode_size_bytes
-    self.descs += [self.desc(blob, imu_i_off, imu_i_sz, am.GFX_FW_TYPE_IMU_I), self.desc(blob, imu_i_off + imu_i_sz, imu_d_sz, am.GFX_FW_TYPE_IMU_D)]
+    if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0):
+      blob, hdr = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_imu.bin", am.struct_imu_firmware_header_v1_0)
+      imu_i_off, imu_i_sz, imu_d_sz = hdr.header.ucode_array_offset_bytes, hdr.imu_iram_ucode_size_bytes, hdr.imu_dram_ucode_size_bytes
+      self.descs += [self.desc(blob, imu_i_off, imu_i_sz, am.GFX_FW_TYPE_IMU_I), self.desc(blob, imu_i_off+imu_i_sz, imu_d_sz, am.GFX_FW_TYPE_IMU_D)]
 
     # RLC firmware
-    blob, hdr0, _hdr1, hdr2, hdr3 = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_rlc.bin", am.struct_rlc_firmware_header_v2_0,
+    blob, hdr0, hdr1, hdr2, hdr3 = self.load_fw(f"gc_{fmt_ver(am.GC_HWIP)}_rlc.bin", am.struct_rlc_firmware_header_v2_0,
       am.struct_rlc_firmware_header_v2_1, am.struct_rlc_firmware_header_v2_2, am.struct_rlc_firmware_header_v2_3)
 
-    for mem,fmem in [('IRAM', 'iram'), ('DRAM_BOOT', 'dram')]:
-      off, sz = getattr(hdr2, f'rlc_{fmem}_ucode_offset_bytes'), getattr(hdr2, f'rlc_{fmem}_ucode_size_bytes')
-      self.descs += [self.desc(blob, off, sz, getattr(am, f'GFX_FW_TYPE_RLC_{mem}'))]
+    if hdr0.header.header_version_minor == 1:
+      for mem,fmem in [('LIST_SRM_CNTL', 'list_cntl'), ('LIST_GPM_MEM', 'list_gpm'), ('LIST_SRM_MEM', 'list_srm')]:
+        off, sz = getattr(hdr1, f'save_restore_{fmem}_offset_bytes'), getattr(hdr1, f'save_restore_{fmem}_size_bytes')
+        self.descs += [self.desc(blob, off, sz, getattr(am, f'GFX_FW_TYPE_RLC_RESTORE_{mem}'))]
+
+    if hdr0.header.header_version_minor >= 2:
+      for mem,fmem in [('IRAM', 'iram'), ('DRAM_BOOT', 'dram')]:
+        off, sz = getattr(hdr2, f'rlc_{fmem}_ucode_offset_bytes'), getattr(hdr2, f'rlc_{fmem}_ucode_size_bytes')
+        self.descs += [self.desc(blob, off, sz, getattr(am, f'GFX_FW_TYPE_RLC_{mem}'))]
 
     if hdr0.header.header_version_minor == 3:
       for mem in ['P', 'V']:
@@ -107,7 +130,7 @@ class AMPageTableEntry:
   def address(self, entry_id:int) -> int:
     assert self.entries[entry_id] & am.AMDGPU_PTE_SYSTEM == 0, "should not be system address"
     return self.adev.xgmi2paddr(self.entries[entry_id] & 0x0000FFFFFFFFF000)
-  def is_page(self, entry_id:int) -> bool: return self.lv == am.AMDGPU_VM_PTB or self.adev.gmc.is_pte_huge_page(self.entries[entry_id])
+  def is_page(self, entry_id:int) -> bool: return self.lv == am.AMDGPU_VM_PTB or self.adev.gmc.is_pte_huge_page(self.lv, self.entries[entry_id])
   def supports_huge_page(self, paddr:int): return self.lv >= am.AMDGPU_VM_PDB2
 
 class AMMemoryManager(MemoryManager):
@@ -121,7 +144,7 @@ class AMMemoryManager(MemoryManager):
 class AMDev(PCIDevImplBase):
   Version = 0xA0000006
 
-  def __init__(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
+  def __init__(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None, reset_mode=False):
     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')
 
@@ -147,7 +170,11 @@ class AMDev(PCIDevImplBase):
 
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
-      if self.psp.is_sos_alive() and self.smu.is_smu_alive(): self.smu.mode1_reset()
+      if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+        if self.gmc.xgmi_seg_sz > 0:
+          if reset_mode: return # in reset mode, do not raise
+          raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
+        self.smu.mode1_reset()
       for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]:
         ip.init_hw()
         if DEBUG >= 2: print(f"am {self.devfmt}: {ip.__class__.__name__} initialized")
@@ -199,14 +226,14 @@ class AMDev(PCIDevImplBase):
   def reg(self, reg:str) -> AMRegister: return self.__dict__[reg]
 
   def rreg(self, reg:int) -> int:
-    val = self.indirect_rreg(reg * 4) if reg > len(self.mmio) else self.mmio[reg]
+    val = self.indirect_rreg(reg) if reg > len(self.mmio) else self.mmio[reg]
     if AM_DEBUG >= 4 and getattr(self, '_prev_rreg', None) != (reg, val): print(f"am {self.devfmt}: Reading register {reg:#x} with value {val:#x}")
     self._prev_rreg = (reg, val)
     return val
 
   def wreg(self, reg:int, val:int):
     if AM_DEBUG >= 4: print(f"am {self.devfmt}: Writing register {reg:#x} with value {val:#x}")
-    if reg > len(self.mmio): self.indirect_wreg(reg * 4, val)
+    if reg > len(self.mmio): self.indirect_wreg(reg, val)
     else: self.mmio[reg] = val
 
   def wreg_pair(self, reg_base:str, lo_suffix:str, hi_suffix:str, val:int, inst:int=0):
@@ -214,13 +241,19 @@ class AMDev(PCIDevImplBase):
     self.reg(f"{reg_base}{hi_suffix}").write(val >> 32, inst=inst)
 
   def indirect_rreg(self, reg:int) -> int:
-    self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg)
+    self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg * 4)
     return self.reg("regBIF_BX_PF0_RSMU_DATA").read()
 
   def indirect_wreg(self, reg:int, val:int):
-    self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg)
+    self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg * 4)
     self.reg("regBIF_BX_PF0_RSMU_DATA").write(val)
 
+  def indirect_wreg_pcie(self, reg:int, val:int, aid:int=0):
+    self.reg("regBIF_BX0_PCIE_INDEX2").write(reg * 4 + ((((aid & 0b11) << 32) | (1 << 34)) if aid > 0 else 0))
+    self.reg("regBIF_BX0_PCIE_INDEX2").read()
+    self.reg("regBIF_BX0_PCIE_DATA2").write(val)
+    self.reg("regBIF_BX0_PCIE_DATA2").read()
+
   def _read_vram(self, addr, size) -> bytes:
     assert addr % 4 == 0 and size % 4 == 0, f"Invalid address {addr:#x} or size {size:#x}"
     res = []
@@ -268,6 +301,7 @@ class AMDev(PCIDevImplBase):
   def _build_regs(self):
     mods = [("mp", am.MP0_HWIP), ("hdp", am.HDP_HWIP), ("gc", am.GC_HWIP), ("mmhub", am.MMHUB_HWIP), ("osssys", am.OSSSYS_HWIP),
       ("nbio" if self.ip_ver[am.GC_HWIP] < (12,0,0) else "nbif", am.NBIO_HWIP)]
+    if self.ip_ver[am.SDMA0_HWIP] in {(4,4,2), (4,4,4)}: mods += [("sdma", am.SDMA0_HWIP)]
 
     for prefix, hwip in mods:
       self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))

tinygrad/runtime/support/am/ip.py

@@ -15,9 +15,10 @@ class AM_SOC(AM_IP):
   def init_sw(self): self.module = import_soc(self.adev.ip_ver[am.GC_HWIP])
 
   def init_hw(self):
-    if self.adev.ip_ver[am.NBIO_HWIP] == (7,9,0):
+    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
       self.adev.regXCC_DOORBELL_FENCE.write(0x0)
       self.adev.regBIFC_GFX_INT_MONITOR_MASK.write(0x7ff)
+      self.adev.regBIFC_DOORBELL_ACCESS_EN_PF.write(0xfffff)
     else: self.adev.regRCC_DEV0_EPF2_STRAP2.update(strap_no_soft_reset_dev0_f2=0x0)
     self.adev.regRCC_DEV0_EPF0_RCC_DOORBELL_APER_EN.write(0x1)
   def set_clockgating_state(self):
@@ -28,7 +29,7 @@ class AM_SOC(AM_IP):
     val = reg.encode(**{f"s2a_doorbell_port{port}_enable":1, f"s2a_doorbell_port{port}_awid":awid,  f"s2a_doorbell_port{port}_range_size":size,
       f"s2a_doorbell_port{port}_awaddr_31_28_value":awaddr_31_28_value, f"s2a_doorbell_port{port}_range_offset":offset})
 
-    if self.adev.ip_ver[am.NBIO_HWIP] == (7,9,0): self.adev.indirect_wreg_pcie(reg.addr[0], val)
+    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}: self.adev.indirect_wreg_pcie(reg.addr[0], val)
     else: reg.write(val)
 
 class AM_GMC(AM_IP):
@@ -36,10 +37,10 @@ class AM_GMC(AM_IP):
     self.vmhubs = len(self.adev.regs_offset[am.MMHUB_HWIP])
 
     # XGMI (for supported systems)
-    xgmi_phys_id = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields()['pf_lfb_region'] if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else 0
-    xgmi_seg_sz = (self.adev.regMMMC_VM_XGMI_LFB_SIZE.read_bitfields()['pf_lfb_size'] << 24) if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_SIZE') else 0
+    self.xgmi_phys_id = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields()['pf_lfb_region'] if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else 0
+    self.xgmi_seg_sz = self.adev.regMMMC_VM_XGMI_LFB_SIZE.read_bitfields()['pf_lfb_size']<<24 if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_SIZE') else 0
 
-    self.paddr_base = xgmi_phys_id * xgmi_seg_sz
+    self.paddr_base = self.xgmi_phys_id * self.xgmi_seg_sz
 
     self.fb_base = (self.adev.regMMMC_VM_FB_LOCATION_BASE.read() & 0xFFFFFF) << 24
     self.fb_end = (self.adev.regMMMC_VM_FB_LOCATION_TOP.read() & 0xFFFFFF) << 24
@@ -51,12 +52,16 @@ class AM_GMC(AM_IP):
     self.vm_base = self.adev.mm.va_base
     self.vm_end = min(self.vm_base + (1 << self.adev.mm.va_bits) - 1, 0x7fffffffffff)
 
+    self.trans_futher = self.adev.ip_ver[am.GC_HWIP] < (10, 0, 0)
+
     # GFX11/GFX12 has 44-bit address space
     self.address_space_mask = (1 << 44) - 1
 
     self.memscratch_xgmi_paddr = self.adev.paddr2xgmi(self.adev.mm.palloc(0x1000, zero=False, boot=True))
     self.dummy_page_xgmi_paddr = self.adev.paddr2xgmi(self.adev.mm.palloc(0x1000, zero=False, boot=True))
-    self.hub_initted = {"MM": False, "GC": False}
+
+    # MM hub is inited before any tlb flushes and is still valid during partial_boot, so set it to true
+    self.hub_initted = {"MM": True, "GC": False}
 
     self.pf_status_reg = lambda ip: f"reg{ip}VM_L2_PROTECTION_FAULT_STATUS{'_LO32' if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else ''}"
 
@@ -69,32 +74,30 @@ class AM_GMC(AM_IP):
     # Can't issue TLB invalidation if the hub isn't initialized.
     if not self.hub_initted[ip]: return
 
-    if ip == "MM": wait_cond(lambda: self.adev.regMMVM_INVALIDATE_ENG17_SEM.read() & 0x1, value=1, msg="mm flush_tlb timeout")
+    for inst in range(self.adev.gmc.vmhubs if ip == "MM" else self.adev.gfx.xccs):
+      if ip == "MM": wait_cond(lambda: self.adev.regMMVM_INVALIDATE_ENG17_SEM.read(inst=inst) & 0x1, value=1, msg="mm flush_tlb timeout")
 
-    self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_REQ").write(flush_type=flush_type, per_vmid_invalidate_req=(1 << vmid), invalidate_l2_ptes=1,
-      invalidate_l2_pde0=1, invalidate_l2_pde1=1, invalidate_l2_pde2=1, invalidate_l1_ptes=1, clear_protection_fault_status_addr=0)
+      self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_REQ").write(flush_type=flush_type, per_vmid_invalidate_req=(1 << vmid), invalidate_l2_ptes=1,
+        invalidate_l2_pde0=1, invalidate_l2_pde1=1, invalidate_l2_pde2=1, invalidate_l1_ptes=1, clear_protection_fault_status_addr=0, inst=inst)
 
-    wait_cond(lambda: self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_ACK").read() & (1 << vmid), value=(1 << vmid), msg="flush_tlb timeout")
+      wait_cond(lambda: self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_ACK").read(inst=inst) & (1 << vmid), value=(1 << vmid), msg="flush_tlb timeout")
 
-    if ip == "MM":
-      self.adev.regMMVM_INVALIDATE_ENG17_SEM.write(0x0)
-      self.adev.regMMVM_L2_BANK_SELECT_RESERVED_CID2.update(reserved_cache_private_invalidation=1)
+      if ip == "MM": self.adev.regMMVM_INVALIDATE_ENG17_SEM.write(0x0, inst=inst)
+      if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0) and ip == "MM":
+        self.adev.regMMVM_L2_BANK_SELECT_RESERVED_CID2.update(reserved_cache_private_invalidation=1, inst=inst)
 
-      # Read back the register to ensure the invalidation is complete
-      self.adev.regMMVM_L2_BANK_SELECT_RESERVED_CID2.read()
+        # Read back the register to ensure the invalidation is complete
+        self.adev.regMMVM_L2_BANK_SELECT_RESERVED_CID2.read(inst=inst)
 
   def enable_vm_addressing(self, page_table, ip:Literal["MM", "GC"], vmid, inst):
     self.adev.wreg_pair(f"reg{ip}VM_CONTEXT{vmid}_PAGE_TABLE_START_ADDR", "_LO32", "_HI32", self.vm_base >> 12, inst=inst)
     self.adev.wreg_pair(f"reg{ip}VM_CONTEXT{vmid}_PAGE_TABLE_END_ADDR", "_LO32", "_HI32", self.vm_end >> 12, inst=inst)
     self.adev.wreg_pair(f"reg{ip}VM_CONTEXT{vmid}_PAGE_TABLE_BASE_ADDR", "_LO32", "_HI32", self.adev.paddr2xgmi(page_table.paddr) | 1, inst=inst)
-    self.adev.reg(f"reg{ip}VM_CONTEXT{vmid}_CNTL").write(0x1800000, pde0_protection_fault_enable_interrupt=1, pde0_protection_fault_enable_default=1,
-                                                         dummy_page_protection_fault_enable_interrupt=1, dummy_page_protection_fault_enable_default=1,
-                                                         range_protection_fault_enable_interrupt=1, range_protection_fault_enable_default=1,
-                                                         valid_protection_fault_enable_interrupt=1, valid_protection_fault_enable_default=1,
-                                                         read_protection_fault_enable_interrupt=1, read_protection_fault_enable_default=1,
-                                                         write_protection_fault_enable_interrupt=1, write_protection_fault_enable_default=1,
-                                                         execute_protection_fault_enable_interrupt=1, execute_protection_fault_enable_default=1,
-                                                         enable_context=1, page_table_depth=(3 - page_table.lv), inst=inst)
+
+    fault_flags = {f'{x}_protection_fault_enable_interrupt':1 for x in ['pde0', 'dummy_page', 'range', 'valid', 'read', 'write', 'execute']}
+    en_def_flags = {f'{x}_protection_fault_enable_default':1 for x in ['pde0', 'dummy_page', 'range', 'valid', 'read', 'write', 'execute']}
+    self.adev.reg(f"reg{ip}VM_CONTEXT{vmid}_CNTL").write(0x1800000, **fault_flags, **en_def_flags, enable_context=1,
+      page_table_depth=((2 if self.trans_futher else 3) - page_table.lv), page_table_block_size=9 if self.trans_futher else 0, inst=inst)
 
   def init_hub(self, ip:Literal["MM", "GC"], inst_cnt:int):
     # Init system apertures
@@ -112,15 +115,15 @@ class AM_GMC(AM_IP):
 
       # Init TLB and cache
       self.adev.reg(f"reg{ip}MC_VM_MX_L1_TLB_CNTL").update(enable_l1_tlb=1, system_access_mode=3, enable_advanced_driver_model=1,
-        system_aperture_unmapped_access=0, eco_bits=0, mtype=self.adev.soc.module.MTYPE_UC, inst=inst)
+        system_aperture_unmapped_access=0, mtype=self.adev.soc.module.MTYPE_UC, inst=inst)
 
       self.adev.reg(f"reg{ip}VM_L2_CNTL").update(enable_l2_cache=1, enable_l2_fragment_processing=0, enable_default_page_out_to_system_memory=1,
         l2_pde0_cache_tag_generation_mode=0, pde_fault_classification=0, context1_identity_access_mode=1, identity_mode_fragment_size=0, inst=inst)
       self.adev.reg(f"reg{ip}VM_L2_CNTL2").update(invalidate_all_l1_tlbs=1, invalidate_l2_cache=1, inst=inst)
-      self.adev.reg(f"reg{ip}VM_L2_CNTL3").write(bank_select=9, l2_cache_bigk_fragment_size=6,l2_cache_4k_associativity=1,
-        l2_cache_bigk_associativity=1, inst=inst)
+      self.adev.reg(f"reg{ip}VM_L2_CNTL3").write(l2_cache_4k_associativity=1, l2_cache_bigk_associativity=1,
+        bank_select=12 if self.trans_futher else 9, l2_cache_bigk_fragment_size=9 if self.trans_futher else 6, inst=inst)
       self.adev.reg(f"reg{ip}VM_L2_CNTL4").write(l2_cache_4k_partition_count=1, inst=inst)
-      self.adev.reg(f"reg{ip}VM_L2_CNTL5").write(walker_priority_client_id=0x1ff, inst=inst)
+      if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0): self.adev.reg(f"reg{ip}VM_L2_CNTL5").write(walker_priority_client_id=0x1ff, inst=inst)
 
       self.enable_vm_addressing(self.adev.mm.root_page_table, ip, vmid=0, inst=inst)
 
@@ -139,11 +142,18 @@ class AM_GMC(AM_IP):
     if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0):
       extra |= am.AMDGPU_PTE_MTYPE_GFX12(0, self.adev.soc.module.MTYPE_UC if uncached else 0)
       extra |= (am.AMDGPU_PDE_PTE_GFX12 if not is_table and pte_lv != am.AMDGPU_VM_PTB else (am.AMDGPU_PTE_IS_PTE if not is_table else 0))
-    else:
+    elif self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
       extra |= am.AMDGPU_PTE_MTYPE_NV10(0, self.adev.soc.module.MTYPE_UC if uncached else 0)
       extra |= (am.AMDGPU_PDE_PTE if not is_table and pte_lv != am.AMDGPU_VM_PTB else 0)
+    else:
+      extra |= am.AMDGPU_PTE_MTYPE_VG10(0, self.adev.soc.module.MTYPE_UC if uncached else 0)
+      if is_table and pte_lv == am.AMDGPU_VM_PDB1: extra |= am.AMDGPU_PDE_BFS(0x9)
+      if is_table and pte_lv == am.AMDGPU_VM_PDB0: extra |= am.AMDGPU_PTE_TF
+      if not is_table and pte_lv not in {am.AMDGPU_VM_PTB, am.AMDGPU_VM_PDB0}: extra |= am.AMDGPU_PDE_PTE
     return extra
-  def is_pte_huge_page(self, pte): return pte & (am.AMDGPU_PDE_PTE_GFX12 if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else am.AMDGPU_PDE_PTE)
+  def is_pte_huge_page(self, pte_lv, pte):
+    if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): return (pte & am.AMDGPU_PDE_PTE) if pte_lv != am.AMDGPU_VM_PDB0 else not (pte & am.AMDGPU_PTE_TF)
+    return pte & (am.AMDGPU_PDE_PTE_GFX12 if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else am.AMDGPU_PDE_PTE)
 
   def on_interrupt(self):
     for ip in ["MM", "GC"]:
@@ -168,6 +178,7 @@ class AM_SMU(AM_IP):
   def mode1_reset(self):
     if DEBUG >= 2: print(f"am {self.adev.devfmt}: mode1 reset")
     if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0): self._send_msg(__DEBUGSMC_MSG_Mode1Reset:=2, 0, debug=True)
+    elif self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12)}: self._send_msg(self.smu_mod.PPSMC_MSG_GfxDriverReset, 1)
     else: self._send_msg(self.smu_mod.PPSMC_MSG_Mode1Reset, 0)
     time.sleep(0.5) # 500ms
 
@@ -177,6 +188,8 @@ class AM_SMU(AM_IP):
   def read_metrics(self): return self.read_table(self.smu_mod.SmuMetricsExternal_t, self.smu_mod.TABLE_SMU_METRICS)
 
   def set_clocks(self, level):
+    if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12)}: return # TODO
+
     if not hasattr(self, 'clcks'):
       self.clcks = {}
       for clck in [self.smu_mod.PPCLK_GFXCLK, self.smu_mod.PPCLK_UCLK, self.smu_mod.PPCLK_FCLK, self.smu_mod.PPCLK_SOCCLK]:
@@ -211,18 +224,26 @@ class AM_GFX(AM_IP):
     # NOTE: Golden reg for gfx11. No values for this reg provided. The kernel just ors 0x20000000 to this reg.
     for xcc in range(self.xccs): self.adev.regTCP_CNTL.write(self.adev.regTCP_CNTL.read() | 0x20000000, inst=xcc)
 
+    for xcc in range(self.xccs): self.adev.regRLC_CNTL.write(0x1, inst=xcc)
+
     for xcc in range(self.xccs): self.adev.regRLC_SRM_CNTL.update(srm_enable=1, auto_incr_addr=1, inst=xcc)
 
-    if self.adev.ip_ver[am.NBIO_HWIP] != (7,9,0):
+    for xcc in range(self.xccs): self.adev.regRLC_SPM_MC_CNTL.write(0xf, inst=xcc)
+
+    if self.adev.ip_ver[am.NBIO_HWIP][:2] != (7,9):
       self.adev.soc.doorbell_enable(port=0, awid=0x3, awaddr_31_28_value=0x3)
       self.adev.soc.doorbell_enable(port=3, awid=0x6, awaddr_31_28_value=0x3)
 
     for xcc in range(self.xccs):
+      if self.adev.ip_ver[am.GC_HWIP] in {(9,4,3), (9,5,0)}:
+        self.adev.regGB_ADDR_CONFIG.write(0x2a114042, inst=xcc) # Golden value for mi300/mi350
+        self.adev.regTCP_UTCL1_CNTL2.update(spare=1, inst=xcc)
+
       self.adev.regGRBM_CNTL.update(read_timeout=0xff, inst=xcc)
       for i in range(0, 16):
         self._grbm_select(vmid=i, inst=xcc)
-        self.adev.regSH_MEM_CONFIG.write(address_mode=self.adev.soc.module.SH_MEM_ADDRESS_MODE_64,
-                                         alignment_mode=self.adev.soc.module.SH_MEM_ALIGNMENT_MODE_UNALIGNED, initial_inst_prefetch=3, inst=xcc)
+        self.adev.regSH_MEM_CONFIG.write(**({'initial_inst_prefetch':3} if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else {}),
+          address_mode=self.adev.soc.module.SH_MEM_ADDRESS_MODE_64, alignment_mode=self.adev.soc.module.SH_MEM_ALIGNMENT_MODE_UNALIGNED, inst=xcc)
 
         # Configure apertures:
         # LDS:         0x10000000'00000000 - 0x10000001'00000000 (4GB)
@@ -231,14 +252,18 @@ class AM_GFX(AM_IP):
       self._grbm_select(inst=xcc)
 
       # Configure MEC doorbell range
-      self.adev.regCP_MEC_DOORBELL_RANGE_LOWER.write(0x0, inst=xcc)
-      self.adev.regCP_MEC_DOORBELL_RANGE_UPPER.write(0x450, inst=xcc)
+      self.adev.regCP_MEC_DOORBELL_RANGE_LOWER.write(0x100 * xcc, inst=xcc)
+      self.adev.regCP_MEC_DOORBELL_RANGE_UPPER.write(0x100 * xcc + 0xf8, inst=xcc)
 
       # Enable MEC
-      self.adev.regCP_MEC_RS64_CNTL.update(mec_invalidate_icache=0, mec_pipe0_reset=0, mec_pipe0_active=1, mec_halt=0, inst=xcc)
+      if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): self.adev.regCP_MEC_CNTL.write(0x0, inst=xcc)
+      else: self.adev.regCP_MEC_RS64_CNTL.update(mec_invalidate_icache=0, mec_pipe0_reset=0, mec_pipe0_active=1, mec_halt=0, inst=xcc)
     # NOTE: Wait for MEC to be ready. The kernel does udelay here as well.
     time.sleep(0.05)
 
+    # Set 1 partition
+    if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)
+
   def fini_hw(self):
     for xcc in range(self.xccs):
       self._grbm_select(me=1, pipe=0, queue=0, inst=xcc)
@@ -247,11 +272,11 @@ class AM_GFX(AM_IP):
     for xcc in range(self.xccs): self.adev.regGCVM_CONTEXT0_CNTL.write(0, inst=xcc)
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, doorbell:int, pipe:int, queue:int,
-                 aql:bool):
+                 aql:bool) -> int:
     for xcc in range(self.xccs if aql else 1):
       mqd = self.adev.mm.valloc(0x1000, uncached=True, contiguous=True)
 
-      struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}_compute_mqd")
+      struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}{'_compute' if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else ''}_mqd")
       mqd_struct = struct_t(header=0xC0310800, cp_mqd_base_addr_lo=lo32(mqd.va_addr), cp_mqd_base_addr_hi=hi32(mqd.va_addr),
         cp_hqd_persistent_state=self.adev.regCP_HQD_PERSISTENT_STATE.encode(preload_size=0x55, preload_req=1),
         cp_hqd_pipe_priority=0x2, cp_hqd_queue_priority=0xf, cp_hqd_quantum=0x111,
@@ -260,12 +285,13 @@ class AM_GFX(AM_IP):
         cp_hqd_pq_wptr_poll_addr_lo=lo32(wptr_addr), cp_hqd_pq_wptr_poll_addr_hi=hi32(wptr_addr),
         cp_hqd_pq_doorbell_control=self.adev.regCP_HQD_PQ_DOORBELL_CONTROL.encode(doorbell_offset=doorbell*2, doorbell_en=1),
         cp_hqd_pq_control=self.adev.regCP_HQD_PQ_CONTROL.encode(rptr_block_size=5, unord_dispatch=0, queue_size=(ring_size//4).bit_length()-2,
-          **({'queue_full_en':1, 'slot_based_wptr':2, 'no_update_rptr':1} if aql else {})),
+          **({'queue_full_en':1, 'slot_based_wptr':2, 'no_update_rptr':xcc==0} if aql else {})),
         cp_hqd_ib_control=self.adev.regCP_HQD_IB_CONTROL.encode(min_ib_avail_size=0x3), cp_hqd_hq_status0=0x20004000,
         cp_mqd_control=self.adev.regCP_MQD_CONTROL.encode(priv_state=1), cp_hqd_vmid=0, cp_hqd_aql_control=int(aql),
         cp_hqd_eop_base_addr_lo=lo32(eop_addr>>8), cp_hqd_eop_base_addr_hi=hi32(eop_addr>>8),
-        cp_hqd_eop_control=self.adev.regCP_HQD_EOP_CONTROL.encode(eop_size=(eop_size//4).bit_length()-2))
-      for se in range(8): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)
+        cp_hqd_eop_control=self.adev.regCP_HQD_EOP_CONTROL.encode(eop_size=(eop_size//4).bit_length()-2),
+        **({'compute_tg_chunk_size':1, 'compute_current_logic_xcc_id':xcc} if aql and self.xccs > 1 else {}))
+      for se in range(8 if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else 4): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)
 
       # Copy mqd into memory
       self.adev.vram.view(mqd.paddrs[0][0], ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
@@ -281,6 +307,7 @@ class AM_GFX(AM_IP):
       self._grbm_select(inst=xcc)
 
       self.adev.reg(f"regCP_ME1_PIPE{pipe}_INT_CNTL").update(time_stamp_int_enable=1, generic0_int_enable=1, inst=xcc)
+    return 0
 
   def set_clockgating_state(self):
     if hasattr(self.adev, 'regMM_ATC_L2_MISC_CG'): self.adev.regMM_ATC_L2_MISC_CG.write(enable=1, mem_ls_enable=1)
@@ -293,11 +320,13 @@ class AM_GFX(AM_IP):
 
       self.adev.regCP_RB_WPTR_POLL_CNTL.update(poll_frequency=0x100, idle_poll_count=0x90, inst=xcc)
       self.adev.regCP_INT_CNTL.update(cntx_busy_int_enable=1, cntx_empty_int_enable=1, cmp_busy_int_enable=1, gfx_idle_int_enable=1, inst=xcc)
-      self.adev.regSDMA0_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
-      self.adev.regSDMA1_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
+      if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
+        self.adev.regSDMA0_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
+        self.adev.regSDMA1_RLC_CGCG_CTRL.update(cgcg_int_enable=1, inst=xcc)
 
-      self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(perfmon_clock_state=1, gfxip_fgcg_override=0, gfxip_repeater_fgcg_override=0,
-        grbm_cgtt_sclk_override=0, rlc_cgtt_sclk_override=0, gfxip_mgcg_override=0, gfxip_cgls_override=0, gfxip_cgcg_override=0, inst=xcc)
+      feats_gfx11 = {'perfmon_clock_state':1, 'gfxip_repeater_fgcg_override':0} if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0) else {}
+      self.adev.regRLC_CGTT_MGCG_OVERRIDE.update(**feats_gfx11, gfxip_fgcg_override=0, grbm_cgtt_sclk_override=0, rlc_cgtt_sclk_override=0,
+        gfxip_mgcg_override=0, gfxip_cgls_override=0, gfxip_cgcg_override=0, inst=xcc)
 
       self.adev.regRLC_SAFE_MODE.write(message=0, cmd=1, inst=xcc)
 
@@ -314,10 +343,13 @@ class AM_GFX(AM_IP):
       self.adev.reg(f"regCP_{cntl_reg}_CNTL").update(**{f"{eng_name.lower()}_pipe{pipe}_reset": 0 for pipe in range(pipe_cnt)}, inst=xcc)
 
     for xcc in range(self.adev.gfx.xccs):
+      if self.adev.ip_ver[am.GC_HWIP] < (10,0,0):
+        self.adev.regCP_MEC_CNTL.update(mec_invalidate_icache=1, mec_me1_pipe0_reset=1, mec_me2_pipe0_reset=1, mec_me1_halt=1,mec_me2_halt=1,inst=xcc)
       if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0):
         _config_helper(eng_name="PFP", cntl_reg="ME", eng_reg="PFP", pipe_cnt=1, xcc=xcc)
         _config_helper(eng_name="ME", cntl_reg="ME", eng_reg="ME", pipe_cnt=1, xcc=xcc)
-      _config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)
+      if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
+        _config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)
 
 class AM_IH(AM_IP):
   def init_sw(self):
@@ -339,15 +371,16 @@ class AM_IH(AM_IP):
 
       self.adev.reg(f"regIH_DOORBELL_RPTR{suf}").write(offset=(am.AMDGPU_NAVI10_DOORBELL_IH + ring_id) * 2, enable=1)
 
-    self.adev.regIH_STORM_CLIENT_LIST_CNTL.update(client18_is_storm_client=1)
-    self.adev.regIH_INT_FLOOD_CNTL.update(flood_cntl_enable=1)
-    self.adev.regIH_MSI_STORM_CTRL.update(delay=3)
+    if self.adev.ip_ver[am.OSSSYS_HWIP] != (4,4,2):
+      self.adev.regIH_STORM_CLIENT_LIST_CNTL.update(client18_is_storm_client=1)
+      self.adev.regIH_INT_FLOOD_CNTL.update(flood_cntl_enable=1)
+      self.adev.regIH_MSI_STORM_CTRL.update(delay=3)
 
     # toggle interrupts
     for _, rwptr_vm, suf, ring_id in self.rings:
       self.adev.reg(f"regIH_RB_CNTL{suf}").update(rb_enable=1, **({'enable_intr': 1} if ring_id == 0 else {}))
 
-    if self.adev.ip_ver[am.NBIO_HWIP] != (7,9,0):
+    if self.adev.ip_ver[am.NBIO_HWIP][:2] != (7,9):
       self.adev.soc.doorbell_enable(port=1, awid=0x0, awaddr_31_28_value=0x0, offset=am.AMDGPU_NAVI10_DOORBELL_IH*2, size=2)
 
   def interrupt_handler(self):
@@ -363,42 +396,53 @@ class AM_IH(AM_IP):
 class AM_SDMA(AM_IP):
   def init_sw(self): self.sdma_name = "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
   def init_hw(self):
-    for pipe in range(2):
-      self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
-      self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
+    for pipe_id in range(1):
+      pipe = "" if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else str(pipe_id)
 
-      # rd=noa, wr=bypass
-      self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=0x2, wr_l2_policy=0x3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}))
-      self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0})
-      self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1)
+      if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
+        self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
+        self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
 
-    if self.adev.ip_ver[am.NBIO_HWIP] == (7,9,0):
+        # rd=noa, wr=bypass
+        self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=2, wr_l2_policy=3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}))
+        self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0})
+
+      self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1,
+        **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}))
+
+    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
       self.adev.regDOORBELL0_CTRL_ENTRY_1.write(bif_doorbell1_range_offset_entry=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2,
         bif_doorbell1_range_size_entry=4)
       self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x1, offset=0xe, size=4)
     else: self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x3, offset=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2, size=4)
 
   def fini_hw(self):
-    self.adev.regSDMA0_QUEUE0_RB_CNTL.update(rb_enable=0)
-    self.adev.regSDMA0_QUEUE0_IB_CNTL.update(ib_enable=0)
-    self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
-    time.sleep(0.01)
-    self.adev.regGRBM_SOFT_RESET.write(0x0)
+    reg, inst = ("regSDMA_GFX", 0) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else ("regSDMA0_QUEUE0", 0)
 
-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, doorbell:int, pipe:int, queue:int):
+    self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
+    self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
+    if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
+      self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
+      time.sleep(0.01)
+      self.adev.regGRBM_SOFT_RESET.write(0x0)
+
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, doorbell:int, pipe:int, queue:int) -> int:
     # Setup the ring
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_MINOR_PTR_UPDATE").write(0x1)
-    self.adev.wreg_pair(f"regSDMA{pipe}_QUEUE{queue}_RB_RPTR", "", "_HI", 0)
-    self.adev.wreg_pair(f"regSDMA{pipe}_QUEUE{queue}_RB_WPTR", "", "_HI", 0)
-    self.adev.wreg_pair(f"regSDMA{pipe}_QUEUE{queue}_RB_BASE", "", "_HI", ring_addr >> 8)
-    self.adev.wreg_pair(f"regSDMA{pipe}_QUEUE{queue}_RB_RPTR_ADDR", "_LO", "_HI", rptr_addr)
-    self.adev.wreg_pair(f"regSDMA{pipe}_QUEUE{queue}_RB_WPTR_POLL_ADDR", "_LO", "_HI", wptr_addr)
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_DOORBELL_OFFSET").update(offset=doorbell * 2)
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_DOORBELL").update(enable=1)
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_MINOR_PTR_UPDATE").write(0x0)
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_RB_CNTL").write(rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4,
-      **{f'{self.sdma_name.lower()}_wptr_poll_enable':1}, rb_size=(ring_size//4).bit_length()-1, rb_enable=1, rb_priv=1)
-    self.adev.reg(f"regSDMA{pipe}_QUEUE{queue}_IB_CNTL").update(ib_enable=1)
+    reg, inst = ("regSDMA_GFX", pipe*4+queue) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
+
+    self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x1, inst=inst)
+    self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
+    self.adev.wreg_pair(f"{reg}_RB_WPTR", "", "_HI", 0, inst=inst)
+    self.adev.wreg_pair(f"{reg}_RB_BASE", "", "_HI", ring_addr >> 8, inst=inst)
+    self.adev.wreg_pair(f"{reg}_RB_RPTR_ADDR", "_LO", "_HI", rptr_addr, inst=inst)
+    self.adev.wreg_pair(f"{reg}_RB_WPTR_POLL_ADDR", "_LO", "_HI", wptr_addr, inst=inst)
+    self.adev.reg(f"{reg}_DOORBELL_OFFSET").update(offset=doorbell * 2, inst=inst)
+    self.adev.reg(f"{reg}_DOORBELL").update(enable=1, inst=inst)
+    self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x0, inst=inst)
+    self.adev.reg(f"{reg}_RB_CNTL").write(**({f'{self.sdma_name.lower()}_wptr_poll_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP][:2]!=(4,4) else {}),
+      rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1, rb_priv=1, rb_size=(ring_size//4).bit_length()-1, inst=inst)
+    self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=1, inst=inst)
+    return self.adev.reg(f"{reg}_RB_WPTR").read() | (self.adev.reg(f"{reg}_RB_WPTR_HI").read() << 32)
 
 class AM_PSP(AM_IP):
   def init_sw(self):
@@ -435,14 +479,16 @@ class AM_PSP(AM_IP):
       while not self.is_sos_alive(): time.sleep(0.01)
 
     self._ring_create()
-    self._tmr_init()
+    if am.PSP_FW_TYPE_PSP_TOC in self.adev.fw.sos_fw: self._tmr_init()
 
     # SMU fw should be loaded before TMR.
     if hasattr(self.adev.fw, 'smu_psp_desc'): self._load_ip_fw_cmd(*self.adev.fw.smu_psp_desc)
     if not self.boot_time_tmr or not self.autoload_tmr: self._tmr_load_cmd()
 
     for psp_desc in self.adev.fw.descs: self._load_ip_fw_cmd(*psp_desc)
-    self._rlc_autoload_cmd()
+
+    if self.adev.ip_ver[am.GC_HWIP] >= (11,0,0): self._rlc_autoload_cmd()
+    else: self._load_ip_fw_cmd([am.GFX_FW_TYPE_REG_LIST], self.adev.fw.sos_fw[am.PSP_FW_TYPE_PSP_RL])
 
   def is_sos_alive(self): return self.adev.reg(f"{self.reg_pref}_81").read() != 0x0
 
@@ -537,4 +583,9 @@ class AM_PSP(AM_IP):
     cmd.cmd.cmd_load_toc.toc_size = toc_size
     return self._ring_submit(cmd)
 
+  def _spatial_partition_cmd(self, mode):
+    cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_SRIOV_SPATIAL_PART)
+    cmd.cmd.cmd_spatial_part.mode = mode
+    return self._ring_submit(cmd)
+
   def _rlc_autoload_cmd(self): return self._ring_submit(am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_AUTOLOAD_RLC))

tinygrad/runtime/support/amd.py

@@ -38,8 +38,9 @@ def fixup_ip_version(ip:str, version:tuple[int, ...]) -> list[tuple[int, ...]]:
     return version
 
   if ip in ['nbio', 'nbif']: version = _apply_ovrd({(3,3): (2,3,0), (7,3): (7,2,0)})
-  elif ip in ['mp', 'smu']: version = _apply_ovrd({(14,0,3): (14,0,2)})
+  elif ip in ['mp', 'smu']: version = _apply_ovrd({(14,0,3): (14,0,2), (13,0,12): (13,0,6)})
   elif ip in ['gc']: version = _apply_ovrd({(9,5,0): (9,4,3)})
+  elif ip in ['sdma']: version = _apply_ovrd({(4,4,4): (4,4,2)})
 
   return [version, version[:2], version[:2]+(0,), version[:1]+(0, 0)]
 
@@ -79,7 +80,12 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]:
 def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
   def _split_name(name): return name[:(pos:=next((i for i,c in enumerate(name) if c.isupper()), len(name)))], name[pos:]
   def _extract_regs(txt):
-    return {m.group(1): int(m.group(2), 0) for line in txt.splitlines() if (m:=re.match(r'#define\s+(\S+)\s+(0x[\da-fA-F]+|\d+)', line))}
+    x = {}
+    for k,v in {m.group(1): int(m.group(2), 0) for line in txt.splitlines() if (m:=re.match(r'#define\s+(\S+)\s+(0x[\da-fA-F]+|\d+)', line))}.items():
+      if k.startswith('VM_') or k.startswith('MC_'): x[prefix.upper()[:2]+k] = v
+      elif k.startswith('regVM_') or k.startswith('regMC_'): x["reg"+prefix.upper()[:2]+k[3:]] = v
+      else: x[k] = v
+    return x
   def _download_file(ver, suff) -> str:
     dir_prefix = {"osssys": "oss"}.get(prefix, prefix)
     fetch_name, file_name = f"{prefix}_{'_'.join(map(str, ver))}_{suff}.h", f"{prefix}_{'_'.join(map(str, version))}_{suff}.h"
@@ -98,6 +104,7 @@ def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[st
     for field_name, field_mask in sh_masks.items():
       if not ('__' in field_name and field_name.endswith('_MASK')): continue
       reg_name, reg_field_name = field_name[:-len('_MASK')].split('__')
+      if reg_name.startswith('MC_') or reg_name.startswith('VM_'): reg_name = f"{prefix.upper()[:2]}{reg_name}"
       fields[reg_name][reg_field_name.lower()] = ((field_mask & -field_mask).bit_length()-1, field_mask.bit_length()-1)
 
     # NOTE: Some registers like regGFX_IMU_FUSESTRAP in gc_11_0_0 are missing base idx, just skip them

tinygrad/runtime/support/autogen.py

@@ -1,5 +1,5 @@
 import ctypes, itertools, re, functools, os
-from tinygrad.helpers import flatten, unwrap
+from tinygrad.helpers import unwrap
 from tinygrad.runtime.autogen import libclang as clang # use REGEN=1 to regenerate libclang bindings
 
 def unwrap_cursor(c: clang.CXCursor) -> clang.CXCursor:
@@ -91,7 +91,7 @@ fns, specs = (clang.CXType_FunctionProto, clang.CXType_FunctionNoProto), (clang.
 # https://clang.llvm.org/docs/AutomaticReferenceCounting.html#arc-method-families
 arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
 
-def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_errno=False, anon_names={}, types={}, parse_macros=True):
+def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, parse_macros=True, paths=[]):
   macros, lines, anoncnt, types, objc = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False
   def tname(t, suggested_name=None, typedef=None) -> str:
     suggested_name = anon_names.get(f"{loc_file(loc(decl:=clang.clang_getTypeDeclaration(t)))}:{loc_line(loc(decl))}", suggested_name)
@@ -257,11 +257,9 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
         lines, types = rollback
     clang.clang_disposeTranslationUnit(tu)
     clang.clang_disposeIndex(idx)
-  main = (f"# mypy: ignore-errors\nimport ctypes{', os' if any('os' in s for s in dll) else ''}\n"
-    "from tinygrad.helpers import unwrap\nfrom tinygrad.runtime.support.c import Struct, CEnum, _IO, _IOW, _IOR, _IOWR\n" + '\n'.join([*prolog,
-      *(["from ctypes.util import find_library"]*any('find_library' in s for s in dll)), *(["from tinygrad.runtime.support import objc"]*objc),
-      *(["def dll():",*flatten([[f"  try: return ctypes.CDLL(unwrap({d}){', use_errno=True' if use_errno else ''})",'  except: pass'] for d in dll]),
-         "  return None", "dll = dll()\n"]*bool(dll)), *lines]) + '\n')
+  main = '\n'.join(["# mypy: ignore-errors", "import ctypes", "from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR",
+                    *prolog, *(["from tinygrad.runtime.support import objc"]*objc),
+                    *([f"dll = DLL('{name}', {dll}{f', {paths}'*bool(paths)}{', use_errno=True'*errno})"] if dll else []), *lines]) + '\n'
   macros = [r for m in macros if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]
   while True:
     try:

tinygrad/runtime/support/c.py

@@ -1,6 +1,6 @@
-import ctypes, functools, sys
+import ctypes, functools, os, pathlib, re, sys, sysconfig
 from typing import TYPE_CHECKING
-from tinygrad.helpers import flatten, WIN
+from tinygrad.helpers import flatten, getenv, DEBUG, OSX, WIN
 from _ctypes import _SimpleCData
 
 def _do_ioctl(__idir, __base, __nr, __struct, __fd, *args, __payload=None, **kwargs):
@@ -37,6 +37,42 @@ def CEnum(typ: type[ctypes._SimpleCData]):
 
   return _CEnum
 
+class DLL(ctypes.CDLL):
+  @staticmethod
+  def findlib(nm:str, paths:list[str], extra_paths=[]):
+    if nm == 'libc' and OSX: return '/usr/lib/libc.dylib'
+    if pathlib.Path(path:=getenv(nm.replace('-', '_').upper()+"_PATH", '')).is_file(): return path
+    for p in paths:
+      libpaths = {"posix": ["/usr/lib", "/usr/local/lib"], "nt": os.environ['PATH'].split(os.pathsep),
+                  "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework"],
+                  'linux': ['/lib', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
+      if (pth:=pathlib.Path(p)).is_absolute():
+        if pth.is_file(): return p
+        else: continue
+      for pre in (pathlib.Path(pre) for pre in libpaths.get(os.name, []) + libpaths.get(sys.platform, []) + extra_paths):
+        if not pre.is_dir(): continue
+        if WIN or OSX:
+          for base in ([f"lib{p}.dylib", f"{p}.dylib", str(p)] if OSX else [f"{p}.dll"]):
+            if (l:=pre / base).is_file() or (OSX and 'framework' in str(l) and l.is_symlink()): return str(l)
+        else:
+          for l in (l for l in pre.iterdir() if l.is_file() and re.fullmatch(f"lib{p}\\.so\\.?[0-9]*", l.name)):
+            # filter out linker scripts
+            with open(l, 'rb') as f:
+              if f.read(4) == b'\x7FELF': return str(l)
+
+  def __init__(self, nm:str, paths:str|list[str], extra_paths=[], emsg="", **kwargs):
+    self.nm, self.emsg, self.loaded = nm, emsg, False
+    if (path:= DLL.findlib(nm, paths if isinstance(paths, list) else [paths], extra_paths if isinstance(extra_paths, list) else [extra_paths])):
+      if DEBUG >= 3: print(f"loading {nm} from {path}")
+      try:
+        super().__init__(path, **kwargs)
+        self.loaded = True
+      except OSError as e: self.emsg = str(e)
+
+  def __getattr__(self, nm):
+    if not self.loaded: raise AttributeError(f"failed to load library {self.nm}: " + (self.emsg or f"try setting {self.nm.upper()+'_PATH'}?"))
+    return super().__getattr__(nm)
+
 # supports gcc (C11) __attribute__((packed))
 if TYPE_CHECKING: Struct = ctypes.Structure
 else:

tinygrad/runtime/support/ib.py

@@ -1,173 +0,0 @@
-from __future__ import annotations
-import resource, ctypes, weakref, functools, itertools
-from tinygrad.runtime.autogen import ib
-from typing import Iterator
-from dataclasses import dataclass
-from weakref import WeakKeyDictionary
-from tinygrad.device import Buffer, DMACPURef, DMAFdRef
-from tinygrad.helpers import getenv, round_up, DEBUG
-
-DEFAULT_PORT, DEFAULT_GID = getenv("DEFAULT_PORT", 1), getenv("DEFAULT_GID", 3) # DEFAULT_GID=0 for RXE
-IOVA_ALIGN = resource.getpagesize()
-
-def checkz(x, ret=None):
-  if x != 0: raise RuntimeError(f'{x} != 0 (errno {ctypes.get_errno()})')
-  return ret
-
-@dataclass(frozen=True)
-class SGE:
-  dst_iova: int
-  dst_key: int
-  src_iova: int
-  src_key: int
-  size: int
-
-class IBCtx:
-  def __init__(self, idx:int):
-    # Open the device (aka Host Channel Adapter in ib-speak)
-    devs = ib.ibv_get_device_list(ctypes.byref(ndevs:=ctypes.c_int32()))
-    if idx >= ndevs.value: raise IndexError(f"{idx} > {ndevs.value}")
-    self.ctx = ib.ibv_open_device(devs[idx])
-    ib.ibv_free_device_list(devs)
-
-    # HACK: remove this (and all usage of `ctx.contents.ops`) when clang2py can deal with `static inline` wrapper-functions
-    self.vctx = ctypes.cast(ctypes.addressof(self.ctx.contents) - ib.struct_verbs_context.context.offset, ctypes.POINTER(ib.struct_verbs_context))
-
-    # Get attributes. Something like port_attr.max_msg_sz sound like it might requre taking the min of host's and remote's attributes if they differ
-    self.device_attr = checkz(ib.ibv_query_device(self.ctx, ctypes.byref(da:=ib.struct_ibv_device_attr())), da)
-    self.port_attr = checkz(self.vctx.contents.query_port(self.ctx, DEFAULT_PORT, ctypes.byref(pa:=ib.struct_ibv_port_attr()), ctypes.sizeof(pa)), pa)
-    self.gid_attr = checkz(ib.ibv_query_gid(self.ctx, DEFAULT_PORT, DEFAULT_GID, ctypes.byref(ga:=ib.union_ibv_gid())), ga)
-
-    # Allocate protection domain
-    self.pd = ib.ibv_alloc_pd(self.ctx)
-    self.next_iova: int = IOVA_ALIGN # don't start at zero (nullptr)
-
-    # weakref(buf) => (iova, mr, mr_dealloc). mr_dealloc is kept here to avoid double freeing mrs that are deallocated in __del__
-    self.mrs: WeakKeyDictionary[Buffer, tuple[int, ctypes._Pointer[ib.struct_ibv_mr], weakref.finalize]] = WeakKeyDictionary()
-
-    # Default soft fd limit is 1024, which is not enough, set soft to hard (maximum allowed by the os)
-    IBCtx.rlimit_fix()
-
-  def __del__(self):
-    # must deallocate all mrs in protection domain before deallocating the protection domain
-    if hasattr(self, "mrs"): [fin() for _,_,fin in self.mrs.values()]
-    if hasattr(self, "pd"): ib.ibv_dealloc_pd(self.pd)
-    if hasattr(self, "ctx"): ib.ibv_close_device(self.ctx)
-
-  @functools.cache # run once
-  @staticmethod
-  def rlimit_fix():
-    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-    resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
-    if DEBUG>=2: print(f"IB: Increased fd limit from {soft} to {hard}")
-
-  def alloc_iova(self, size:int, required_offset:int):
-    iova = round_up(self.next_iova - required_offset, IOVA_ALIGN) + required_offset
-    self.next_iova = iova + size
-    return iova
-
-  def reg(self, buf:Buffer) -> tuple[int, ctypes._Pointer[ib.struct_ibv_mr]]:
-    buf = buf.base
-    if buf not in self.mrs:
-      if buf.nbytes > self.device_attr.max_mr_size: raise RuntimeError(f"Buffer too big: {buf.nbytes:#x} > {self.device_attr.max_mr_size:#x}")
-      if len(self.mrs) >= self.device_attr.max_mr: raise RuntimeError(f"Out of memory region cap: {len(self.mrs)} >= {self.device_attr.max_mr}")
-      # Local read is implied (but still have to create the memory region, except for short sends/writes with IBV_SEND_INLINE that are inlined by cpu)
-      mr_flags = ib.IBV_ACCESS_LOCAL_WRITE | ib.IBV_ACCESS_REMOTE_READ | ib.IBV_ACCESS_REMOTE_WRITE
-      match (dmaref:=buf.as_dmaref()):
-        case DMACPURef():
-          iova = self.alloc_iova(dmaref.size, dmaref.addr % IOVA_ALIGN)
-          mr = ib.ibv_reg_mr_iova2(self.pd, ctypes.c_void_p(dmaref.addr), dmaref.size, iova, mr_flags)
-        case DMAFdRef():
-          iova = self.alloc_iova(dmaref.size, dmaref.offset % IOVA_ALIGN)
-          mr = ib.ibv_reg_dmabuf_mr(self.pd, dmaref.offset, dmaref.size, iova, dmaref.fd, mr_flags)
-        case _: raise RuntimeError(f"Unknown type of dma ref: {dmaref}")
-      if not mr: raise RuntimeError(f"Couldn't register memory region for {buf} {dmaref} (errno={ctypes.get_errno()})")
-      self.mrs[buf] = (iova, mr, weakref.finalize(buf, ib.ibv_dereg_mr, mr))
-    return self.mrs[buf][0:2]
-
-class IBConn:
-  def __init__(self, ctx:IBCtx):
-    self.ctx = ctx
-
-    # Create Completion Channel. It is a file descriptor that kernel sends notifications through, not a thing in infiniband spec, just linux-ism
-    self.comp_channel = ib.ibv_create_comp_channel(self.ctx.ctx)
-    # Create Completion Queue. When a Work Request with signaled flag is completed a Completion Queue Entry is pushed onto this queue
-    self.cq = ib.ibv_create_cq(self.ctx.ctx, _capacity:=256, _cq_context:=None, self.comp_channel, _comp_vector:=0)
-    self.pending_wrids: set[int] = set()
-    self.wrid_num: Iterator[int] = itertools.count(0) # wc_id is uint64, this will never overflow
-
-    # Create Queue Pair. It's the closest thing to a socket in infiniband with QP num being the closest thing to a port, except it's allocated by hca
-    qp_init_attrs_cap = ib.struct_ibv_qp_cap(max_send_wr=1024, max_recv_wr=64, max_send_sge=8, max_recv_sge=8, max_inline_data=64)
-    qp_init_attrs = ib.struct_ibv_qp_init_attr(send_cq=self.cq, recv_cq=self.cq, cap=qp_init_attrs_cap, qp_type=ib.IBV_QPT_RC) # Reliable Connection
-    self.qp = ib.ibv_create_qp(self.ctx.pd, ctypes.byref(qp_init_attrs))
-    self.qp_cap = qp_init_attrs.cap
-
-    # The most important thing about QPs is their state, when a new QP is created it's in the RESET state, before it can be properly used it has to go
-    # through Init, Ready To Receive, Ready To Send. A good docs on QP state machine: https://www.rdmamojo.com/2012/05/05/qp-state-machine/
-
-    # INIT
-    qp_access_flags = ib.IBV_ACCESS_REMOTE_WRITE | ib.IBV_ACCESS_REMOTE_READ
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_INIT, port_num=DEFAULT_PORT, qp_access_flags=qp_access_flags)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PORT | ib.IBV_QP_ACCESS_FLAGS | ib.IBV_QP_PKEY_INDEX))
-
-    self.gid, self.qp_num = bytes(self.ctx.gid_attr.raw), self.qp.contents.qp_num
-
-  # Exchange GID and QP num with remote. At least in RoCEv2 gid can be guessed from remote's ip, QP num can't.
-
-  def connect(self, remote_gid:bytes, remote_qp_num:int):
-    # RTR
-    qp_ah_attr_grh = ib.struct_ibv_global_route(hop_limit=1, dgid=ib.union_ibv_gid(raw=(ctypes.c_ubyte * 16)(*remote_gid)), sgid_index=DEFAULT_GID)
-    qp_ah_attr = ib.struct_ibv_ah_attr(is_global=1, port_num=DEFAULT_PORT, grh=qp_ah_attr_grh)
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTR, path_mtu=ib.IBV_MTU_4096, dest_qp_num=remote_qp_num, rq_psn=0, max_dest_rd_atomic=1,
-                                min_rnr_timer=12, ah_attr=qp_ah_attr)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PATH_MTU | ib.IBV_QP_DEST_QPN | ib.IBV_QP_RQ_PSN | \
-                                          ib.IBV_QP_MAX_DEST_RD_ATOMIC | ib.IBV_QP_MIN_RNR_TIMER | ib.IBV_QP_AV))
-
-    # RTS
-    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTS, timeout=14, retry_cnt=7, rnr_retry=7, sq_psn=0, max_rd_atomic=1)
-    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_TIMEOUT | ib.IBV_QP_RETRY_CNT | ib.IBV_QP_RNR_RETRY | ib.IBV_QP_SQ_PSN | \
-                                          ib.IBV_QP_MAX_QP_RD_ATOMIC))
-
-  def __del__(self):
-    self.wait_cq() # need to wait for **everything** to complete before it's safe to dealloc queues and stuff
-    ib.ibv_destroy_qp(self.qp)
-    ib.ibv_destroy_cq(self.cq)
-    ib.ibv_destroy_comp_channel(self.comp_channel)
-
-  def next_wrid(self):
-    self.pending_wrids.add(wrid:=next(self.wrid_num))
-    return wrid
-
-  def wait_cq(self, wr_id: int|None=None):
-    while (wr_id in self.pending_wrids) if wr_id is not None else self.pending_wrids:
-      if self.ctx.ctx.contents.ops.poll_cq(self.cq, _num_entries:=1, ctypes.byref(wc:=ib.struct_ibv_wc())):
-        if wc.status != ib.IBV_WC_SUCCESS:
-          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.enum_ibv_wc_status.get(wc.status, wc.status)}')
-        self.pending_wrids.remove(wc.wr_id)
-
-  def rdma_write(self, sgl:list[SGE]):
-    swr: ctypes._Pointer[ib.struct_ibv_send_wr]|None = None
-    swr_cnt, wr_id = 0, self.next_wrid()
-    def _post():
-      nonlocal swr, swr_cnt, wr_id
-      if swr is not None:
-        # The swr can be freed when this returns, the memory that sge points to can be unmapped after work completion is retrieved from cq
-        checkz(self.ctx.ctx.contents.ops.post_send(self.qp, swr, ctypes.byref(_bad_wr:=ctypes.POINTER(ib.struct_ibv_send_wr)())))
-        # TODO: async
-        self.wait_cq(wr_id)
-        swr, swr_cnt, wr_id = None, 0, self.next_wrid()
-    # Everything is in reverse for elegant chaining
-    for sg in reversed(sgl):
-      # Message size limit (max 2GB per ib spec, 1GB on tinybox mellanoxes) applies to both scatter-gather entries and entire wrs
-      for off in reversed(range(0, sg.size, self.ctx.port_attr.max_msg_sz)):
-        # Scatter-Gather Entry for local memory
-        sge = ctypes.pointer(ib.struct_ibv_sge(addr=sg.src_iova+off, length=min(sg.size-off, self.ctx.port_attr.max_msg_sz), lkey=sg.src_key))
-        # RDMA struct for remote memory
-        wr = ib.struct_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_wr_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
-        # Signal (with chosen work request id) if it's the last wr (first in the loop since it's reversed)
-        wid, flags = (wr_id, ib.IBV_SEND_SIGNALED) if swr is None else (0, 0)
-        # Create Send Request
-        swr = ctypes.pointer(ib.struct_ibv_send_wr(opcode=ib.IBV_WR_RDMA_WRITE, sg_list=sge, num_sge=1, wr=wr, wr_id=wid, send_flags=flags, next=swr))
-        # Flush if queue is being overrun
-        if (swr_cnt:=swr_cnt + 1) >= self.qp_cap.max_send_wr: _post()
-    _post()

tinygrad/runtime/support/llvm.py

@@ -1,25 +0,0 @@
-import ctypes.util, os, sys
-from tinygrad.helpers import DEBUG, OSX, getenv, system
-
-if sys.platform == 'win32':
-  # Windows llvm distribution doesn't seem to add itself to PATH or anywhere else where it can be easily retrieved from.
-  # winget also doesn't have something like `brew --prefix llvm` so just hardcode default installation path with an option to override
-  LLVM_PATH = getenv('LLVM_PATH', 'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll')
-  if not os.path.exists(LLVM_PATH):
-    raise FileNotFoundError('LLVM not found, you can install it with `winget install LLVM.LLVM` or point at a custom dll with LLVM_PATH')
-elif OSX:
-  # Will raise FileNotFoundError if brew is not installed
-  # `brew --prefix` will return even if formula is not installed
-  if not os.path.exists(brew_prefix:=system("brew --prefix llvm@20")):
-    raise FileNotFoundError('LLVM not found, you can install it with `brew install llvm@20`')
-  LLVM_PATH: str|None = os.path.join(brew_prefix, 'lib', 'libLLVM.dylib')
-else:
-  LLVM_PATH = ctypes.util.find_library('LLVM')
-  # use newer LLVM if possible
-  for ver in reversed(range(14, 21+1)):
-    if LLVM_PATH is not None: break
-    LLVM_PATH = ctypes.util.find_library(f'LLVM-{ver}')
-  if LLVM_PATH is None:
-    raise FileNotFoundError("No LLVM library found on the system. Install it via your distro's package manager and ensure it's findable as 'LLVM'")
-
-if DEBUG>=3: print(f'Using LLVM at {repr(LLVM_PATH)}')

tinygrad/runtime/support/webgpu.py

@@ -1,18 +0,0 @@
-import ctypes.util, os, platform, sysconfig
-from tinygrad.helpers import system, OSX
-
-WEBGPU_PATH: str | None
-
-if OSX:
-  if not os.path.exists(brew_prefix:=system("brew --prefix dawn")):
-    raise FileNotFoundError('dawn library not found. Install it with `brew tap wpmed92/dawn && brew install dawn`')
-  WEBGPU_PATH = os.path.join(brew_prefix, 'lib', 'libwebgpu_dawn.dylib')
-elif platform.system() == "Windows":
-  if not os.path.exists(pydawn_path:=os.path.join(sysconfig.get_paths()["purelib"], "pydawn")):
-    raise FileNotFoundError("dawn library not found. Install it with `pip install dawn-python`")
-  WEBGPU_PATH = os.path.join(pydawn_path, "lib", "libwebgpu_dawn.dll")
-else:
-  if (WEBGPU_PATH:=ctypes.util.find_library('webgpu_dawn')) is None:
-    raise FileNotFoundError("dawn library not found. " +
-    "Install it with `sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.3.0/" +
-    f"libwebgpu_dawn_{platform.machine()}.so -o /usr/lib/libwebgpu_dawn.so`")

tinygrad/schedule/multi.py

@@ -218,6 +218,11 @@ multi_pm = PatternMatcher([
     lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
   (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
         src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
+  # multi supports custom kernels with CUSTOM_KERNEL + AFTER
+  (UPat(Ops.CUSTOM_KERNEL, src=UPat(Ops.MULTI), name="ck"),
+    lambda ck: ck.replace(src=tuple(m.src[0] for m in ck.src))),
+  (UPat(Ops.AFTER, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.CUSTOM_KERNEL)), name="a"),
+    lambda multi,a: a.replace(src=(multi.src[0],)+a.src[1:]).multi(multi.axis))
 ])+replace_allreduce
 
 def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:

tinygrad/schedule/rangeify.py

@@ -63,10 +63,17 @@ mop_cleanup = PatternMatcher([
    lambda x,x2: x.replace(src=(x2.src[0], x.src[1])) if x.tag is None and x2.tag is None else None),
 ])
 
+def resolve_custom_kernel(ck:UOp) -> UOp:
+  placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(ck.src)]
+  return UOp(Ops.KERNEL, src=ck.src, arg=Kernel(ck.arg.fxn(*placeholders)))
+
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 
+  # resolve custom kernels
+  (UPat(Ops.CUSTOM_KERNEL, name="ck"), resolve_custom_kernel),
+
   # remove CONTIGUOUS if the BUFFER is already contiguous
   (UPat(Ops.BUFFER).f(Ops.RESHAPE, allow_any_len=True, name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
 

tinygrad/tensor.py

@@ -190,8 +190,9 @@ class Tensor(OpMixin):
     lhs,rhs = self._broadcasted(x, reverse)
     return lhs._apply_uop(fxn, rhs)
 
-  # _binop is used by MathTrait
+  # _binop and alu are used by MathMixin
   def _binop(self, op, x, reverse): return self._apply_broadcasted_uop(lambda *u: UOp.alu(u[0], op, *u[1:]), x, reverse)
+  def alu(self, op: Ops, *src: Tensor) -> Tensor: return self._apply_uop(lambda *u: u[0].alu(op, *u[1:]), *src)
 
   def requires_grad_(self, requires_grad=True) -> Tensor:
     self.requires_grad = requires_grad
@@ -395,7 +396,7 @@ class Tensor(OpMixin):
     print(t.shard((t.device, t.device), axis=1).uop)
     ```
     """
-    assert isinstance(self.device, str), "can't shard a MultiLazyBuffer"
+    if not isinstance(self.device, str): raise RuntimeError("can't shard a MultiLazyBuffer")
     if len(devices) == 1: return self.to(devices[0])
     devices = tuple(canonicalize_device(x) for x in devices)
     mlb = self.uop.shard(devices, self._resolve_dim(axis)) if axis is not None else self.uop.copy_to_device(devices)
@@ -407,6 +408,13 @@ class Tensor(OpMixin):
     """
     return self.replace(self.shard(devices, axis))
 
+  def shard_like(self, y:Tensor) -> Tensor:
+    """
+    Shards the tensor the same way as `y` (same devices and axis).
+    """
+    if isinstance(y.device, str): return self.to(y.device)
+    return self if isinstance(self.device, tuple) and (y.device, y.uop.axis) == (self.device, self.uop.axis) else self.shard(y.device, y.uop.axis)
+
   CHUNK_SIZE = 2**20
   def fs_load(self, size:int) -> Tensor:
     """
@@ -1334,10 +1342,11 @@ class Tensor(OpMixin):
     print("\\n".join([repr(x.numpy()) for x in split]))
     ```
     """
-    assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
     dim = self._resolve_dim(dim)
-    if isinstance(sizes, int): sizes = [min(sizes, self.shape[dim]-i) for i in range(0, max(1, self.shape[dim]), max(1, sizes))]
-    assert sum(sizes) == self.shape[dim], f"expect sizes to sum exactly to {self.shape[dim]}, but got {sum(sizes)}"
+    dim_sz = self.shape[dim]
+    assert isinstance(dim_sz, int), f"does not support symbolic shape in split dimension {dim}: {self.shape}"
+    if isinstance(sizes, int): sizes = [min(sizes, dim_sz-i) for i in range(0, max(1, dim_sz), max(1, sizes))]
+    assert sum(sizes) == dim_sz, f"expect sizes to sum exactly to {dim_sz}, but got {sum(sizes)}"
     return tuple(self[sl] for sl in [tuple([slice(None)]*dim + [slice(sum(sizes[:i]), sum(sizes[:i + 1]))]) for i in range(len(sizes))])
 
   def chunk(self, chunks:int, dim:int=0) -> list[Tensor]:
@@ -1359,10 +1368,11 @@ class Tensor(OpMixin):
     print("\\n".join([repr(x.numpy()) for x in chunked]))
     ```
     """
-    assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
-    assert chunks > 0, f"expect chunks to be greater than 0, got: {chunks}"
     dim = self._resolve_dim(dim)
-    return list(self.split(ceildiv(self.shape[dim], chunks) if self.shape[dim] else [0]*chunks, dim=dim))
+    dim_sz = self.shape[dim]
+    assert isinstance(dim_sz, int), f"does not support symbolic shape in split dimension {dim}: {self.shape}"
+    assert chunks > 0, f"expect chunks to be greater than 0, got: {chunks}"
+    return list(self.split(ceildiv(dim_sz, chunks) if dim_sz else [0]*chunks, dim=dim))
 
   def unfold(self, dim:int, size:sint, step:int) -> Tensor:
     """
@@ -2788,29 +2798,6 @@ class Tensor(OpMixin):
     """
     return self.cast(least_upper_float(self.dtype))._apply_uop(UOp.exp2)
 
-  def relu(self) -> Tensor:
-    """
-    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) -> Tensor:
-    """
-    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 logsigmoid(self) -> Tensor:
     """
     Applies the LogSigmoid function element-wise.
@@ -2823,19 +2810,6 @@ class Tensor(OpMixin):
     """
     return -(-self).softplus()
 
-  def hardsigmoid(self, alpha:float=1/6, beta:float=0.5) -> Tensor:
-    """
-    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 sqrt(self) -> Tensor:
     """
     Computes the square root of the tensor element-wise.
@@ -2846,16 +2820,6 @@ class Tensor(OpMixin):
     """
     return self.cast(least_upper_float(self.dtype))._apply_uop(UOp.sqrt)
 
-  def rsqrt(self) -> Tensor:
-    """
-    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()
-
   def sin(self) -> Tensor:
     """
     Computes the sine of the tensor element-wise.
@@ -2922,36 +2886,6 @@ class Tensor(OpMixin):
 
   # ***** math functions *****
 
-  def trunc(self: Tensor) -> Tensor:
-    """
-    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._apply_uop(UOp.trunc)
-
-  def ceil(self: Tensor) -> Tensor:
-    """
-    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: Tensor) -> Tensor:
-    """
-    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 round(self: Tensor) -> Tensor:
     """
     Rounds the tensor element-wise with rounding half to even.
@@ -2962,36 +2896,6 @@ class Tensor(OpMixin):
     """
     return ((self > 0) == ((b := self.trunc() / 2.0).trunc() == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
 
-  def isinf(self:Tensor, detect_positive:bool=True, detect_negative:bool=True) -> Tensor:
-    """
-    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 == float("inf")) * detect_positive + (self == float("-inf")) * detect_negative
-
-  def isnan(self:Tensor) -> Tensor:
-    """
-    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 isfinite(self:Tensor) -> Tensor:
-    """
-    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 lerp(self, end:Tensor, weight:Tensor|float) -> Tensor:
     """
     Linearly interpolates between `self` and `end` by `weight`.
@@ -3005,36 +2909,6 @@ class Tensor(OpMixin):
       return (self+(((end - self).cast(dtypes.int8) * w_i + (1<<W_PREC-1)).cast(dtypes.uint16) >> W_PREC)).cast(dtypes.uint8)
     return self + (end - self) * weight
 
-  def square(self) -> Tensor:
-    """
-    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) -> Tensor:
-    """
-    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) -> Tensor:
-    """
-    Alias for `Tensor.clamp`.
-    """
-    return self.clamp(min_, max_)
-
   def sign(self) -> Tensor:
     """
     Returns the sign of the tensor element-wise.
@@ -3103,66 +2977,6 @@ class Tensor(OpMixin):
     """
     return gamma * (self >= 0).detach().where(self, alpha * (self.exp() - 1))
 
-  def swish(self) -> Tensor:
-    """
-    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) -> Tensor:
-    """
-    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 relu6(self) -> Tensor:
-    """
-    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) -> Tensor:
-    """
-    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 tanh(self) -> Tensor:
-    """
-    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 sinh(self) -> Tensor:
     """
     Applies the Hyperbolic Sine (sinh) function element-wise.
@@ -3223,16 +3037,6 @@ class Tensor(OpMixin):
     """
     return (self + (self.square() - 1).sqrt()).log()
 
-  def hardtanh(self, min_val=-1, max_val=1) -> Tensor:
-    """
-    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 erf(self) -> Tensor:
     """
     Applies error function element-wise.
@@ -3247,41 +3051,6 @@ class Tensor(OpMixin):
     t = 1.0 / (1.0 + 0.3275911 * self.abs())
     return self.sign() * (1.0 - t * polyN(t, [1.061405429, -1.453152027, 1.421413741, -0.284496736, 0.254829592]) * (-self.square()).exp())
 
-  def gelu(self) -> Tensor:
-    """
-    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 quick_gelu(self) -> Tensor:
-    """
-    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 leaky_relu(self, neg_slope=0.01) -> Tensor:
-    """
-    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 mish(self) -> Tensor:
     """
     Applies the Mish function element-wise.

tinygrad/uop/__init__.py

@@ -75,6 +75,7 @@ class Ops(FastEnum):
 
   # tensor graph ops
   UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); ASSIGN = auto()
+  CUSTOM_KERNEL = auto()
 
   # local unique
   LUNIQUE = auto()

tinygrad/uop/decompositions.py

@@ -80,10 +80,12 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
   intermediate_dtype = dtypes.float32.vec(d.dtype.count) if d.dtype.base.scalar() == dtypes.float16 else d.dtype
 
   f, e = frexp(d)
-  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(dtypes.uint64)
+  vc = d.dtype.count
+  u64, u32, i32 = dtypes.uint64.vec(vc), dtypes.uint32.vec(vc), dtypes.int32.vec(vc)
+  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(u64)
   # extract 96 relevant bits of 2/pi based on magnitude of argument
-  i = shr(e.cast(dtypes.uint64), 5)
-  e = e.cast(dtypes.int32) & 31
+  i = shr(e.cast(u64), 5)
+  e = e.cast(i32) & 31
   offset = 32 - e
 
   def _take(an:UOp, offset:int, count:int=0) -> UOp:
@@ -91,8 +93,8 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
     if count+offset < len(two_over_pi_f) - 1:
       an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
     return an
-  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
-  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(u64) * pow2if(y, d.dtype).cast(u64)).cast(u32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(u64) // pow2if(y, d.dtype).cast(u64)).cast(u32)
 
   a = [_take(UOp.const(dtypes.uint32.vec(d.dtype.count), 0), i) for i in range(4)]
   #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
@@ -101,12 +103,12 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
   mi = _shl_lazy(a[1], e) | _shr_lazy(a[2], offset)
   lo = _shl_lazy(a[2], e) | _shr_lazy(a[3], offset)
 
-  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(dtypes.uint64) * y.cast(dtypes.uint64)
+  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(u64) * y.cast(u64)
   # compute x * 2/pi
   p = shl(_hp_mul(ia, hi), 32) + _hp_mul(ia, mi) + shr(_hp_mul(ia, lo), 32)
 
   # round quotient to nearest
-  q = shr(p, 62).cast(dtypes.int32)
+  q = shr(p, 62).cast(i32)
   p = p & 0x3fffffffffffffff
   r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)
 
@@ -133,7 +135,8 @@ def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
       d = (qdh + q) * -PI_D + d
     elif x.dtype.scalar() == dtypes.float16:
       # [FIXME] when reducing `d`, FP16 needs FP32 precision to achieve 1.0 ULP precision.
-      d = _reduce_d(x.cast(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
+      f32 = dtypes.float32.vec(x.dtype.count)
+      d = _reduce_d(x.cast(f32), q.cast(f32)).cast(x.dtype)
     else:
       # https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
       d = q * -3.1414794921875 + x
@@ -143,9 +146,9 @@ def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
     return d
 
   m_1_pi = 0.318309886183790671537767526745028724
-  qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64).cast(d.dtype) * (2.0**24)
+  qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64.vec(d.dtype.count)).cast(d.dtype) * (2.0**24)
   quadrant = rintk(d * m_1_pi -qdh) if d.dtype.base.scalar() == dtypes.float64 else rintk(d * m_1_pi)
-  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32)
+  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32.vec(d.dtype.count))
 
 # *** approximate sine on small angle. ***
 def trig_poly(d:UOp, coeff32, coeff64): return d * (polyN(d*d, coeff64) if d.dtype.scalar() == dtypes.float64 else polyN(d*d, coeff32))
@@ -224,7 +227,8 @@ def xlog2(d:UOp) -> UOp:
   """
   assert d.dtype.scalar() in TRANSCENDENTAL_DTYPES
   # TODO: float16 denormal need float32 to achieve precision
-  if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast(dtypes.float32)).cast(dtypes.float16)
+  if d.dtype.scalar() == dtypes.float16:
+    return xlog2(d.cast(dtypes.float32.vec(d.dtype.count))).cast(d.dtype)
   FLT_MIN = d.const_like(1e-6 if d.dtype.scalar() == dtypes.float16 else 1e-4)
   is_denormal = d<FLT_MIN
   a = is_denormal.where(d * (2 ** 64), d)
@@ -261,9 +265,10 @@ def xpow(base:UOp, exponent:UOp) -> UOp:
   # start with b ** e = exp2(e * log2(b))
   ret = (base < 0).where(-base, base).log2().mul(exponent).exp2()
   # negative base adjustment: nan for non-integer exponent and -1 for odd exponent
-  non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype)
+  int32_dtype = dtypes.int32.vec(exponent.dtype.count)
+  non_int = exponent != exponent.cast(int32_dtype).cast(exponent.dtype)
   adj = non_int.where(ret.const_like(math.nan),
-    (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool).where(ret.const_like(-1), ret.const_like(1)))
+    (exponent < 0).where(-exponent, exponent).cast(int32_dtype).mod(2).cast(dtypes.bool.vec(exponent.dtype.count)).where(ret.const_like(-1), ret.const_like(1)))
   # fix 0 ** 0 = 1
   return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), ret * (base < 0).where(adj, ret.const_like(1)))
 

tinygrad/uop/ops.py

@@ -7,7 +7,7 @@ from tinygrad.uop import Ops, GroupOp
 from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType, AddrSpace
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CI
-from tinygrad.helpers import strip_parens, colored, ansilen, printable
+from tinygrad.helpers import strip_parens, colored, ansilen, printable, panic
 if TYPE_CHECKING:
   from tinygrad.device import Buffer, MultiBuffer
 
@@ -218,7 +218,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT:
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL:
         return None
 
       case Ops.INDEX:
@@ -402,8 +402,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if count == 1: return self
     return UOp(Ops.VECTORIZE, self.dtype.vec(count), (self,)*count)
   def cast(self, dtype:DType):
-    # TODO: we shouldn't have to check for dtype.count == 1 here, but CAST is misused in AMD LLVM
-    if dtype.count == 1 and dtype.count != self.dtype.count: dtype = dtype.vec(self.dtype.count)
     if self.dtype == dtype: return self
     return UOp(Ops.CAST, dtype, (self,))
   def bitcast(self, dtype:DType): return UOp(Ops.BITCAST, dtype, (self,))
@@ -470,7 +468,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   def is_contiguous(self):
     # TODO: this is is_realized
-    if self.op is Ops.RESHAPE: return self.src[0].is_contiguous()
+    if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].is_contiguous()
     return self.op is Ops.BUFFER
 
   def contiguous(self, *args, **kwargs):
@@ -616,7 +614,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op is Ops.BUFFERIZE: return self.arg.device
     if self.op is Ops.AFTER: return self.src[0]._device
     if self.op is Ops.MSELECT:
-      assert isinstance(self.src[0].device, tuple), "mselect must be on tuple device"
+      assert isinstance(self.src[0].device, tuple), f"mselect must be on tuple device, getting {self.src[0].device}"
       return self.src[0].device[self.arg]
     if self.op is Ops.MSTACK: return tuple(cast(str, x.device) for x in self.src)
     if self.op in {Ops.COPY, Ops.BUFFER, Ops.ALLREDUCE}: return self.src[1].device
@@ -840,9 +838,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return self.src[0].after(self.store(val).end(*argfix(end)))
 
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
     contig_srcs = tuple(x.contiguous() for x in srcs)
-    kernel = UOp(Ops.KERNEL, src=tuple(x.base for x in contig_srcs), arg=Kernel(fxn(*placeholders), grad_fxn=grad_fxn))
+    kernel = UOp(Ops.CUSTOM_KERNEL, src=contig_srcs, arg=CustomKernel(fxn=fxn, grad_fxn=grad_fxn))
     return [s.after(kernel) for s in contig_srcs]
 
 @dataclass(frozen=True)
@@ -855,6 +852,14 @@ class KernelInfo:
   @property
   def function_name(self): return to_function_name(self.name)
 
+@dataclass(frozen=True)
+class CustomKernel:
+  fxn: Callable
+  grad_fxn: Callable|None = None
+  # sadly CustomKernel can't be pickled or reconstructed as a str
+  def __reduce__(self): return (CustomKernel, (panic,))
+  def __repr__(self): return "CustomKernel(panic)"
+
 @dataclass(frozen=True)
 class Kernel:
   ast: UOp
@@ -1054,6 +1059,7 @@ class PatternMatcher:
 # *** tracking pattern matcher ***
 
 TRACK_MATCH_STATS = ContextVar("TRACK_MATCH_STATS", 2 if VIZ else 0)
+REWRITE_STACK_LIMIT = ContextVar("REWRITE_STACK_LIMIT", 250000)
 match_stats:dict[UPat, list[int|float]] = dict()
 
 # TRACK_MATCH_STATS>=2 or VIZ=1 saves all matches
@@ -1210,7 +1216,6 @@ class RewriteContext:
   def unified_rewrite(self, root:UOp) -> UOp:
     stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
     on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
-    REWRITE_STACK_LIMIT = getenv("REWRITE_STACK_LIMIT", 250000)
     while stack:
       if len(stack) > REWRITE_STACK_LIMIT: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
       n, stage, new_n = stack.pop()

tinygrad/uop/spec.py

@@ -1,8 +1,8 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel, CustomKernel
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata
+from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic
 from tinygrad.uop.validate import validate_index
 
 # four specs:
@@ -54,7 +54,10 @@ movement_ops = PatternMatcher([
   (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
 
   # AFTER on Movement Op
-  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement),), allow_any_len=True), lambda: True),
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI})),), allow_any_len=True), lambda: True),
+
+  # custom kernels allowed here
+  (UPat(Ops.CUSTOM_KERNEL), lambda: True),
 ])
 
 _tensor_spec = PatternMatcher([
@@ -274,8 +277,8 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.schedule.rangeify import BufferizeOpts
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel, "Metadata": Metadata,
-                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
-                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}
+                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid, "CustomKernel": CustomKernel,
+                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic}
 def eval_pyrender(code:str) -> UOp:
   lcls:dict[str, Any] = {}
   exec(code, glbls, lcls)

tinygrad/viz/js/index.js

@@ -698,6 +698,29 @@ const toggleLabel = d3.create("label").text("Show indexing (r)").node();
 const toggle = d3.create("input").attr("type", "checkbox").attr("id", "show-indexing").property("checked", true).node();
 toggleLabel.prepend(toggle);
 
+function appendSteps(root, idx, steps) {
+  const stack = [];
+  for (const [j,u] of steps.entries()) {
+    while (stack.length && stack.at(-1).depth >= u.depth) stack.pop();
+    const list = stack.length > 0 ? stack.at(-1).li : root;
+    u.li = list.appendChild(document.createElement("ul"));
+    u.li.id = `step-${idx}-${j}`
+    const p = u.li.appendChild(document.createElement("p"));
+    p.appendChild(colored(`${u.name}`+(u.match_count ? ` - ${u.match_count}` : '')));
+    p.onclick = (e) => {
+      e.stopPropagation();
+      const subrewrites = getSubrewrites(e.currentTarget.parentElement);
+      if (subrewrites.length) { e.currentTarget.parentElement.classList.toggle("expanded"); }
+      setState({ currentStep:j, currentCtx:idx, currentRewrite:0 });
+    }
+    stack.push(u);
+  }
+  for (const l of root.querySelectorAll("ul > ul > p")) {
+    const subrewrites = getSubrewrites(l.parentElement);
+    if (subrewrites.length > 0) { l.appendChild(d3.create("span").text(` (${subrewrites.length})`).node()); l.parentElement.classList.add("has-children"); }
+  }
+}
+
 async function main() {
   // ** left sidebar context list
   if (ctxs == null) {
@@ -712,26 +735,7 @@ async function main() {
       p.onclick = () => {
         setState(i === state.currentCtx ? { expandSteps:!state.expandSteps } : { expandSteps:true, currentCtx:i, currentStep:0, currentRewrite:0 });
       }
-      const stack = []; let list = ul;
-      for (const [j,u] of steps.entries()) {
-        while (stack.length && stack.at(-1).depth >= u.depth) stack.pop();
-        const list = stack.length > 0 ? stack.at(-1).li : ul;
-        u.li = list.appendChild(document.createElement("ul"));
-        u.li.id = `step-${i}-${j}`
-        const p = u.li.appendChild(document.createElement("p"));
-        p.appendChild(colored(`${u.name}`+(u.match_count ? ` - ${u.match_count}` : '')));
-        p.onclick = (e) => {
-          e.stopPropagation();
-          const subrewrites = getSubrewrites(e.currentTarget.parentElement);
-          if (subrewrites.length) { e.currentTarget.parentElement.classList.toggle("expanded"); }
-          setState({ currentStep:j, currentCtx:i, currentRewrite:0 });
-        }
-        stack.push(u);
-      }
-      for (const l of ul.querySelectorAll("ul > ul > p")) {
-        const subrewrites = getSubrewrites(l.parentElement);
-        if (subrewrites.length > 0) { l.appendChild(d3.create("span").text(` (${subrewrites.length})`).node()); l.parentElement.classList.add("has-children"); }
-      }
+      appendSteps(ul, i, steps);
     }
     return setState({ currentCtx:-1 });
   }
@@ -756,6 +760,15 @@ async function main() {
   // ** Disassembly view
   if (!ckey.startsWith("/rewrites")) {
     if (!(ckey in cache)) cache[ckey] = ret = await fetchValue(ckey);
+    if (ret.steps?.length > 0) {
+      const el = select(state.currentCtx, state.currentStep);
+      if (el.step.querySelectorAll("ul").length === ret.steps.length) return;
+      // re render the list with new items
+      ctx.steps.push(...ret.steps);
+      while (el.ctx.children.length > 1) el.ctx.children[1].remove();
+      appendSteps(el.ctx, state.currentCtx, ctx.steps);
+      return setState({ currentStep:state.currentStep+1, expandSteps:true });
+    }
     // cycles on the x axis
     if (ret instanceof ArrayBuffer) {
       opts = {heightScale:0.5, hideLabels:true, levelKey:(e) => parseInt(e.name.split(" ")[1].split(":")[1])};
@@ -817,7 +830,7 @@ async function main() {
     const eventSource = new EventSource(ckey);
     evtSources.push(eventSource);
     eventSource.onmessage = (e) => {
-      if (e.data === "END") return eventSource.close();
+      if (e.data === "[DONE]") return eventSource.close();
       const chunk = JSON.parse(e.data);
       ret.push(chunk);
       // if it's the first one render this new rgaph

tinygrad/viz/serve.py

@@ -9,14 +9,14 @@ from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA
 from tinygrad.helpers import printable, system, TCPServerWithReuse, HTTPRequestHandler
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, pyrender
 from tinygrad.uop.ops import print_uops, range_start, multirange_str
-from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device, ProfileProgramEvent
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.DEFINE_GLOBAL:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
-               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
+               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.CUSTOM_KERNEL: "#3ebf55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.ENCDEC: "#bf71b6",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
@@ -234,70 +234,66 @@ def unpack_pmc(e) -> dict:
     rows.append(row)
   return {"rows":rows, "cols":agg_cols}
 
-@soft_err(lambda err: ctxs.append({"name":"ERR", "steps":[create_step("Loader error", ("render",len(ctxs),0), err)]}))
-def load_sqtt(profile:list[ProfileEvent]) -> None:
+# ** on startup, list all the performance counter traces
+
+def load_counters(profile:list[ProfileEvent]) -> None:
   from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
-  counter_events:dict[tuple[str, int], list[ProfileSQTTEvent|ProfilePMCEvent]] = {}
+  counter_events:dict[tuple[str, int], dict] = {}
   durations:dict[str, list[float]] = {}
+  prg_events:dict[str, ProfileProgramEvent] = {}
+  dev_events:dict[str, ProfileDeviceEvent] = {}
   for e in profile:
-    if isinstance(e, (ProfilePMCEvent, ProfileSQTTEvent)): counter_events.setdefault((e.kern, e.exec_tag), []).append(e)
+    if isinstance(e, (ProfilePMCEvent, ProfileSQTTEvent)): counter_events.setdefault((e.kern, e.exec_tag), {}).setdefault(type(e), []).append(e)
     if isinstance(e, ProfileRangeEvent) and e.device.startswith("AMD") and e.en is not None:
       durations.setdefault(str(e.name), []).append(float(e.en-e.st))
-  if not counter_events: return
-  # ** init decoder
+    if isinstance(e, ProfileProgramEvent): prg_events[str(e.name)] = e
+    if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
+  ctxs.append({"name":"All Counters", "steps":[create_step("PMC", ("/all-pmc", len(ctxs), 0), \
+      (durations, {k:v[ProfilePMCEvent][0] for k,v in counter_events.items()}))]})
+  run_number = {n:0 for n,_ in counter_events}
+  for k,v in counter_events.items():
+    prg = trace.keys[r].ret if (r:=ref_map.get(k[0])) else None
+    name = prg.name if prg is not None else k[0]
+    run_number[k[0]] += 1
+    steps:list[dict] = []
+    if (pmc:=v.get(ProfilePMCEvent)): steps.append(create_step("PMC", ("/prg-pmc", len(ctxs), len(steps)), pmc))
+    if (sqtt:=v.get(ProfileSQTTEvent)):
+      # to decode a SQTT trace, we need the raw stream, program binary and device properties
+      steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), (k, [*sqtt, prg_events[k[0]], dev_events[sqtt[0].device]])))
+      if getenv("SQTT_PARSE"):
+        # run our decoder on startup, we don't use this since it only works on gfx11
+        from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+        for e in sqtt: parse_sqtt_print_packets(e.blob)
+    ctxs.append({"name":f"Exec {name} n{run_number[k[0]]}", "steps":steps})
+
+# ** SQTT OCC only unpacks wave start, end time and SIMD location
+
+def unpack_sqtt(key:tuple[str, int], profile:list[ProfileEvent]) -> tuple[dict[str, list[ProfileEvent]], list[str], dict[str, dict[str, dict]]]:
+  # * init decoder
   from extra.sqtt.roc import decode
   rctx = decode(profile)
-  if getenv("SQTT_PARSE"):
-    from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-    for counters in counter_events.values():
-      for e in counters:
-        if isinstance(e, ProfileSQTTEvent): parse_sqtt_print_packets(e.blob)
-  # ** decode traces for each run
-  all_runs:dict = {"cols":{}, "rows":[]}
-  steps:list[dict] = [create_step("All", ("/all", len(ctxs), 0), data=all_runs)]
-  for key,counters in counter_events.items():
-    # ** Run summary
-    program = trace.keys[r].ret if (r:=ref_map.get(key[0])) else None
-    summary = [f"{program.global_size=} {program.local_size=}"] if program else [repr(key)]
-    # ** SQTT events
-    disasm = rctx.disasms[key[0]]
-    cu_events:dict[str, list[ProfileEvent]] = {}
-    # * INST waves
-    wave_insts:dict[str, dict[str, dict]] = {}
-    inst_units:dict[str, itertools.count] = {}
-    for w in rctx.inst_execs.get(key, []):
-      if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
-      n = next(inst_units[u])
-      if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
-      events.append(ProfileRangeEvent(w.simd_loc, loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
-      wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "run_number":n, "loc":loc}
-    # * OCC waves
-    units:dict[str, itertools.count] = {}
-    wave_start:dict[str, int] = {}
-    for occ in rctx.occ_events.get(key, []):
-      if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
-      if u in inst_units: continue
-      if occ.start: wave_start[u] = occ.time
-      else:
-        if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
-        events.append(ProfileRangeEvent(occ.simd_loc, f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)), Decimal(occ.time)))
-    prg_cu = sorted(cu_events, key=row_tuple)
-    if cu_events: summary.append(f"Scheduled on {len(prg_cu)} CUs")
-    steps.append(create_step(program.name if program else key[0], ("/prg-run", len(ctxs), len(steps)), {"src":"\n\n".join(summary)}, depth=0))
-    # ** PMC events
-    if (pmc_event:=next((e for e in counters if isinstance(e, ProfilePMCEvent)), None)) is not None:
-      steps.append(create_step("PMC", ("/pmc", len(ctxs), len(steps)), pmc_table:=unpack_pmc(pmc_event), depth=1))
-      all_runs["cols"].update([(r[0], None) for r in pmc_table["rows"]])
-      all_runs["rows"].append((key[0], durations[key[0]].pop(0), *[r[1] for r in pmc_table["rows"]]))
-    for cu in prg_cu:
-      events = [ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]
-      steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/counters", len(ctxs), len(steps)),
-                               {"value":get_profile(events, sort_fn=row_tuple), "content_type":"application/octet-stream"}, depth=1))
-      for k in sorted(wave_insts.get(cu, []), key=row_tuple):
-        data = wave_insts[cu][k]
-        steps.append(create_step(k.replace(cu, ""), ("/sqtt-insts", len(ctxs), len(steps)), data, loc=data["loc"], depth=2))
-  all_runs["cols"] = ["Kernel", "Duration", *all_runs["cols"]]
-  ctxs.append({"name":"Counters", "steps":steps})
+  disasm = rctx.disasms[key[0]]
+  cu_events:dict[str, list[ProfileEvent]] = {}
+  # * INST waves
+  wave_insts:dict[str, dict[str, dict]] = {}
+  inst_units:dict[str, itertools.count] = {}
+  for w in rctx.inst_execs.get(key, []):
+    if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
+    n = next(inst_units[u])
+    if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
+    events.append(ProfileRangeEvent(w.simd_loc, loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
+    wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "run_number":n, "loc":loc}
+  # * OCC waves
+  units:dict[str, itertools.count] = {}
+  wave_start:dict[str, int] = {}
+  for occ in rctx.occ_events.get(key, []):
+    if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
+    if u in inst_units: continue
+    if occ.start: wave_start[u] = occ.time
+    else:
+      if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
+      events.append(ProfileRangeEvent(occ.simd_loc, f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)), Decimal(occ.time)))
+  return cu_events, list(units), wave_insts
 
 def device_sort_fn(k:str) -> tuple[int, str, int]:
   order = {"GC": 0, "USER": 1, "TINY": 2, "DISK": 999}
@@ -307,13 +303,12 @@ def device_sort_fn(k:str) -> tuple[int, str, int]:
 
 def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_sort_fn) -> bytes|None:
   # start by getting the time diffs
-  for ev in profile:
-    if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
-  # load device specific counters
   device_decoders:dict[str, Callable[[list[ProfileEvent]], None]] = {}
-  for device in device_ts_diffs:
-    d = device.split(":")[0]
-    if d == "AMD": device_decoders[d] = load_sqtt
+  for ev in profile:
+    if isinstance(ev, ProfileDeviceEvent):
+      device_ts_diffs[ev.device] = (ev.comp_tdiff,ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
+      if (d:=ev.device.split(":")[0]) == "AMD": device_decoders[d] = load_counters
+  # load device specific counters
   for fxn in device_decoders.values(): fxn(profile)
   # map events per device
   dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
@@ -381,6 +376,8 @@ def amd_readelf(lib:bytes) -> list[dict]:
           ".group_segment_fixed_size":"LDS size", ".private_segment_fixed_size":"Scratch size"}
   return [{"label":label, "value":v} for k,label in keys.items() if (v:=notes["amdhsa.kernels"][0][k]) > 0]
 
+# ** Main render function to get the complete details about a trace event
+
 def get_render(i:int, j:int, fmt:str) -> dict:
   data = ctxs[i]["steps"][j]["data"]
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data.uops or [])), "lang":"txt"}
@@ -397,6 +394,28 @@ def get_render(i:int, j:int, fmt:str) -> dict:
       with soft_err(lambda err: metadata.append(err)):
         metadata.append(amd_readelf(compiler.compile(data.src)))
     return {"src":disasm_str, "lang":"amdgpu" if data.device.startswith("AMD") else None, "metadata":metadata}
+  if fmt == "all-pmc":
+    durations, pmc = data
+    ret:dict = {"cols":{}, "rows":[]}
+    for (prg,_),events in pmc.items():
+      pmc_table = unpack_pmc(events)
+      ret["cols"].update([(r[0], None) for r in pmc_table["rows"]])
+      ret["rows"].append((prg, durations[prg].pop(0), *[r[1] for r in pmc_table["rows"]]))
+    ret["cols"] = ["Kernel", "Duration", *ret["cols"]]
+    return ret
+  if fmt == "prg-pmc": return unpack_pmc(data[0])
+  if fmt == "prg-sqtt":
+    ret = {}
+    if len((steps:=ctxs[i]["steps"])[j+1:]) == 0:
+      with soft_err(lambda err: ret.update(err)):
+        cu_events, units, wave_insts = unpack_sqtt(*data)
+        for cu in sorted(cu_events, key=row_tuple):
+          steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
+                                   data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
+          for k in sorted(wave_insts.get(cu, []), key=row_tuple):
+            steps.append(create_step(k.replace(cu, ""), ("/sqtt-insts", i, len(steps)), loc=(data:=wave_insts[cu][k])["loc"], depth=2, data=data))
+    return {**ret, "steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
+  if fmt == "cu-sqtt": return {"value":get_profile(data, sort_fn=row_tuple), "content_type":"application/octet-stream"}
   if fmt == "sqtt-insts":
     columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
     inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]