something

* add tiny markdown parser to llm

* disable on generating

* regex slop

* more markdown

* okay, real markdown lib, reusing viz mech

* min diff

* simpler css

* rm that

.github/workflows/benchmark.yml

@@ -82,19 +82,22 @@ jobs:
   #      pytest -nauto --durations=20
 
   llmbenchmark:
-    name: LLM (DEV=${{ matrix.dev }})
+    name: Benchmark ${{ matrix.model }} (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
     strategy:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 30
+        model: ['llama3.2:3b-f16', 'qwen3.8:27b', 'olmoe']
+        # qwen3.8:27b doesn't fit on mac
+        exclude: [{ dev: 'METAL', model: 'qwen3.8:27b' }, { dev: 'AMD', model: 'olmoe' }, { dev: 'NV', model: 'olmoe' }]
+    timeout-minutes: 15
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -114,16 +117,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: python3 test/external/process_replay/reset.py
-    - name: Run llama3.2
-      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.8
-      # qwen3.8:27b doesn't fit on mac
-      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
-    - name: Run olmoe
-      # just metal for now
-      if: ${{ matrix.dev == 'METAL' }}
-      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run ${{ matrix.model }}
+      run: |
+        MODEL=${{ matrix.model }}
+        BENCHMARK_LOG=${MODEL//./} JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m $MODEL --benchmark --warmup
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -140,7 +137,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -188,7 +185,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -230,7 +227,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -275,7 +272,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -322,7 +319,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 11
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -439,13 +436,7 @@ jobs:
     - name: UsbGPU tiny tests
       run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    - name: UsbGPU (USB4/TB) install script
-      run: sh extra/setup_tinygpu_osx.sh
-    - name: UsbGPU (USB4/TB) boot time
-      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py
 
   testcomma:
     strategy:
@@ -565,7 +556,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -632,16 +623,6 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
-        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
-          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
-        fi
-        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 

.github/workflows/platform.yml

@@ -76,11 +76,12 @@ jobs:
     - name: Run pytest (ptx)
       env:
         DEV: "MOCK+NV:PTX"
+        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
+        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py \
           test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay

.github/workflows/test.yml

@@ -544,14 +544,6 @@ jobs:
         run: python -m pytest test/device/test_hcq2.py
       - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
-      - name: Run test_tiny on MOCKPCI Remote
-        env:
-          HCQ2: 0
-        run: |
-          python extra/remote/serve.py 6667 &
-          sleep 2
-          REMOTE=127.0.0.1:6667 python test/test_tiny.py
-          REMOTE=127.0.0.1:6667 python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk; kill %1
 
   testamd:
     strategy:
@@ -619,7 +611,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1\nHCQ_RUNTIME_DEV=PYTHON' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"

AGENTS.md

@@ -5,3 +5,4 @@
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
 - Do not do amend commits. Always do a new commit if a force push to origin would be required.
+- tinygrad has user space PCI drivers for AMD and NVIDIA GPUs. Do not insert the unneeded kernel modules.

conftest.py

@@ -2,7 +2,7 @@ import os, pytest, signal, threading
 
 @pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 120)), os.kill, args=(os.getpid(), signal.SIGABRT))
   t.start()
   try: yield
   finally:

examples/mlperf/models/gpt_oss.py

@@ -12,9 +12,9 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
 from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine
+from extra.gemm.moe_routing import route, dispatch, combine, router_mfma
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -41,7 +41,7 @@ def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 def quant_dequant_mx(x:Tensor) -> Tensor:
   fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
-  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
+  return Tensor(fxn.uop.call_with_output(x.uop, grad_fxn=_quant_dequant_bwd))
 
 def _mx_scale(e8:Tensor) -> Tensor:
   return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
@@ -58,8 +58,7 @@ def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
 
 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
   fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
-  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
-  return Tensor(call.gettuple(0))
+  return Tensor(fxn.uop.call_with_output(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd))
 
 def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
   if isinstance(x, tuple):
@@ -247,7 +246,7 @@ class GPTOSS:
       x_normed, rrms = rmsnorm(x, self.norm_eps)
       inp = x_normed * ffn_norm
 
-    logits = inp.float() @ gate.float().T + gate_bias.float()
+    logits = router_mfma(inp, gate, gate_bias) if getenv("ROUTER_MFMA", 0) else inp.float() @ gate.float().T + gate_bias.float()
     dim, inter = self.dim, self.intermediate_size
 
     if getenv("GROUPED_MOE", 0):
@@ -306,10 +305,25 @@ class GPTOSS:
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     h_normed = self.norm(h)
-    pad = (-self.dim) % 256
-    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
-    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
-    else: logits = h_normed @ self.output.T
+
+    if getenv("FP8_LMHEAD", 0) and ASM_GEMM:
+      pad = (-self.dim) % 256
+      h2 = h_normed.reshape(-1, self.dim).pad(((0, 0), (0, pad)))
+      w2 = self.output.pad(((0, 0), (0, pad)))
+      hq, he8, hsi = quantize_mxfp8(h2)
+      oq, oe8, _ = quantize_mxfp8(w2)
+      if hsi is not None and can_use_asm_gemm(hq, oq.T):
+        logits = asm_gemm(hq, oq.T, mx=True, mx_scales=(hsi, he8, mx_pack(oe8), oe8), mx_w_stored=False)
+        logits = logits.reshape(bsz, seqlen, self.vocab_size).cast(dtypes.bfloat16)
+      else:
+        logits = h_normed @ self.output.T
+    elif ASM_GEMM:
+      pad = (-self.dim) % 256
+      h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+      logits = asm_gemm(h_padded, w_padded.T) if can_use_asm_gemm(h_padded, w_padded.T) and getenv("VOCAB_ASM", 1) else h_normed @ self.output.T
+    else:
+      logits = h_normed @ self.output.T
+
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

extra/amdflash/README.md

@@ -0,0 +1,116 @@
+# Navi31 flash tools
+
+Utilities for reading and recovering the 2 MiB SPI flash on Navi31 boards.
+Run them from the tinygrad repository root. No image is bundled; keep a verified
+full-ROM backup before performing any write.
+
+`fw_live.py` accesses BAR5 through tinygrad's `PCIDevice.map_bar()` abstraction
+and supports either the custom ASM24 USB-PCIe bridge or native PCIe. Select the
+transport before the subcommand:
+
+```sh
+python3 extra/amdflash/fw_live.py --transport usb probe
+python3 extra/amdflash/fw_live.py --transport pci probe
+```
+
+The default, `--transport auto`, considers USB devices first and then native
+PCI devices. Native PCI access requires the usual tinygrad PCI permissions and
+an unbound kernel driver.
+
+## Access paths and hardware state
+
+The paths are state-dependent and are not interchangeable:
+
+* **`romless.py`** drives SMUIO `ROM_SW_*` directly through the ASM24 bridge.
+  Use it only when an empty or corrupt flash has stalled the PSP PBL. Healthy
+  autonomous boot gates this engine; the usual gated status is
+  `ROM_SW_STATUS=0x04000800`.
+* **`fw_live.py probe`** queries the early PSP boot-firmware mailbox.
+* Firmware-mediated write commands are retained for protocol documentation but
+  are disabled because an exact stock reflash did not validate safely.
+* **`fw_live.py dump`** reads an exact 2 MiB raw image through
+  `ROM_INDEX/ROM_DATA`. It refuses devices where the raw SMUIO controller is
+  unavailable; the NBIO SOC15 function-ROM aperture is not a physical SPI
+  mapping and is deliberately not used as a fallback.
+
+The tools do not reset or power-cycle the board.
+
+## Raw ROM_SW recovery
+
+Identification and read-only operations:
+
+```sh
+python3 extra/amdflash/romless.py info
+python3 extra/amdflash/romless.py read 0 0x40
+python3 extra/amdflash/romless.py dump spi.bin
+python3 extra/amdflash/romless.py verify known-good.bin
+```
+
+Restore an exact 2 MiB image:
+
+```sh
+python3 extra/amdflash/romless.py flash known-good.bin --yes
+```
+
+If GD25 status-register bit `SR2.CMP` protects the complete array, clearing it
+requires separate authorization:
+
+```sh
+python3 extra/amdflash/romless.py flash known-good.bin --clear-cmp --yes
+```
+
+Programming is sector-granular. Every written 4 KiB sector is immediately read
+back and compared with the input. A range can be resumed independently:
+
+```sh
+python3 extra/amdflash/romless.py flash known-good.bin \
+  --start-sector 128 --sector-count 64 --yes
+```
+
+Navi31 ROM_SW details used by the implementation:
+
+* `ROM_SW_COMMAND = (address << 8) | opcode`
+* TX data uses big-endian stream dwords
+* `RETURN_DATA_EN` (bit 19) is clear for TX and set for RX
+* the RX window exposes the preceding transaction, so reads are primed once
+
+## Firmware-mediated access
+
+The read-only commands are:
+
+```sh
+python3 extra/amdflash/fw_live.py probe
+python3 extra/amdflash/fw_live.py dump current-spi.bin
+```
+
+`dump` produces exactly `0x200000` bytes, requires the raw IFWI magic at offset
+zero, rejects mirrored 1 MiB apertures, and restores the ROM controller/index
+state before writing output.
+
+The validated early-firmware sequence is available as:
+
+```sh
+python3 extra/amdflash/fw_live.py --transport usb ifwi-all full-ifwi.bin --yes
+```
+
+It resolves at most Navi31's configured 19 items, streams the item associated
+with terminal phase `0x2xx`, and then stops. PSP selects the destination
+partition; item `0x08` always comes from the payload referenced by the first
+ISH descriptor, matching AMDVBFlash. A hard power cycle is required afterward.
+
+A successful PSP update is not a byte-identical raw rewrite. On the validated
+stock test, both A/B payloads matched the source exactly, PSP selected and
+booted the updated B partition, and firmware changed only its update cookie,
+B descriptor counter/checksum, and generated metadata near `0x1ef000`.
+
+The `stream`, `ifwi-step`, and `live-flash` commands remain disabled. Testing
+showed that the PSP live path parses a raw stock IFWI but fails with status
+`0xC` (`PSP Write To SPI Error`) after writing an `$AMDVBFL` cookie. Use the
+verified ROM_SW path for recovery.
+
+## Safety
+
+ROM_SW erase/program and `ifwi-all` commands require `--yes`; other
+firmware-streaming commands are disabled. Read-only commands still touch controller and mailbox registers but
+do not issue SPI program/erase or PSP transfer-start commands. Preserve a
+known-good full dump outside the repository.

extra/amdflash/common.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+import struct, sys, time
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT))
+from tinygrad.runtime.support.usb import USB3
+from tinygrad.runtime.support.system import PCIDevice, System, USBPCIDevice
+
+USB_IDS = ((0x3801, 0x0001), (0xADD1, 0x0001))
+NAVI31_DEVICES = ((0xffff, (0x744c,)),)
+
+
+def open_gpu(index: int = 0, transport: str = 'auto') -> PCIDevice:
+  """Open an AMD GPU through tinygrad's transport-independent PCI interface."""
+  if transport not in ('auto', 'usb', 'pci'): raise ValueError(f"unsupported transport {transport!r}")
+  candidates = []
+  if transport in ('auto', 'usb'):
+    for vendor, product in USB_IDS:
+      candidates += [(USBPCIDevice, dev) for dev in USB3.list_devices(vendor, product)]
+  if transport in ('auto', 'pci'):
+    candidates += System.list_devices(0x1002, NAVI31_DEVICES)
+  if not candidates: raise RuntimeError(f"no supported {transport} AMD GPU found")
+  if not 0 <= index < len(candidates): raise RuntimeError(f"device index {index} out of range (found {len(candidates)})")
+  cls, descriptor = candidates[index]
+  return cls("AM", *descriptor) if cls is USBPCIDevice else cls("AM", descriptor)
+
+
+class MMIO:
+  """Transport-independent byte view of BAR5."""
+  def __init__(self, pci_dev: PCIDevice): self.bar = pci_dev.map_bar(5, fmt='B')
+
+  def read32(self, offset: int) -> int:
+    return struct.unpack('<I', bytes(self.bar[offset:offset+4]))[0]
+
+  def write32(self, offset: int, value: int):
+    self.write(offset, struct.pack('<I', value & 0xffffffff))
+
+  def read(self, offset: int, size: int) -> bytes:
+    return bytes(self.bar[offset:offset+size])
+
+  def write(self, offset: int, data: bytes):
+    self.bar[offset:offset+len(data)] = data
+
+
+def wait_until(fn, timeout: float, message: str, interval: float = 0.001):
+  if timeout <= 0 or timeout > 60: raise ValueError("timeout must be in (0, 60] seconds")
+  end = time.monotonic() + timeout
+  while True:
+    value = fn()
+    if value: return value
+    if time.monotonic() >= end: raise TimeoutError(message)
+    time.sleep(interval)

extra/amdflash/fw_live.py

@@ -0,0 +1,292 @@
+#!/usr/bin/env python3
+"""Navi31 firmware-mediated flash access and ROM aperture dumping.
+
+Early item streaming must run after autonomous PSP boot but before a host
+driver or AMDev loads SOS. A fully initialized SOS rejects those commands.
+"""
+from __future__ import annotations
+import argparse, struct, sys, time
+from pathlib import Path
+from common import MMIO, open_gpu, wait_until
+
+ROM_CNTL, ROM_INDEX, ROM_DATA = 0x5A380, 0x5A390, 0x5A394
+FLASH_SIZE, INDEX_PAGE = 0x200000, 0x10000
+
+def bswap32(value: int) -> int: return int.from_bytes(value.to_bytes(4, 'little'), 'big')
+
+COMMAND_DATA, COMMAND, DOORBELL = 0x582D0, 0x582CC, 0x58224
+GET_BOOT_PARTITION, GET_FB_STATE, GET_TRANSFER_TYPE = 0x01, 0x06, 0x07
+START_TRANSFER, DATA_TRANSFER, END_TRANSFER = 0x08, 0x09, 0x0A
+SPI_GET_MODEL_ID = 0x0B
+LIVE_ADDR_LO, LIVE_ADDR_HI, LIVE_UPDATE = 0x02, 0x03, 0x04
+PSP_ERRORS = {
+  0x01: "generic error", 0x02: "out of bounds", 0x03: "invalid parameter",
+  0x04: "off-chip boot error", 0x05: "address not set", 0x06: "parse off-chip error",
+  0x07: "address map error", 0x08: "parse on-chip error", 0x09: "full update error",
+  0x0A: "partition update error", 0x0B: "map on-chip error", 0x0C: "write to SPI error",
+  0x0D: "signature validation error", 0x0E: "invalid command", 0x0F: "signature not found",
+  0x10: "state machine not initialized", 0x11: "state machine transfer error",
+  0x12: "initialization error",
+}
+
+
+class PSPFlashMailbox:
+  def __init__(self, pci_dev): self.mmio = MMIO(pci_dev)
+
+  def command(self, command: int, data: int | None = None, *, timeout: float = 10.0) -> tuple[int, int]:
+    status = self.mmio.read32(COMMAND)
+    if not status & 0x80000000:
+      raise RuntimeError(f"PSP mailbox is not ready before command {command:#x}: status={status:#010x}")
+    if data is not None: self.mmio.write32(COMMAND_DATA, data)
+    self.mmio.write32(COMMAND, command << 16)
+    self.mmio.write32(DOORBELL, 1)
+    wait_until(lambda: self.mmio.read32(COMMAND) & 0x80000000, timeout,
+               f"PSP mailbox command {command:#x} timed out")
+    value = self.mmio.read32(COMMAND)
+    return value & 0xffff, self.mmio.read32(COMMAND_DATA)
+
+  def require(self, command: int, data: int | None = None, *, timeout: float = 10.0, name: str = '') -> int:
+    error, response = self.command(command, data, timeout=timeout)
+    if error:
+      detail = PSP_ERRORS.get(error, "unknown error")
+      raise RuntimeError(f"PSP {name or hex(command)} failed: error={error:#x} ({detail})")
+    return response
+
+  def probe(self) -> dict[str, tuple[int, int]]:
+    result = {}
+    for name, command in (("boot_partition", GET_BOOT_PARTITION), ("fb_state", GET_FB_STATE),
+                          ("model_id", SPI_GET_MODEL_ID), ("transfer_type", GET_TRANSFER_TYPE)):
+      result[name] = self.command(command)
+    return result
+
+  def stream(self, payload: bytes, item_type: int, transfer_type: int | None = None):
+    if not payload: raise ValueError("payload is empty")
+    if len(payload) > 0xFFFFFF: raise ValueError("payload exceeds the mailbox's 24-bit size field")
+    if len(payload) & 3: raise ValueError("payload size must be divisible by four")
+    if not 0 <= item_type <= 0xff: raise ValueError("item type must fit in eight bits")
+    if transfer_type is None: transfer_type = self.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+    requested = transfer_type & 0xff
+    print(f"firmware transfer_type={transfer_type:#x}", flush=True)
+    if requested != item_type:
+      raise RuntimeError(f"firmware requests item {requested:#x}, not {item_type:#x}")
+    self.require(START_TRANSFER, (len(payload) << 8) | item_type, name="START_TRANSFER")
+    sent, started = 0, time.monotonic()
+    try:
+      for offset in range(0, len(payload), 4):
+        word = struct.unpack_from('<I', payload, offset)[0]
+        self.require(DATA_TRANSFER, word, name=f"DATA_TRANSFER@{offset:#x}")
+        sent = offset + 4
+        if sent % 0x1000 == 0:
+          print(f"{sent:#x}/{len(payload):#x} ({sent/(time.monotonic()-started)/1024:.1f} KiB/s)", flush=True)
+      self.require(END_TRANSFER, (sent << 8) | item_type, timeout=60.0, name="END_TRANSFER")
+    except BaseException:
+      # Give firmware a chance to terminate an interrupted partial session. Do
+      # not submit END_TRANSFER twice if firmware rejected the original END.
+      if sent != len(payload):
+        try: self.command(END_TRANSFER, (sent << 8) | item_type, timeout=10.0)
+        except Exception: pass
+      raise
+    print(f"stream complete: type={item_type:#x} size={sent:#x} elapsed={time.monotonic()-started:.1f}s")
+
+
+def resolve_ifwi_item(image: bytes, item_type: int) -> tuple[int, bytes]:
+  """Resolve AMDVBFlash recovery-layout item types to exact IFWI bytes."""
+  if item_type == 0x01: offset, size = 0, 0x54
+  elif item_type in (0x02, 0x03):
+    offset = 0x2000 if item_type == 0x02 else 0x3000
+    if image[offset:offset+4] != b'$PSP': raise ValueError(f"invalid PSP directory at {offset:#x}")
+    size = (struct.unpack_from('<I', image, offset + 8)[0] + 1) * 0x10
+  elif item_type == 0x04: offset, size = 0x10000, 0x1000
+  elif item_type == 0x05: offset, size = 0x11000, 0x1000
+  elif item_type == 0x06: offset, size = 0x12000, 0x20
+  elif item_type == 0x07: offset, size = 0x13000, 0x20
+  elif item_type == 0x80: offset, size = 0x1000, 4
+  elif item_type == 0x81:
+    offset = struct.unpack_from('<I', image, 0x1000)[0]
+    if image[offset:offset+4] != b'$SGN': raise ValueError("invalid $SGN table pointer")
+    size = (struct.unpack_from('<I', image, offset + 8)[0] + 1) * 0x10
+  elif 0x82 <= item_type <= 0x88:
+    table = struct.unpack_from('<I', image, 0x1000)[0]
+    if image[table:table+4] != b'$SGN': raise ValueError("invalid $SGN table pointer")
+    wanted = item_type - 0x81  # 82h..88h map to SIGN_TYPE 1..7
+    count = struct.unpack_from('<I', image, table + 8)[0]
+    entries = [struct.unpack_from('<IIII', image, table + 0x10 + i*0x10) for i in range(count)]
+    match = [entry for entry in entries if entry[0] == wanted]
+    if len(match) != 1: raise ValueError(f"missing $SGN type {wanted}")
+    _, _, size, offset = match[0]
+  elif item_type == 0x89: offset, size = 0x1f0000, 0x100
+  elif item_type == 0x08:
+    # AMDVBFlash's GetPartitionDetails follows the first ISH entry (firmware ID
+    # 0x13c) and streams its payload. PSP, not the host resolver, selects the
+    # destination partition.
+    offset = struct.unpack_from('<I', image, 0x12000 + 0x10)[0]
+    size = struct.unpack_from('<I', image, 0x12000 + 0x18)[0]
+  else:
+    raise ValueError(f"IFWI resolver does not yet support requested item {item_type:#x}")
+  payload = image[offset:offset+size]
+  if len(payload) != size: raise ValueError(f"item {item_type:#x} extends beyond IFWI")
+  print(f"resolved requested item {item_type:#x}: offset={offset:#x} size={size:#x}")
+  return offset, payload
+
+
+class LivePSPFlash:
+  """Linux psp_v13_0_update_spirom protocol, used with SOS and trained VRAM."""
+  def __init__(self, pci_dev): self.mailbox = PSPFlashMailbox(pci_dev)
+
+  def command(self, command: int, data: int | None = None, timeout: float = 10.0):
+    # Same C2PMSG registers, but the live PSP command set uses IDs 2/3/4.
+    return self.mailbox.require(command, data, timeout=timeout, name=f"LIVE_SPI_{command:#x}")
+
+  def update(self, mc_address: int):
+    status = self.mailbox.mmio.read32(COMMAND)
+    if not status & 0x80000000: raise RuntimeError(f"live PSP mailbox is not ready: {status:#x}")
+    self.command(LIVE_ADDR_LO, mc_address & 0xffffffff)
+    self.command(LIVE_ADDR_HI, mc_address >> 32)
+    self.command(LIVE_UPDATE, timeout=60.0)
+
+
+def open_mailbox(args): return PSPFlashMailbox(open_gpu(args.device, args.transport))
+
+
+def reject_unvalidated_firmware_write():
+  raise RuntimeError("firmware writes are disabled: stock reflash validation failed; use romless.py for recovery")
+
+
+def cmd_probe(args):
+  result = open_mailbox(args).probe()
+  for name, (error, response) in result.items(): print(f"{name}: error={error:#x} response={response:#x}")
+  if result['transfer_type'][0] == 0xA: print("update commands gated: reset card and do not initialize AMDev/SOS", file=sys.stderr)
+
+
+def cmd_stream(args):
+  if not args.yes: raise RuntimeError("refusing to stream without --yes")
+  reject_unvalidated_firmware_write()
+  payload = Path(args.image).read_bytes()
+  open_mailbox(args).stream(payload, args.item_type)
+
+
+def cmd_ifwi_step(args):
+  if not args.yes: raise RuntimeError("refusing to stream without --yes")
+  reject_unvalidated_firmware_write()
+  image = Path(args.ifwi).read_bytes()
+  if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
+  mailbox = open_mailbox(args)
+  state = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+  request = state & 0xff
+  _, payload = resolve_ifwi_item(image, request)
+  mailbox.stream(payload, request, transfer_type=state)
+  next_request = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+  print(f"next firmware transfer_type={next_request:#x}")
+
+
+def cmd_ifwi_all(args):
+  if not args.yes: raise RuntimeError("refusing to stream without --yes")
+  image = Path(args.ifwi).read_bytes()
+  if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
+  mailbox = open_mailbox(args)
+  current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+  for step in range(19):  # Navi31 ROMItemCount from AMDVBFlash ASICDetails.xml
+    request, phase = current & 0xff, current >> 8
+    print(f"IFWI step {step}: state={current:#x} item={request:#x} phase={phase}", flush=True)
+    _, payload = resolve_ifwi_item(image, request)
+    mailbox.stream(payload, request, transfer_type=current)
+    # AMDVBFlash tests the high byte belonging to the item just streamed. Phase
+    # 2 terminates the loop only after that item has completed successfully.
+    if phase == 2:
+      print(f"IFWI stream complete after terminal state {current:#x}; hard power cycle required")
+      return
+    current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+  raise RuntimeError(f"IFWI stream did not reach terminal phase after 19 items (state={current:#x})")
+
+
+def cmd_live_flash(args):
+  if not args.yes: raise RuntimeError("refusing to flash without --yes")
+  reject_unvalidated_firmware_write()
+  image = Path(args.ifwi).read_bytes()
+  if not image or len(image) > 16 * 1024 * 1024 or len(image) & 3:
+    raise ValueError("live PSP image must be non-empty, 4-byte aligned, and at most 16 MiB")
+  pci_dev = open_gpu(args.device, args.transport)
+  from tinygrad.runtime.support.am.amdev import AMDev
+  started = time.monotonic()
+  adev = AMDev(pci_dev)
+  print(f"AMDev booted, SOS alive={adev.psp.is_sos_alive()}", flush=True)
+  paddr = adev.mm.palloc(len(image), align=0x1000, zero=False)
+  try:
+    adev.vram.view(paddr, len(image), 'B')[:] = image
+    adev.gmc.flush_hdp()
+    mc_address = adev.paddr2mc(paddr)
+    print(f"staged IFWI at VRAM paddr={paddr:#x} mc={mc_address:#x}", flush=True)
+    LivePSPFlash(pci_dev).update(mc_address)
+    print(f"live PSP flash update complete in {time.monotonic()-started:.1f}s")
+  finally:
+    adev.mm.pfree(paddr)
+
+
+def cmd_dump(args):
+  import hashlib
+  pci_dev = open_gpu(args.device, args.transport)
+  mmio, output, started = MMIO(pci_dev), bytearray(), time.monotonic()
+  original_cntl, original_index = mmio.read32(ROM_CNTL), mmio.read32(ROM_INDEX)
+  if original_cntl == 0xFFFFFFFF:
+    raise RuntimeError("raw SMUIO ROM controller is unavailable; the SOC15 function-ROM aperture is not a raw SPI dump")
+  try:
+    # ROM_DATA must be read one dword at a time; a block read increments MMIO
+    # addresses rather than repeatedly reading the flash aperture register.
+    mmio.write32(ROM_CNTL, bswap32(original_cntl | (1 << 29)))
+    for page in range(0, FLASH_SIZE, INDEX_PAGE):
+      mmio.write32(ROM_INDEX, bswap32(page >> 8))
+      for _ in range(INDEX_PAGE // 4): output += struct.pack('<I', mmio.read32(ROM_DATA))
+      print(f"{page+INDEX_PAGE:#08x}/{FLASH_SIZE:#08x}", flush=True)
+  finally:
+    mmio.write32(ROM_INDEX, bswap32(original_index))
+    mmio.write32(ROM_CNTL, bswap32(original_cntl))
+  if len(output) != FLASH_SIZE or output[:4] != b'\xaa\x55\xaa\x55':
+    raise RuntimeError(f"invalid raw flash dump: size={len(output):#x} magic={output[:4].hex()}")
+  if output[:FLASH_SIZE//2] == output[FLASH_SIZE//2:]:
+    raise RuntimeError("ROM aperture contains mirrored 1 MiB halves; refusing to write a non-raw 2 MiB dump")
+  Path(args.output).write_bytes(output)
+  print(f"dumped {len(output):#x} bytes in {time.monotonic()-started:.1f}s sha256={hashlib.sha256(output).hexdigest()}")
+
+
+def parser():
+  p = argparse.ArgumentParser(description=__doc__)
+  p.add_argument('--device', type=int, default=0, help='device index for the selected transport')
+  p.add_argument('--transport', choices=('auto', 'usb', 'pci'), default='auto', help='PCIe transport (default: USB first, then native PCI)')
+  sub = p.add_subparsers(dest='command', required=True)
+  sub.add_parser('probe', help='query firmware mailbox state without writing').set_defaults(func=cmd_probe)
+
+  s = sub.add_parser('stream', help='stream one exact PSP ROM-item payload')
+  s.add_argument('item_type', type=lambda x:int(x, 0))
+  s.add_argument('image')
+  s.add_argument('--yes', action='store_true')
+  s.set_defaults(func=cmd_stream)
+
+  v = sub.add_parser('ifwi-step', help='resolve and stream the next early-firmware-requested item from a 2 MiB IFWI')
+  v.add_argument('ifwi')
+  v.add_argument('--yes', action='store_true')
+  v.set_defaults(func=cmd_ifwi_step)
+
+  a = sub.add_parser('ifwi-all', help='stream requested IFWI items until firmware reports completion')
+  a.add_argument('ifwi')
+  a.add_argument('--yes', action='store_true')
+  a.set_defaults(func=cmd_ifwi_all)
+
+  l = sub.add_parser('live-flash', help='stage an image in VRAM and invoke the PSP v13 live-update command')
+  l.add_argument('ifwi')
+  l.add_argument('--yes', action='store_true')
+  l.set_defaults(func=cmd_live_flash)
+
+  d = sub.add_parser('dump', help='dump the exact 2 MiB flash through ROM_INDEX/ROM_DATA')
+  d.add_argument('output')
+  d.set_defaults(func=cmd_dump)
+  return p
+
+
+def main():
+  args = parser().parse_args()
+  try: args.func(args)
+  except (RuntimeError, TimeoutError, ValueError, OSError) as error:
+    print(f"error: {error}", file=sys.stderr)
+    raise SystemExit(1)
+
+if __name__ == '__main__': main()

extra/amdflash/romless.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+"""Direct Navi31 ROM_SW access for GD25LQ16E-class 2 MiB SPI flash."""
+from __future__ import annotations
+import argparse, hashlib, sys, time
+from pathlib import Path
+from common import MMIO, open_gpu, wait_until
+
+FLASH_SIZE, SECTOR_SIZE, PAGE_SIZE, MAX_DATA = 0x200000, 0x1000, 0x100, 0x100
+ROM_CNTL, PAGE_MIRROR_CNTL = 0x5A380, 0x5A384
+ROM_SW_CNTL, ROM_SW_STATUS, ROM_SW_COMMAND, ROM_SW_DATA = 0x5A3A0, 0x5A3A4, 0x5A3A8, 0x5A3B0
+GPIO_PAD_MASK, GPIO_PAD_A, GPIO_PAD_EN = 0x5A504, 0x5A508, 0x5A510
+SPI_GPIO_BITS, RETURN_DATA_EN = 0x780, 0x80000
+EXPECTED_JEDEC = b'\xc8\x60\x15'
+
+
+class Navi31SPI:
+  def __init__(self, pci_dev, prescale: int = 8):
+    if not 0 <= prescale <= 15: raise ValueError("prescale must be 0..15")
+    self.mmio = MMIO(pci_dev)
+    rc = self.mmio.read32(ROM_CNTL)
+    # Select the prescaler instead of inheriting a potentially unusable BL value.
+    self.mmio.write32(ROM_CNTL, (rc & 0xE0FFFFFF) | (1 << 28) | (prescale << 24) | 1)
+
+  def transfer(self, opcode: int, *, address: int = 0, address_len: int = 0,
+               data_out: bytes = b'', data_in: int = 0, timeout: float = 2.0) -> bytes:
+    if data_out and data_in: raise ValueError("simultaneous TX and RX is unsupported")
+    if not 0 <= address_len <= 3: raise ValueError("address_len must be 0..3")
+    count = len(data_out) if data_out else data_in
+    if not 0 <= count <= MAX_DATA: raise ValueError(f"transfer data must be <= {MAX_DATA} bytes")
+    ncmd = 1 + address_len
+    m = self.mmio
+    gpio_mask, gpio_a, gpio_en = m.read32(GPIO_PAD_MASK), m.read32(GPIO_PAD_A), m.read32(GPIO_PAD_EN)
+    page_mirror, rom_cntl = m.read32(PAGE_MIRROR_CNTL), m.read32(ROM_CNTL)
+    try:
+      m.write32(GPIO_PAD_MASK, gpio_mask & ~SPI_GPIO_BITS)
+      m.write32(GPIO_PAD_A, gpio_a & ~SPI_GPIO_BITS)
+      m.write32(GPIO_PAD_EN, gpio_en & ~SPI_GPIO_BITS)
+      m.write32(PAGE_MIRROR_CNTL, (page_mirror & 0xF1FFFFFF) | 0x06000000)
+      m.write32(ROM_CNTL, (rom_cntl & ~0xF) | 8)
+      m.write32(ROM_SW_CNTL, 0)
+      m.write32(ROM_SW_STATUS, 0)
+      if m.read32(ROM_SW_STATUS) != 0: raise RuntimeError("ROM_SW_STATUS did not clear")
+
+      # Navi31 serializes the low instruction byte first, followed by ADDRESS[23:0].
+      m.write32(ROM_SW_COMMAND, ((address & 0xFFFFFF) << 8) | (opcode & 0xFF))
+      for offset in range(0, len(data_out), 4):
+        word = data_out[offset:offset+4].ljust(4, b'\0')
+        m.write32(ROM_SW_DATA + offset, int.from_bytes(word, 'big'))
+
+      control = ((ncmd - 1) << 16) | (RETURN_DATA_EN if data_in else 0) | count
+      m.write32(ROM_SW_CNTL, control)
+      m.read32(ROM_SW_CNTL)  # posted-write flush
+      wait_until(lambda: m.read32(ROM_SW_STATUS) & 1, timeout,
+                 f"ROM_SW transaction timeout (status={m.read32(ROM_SW_STATUS):#x}); engine may be gated after SOS boot")
+      return m.read(ROM_SW_DATA, (data_in + 3) & ~3)[:data_in] if data_in else b''
+    finally:
+      m.write32(ROM_SW_CNTL, 0)
+      m.write32(ROM_SW_STATUS, 0)
+      m.write32(ROM_CNTL, rom_cntl)
+      m.write32(PAGE_MIRROR_CNTL, page_mirror)
+      m.write32(GPIO_PAD_A, gpio_a)
+      m.write32(GPIO_PAD_EN, gpio_en)
+      m.write32(GPIO_PAD_MASK, gpio_mask)
+
+
+class GD25LQ16E:
+  def __init__(self, spi: Navi31SPI): self.spi = spi
+
+  def read_register(self, opcode: int, count: int = 1) -> bytes:
+    # Navi31 exposes the preceding transaction's RX capture. Prime identically.
+    self.spi.transfer(opcode, data_in=max(2, count))
+    return self.spi.transfer(opcode, data_in=count)
+
+  def status(self, opcode: int = 0x05) -> int: return self.read_register(opcode)[0]
+  def rdid(self) -> bytes: return self.read_register(0x9F, 4)
+
+  def sfdp(self, count: int = 20) -> bytes:
+    # 5Ah has one dummy byte after its 24-bit address; retain it for diagnostics.
+    self.spi.transfer(0x5A, address_len=3, data_in=count)
+    return self.spi.transfer(0x5A, address_len=3, data_in=count)
+
+  def wait_idle(self, timeout: float = 2.0) -> int:
+    end = time.monotonic() + timeout
+    while time.monotonic() < end:
+      sr1 = self.status()
+      if not sr1 & 1: return sr1
+      time.sleep(0.002)
+    raise TimeoutError(f"flash remained busy for {timeout}s")
+
+  def write_enable(self):
+    self.spi.transfer(0x06)
+    sr1 = self.status()
+    if not sr1 & 2: raise RuntimeError(f"WREN failed (SR1={sr1:#04x})")
+
+  def clear_cmp(self):
+    sr1, sr2 = self.status(), self.status(0x35)
+    if not sr2 & 0x40: return False
+    self.write_enable()
+    # BUSY/WEL are not writable; preserve all protection/QE fields except CMP.
+    self.spi.transfer(0x01, data_out=bytes((sr1 & 0xFC, sr2 & ~0x40)))
+    self.wait_idle(1.0)
+    new_sr2 = self.status(0x35)
+    if new_sr2 & 0x40: raise RuntimeError(f"failed to clear CMP (SR2={new_sr2:#04x})")
+    return True
+
+  def erase_sector(self, address: int):
+    if address & (SECTOR_SIZE - 1): raise ValueError("sector address is not 4 KiB aligned")
+    self.write_enable()
+    self.spi.transfer(0x20, address=address, address_len=3)
+    self.wait_idle(2.0)
+
+  def program_page(self, address: int, data: bytes):
+    if not data or len(data) > PAGE_SIZE or (address & 0xFF) + len(data) > PAGE_SIZE:
+      raise ValueError("page program crosses a 256-byte boundary")
+    self.write_enable()
+    self.spi.transfer(0x02, address=address, address_len=3, data_out=data)
+    self.wait_idle(1.0)
+
+  def read(self, address: int, count: int) -> bytes:
+    if address < 0 or count < 0 or address + count > FLASH_SIZE: raise ValueError("read outside 2 MiB flash")
+    output = bytearray()
+    while count:
+      size = min(count, MAX_DATA)
+      self.spi.transfer(0x03, address=address, address_len=3, data_in=size)
+      output += self.spi.transfer(0x03, address=address, address_len=3, data_in=size)
+      address, count = address + size, count - size
+    return bytes(output)
+
+
+def has_jedec(raw: bytes) -> bool:
+  return EXPECTED_JEDEC in raw + raw[:2]
+
+
+def open_flash(args) -> GD25LQ16E:
+  flash = GD25LQ16E(Navi31SPI(open_gpu(args.device, 'usb'), args.prescale))
+  raw = flash.rdid()
+  if not has_jedec(raw): raise RuntimeError(f"unexpected GD25LQ16E JEDEC capture: {raw.hex()}")
+  return flash
+
+
+def cmd_info(args):
+  f = open_flash(args)
+  sr1, sr2, sr3 = f.status(), f.status(0x35), f.status(0x15)
+  sfdp = f.sfdp(24)
+  pos = sfdp.find(b'SFDP')
+  print(f"JEDEC capture: {f.rdid().hex()} (C8 60 15 detected)")
+  print(f"SR1/SR2/SR3: {sr1:02x}/{sr2:02x}/{sr3:02x}  CMP={'set' if sr2 & 0x40 else 'clear'}")
+  print(f"SFDP capture: {sfdp.hex()}  signature_offset={pos}")
+
+
+def cmd_read(args):
+  data = open_flash(args).read(args.address, args.size)
+  if args.output: Path(args.output).write_bytes(data)
+  else: print(data.hex())
+
+
+def cmd_dump(args):
+  f = open_flash(args)
+  out = Path(args.output)
+  digest = hashlib.sha256()
+  with out.open('wb') as file:
+    for address in range(0, FLASH_SIZE, SECTOR_SIZE):
+      data = f.read(address, SECTOR_SIZE)
+      file.write(data)
+      digest.update(data)
+      if not (address & 0xFFFF): print(f"{address + SECTOR_SIZE:#08x}/{FLASH_SIZE:#08x}", flush=True)
+  print(f"wrote {out} sha256={digest.hexdigest()}")
+
+
+def cmd_verify(args):
+  expected = Path(args.image).read_bytes()
+  if len(expected) != FLASH_SIZE: raise ValueError(f"image must be exactly {FLASH_SIZE:#x} bytes")
+  f = open_flash(args)
+  digest = hashlib.sha256()
+  for address in range(0, FLASH_SIZE, SECTOR_SIZE):
+    got, wanted = f.read(address, SECTOR_SIZE), expected[address:address+SECTOR_SIZE]
+    digest.update(got)
+    if got != wanted:
+      index = next(i for i, (a, b) in enumerate(zip(got, wanted)) if a != b)
+      raise RuntimeError(f"verify mismatch at {address+index:#x}: flash={got[index]:02x} image={wanted[index]:02x}")
+  print(f"verified {FLASH_SIZE:#x} bytes sha256={digest.hexdigest()}")
+
+
+def cmd_flash(args):
+  if not args.yes: raise RuntimeError("refusing to write without --yes")
+  image = Path(args.image).read_bytes()
+  if len(image) != FLASH_SIZE: raise ValueError(f"image must be exactly {FLASH_SIZE:#x} bytes")
+  total_sectors = FLASH_SIZE // SECTOR_SIZE
+  start, count = args.start_sector, args.sector_count if args.sector_count is not None else total_sectors - args.start_sector
+  if not 0 <= start < total_sectors or not 1 <= count <= total_sectors - start: raise ValueError("invalid sector range")
+  f = open_flash(args)
+  if f.status(0x35) & 0x40:
+    if not args.clear_cmp: raise RuntimeError("CMP protects the full array; rerun with --clear-cmp")
+    f.clear_cmp()
+    print("cleared SR2.CMP", flush=True)
+  begin = time.monotonic()
+  for sector in range(start, start + count):
+    address = sector * SECTOR_SIZE
+    wanted = image[address:address+SECTOR_SIZE]
+    f.erase_sector(address)
+    for offset in range(0, SECTOR_SIZE, PAGE_SIZE):
+      page = wanted[offset:offset+PAGE_SIZE]
+      if page != b'\xff' * PAGE_SIZE: f.program_page(address + offset, page)
+    got = f.read(address, SECTOR_SIZE)
+    if got != wanted:
+      index = next(i for i, (a, b) in enumerate(zip(got, wanted)) if a != b)
+      raise RuntimeError(f"verify mismatch at {address+index:#x}: flash={got[index]:02x} image={wanted[index]:02x}")
+    print(f"OK sector {sector:03d}/{total_sectors-1} @{address:#07x} elapsed={time.monotonic()-begin:.1f}s", flush=True)
+
+
+def parser():
+  p = argparse.ArgumentParser(description=__doc__)
+  p.add_argument('--device', type=int, default=0, help='USB bridge device index')
+  p.add_argument('--prescale', type=int, default=8, help='SCK prescaler 0..15 (default: 8)')
+  sub = p.add_subparsers(dest='command', required=True)
+  sub.add_parser('info', help='read JEDEC, status and SFDP').set_defaults(func=cmd_info)
+
+  r = sub.add_parser('read', help='read a flash range')
+  r.add_argument('address', type=lambda x:int(x, 0))
+  r.add_argument('size', type=lambda x:int(x, 0))
+  r.add_argument('-o', '--output')
+  r.set_defaults(func=cmd_read)
+
+  d = sub.add_parser('dump', help='dump the complete 2 MiB flash')
+  d.add_argument('output')
+  d.set_defaults(func=cmd_dump)
+
+  v = sub.add_parser('verify', help='compare the complete flash with an image')
+  v.add_argument('image')
+  v.set_defaults(func=cmd_verify)
+
+  w = sub.add_parser('flash', help='erase, program, and verify one or more sectors')
+  w.add_argument('image')
+  w.add_argument('--start-sector', type=lambda x:int(x, 0), default=0)
+  w.add_argument('--sector-count', type=lambda x:int(x, 0))
+  w.add_argument('--clear-cmp', action='store_true')
+  w.add_argument('--yes', action='store_true')
+  w.set_defaults(func=cmd_flash)
+  return p
+
+
+def main():
+  args = parser().parse_args()
+  try: args.func(args)
+  except (RuntimeError, TimeoutError, ValueError, OSError) as error:
+    print(f"error: {error}", file=sys.stderr)
+    raise SystemExit(1)
+
+if __name__ == '__main__': main()

extra/amdpci/am_smi.py

@@ -66,6 +66,8 @@ class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
     self.pcibus, self.devfmt = pcibus, pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
+    self.is_vf = bool(self.mmio[am.mmRCC_IOV_FUNC_IDENTIFIER] & 1)
+    self.vf_rlc_gated:list[tuple[int, int]] = []
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 

extra/gemm/moe_routing.py

@@ -1,8 +1,53 @@
+import functools, math, pathlib
 from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 BLOCK_ROW = 256
 
+@functools.cache
+def _router_mfma_fwd(out:UOp, x:UOp, weight:UOp, bias:UOp, *, dname:str) -> UOp:
+  *lead, K = x.shape
+  M = math.prod(lead)
+  E = weight.shape[0]
+  threads = UOp.special(256, "lidx0")
+  workgroups = UOp.special((M + 63) // 64, "gidx0")
+  sink = UOp.sink(out.base, x.base, weight.base, bias.base, threads, workgroups,
+                  arg=KernelInfo(f"moe_router_mfma_{M}_{K}_{E}", estimates=Estimates(ops=2*M*E*K, mem=(M*K+E*K+E)*2+M*E*4)))
+  amd = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (amd/"moe_router_mfma.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(amd/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
+                                 f"-DROUTER_M={M}", f"-DROUTER_K={K}", f"-DROUTER_E={E}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def _router_mfma_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, x_u, weight_u, bias_u = kernel.src[1:5]
+  x, weight, bias = (Tensor(u, device=u.device) for u in (x_u, weight_u, bias_u))
+  reference = x.float() @ weight.float().T + bias.float()
+  grad_x, grad_weight, grad_bias = reference.gradient(x, weight, bias, gradient=Tensor(gradient, device=x_u.device))
+  return None, grad_x.uop, grad_weight.uop, grad_bias.uop
+
+def router_mfma(x:Tensor, weight:Tensor, bias:Tensor) -> Tensor:
+  assert x.ndim >= 2 and weight.ndim == 2 and bias.ndim == 1
+  K = x.shape[-1]
+  E = weight.shape[0]
+  assert weight.shape == (E, K) and bias.shape == (E,)
+  assert x.dtype == weight.dtype == bias.dtype == dtypes.bfloat16
+  assert E == 32 and K % 64 == 0
+  if isinstance(x.device, tuple):
+    assert x.uop.axis == 0, f"router MFMA requires axis-0 sharding, got axis={x.uop.axis}"
+    local_shape = x.uop.shard_shape
+    assert local_shape[-1] == K and math.prod(local_shape[:-1]) % 64 == 0, f"unsupported local router shape {local_shape}"
+  else:
+    assert math.prod(x.shape[:-1]) % 64 == 0
+  x, weight, bias = x.contiguous(), weight.contiguous(), bias.contiguous()
+  out = _sharded_invalids((*x.shape[:-1], E), dtypes.float32, x.device)
+  out, *_ = Tensor.custom_kernel(out, x, weight, bias,
+    fxn=functools.partial(_router_mfma_fwd, dname=str(x.device)), grad_fxn=_router_mfma_bwd)
+  return out
+
 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
   if isinstance(device, tuple):
     per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.engine.realize import compile_linear
+from tinygrad.engine.realize import lower_and_compile
 from tinygrad.codegen.opt import OptOps
 
 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
@@ -39,7 +39,7 @@ if __name__ == "__main__":
     c = a.matmul(b, dtype=acc_dtype).realize()
 
   if getenv("SHOULD_USE_TC"):
-    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
+    linear = lower_and_compile(a.matmul(b, dtype=acc_dtype).schedule_linear())
     call = get_single_element(list(linear.src))
     applied_opts = call.src[0].src[0].arg.applied_opts
     assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"

extra/gptoss_kernels/rmsnorm/__init__.py

@@ -17,7 +17,7 @@ def _rmsnorm_mul_fwd_fxn(x_in_p, w_p, eps, device):
 
 def _rmsnorm_mul_bwd(grad:UOp, call:UOp) -> tuple:
   x = Tensor(call.src[1]).float(); weight = Tensor(call.src[2]).float()
-  rrms = Tensor(call.gettuple(1))
+  rrms = Tensor(call.unbound_outputs[1])
   x_normed = x * rrms                                  # recompute unweighted normed (x is call.src[1])
   d_y = Tensor(grad).float()
   dxn = d_y * weight                                   # d/d(x_normed)
@@ -28,8 +28,8 @@ def _rmsnorm_mul_bwd(grad:UOp, call:UOp) -> tuple:
 
 def rmsnorm_mul(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
   fxn = _rmsnorm_mul_fwd_fxn(x_in.as_param(0).uop, weight.as_param(1).uop, eps, x_in.device)
-  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, weight.uop, grad_fxn=_rmsnorm_mul_bwd)
-  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
+  outs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop), x_in.uop, weight.uop, grad_fxn=_rmsnorm_mul_bwd)
+  return Tensor(outs[0]), Tensor(outs[1])
 
 @functools.cache
 def _custom_rmsnorm_mul_quantize_mxfp8_fwd(q:UOp, e8:UOp, rrms:UOp, x:UOp, weight:UOp, *, dname:str, eps:float) -> UOp:

extra/hcq2/ops_amd2.py

@@ -1,17 +1,17 @@
 from __future__ import annotations
 from typing import cast
-import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
+import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, hashlib, itertools, collections, atexit
 assert sys.platform != 'win32'
-from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name
-from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import BufferSpec, Buffer
+from dataclasses import dataclass, replace
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
+from tinygrad.uop.ops import sint, UOp, ProgramInfo
+from tinygrad.device import BufferSpec, Buffer, Device, Compiled, ProfileProgramEvent
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32, prod, colored
 from tinygrad.helpers import ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
@@ -19,9 +19,10 @@ from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
-from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.ops_amd import SQTT, PMC
-from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
+from tinygrad.runtime.support.memory import AddrSpace
+from tinygrad.runtime.ops_amd import SQTT, PMC, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, AQL_HDR
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent, PMCSample
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ, WAIT_REG_MEM_FUNCTION_EQ
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
 from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
@@ -36,6 +37,14 @@ def _queue_args(hq:HWQueue, q) -> list[UOp]: # the ring and its pointers, tagged
 
 def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
 
+def dispatch_packet(data:AMDProgramData, info:ProgramInfo, kernel_object:UOp=UOp.const(0, dtypes.uint64),
+                    kernarg_address:UOp=UOp.const(0, dtypes.uint64)) -> list: # as words: the grid may be symbolic
+  pkt = bytes(hsa.hsa_kernel_dispatch_packet_t(header=AQL_HDR | (hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE),
+    setup=3 << hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS, private_segment_size=data.private_segment_size,
+    group_segment_size=data.group_segment_size, **{f"workgroup_size_{d}": l for d, l in zip("xyz", info.local_size)}))
+  grid = [(g * l).cast(dtypes.uint32) if isinstance(g, UOp) else g * l for g, l in zip(info.global_size, info.local_size)]
+  return [UOp(Ops.BINARY, arg=pkt[:12]), *grid, UOp(Ops.BINARY, arg=pkt[24:32]), kernel_object, kernarg_address, UOp(Ops.BINARY, arg=pkt[48:])]
+
 class AMDComputeQueue(HWQueue):
   q_rewrite = PatternMatcher([
     (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
@@ -49,6 +58,8 @@ class AMDComputeQueue(HWQueue):
   def __init__(self, ctx, submit):
     super().__init__(ctx, submit)
     self.pm4, self.gc, self.soc, self.nbio, self.target = self.dev.pm4, self.dev.gc, self.dev.soc, self.dev.nbio, self.dev.target
+    self.profiled:list[UOp] = []
+    if self.dev.pmc_enabled: self.pmc_start()
 
   def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, _dw(vals) - 1), *vals)
 
@@ -61,6 +72,20 @@ class AMDComputeQueue(HWQueue):
     else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
     self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
 
+  @contextlib.contextmanager
+  def pred_exec(self, xcc_mask:int): # the count fills in when the block closes
+    if self.dev.xccs > 1: self.pkt3(self.pm4.PACKET3_PRED_EXEC, xcc_mask << 24)
+    start = len(self.blob)
+    yield
+    if self.dev.xccs > 1:
+      cnt, = struct.unpack("I", self.blob[start-4:start])
+      self.blob[start-4:start] = struct.pack("I", cnt | (len(self.blob) - start) // 4)
+
+  def set_grbm(self, instance=None, se=None, sh=None, wgp=None):
+    instance_val = (wgp << 2 | (instance or 0)) if wgp is not None else instance
+    self.wreg(self.gc.regGRBM_GFX_INDEX, **{(f'{key}_broadcast_writes' if val is None else f'{key}_index'): (1 if val is None else val)
+      for key, val in [('instance', instance_val), ('se', se), ('sh' if self.target[0] == 9 else 'sa', sh)]})
+
   def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
     wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
                 | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
@@ -106,16 +131,232 @@ class AMDComputeQueue(HWQueue):
                       reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
     self.acquire_mem()
 
+  def spi_config(self, tracing:bool):
+    self.wreg(self.gc.regSPI_CONFIG_CNTL, ps_pkr_priority_cntl=3, exp_priority_order=3, gpr_write_priority=0x2c688,
+              enable_sqg_bop_events=int(tracing), enable_sqg_top_events=int(tracing))
+
+  ### profiling: a kernel's slot holds its counters and trace until a synchronize reads them back
+
+  def prof_buf(self, name:str) -> UOp:
+    return UOp.placeholder((getattr(self.dev, name).size,), getattr(self.dev, name).dtype, 0, device=self.devs, tag=name)
+
+  def prof_start(self, data:AMDProgramData, info:ProgramInfo, lib:UOp) -> UOp|None:
+    if not (self.dev.pmc_enabled or self.dev.sqtt_enabled): return None
+    slot = (self.prof_buf("prof_log").index(0).load() + len(self.profiled)) % self.dev.prof_slots
+    tag = UOp.const(unwrap_view(lib)[0].arg.slot, dtypes.uint64)
+    self.profiled.append(self.prof_buf("prof_log").index(1 + slot.cast(dtypes.int)).store(tag))
+    if self.dev.sqtt_enabled:
+      self.sqtt_start(slot)
+      self.sqtt_setup_exec(data, info)
+    return slot
+
+  def prof_stop(self, slot:UOp|None):
+    if slot is None: return
+    if self.dev.pmc_enabled: self.pmc_read(slot)
+    if self.dev.sqtt_enabled: self.sqtt_stop(slot)
+
+  def prof_bump(self, cmdbuf:UOp) -> UOp:
+    if not self.profiled: return cmdbuf
+    log = self.prof_buf("prof_log")
+    return cmdbuf.after(log.after(cmdbuf, *self.profiled).index(0).store(log.index(0).load() + len(self.profiled)))
+
+  ### PMC
+
+  def pmc_reset_counters(self, en=True):
+    self.set_grbm()
+    self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=0)
+    if en: self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1)
+
+  def pmc_start(self): # every submit
+    self.pmc_reset_counters(en=False)
+    self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL, cs_en=1, ps_en=1, gs_en=1, hs_en=1, **({'vmid_mask':0xffff} if (gfx9:=self.target[0] == 9) else {}))
+    if not gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_CTRL2, force_en=1, vmid_en=0xffff)
+
+    end_off, sched = 0, []
+    block2pid:dict[str, itertools.count] = collections.defaultdict(lambda: itertools.count())
+    for name in self.dev.pmc_names:
+      block, idx = self.dev.pmc_counters[name]
+      # sq block on gfx11+ goes down to wgps
+      inst_cnt, se_cnt, sa_cnt, wgp_cnt = {"GRBM": (1, 1, 1, 1), "GL2C": (32, 1, 1, 1), "TCC": (16, 1, 1, 1),
+        "SQ": (1, self.dev.se_cnt) + ((1, 1) if gfx9 else (2, self.dev.iface.props['cu_per_simd_array'] // 2))}[block]
+      end_off += (rec_size:=prod((self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt)) * 8)
+
+      # gfx11+ and later require even-numbered SQ *_SELECT registers
+      regsample = f'reg{block}_PERFCOUNTER{(pcid:=next(block2pid[block]))}'
+      if (regsel:=getattr(self.gc, (f'reg{block}_PERFCOUNTER{(pcid*2) if not gfx9 and block=="SQ" else pcid}_SELECT'), None)) is None:
+        raise RuntimeError(f'{block} is out of perfcounter registers: ({regsample} is not found)')
+
+      self.wreg(regsel, perf_sel=idx, **({'simd_mask':0xf, 'sqc_bank_mask':0xf, 'sqc_client_mask':0xf} if gfx9 and block == "SQ" else {}))
+      sched.append(PMCSample(name, block, self.dev.xccs, inst_cnt, se_cnt, sa_cnt, wgp_cnt, end_off-rec_size, rec_size, regsample))
+    self.dev.pmc_sched = sched
+
+    if gfx9: self.wreg(self.gc.regSQ_PERFCOUNTER_MASK, sh0_mask=0xffff, sh1_mask=0xffff)
+    self.wreg(self.gc.regCOMPUTE_PERFCOUNT_ENABLE, 1)
+    self.pmc_reset_counters(en=True)
+
+  def pmc_read(self, slot:UOp):
+    buf = rt_addr(self.prof_buf("pmc_buf"), self.devs) + slot * self.dev.pmc_size
+    self.set_grbm()
+    self.wreg(self.gc.regCP_PERFMON_CNTL if self.target[0] <= 11 else self.gc.regCP_PERFMON_CNTL_1, perfmon_state=1, perfmon_sample_enable=1)
+
+    for smp in self.dev.pmc_sched:
+      offset = itertools.count(smp.off, step=8)
+
+      for xcc in range(smp.xcc):
+        with self.pred_exec(xcc_mask=1 << xcc):
+          for inst, se_idx, sa_idx, wgp_idx in itertools.product(range(smp.inst), range(smp.se), range(smp.sa), range(smp.wgp)):
+            loff = next(offset)
+            if smp.wgp > 1 and not self.dev.iface.is_wgp_active(xcc, se_idx, sa_idx, wgp_idx): continue
+            self.set_grbm(**({'instance':inst} if smp.inst > 1 else ({'se':se_idx}|({'sh':sa_idx, 'wgp':wgp_idx} if self.target[0] != 9 else {}))))
+
+            # Copy counter to memory (src_sel = perf, dst_sel = tc_l2)
+            lo, hi = getattr(self.gc, f'{smp.regsample}_LO'), getattr(self.gc, f'{smp.regsample}_HI', None)
+            self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, lo.addr[0], 0, buf + loff)
+            if hi is not None: self.pkt3(self.pm4.PACKET3_COPY_DATA, (2 << 8) | 4, hi.addr[0], 0, buf + (loff + 4))
+
+    self.pmc_reset_counters(en=True)
+
+  ### SQTT
+
+  def sqtt_userdata(self, data, *extra_dwords):
+    data_ints = [x[0] for x in struct.iter_unpack('<I', bytes(data))] + list(extra_dwords)
+    for i in range(0, len(data_ints), 2):
+      self.wreg(self.gc.regSQ_THREAD_TRACE_USERDATA_2, *data_ints[i:i+2])
+
+  def sqtt_config(self, tracing:bool):
+    trace_ctrl = {'rt_freq': self.soc.SQ_TT_RT_FREQ_4096_CLK} if self.target < (12,0,0) else {}
+    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1,
+      mode=int(tracing), **trace_ctrl)
+
+  def sqtt_setup_exec(self, data:AMDProgramData, info:ProgramInfo):
+    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_pipeline_bind(identifier=sqtt.RGP_SQTT_MARKER_IDENTIFIER_BIND_PIPELINE,
+                                                                 bind_point=(__BIND_POINT_COMPUTE:=1), api_pso_hash=data64_le(data.libhash)))
+    self.sqtt_userdata(sqtt.struct_rgp_sqtt_marker_event(has_thread_dims=1, cmd_id=next(self.dev.sqtt_next_cmd_id)), *info.global_size)
+
+    if SQTT_LIMIT_SE:
+      # Calculate number of CUs per SE to enable based on blocks count. 4 is maximum simd per CU, but on rdna we can trace only 1.
+      cu_per_se = prod([x if isinstance(x, int) else 1 for x in info.global_size]) // ((self.dev.cu_cnt // self.dev.se_cnt) * 4)
+      for xcc in range(self.dev.xccs):
+        with self.pred_exec(xcc_mask=1 << xcc):
+          for i in range(8 if self.target[0] != 9 else 4):
+            if SQTT_LIMIT_SE > 1: mask = 1 if SQTT_ITRACE_SE_MASK.value & (1 << i) else 0 # only run unmasked shader engines
+            else:
+              sa_mask = (1 << (self.dev.iface.props['cu_per_simd_array'] // 2)) - 1
+              cu_mask = (1 << (cu_per_se + (1 if i == 0 else 0))) - 1
+              mask = lo32((cu_mask & sa_mask) | (cu_mask & (sa_mask << 16)) << 16)
+            self.wreg(getattr(self.gc, f'regCOMPUTE_STATIC_THREAD_MGMT_SE{i}'), mask)
+
+  def sqtt_start(self, slot:UOp):
+    self.memory_barrier()
+    win, ses = self.dev.sqtt_win, self.dev.sqtt_ses
+    base = rt_addr(self.prof_buf("sqtt_buf"), self.devs) + slot * win
+    if self.target[0] == 9:
+      self.set_grbm()
+      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, simd_en=0xf, cu_sel=0, sq_stall_en=1, spi_stall_en=1, reg_stall_en=1, vm_id_mask=0)
+      for se in range(ses):
+        mask = (__SQTT_MISC:=1<<0) | (__SQTT_TIME:=1<<1) | (__SQTT_REG:=1<<2) | (__SQTT_WAVE_START:=1<<3) | (__SQTT_WAVE_END:=1<<6) \
+             | (__SQTT_USERDATA:=1<<12) | (__SQTT_REG_CS:=1<<5) | (__SQTT_REG_CS_PRIV:=1<<15)
+        if (SQTT_ITRACE_SE_MASK.value >> se) & 0b1: mask |= (__SQTTINST:=1<<10) | (__SQTT_INST_PC:=1<<11) | (__SQTT_ISSUE:=1<<13)
+
+        buf0_lo, buf0_hi = [((base + se * self.dev.prof_slots * win) >> sh).cast(dtypes.uint32) for sh in (12, 44)]
+        with self.pred_exec(xcc_mask=1<<(se // self.dev.se_cnt)):
+          self.set_grbm(se=se % self.dev.se_cnt, sh=0)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_mask=0xf, token_mask=mask)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK2, inst_mask=0xffffffff)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BASE, buf0_lo)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BASE2, buf0_hi)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_SIZE, size=win >> 12)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, reset_buffer=1)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=1)
+    else:
+      self.spi_config(tracing=True)
+      # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
+      for se in range(ses):
+        self.set_grbm(se=se, sh=0)
+
+        buf0_lo, buf0_hi = [((base + se * self.dev.prof_slots * win) >> sh).cast(dtypes.uint32) for sh in (12, 44)]
+        if self.target >= (12,0,0):
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=win >> 12)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, buf0_lo)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, buf0_hi)
+        else:
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, self.gc.regSQ_THREAD_TRACE_BUF0_SIZE.encode(size=win >> 12) | buf0_hi)
+          self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, buf0_lo)
+        # NOTE: SQTT can only trace instructions on one simd per se, this selects the simd in first wgp in first sa.
+        # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se,
+        # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but
+        # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and
+        # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
+        # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
+        # Use SQTT_SIMD_SEL to select which SIMD to trace (0-3). Memory ops show different InstOp values (0x2x vs 0x5x) based on SIMD.
+        cs_wtype = (1 << 6) if self.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT
+        self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=SQTT_SIMD_SEL.value, wgp_sel=0, sa_sel=0)
+        reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
+                      self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
+        token_exclude = SQTT_TOKEN_EXCLUDE.value | ((1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.target < (12,0,0) else 0)
+
+        # disable instr tracing
+        if not (SQTT_ITRACE_SE_MASK.value >> se) & 0b1:
+          # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
+          token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
+                            1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
+                            1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.target < (12,0,0) else 0x927
+
+        self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1,
+                  **({} if self.target < (12,0,0) else {'exclude_barrier_wait': 1}))
+        self.sqtt_config(tracing=True)
+
+    self.set_grbm()
+    if self.target[0] != 9: self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 1)
+    self.memory_barrier()
+
+  # Magic values from src/amd/common/ac_sqtt.c:ac_sqtt_emit_stop and src/amd/common/ac_sqtt.c:ac_sqtt_emit_wait
+  def sqtt_stop(self, slot:UOp):
+    self.memory_barrier()
+    self.set_grbm()
+    ses = self.dev.sqtt_ses
+    wptrs = rt_addr(self.prof_buf("sqtt_wptrs"), self.devs) + slot * (ses * 4)
+
+    # Start shutting everything down
+    if self.target[0] == 9: self.wreg(self.gc.regSQ_THREAD_TRACE_MODE, mask_cs=1, autoflush_en=1, mode=0)
+    else:
+      self.wreg(self.gc.regCOMPUTE_THREAD_TRACE_ENABLE, 0)
+      self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_FINISH) | self.pm4.EVENT_INDEX(0))
+
+    # For each SE wait for finish to complete and copy regSQ_THREAD_TRACE_WPTR to know where in the buffer trace data ends
+    for se in range(ses):
+      with self.pred_exec(xcc_mask=1<<(se // self.dev.se_cnt)):
+        self.set_grbm(se=se % self.dev.se_cnt, sh=0)
+
+        regstatus = self.gc.regSQ_THREAD_TRACE_STATUS.addr[0] - (self.pm4.PACKET3_SET_UCONFIG_REG_START if self.target[0] == 9 else 0)
+        if self.target[0] != 9:
+          self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+          self.sqtt_config(tracing=False)
+        self.wait_reg_mem(reg=regstatus, mask=self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), op=WAIT_REG_MEM_FUNCTION_EQ, value=0)
+        self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+
+        # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True)
+        self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, wptrs + se * 4)
+
+    self.set_grbm()
+    if self.target[0] != 9: self.spi_config(tracing=False)
+    self.memory_barrier()
+
+  ### exec
+
+  def kernargs(self, call:UOp, prg:UOp, data:AMDProgramData) -> list[UOp]:
+    words = [get_call_arg_uops(call)[gi].getaddr(self.devs) for gi in prg.arg.globals] + \
+            [b.ccast(v.dtype) for v, b in zip(prg.arg.vars, get_call_var_uops(call, prg))] # a bound value is a bare const, the var has the width
+    pad = data.kernargs_segment_size - sum(w.dtype.itemsize for w in words)
+    assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
+    return words + [UOp.const(0, dtypes.uint32)] * (pad // 4) + (dispatch_packet(data, prg.arg) if data.enable_dispatch_ptr else [])
+
   def exec(self, call:UOp, prg:UOp):
     data, lib = amd_build_program(self.dev, prg, self.devs)
     info = prg.arg
 
     # kernargs: a nested blob linear inside a getaddr, packed into the tail of the cmdbuf
-    ka_words = [get_call_arg_uops(call)[gi].getaddr(self.devs) for gi in info.globals] + \
-               [b.ccast(v.dtype) for v, b in zip(info.vars, get_call_var_uops(call, prg))] # a bound value is a bare const, the var has the width
-    pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
-    assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
-    ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4))
+    ka = UOp(Ops.LINEAR, src=tuple(self.kernargs(call, prg, data)))
 
     prog_addr = lib.getaddr(self.devs) + data.entry_point_offset
     scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=self.devs).rtag("scratch").getaddr(self.devs)
@@ -129,41 +370,85 @@ class AMDComputeQueue(HWQueue):
     dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
       **({'cs_w32_en': int(data.wave32)} if self.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
     self.acquire_mem(gli=0, gl2=0)
+    slot = self.prof_start(data, info, lib)
     self.wreg(self.gc.regCOMPUTE_PGM_LO, prog_addr >> 8)
     self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
     self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
     self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
-    for xcc_id in range(self.dev.xccs):
-      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8)
+    for xcc_id in range(self.dev.xccs): # architected flat scratch: each xcc gets its part
+      with self.pred_exec(xcc_mask=1 << xcc_id):
+        self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8)
     self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
     self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
     self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH")))
     self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0)
     self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *info.global_size, dispatch_init)
+    if self.dev.sqtt_enabled: self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.THREAD_TRACE_MARKER) | self.pm4.EVENT_INDEX(0))
     self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+    self.prof_stop(slot)
 
   def wait(self, signal:UOp, value:UOp): self.wait_reg_mem(value.cast(dtypes.uint32), mem=signal.getaddr(self.devs))
 
   def timestamp(self, signal:UOp):
-    self.release_mem(signal.getaddr(self.devs), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
-                     self.pm4.int_sel__mec_release_mem__none)
+    with self.pred_exec(xcc_mask=0b1):
+      self.release_mem(signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+                       self.pm4.int_sel__mec_release_mem__none)
 
   def signal(self, signal:UOp, value:UOp):
-    self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
-                     self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
+    with self.pred_exec(xcc_mask=0b1):
+      self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
+                       self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
 
-  def submit(self, cmdbuf:UOp) -> UOp:
-    q = self.dev.compute_queue
+  def submit(self, cmdbuf:UOp) -> UOp: # the ring gets an indirect buffer packet: 4 dwords, put stays aligned so it never wraps mid packet
+    base, off = unwrap_view(cmdbuf)
+    blob = struct.pack("IIII", self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), 0, 0, cmdbuf.max_numel() // 4 | self.pm4.INDIRECT_BUFFER_VALID)
+    ib = patch(UOp.placeholder((16,), dtypes.uint8, device="CPU", tag=to_name("ib", self.queue)), [(4, base.getaddr(self.devs) + off)], blob)
+    return self.push(self.prof_bump(cmdbuf), ib, self.dev.compute_queue)
 
+  def push(self, cmdbuf:UOp, words:UOp, q, unit:int=4, doorbell_lag:int=0) -> UOp:
     ring, wptr, doorbell, put = _queue_args(self, q)
+    n, p = words.max_numel() // unit, put.index(0).load() # put counts units
+    i = UOp.range(words.max_numel() // 4, 10, dtype=dtypes.int, src=(cmdbuf,))
+    at = ((p * (unit // 4) + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)
+    written = ring.index(at).store(words.bitcast(dtypes.uint32).index(i).load()).end(i)
+    w = wptr.after(written).index(0).store(p + n)
+    return doorbell.after(put.after(w).index(0).store(p + n)).index(0).store(p + n - doorbell_lag)
 
-    size_dw = cmdbuf.max_numel() // 4
-    p = put.after(*self.deps).index(0).load()
-    i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf,))
-    copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
-    next_put = p + size_dw
-    flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
-    return doorbell.after(flush).index(0).store(next_put)
+class AMDComputeAQLQueue(AMDComputeQueue): # the ring holds 64 byte aql packets: a dispatch per kernel, the pm4 between them wrapped as an ib
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
+    self.cmd_addr = UOp.variable("cmdbuf", 0, 2**48, dtypes.uint64) # the packets point into the cmdbuf, its address binds at submit
+    self.pkts:list[UOp] = []
+    self.run_start = 0
+
+  def close_run(self, end:int):
+    if end > self.run_start:
+      hdr = AQL_HDR | (hsa.HSA_PACKET_TYPE_VENDOR_SPECIFIC << hsa.HSA_PACKET_HEADER_TYPE) | (1 << 16)
+      ib = [self.pm4.PACKET3(self.pm4.PACKET3_INDIRECT_BUFFER, 2), self.cmd_addr + self.run_start,
+            (end - self.run_start) // 4 | self.pm4.INDIRECT_BUFFER_VALID]
+      self.pkts += [UOp.const(w, dtypes.uint32) if isinstance(w, int) else w for w in [hdr, *ib, 10, *[0] * 10]]
+    self.run_start = end
+
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = amd_build_program(self.dev, prg, self.devs)
+    self.dev.scratch_buffer(data.private_segment_size) # the queue descriptor holds the scratch
+    slot = self.prof_start(data, prg.arg, lib)
+    self.close_run(len(self.blob))
+    self.blob += bytes(-len(self.blob) % 16)
+    kernarg_address = self.cmd_addr + len(self.blob) # the kernargs go inline in the cmdbuf: the runs skip them
+    self.q(*self.kernargs(call, prg, data))
+    self.pkts += [UOp.const(w, dtypes.uint32) if isinstance(w, int) else w
+                  for w in dispatch_packet(data, prg.arg, lib.getaddr(self.devs) + data.desc_offset, kernarg_address)]
+    self.run_start = len(self.blob)
+    self.prof_stop(slot)
+
+  def submit(self, cmdbuf:UOp) -> UOp: # the doorbell is the last packet's index
+    self.close_run(cmdbuf.max_numel())
+    base, off = unwrap_view(cmdbuf)
+    self.blob, self.patches = bytearray(), [] # q again, for the aql stream
+    self.q(*UOp.sink(*self.pkts).substitute({self.cmd_addr: base.getaddr(self.devs) + off}).src)
+    aql = UOp.placeholder((len(self.blob),), dtypes.uint8, device="CPU", tag=to_name("aql", self.queue))
+    return self.push(self.prof_bump(cmdbuf), patch(aql, self.patches, bytes(self.blob)), self.dev.compute_queue, unit=64, doorbell_lag=1)
 
 # *****************
 # SDMA
@@ -197,7 +482,7 @@ class AMDSDMAQueue(HWQueue):
 
   def timestamp(self, signal:UOp):
     self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
-           signal.getaddr(self.devs))
+           signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64))
 
   def signal(self, signal:UOp, value:UOp): # a fence packet then a trap
     op = self.sdma.SDMA_OP_FENCE | (self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.target[0] != 9 else 0)
@@ -208,9 +493,11 @@ class AMDSDMAQueue(HWQueue):
     q = unwrap(self.dev.sdma_queue(int(self.queue.split(":")[1])))
 
     ring, wptr, doorbell, put = _queue_args(self, q)
+    base = unwrap_view(cmdbuf)[0] # in host memory: streamed into the ring, the device never reads it
+    cmdbuf = cmdbuf.substitute({base: base.replace(arg=replace(base.arg, device="CPU"))})
 
     rs, size_dw = q.ring.size, cmdbuf.max_numel() // 4
-    put_b = put.after(*self.deps).index(0).load()
+    put_b = put.index(0).load()
     tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
     fits = (size_dw <= rs - tail).cast(dtypes.int)
     start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
@@ -222,19 +509,24 @@ class AMDSDMAQueue(HWQueue):
     flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
     return doorbell.after(flush).index(0).store(next_put)
 
+def amd_compute_queue(ctx, submit:UOp) -> HWQueue:
+  return (AMDComputeAQLQueue if Device[submit.src[0].arg[0][0]].is_aql else AMDComputeQueue)(ctx, submit)
+
 @dataclass(frozen=True)
 class AMDProgramData:
-  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
-  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
+  desc_offset:int; entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool; libhash:int
+  private_segment_size:int; group_segment_size:int; kernargs_segment_size:int
   enable_dispatch_ptr:int; enable_private_segment_sgpr:int
 
 _amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
+_amd_program_prof:dict[UOp, tuple[str, bytes, bytes]] = {} # placeholder -> (name, lib, key) for its profile event
 def amd_build_program(dev, prg:UOp, devs:tuple[str, ...]) -> tuple[AMDProgramData, UOp]:
   # the image parses once per lib, each device set gets its own program buffer of it
   if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, devs))) is None:
     data, image = _amd_program_image(dev, lib)
     buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
     cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
+    if PROFILE: _amd_program_prof[buf] = (prg.arg.function_name, lib, prg.key)
   return cached
 
 @functools.cache
@@ -249,11 +541,11 @@ def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]:
     raise RuntimeError("Too many resources requested: group_segment_size")
   edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
 
-  data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+  data = AMDProgramData(desc_offset=rodata, entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
     rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
-    rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-    wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
-    kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
+    rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3, wave32=bool(desc.kernel_code_properties & 0x400),
+    libhash=struct.unpack('<Q', hashlib.md5(lib).digest()[:8])[0], private_segment_size=desc.private_segment_fixed_size,
+    group_segment_size=desc.group_segment_fixed_size, kernargs_segment_size=desc.kernarg_size, enable_dispatch_ptr=edp,
     enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
   return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
 
@@ -495,7 +787,8 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      (tl:=d.timeline._buf.cpu_view().view(fmt='Q'))[0] = tl[1]
+      tl = d.timeline._buf.cpu_view().view(fmt='Q')
+      tl[0] = tl[1]
 
   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))):
@@ -541,7 +834,7 @@ class AMDDevice(HCQ2Compiled):
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
   max_scratch_psize = 0
   pm_encode = PatternMatcher([
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_compute", name="submit"), lambda ctx, submit: encode_submit(AMDComputeQueue(ctx, submit))),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_compute", name="submit"), lambda ctx, submit: encode_submit(amd_compute_queue(ctx, submit))),
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_copy", name="submit"), lambda ctx, submit: encode_submit(AMDSDMAQueue(ctx, submit))),
   ])
 
@@ -578,10 +871,6 @@ class AMDDevice(HCQ2Compiled):
                       bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
 
     self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
-    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
-      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
-
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queues:dict = {}
     self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
@@ -590,37 +879,32 @@ class AMDDevice(HCQ2Compiled):
 
     # Scratch setup
     self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel())),
+      (UPat(Ops.PARAM, tag="program", name="b"), lambda ctx, b: ctx.program_buffer(b)),
+    ]) + self.pm_bufferize
 
     if self.is_usb:
       self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
       raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
 
-    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
-    if self.pmc_enabled:
+    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
+    self.pmc_enabled, self.sqtt_enabled = PROFILE > 0 and PMC > 0, PROFILE > 0 and SQTT > 0
+    if self.pmc_enabled or self.sqtt_enabled:
       self.iface.require_profile_mode()
-
-      self.pmc_sched:list[PMCSample] = []
+      self.prof_slots, self.prof_read, self.pmc_sched, self.sqtt_next_cmd_id = getenv("PROF_SLOTS", 32), 0, [], itertools.count(0)
+      self.sqtt_ses, self.sqtt_win = self.se_cnt * self.xccs, (getenv("SQTT_BUFFER_SIZE", 256) << 20) // self.prof_slots # mb, per shader engine
+      self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag=n), lambda ctx, n=n: getattr(ctx, n))
+                                          for n in ("prof_log", "pmc_buf", "sqtt_buf", "sqtt_wptrs")]) + self.pm_bufferize
+    if self.pmc_enabled:
       self.pmc_counters = import_pmc(self.target)
-
       # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses
       l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC")
       pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS"
-      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
+      self.pmc_names = getenv("PMC_COUNTERS", pmc_default).split(",")
+      for k in self.pmc_names:
         if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
 
-      raise NotImplementedError("PMC start not migrated to hcq2 yet")
-
-    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
-    self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
-    if self.sqtt_enabled:
-      self.iface.require_profile_mode()
-
-      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt * self.xccs)]
-      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
-      self.sqtt_next_cmd_id = itertools.count(0)
-
   def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
     ring = Buffer(self.device, ring_size // 4, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
     gart = Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
@@ -630,7 +914,8 @@ class AMDDevice(HCQ2Compiled):
       self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
         read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
         max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1)
-      self.aql_gart._buf.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+      if hasattr(self, 'scratch'): self.aql_scratch()
+      else: self.aql_gart._buf.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
 
     cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE)
     cwsr_buffer = Buffer(self.device, cwsr_buffer_size, dtypes.uint8, preallocate=True) if ctx_save_restore_size else None
@@ -681,24 +966,7 @@ class AMDDevice(HCQ2Compiled):
     num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
 
     tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
-
-    if hasattr(self, 'aql_desc'):
-      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
-      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
-      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-
-      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
-      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
-      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
-        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
-        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
-      self.aql_desc.compute_tmpring_size = tmpring
-      self.aql_gart._buf.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    return tmpring
+    return int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
 
   def scratch_buffer(self, private_segment_size):
     AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
@@ -709,8 +977,79 @@ class AMDDevice(HCQ2Compiled):
       size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
       self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
       self.max_private_segment_size = private_segment_size
+      if hasattr(self, 'aql_desc'): self.aql_scratch()
     return self.scratch
 
+  def aql_scratch(self):
+    gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
+    rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
+            'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+    rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+    rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+
+    base = self.scratch._buf.va_addr
+    self.aql_desc.scratch_backing_memory_location = base
+    self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size
+    self.aql_desc.scratch_resource_descriptor[:] = [lo32(base), int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(base), SWIZZLE_ENABLE=1), 'little'),
+                                                    lo32(self.scratch.nbytes // self.xccs), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
+    self.aql_desc.compute_tmpring_size = self.tmpring_size(self.max_private_segment_size)
+    self.aql_gart._buf.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+  def _prof_buffer(self, size:int, dtype, host:bool=False) -> Buffer:
+    buf = Buffer(self.device, size, dtype, options=BufferSpec(host=host, nolru=True, uncached=True, cpu_access=True), preallocate=True)
+    buf._buf.cpu_view().view(fmt='B')[:buf.nbytes] = bytes(buf.nbytes)
+    return buf
+
+  @functools.cached_property
+  def prof_log(self) -> Buffer: return self._prof_buffer(1 + self.prof_slots, dtypes.uint64, host=True)
+  @property
+  def pmc_size(self) -> int: return self.pmc_sched[-1].off + self.pmc_sched[-1].size
+  @functools.cached_property
+  def pmc_buf(self) -> Buffer: return self._prof_buffer(self.pmc_size * self.prof_slots, dtypes.uint8)
+  @functools.cached_property
+  def sqtt_buf(self) -> Buffer: return self._prof_buffer(self.sqtt_win * self.prof_slots * self.sqtt_ses, dtypes.uint8)
+  @functools.cached_property
+  def sqtt_wptrs(self) -> Buffer: return self._prof_buffer(self.prof_slots * self.sqtt_ses, dtypes.uint32)
+
+  def program_buffer(self, b:UOp) -> Buffer:
+    if b not in self.prog_bufs:
+      buf = self.prog_bufs[b] = Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
+      if PROFILE:
+        name, lib, key = _amd_program_prof[b]
+        Compiled.profile_events.append(ProfileProgramEvent(self.device, name, lib, buf._buf.va_addr, b.arg.slot, key))
+    return self.prog_bufs[b]
+
+  def sqtt_trace(self, slot:int, se:int) -> bytes:
+    off = (se * self.prof_slots + slot) * self.sqtt_win
+    wptr = (self.sqtt_wptrs._buf.cpu_view().view(fmt='I')[slot * self.sqtt_ses + se] & 0x1FFFFFFF) * 32
+    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf.va_addr + off) // 32) & 0x1FFFFFFF) * 32
+    assert 0 <= wptr <= self.sqtt_win, f"{wptr} > {self.sqtt_win}, should never happen"
+    if wptr >= self.sqtt_win - 32: # the wptr stops at the last dword when the window overflows
+      print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {se})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))
+    blob = bytes(self.sqtt_buf._buf.cpu_view()[off:off + wptr])
+    return (struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (se << 24)) + blob) if self.target[0] == 9 else blob
+
+  def _at_profile_finalize(self): # the calibration kernels aren't profiles
+    self.synchronize()
+    super()._at_profile_finalize()
+    if self.pmc_enabled or self.sqtt_enabled: self.prof_read = self.prof_log._buf.cpu_view().view(fmt='Q')[0]
+
+  def collect_prof(self):
+    if self.pmc_enabled or self.sqtt_enabled:
+      log = self.prof_log._buf.cpu_view().view(fmt='Q')
+      if (lost:=log[0] - self.prof_read - self.prof_slots) > 0:
+        print(colored(f"{self.device}: Warning: {lost} kernel profiles were overwritten: synchronize more often or raise PROF_SLOTS", "yellow"))
+      for k in range(max(self.prof_read, log[0] - self.prof_slots), log[0]):
+        slot, tag = k % self.prof_slots, log[1 + k % self.prof_slots]
+        if self.pmc_enabled:
+          blob = bytes(self.pmc_buf._buf.cpu_view()[slot * self.pmc_size:(slot + 1) * self.pmc_size])
+          Compiled.profile_events.append(ProfilePMCEvent(self.device, tag, self.pmc_sched, blob, k))
+        for se in range(self.sqtt_ses if self.sqtt_enabled else 0):
+          itrace = bool((SQTT_ITRACE_SE_MASK.value >> se) & 1)
+          Compiled.profile_events.append(ProfileSQTTEvent(self.device, tag, se, self.sqtt_trace(slot, se), itrace, k))
+      self.prof_read = log[0]
+    super().collect_prof()
+
   def on_device_hang(self): self.iface.on_device_hang()
 
   def device_props(self): return self.iface.props

extra/llama_kernels/__init__.py

@@ -16,7 +16,7 @@ def _local_abs_max_fxn(x_p, device):
 def local_abs_max(x:Tensor) -> Tensor:
   param = x.as_param(0)
   fxn = _local_abs_max_fxn(param.uop, x.device)
-  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+  return Tensor(fxn[0].uop.call_with_output(x.uop))
 
 def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
   s = list(shape)

extra/llama_kernels/rmsnorm/__init__.py

@@ -13,12 +13,13 @@ def _rmsnorm_fwd_fxn(x_in_p, eps, device):
   return rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
 
 def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
-  x_normed = Tensor(call.gettuple(0)).float()
+  outs = call.unbound_outputs
+  x_normed = Tensor(outs[0]).float()
   do_float = Tensor(grad).float()
-  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  d_x = Tensor(outs[1]) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
   return (d_x.cast(call.src[1].dtype).uop,)
 
 def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
   fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
-  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
-  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
+  outs = UOp.call_with_outputs((fxn[0].uop, fxn[1].uop), x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(outs[0]), Tensor(outs[1])

extra/remote/bench.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os, sys, time
-from tinygrad.runtime.support.system import RemotePCIDevice
+from extra.remote.hcq1_remote import RemotePCIDevice
 
 LAT_N_RUNS = 500
 THROUGHPUT_N_RUNS = 8
@@ -18,7 +18,7 @@ if __name__ == "__main__":
   print(f"connected to {os.environ['REMOTE']}, device: {name}\n")
 
   # ping (minimal server round-trip, no device I/O)
-  from tinygrad.runtime.support.system import RemoteCmd
+  from extra.remote.hcq1_remote import RemoteCmd
   sock = pci.sock
   for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
   st = time.perf_counter()

extra/remote/hcq1_remote.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+import os, mmap, array, functools, contextlib, itertools, struct, socket, subprocess, time, enum, atexit
+from tinygrad.helpers import getenv, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
+from tinygrad.runtime.support.system import PCIDevice, System
+
+class RemoteCmd(enum.IntEnum):
+  PROBE,MAP_BAR,MAP_SYSMEM_FD,CFG_READ,CFG_WRITE,RESET,MMIO_READ,MMIO_WRITE,MAP_SYSMEM,SYSMEM_READ,SYSMEM_WRITE,RESIZE_BAR,PING = range(13)
+
+class RemoteMMIOInterface(MMIOInterface):
+  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0, rd_cmd=RemoteCmd.MMIO_READ, wr_cmd=RemoteCmd.MMIO_WRITE):
+    self.dev, self.residx, self.nbytes, self.fmt, self.off, self.el_sz = dev, residx, nbytes, fmt, off, struct.calcsize(fmt)
+    self.rd_cmd, self.wr_cmd = rd_cmd, wr_cmd
+
+  def __getitem__(self, index):
+    sl = index if isinstance(index, slice) else slice(index, index + 1)
+    start, stop = (sl.start or 0) * self.el_sz, (sl.stop or len(self)) * self.el_sz
+    data = self.dev._bulk_read(self.rd_cmd, self.residx, self.off + start, stop - start)
+    result = data if self.fmt == 'B' else list(struct.unpack(f'<{(stop - start) // self.el_sz}{self.fmt}', data))
+    return result if isinstance(index, slice) else result[0]
+
+  def __setitem__(self, index, val):
+    start = (index.start or 0) * self.el_sz if isinstance(index, slice) else index * self.el_sz
+    data = (val if self.fmt == 'B' else struct.pack(f'<{len(val)}{self.fmt}', *val)) if isinstance(index, slice) else struct.pack(f'<{self.fmt}', val)
+    self.dev._bulk_write(self.wr_cmd, self.residx, self.off + start, data)
+
+  def view(self, offset:int=0, size:int|None=None, fmt=None):
+    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset, self.rd_cmd, self.wr_cmd)
+
+class RemotePCIDevice(PCIDevice):
+  _bulk_sent:int = 0
+  _bulk_recv:int = 0
+  _rpc_count:int = 0
+  _start_time:float = 0.0
+
+  @staticmethod
+  @functools.cache
+  def remote_sock(host:str, port:int) -> socket.socket:
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    sock.settimeout(getenv("REMOTE_TIMEOUT", 3))
+    sock.connect((host, port))
+    sock.settimeout(None)
+    if DEBUG >= 1 and RemotePCIDevice._start_time == 0.0:
+      RemotePCIDevice._start_time = time.perf_counter()
+      def _print_stats():
+        dt = time.perf_counter() - RemotePCIDevice._start_time
+        sent_mb, recv_mb = RemotePCIDevice._bulk_sent / 1e6, RemotePCIDevice._bulk_recv / 1e6
+        print(f"remote: sent {sent_mb:,.2f} MB ({sent_mb/dt:,.2f} MB/s), recv {recv_mb:,.2f} MB ({recv_mb/dt:,.2f} MB/s), "
+              f"{RemotePCIDevice._rpc_count:,} roundtrips in {dt:.2f}s")
+      atexit.register(_print_stats)
+    return sock
+
+  @staticmethod
+  @functools.cache
+  def remote_list(vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None) -> list[tuple[socket.socket, str]]:
+    payload = array.array('I', itertools.chain.from_iterable((m, d) for m, ds in devices for d in ds)).tobytes()
+    def q(r:str) -> list[tuple[socket.socket, str]]:
+      sock = RemotePCIDevice.remote_sock((host:=r.strip().split(":")[0]), (port:=int(r.strip().split(":")[1]) if ":" in r else 6667))
+      data_len, _, _, _ = RemotePCIDevice._rpc(sock, 0, RemoteCmd.PROBE, base_class or 0, len(payload), vendor, payload=payload)
+      return [(sock, f"remote:{host}:{port}:{d}") for d in RemotePCIDevice._recvall(sock, data_len).decode().split('\n')]
+    return flatten([q(r) for r in getenv("REMOTE", "").split(",") if r.strip()])
+
+  @staticmethod
+  def _recvall(sock:socket.socket, n:int) -> bytes:
+    data = b''
+    while len(data) < n and (chunk:=sock.recv(n - len(data))): data += chunk
+    if len(data) < n: raise RuntimeError("Connection closed")
+    return data
+
+  @staticmethod
+  def _rpc(sock:socket.socket, dev_id:int, cmd:int, *args:int, bar:int=0, readout_size:int=0, payload:bytes=b'', has_fd=False):
+    sock.sendall(struct.pack('<BIIQQQ', cmd, dev_id, bar, *(*args, 0, 0, 0)[:3]) + payload)
+    if has_fd:
+      msg, anc, _, _ = sock.recvmsg(17, socket.CMSG_LEN(4))
+      fd = struct.unpack('<i', anc[0][2][:4])[0]
+    else: msg, fd = RemotePCIDevice._recvall(sock, 17), None
+    if (resp:=struct.unpack('<BQQ', msg))[0] != 0:
+      raise RuntimeError(f"RPC failed: {RemotePCIDevice._recvall(sock, resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
+    RemotePCIDevice._rpc_count += 1
+    return (resp[1], resp[2]) + ((RemotePCIDevice._recvall(sock, readout_size) if readout_size > 0 else None),) + (fd,)
+
+  def __init__(self, devpref:str, pcibus:str, sock:socket.socket):
+    self.sock, self.pcibus, self.dev_id = sock, pcibus, int(pcibus.split(':')[-1]) if ':' in pcibus else 0
+    self.peer_group = sock.getpeername()[0]
+    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
+
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
+
+  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes:
+    RemotePCIDevice._bulk_recv += size
+    return unwrap(self._rpc(self.sock, self.dev_id, cmd, offset, size, bar=idx, readout_size=size)[2])
+  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes):
+    RemotePCIDevice._bulk_sent += len(data)
+    self.sock.sendall(struct.pack('<BIIQQQ', cmd, self.dev_id, idx, offset, len(data), 0) + data)
+
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    paddrs_len, handle, _, _ = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM, size, int(contiguous))
+    paddrs = list(struct.unpack(f'<{paddrs_len // 8}Q', self._recvall(self.sock, paddrs_len)))
+    return RemoteMMIOInterface(self, handle, size, fmt='B', rd_cmd=RemoteCmd.SYSMEM_READ, wr_cmd=RemoteCmd.SYSMEM_WRITE), paddrs
+
+  def reset(self): self._rpc(self.sock, self.dev_id, RemoteCmd.RESET)
+  def read_config(self, offset:int, size:int): return self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_READ, offset, size)[0]
+  def write_config(self, offset:int, value:int, size:int): self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_WRITE, offset, size, value)
+
+  @functools.cache
+  def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_BAR, bar=bar_idx)[:2]
+  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
+    return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt)
+  def resize_bar(self, bar_idx:int): self._rpc(self.sock, self.dev_id, RemoteCmd.RESIZE_BAR, bar=bar_idx)
+
+class APLRemotePCIDevice(RemotePCIDevice):
+  APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"
+
+  @classmethod
+  def ensure_app(cls):
+    commit = "c0d024f9ff0e1dc8fdf217f255da7101d91e8323"
+    app_name = f"TinyGPU_{commit}.zip"
+    if (_ensure_downloads_dir() / app_name).is_file() and os.path.exists(cls.APP_PATH): return
+    print("Downloading TinyGPU.app...")
+    with contextlib.suppress(RuntimeError): system("pkill -f TinyGPU")
+    system(f"ditto -xk {fetch(f'https://github.com/tinygrad/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications")
+    print(system(f"{cls.APP_PATH} install"))
+
+  def __init__(self, devpref:str, pcibus:str):
+    self.ensure_app()
+    sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    for i in range(100):
+      with contextlib.suppress(ConnectionRefusedError, FileNotFoundError):
+        sock.connect(sock_path)
+        break
+      if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+      time.sleep(0.05)
+    else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
+    super().__init__(devpref, "usb4", sock=sock)
+
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    mapped_size, _, _, fd = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM_FD, size, int(contiguous), has_fd=True)
+    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
+
+    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
+    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
+    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

extra/remote/serve.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import socket, struct, sys
-from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
+from tinygrad.runtime.support.system import PCIDevice, System
+from extra.remote.hcq1_remote import RemoteCmd
 from tinygrad.helpers import DEBUG, OSX
 
 def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)

extra/sqtt/roc.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
-from typing import Any, Generator
+from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.runtime.autogen import rocprof
@@ -37,7 +37,8 @@ class WaveExec(WaveSlot):
     insts_array = (struct*(len(self.insts)//sz)).from_buffer(self.insts)
     for inst in insts_array:
       inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst.category)
-      yield InstExec(inst_typ or "UNKNOWN", inst.pc.address, inst.stall, inst.duration, inst.time)
+      yield InstExec(inst_typ.replace("ROCPROFILER_THREAD_TRACE_DECODER_", "") if inst_typ else "UNKNOWN",
+                     inst.pc.address, inst.stall, inst.duration, inst.time)
 
 @dataclasses.dataclass(frozen=True)
 class OccEvent(WaveSlot):
@@ -127,31 +128,6 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[int, dict[int, Inst]])
     raise exc
   return ROCParseCtx
 
-def unpack_insts(w:WaveExec, pc_to_inst:dict[int, Inst]) -> dict:
-  columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
-  inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
-  # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
-  #           The idle time can be caused by:
-  #             * Arbiter loss
-  #             * Source or destination register dependency
-  #             * Instruction cache miss
-  # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
-  # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
-  prev_instr = w.begin_time
-  start_pc = None
-  rows:dict[int, dict[str, Any]] = {}
-  for pc, inst in pc_to_inst.items():
-    if start_pc is None: start_pc = pc
-    rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
-  for e in w.unpack_insts():
-    if not (row:=rows[e.pc]).get("type"): row["type"] = str(e.typ).split("_")[-1]
-    row["hit_count"] += 1
-    row["dur"] += e.dur
-    row["stall"] += e.stall
-    row["hits"]["rows"].append((row["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
-    prev_instr = max(prev_instr, e.time + e.dur)
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns}
-
 def main() -> None:
   from tabulate import tabulate
   from tinygrad.viz.serve import amd_decode
@@ -185,8 +161,8 @@ def main() -> None:
   for w in itertools.islice(waves, args.n):
     if w.wave_loc not in run_numbers: run_numbers[w.wave_loc] = itertools.count()
     print(f"{w.wave_loc} N:{next(run_numbers[w.wave_loc])} Total Cycles:{w.end_time-w.begin_time}")
-    table = unpack_insts(w, pc_to_inst)
-    print(tabulate([r[:len(table["cols"])] for r in table["rows"]], headers=table["cols"], tablefmt="github"))
+    rows = [(e.time, f"0x{e.pc:x}", pc_to_inst[e.pc], e.typ, e.dur, e.stall) for e in w.unpack_insts()]
+    print(tabulate(rows, headers=("Timestamp", "PC", "Instruction", "Type", "Duration", "Stall"), tablefmt="github"))
 
 if __name__ == "__main__":
   main()

extra/thunder/amd/moe_router_mfma.cpp

@@ -0,0 +1,80 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef ROUTER_M
+#define ROUTER_M 16384
+#endif
+#ifndef ROUTER_K
+#define ROUTER_K 2880
+#endif
+#ifndef ROUTER_E
+#define ROUTER_E 32
+#endif
+
+constexpr int BLOCK_M = 64;
+constexpr int BLOCK_K = 64;
+constexpr int NUM_WARPS = 4;
+constexpr int THREADS = NUM_WARPS * WARP_THREADS;
+
+using G = kittens::group<NUM_WARPS>;
+using XST = st_bf<BLOCK_M, BLOCK_K, st_16x32_s>;
+using WST = st_bf<ROUTER_E, BLOCK_K, st_16x32_s>;
+using XRT = rt_bf<16, BLOCK_K, row_l, rt_16x32_s>;
+using WRT = rt_bf<ROUTER_E, BLOCK_K, row_l, rt_16x32_s>;
+using CRT = rt_fl<16, ROUTER_E, col_l, rt_16x16_s>;
+
+static_assert(ROUTER_M % BLOCK_M == 0, "ROUTER_M must be divisible by 64");
+static_assert(ROUTER_K % BLOCK_K == 0, "ROUTER_K must be divisible by 64");
+static_assert(ROUTER_E == 32, "the small-N tile is specialized for 32 experts");
+
+extern "C" __global__ __launch_bounds__(THREADS, 4) void moe_router_mfma(
+    float *__restrict__ out, bf16 *__restrict__ x_ptr, bf16 *__restrict__ weight_ptr,
+    bf16 *__restrict__ bias) {
+  gl<bf16, 1, 1, ROUTER_M, ROUTER_K> X{x_ptr, nullptr, nullptr, nullptr, nullptr};
+  gl<bf16, 1, 1, ROUTER_E, ROUTER_K> W{weight_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  __shared__ XST Xs;
+  __shared__ WST Ws;
+
+  XRT xr;
+  WRT wr;
+  CRT accum;
+  zero(accum);
+
+  const int block_m = __builtin_amdgcn_workgroup_id_x();
+  const int warp_m = warpid();
+
+  #pragma unroll
+  for (int kk = 0; kk < ROUTER_K / BLOCK_K; kk++) {
+    G::load(Xs, X, {0, 0, block_m, kk});
+    G::load(Ws, W, {0, 0, 0, kk});
+    asm volatile("s_waitcnt vmcnt(0)");
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    __builtin_amdgcn_s_barrier();
+
+    load(xr, subtile_inplace<16, BLOCK_K>(Xs, {warp_m, 0}));
+    load(wr, subtile_inplace<ROUTER_E, BLOCK_K>(Ws, {0, 0}));
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    __builtin_amdgcn_s_setprio(1);
+    mma_ABt(accum, xr, wr, accum);
+    __builtin_amdgcn_s_setprio(0);
+    __builtin_amdgcn_sched_barrier(0);
+    __builtin_amdgcn_s_barrier();
+  }
+
+  // A 16x16 MFMA accumulator is column-layout: each lane owns four consecutive rows
+  // at one column. Store all 64x32 FP32 results directly; no padded or undersized output ABI.
+  const int lane = laneid();
+  const int row0 = block_m * BLOCK_M + warp_m * 16 + 4 * (lane / 16);
+  const int lane_col = lane % 16;
+  #pragma unroll
+  for (int j = 0; j < ROUTER_E / 16; j++) {
+    const int col = j * 16 + lane_col;
+    const float b = (float)bias[col];
+    const float vals[4] = {accum.tiles[0][j].data[0].x, accum.tiles[0][j].data[0].y,
+                           accum.tiles[0][j].data[1].x, accum.tiles[0][j].data[1].y};
+    #pragma unroll
+    for (int r = 0; r < 4; r++) out[(long long)(row0 + r) * ROUTER_E + col] = vals[r] + b;
+  }
+}

spec/tinyspec.tex

@@ -150,6 +150,18 @@ A value \op{Call} is void: its \op{Sink} body stores to output \op{Param}s bound
 
 \smallskip
 Assign is \op{Store} followed by \op{After}: write the value, then return the buffer with an ordering dependency.
+\op{After} orders consumers after its dependencies; it neither declares a write nor snapshots memory.
+In particular, \op{After}$(b, \op{Store}(d,v))$ returns $b$, not $v$, when $b$ and $d$ are disjoint.
+Views may share storage despite having different UOps. Differentiation follows the returned value:
+a matching unconditional full overwrite routes its gradient to the stored value; an unrelated write does not create a gradient path.
+Partial or uncertain aliased mutation gradients may be rejected.
+
+\smallskip
+\textbf{Tensor scheduling contract.} Within a lazy Tensor schedule, reads retain their assignment dependencies.
+A read must follow those dependencies and precede other writes that would destroy the required contents.
+Lowering must preserve these requirements until accesses are ordered, even when arguments share storage.
+Unsatisfiable requirements raise rather than read overwritten contents. This is a frontend requirement, not snapshot semantics for \op{After}.
+An executed \texttt{clone()} preserves data in fresh storage; \texttt{contiguous()} need not allocate.
 
 %% ============================================================
 \subsection*{{\color{elwyellow}Elementwise Ops} \normalfont\small--- all inputs same shape, output same shape, applied per-element}
@@ -236,7 +248,7 @@ Ternary & $(P, A, B)$
 \op{Custom} & (args\ldots) & fmt & Inject custom code string into generated source. \\
 \op{AtomicAdd} & (idx, val) & --- & Atomic read-modify-write: \texttt{buf[idx] += val}. \\[4pt]
 \op{CustomFunction} & (meta\ldots) & name & Opaque device function (e.g.\ HW decode). Via \op{Call}. \\
-\op{Program} & (linear, source, binary) & --- & Compiled kernel: instructions, source, and machine code. \\
+\op{Program} & (sink, \ldots) & metadata? & Kernel through compilation stages. \\
 \op{Source} & () & str & Human-readable rendered source code. \\
 \op{Binary} & () & bytes & Compiled machine code. \\
 \bottomrule
@@ -244,6 +256,11 @@ Ternary & $(P, A, B)$
 
 \smallskip
 These ops are not part of the core specification and are subject to change.
+\op{Program} contains a \op{Sink}, followed progressively by \op{Linear}, \op{Source}, and \op{Binary}.
+Access analysis derives reads and writes from the memory operands of \op{Load}/\op{Store}, resolving \op{Param}s through \op{Call} arguments.
+Compilation records these sets in \texttt{ProgramInfo.ins/outs} as zero-based argument slots; a read-modify-write belongs in both.
+Listing a parameter, returning an \op{After}, or declaring a write does not establish full initialization.
+Opaque code without computed access information is unsupported by assignment scheduling; its effects must not be guessed from its argument list.
 
 %% ============================================================
 \subsection*{Derived Properties}

test/amd/hw/test_ds.py

@@ -7,6 +7,29 @@ Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
 import unittest
 from test.amd.hw.helpers import *
 
+class TestDSSwizzle(unittest.TestCase):
+  def test_modes_and_overlapping_registers(self):
+    for offset in (0x041f, 0x401f, 0x7c1f, 0x00a0, 0x801b, 0xc020, 0xc420, 0xc021, 0xe000, 0xe010, 0xe01f):
+      for dst in (0, 1):
+        with self.subTest(offset=hex(offset), dst=dst):
+          st = run_program([
+            v_add_nc_u32_e32(v[0], 1, v[255]),
+            ds_swizzle_b32(vdst=v[dst], addr=v[0], offset0=offset & 255, offset1=offset >> 8),
+            s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+          ], n_lanes=32)
+          self.assertEqual(sorted(st.vgpr[i][dst] for i in range(32)), [6]*32 if offset == 0x00a0 else list(range(1, 33)))
+
+  def test_inactive_sources_and_destinations(self):
+    st = run_program([
+      v_add_nc_u32_e32(v[0], 1, v[255]),
+      v_mov_b32_e32(v[1], 99),
+      s_mov_b32(EXEC_LO, 0x55555555),
+      ds_swizzle_b32(vdst=v[1], addr=v[0], offset0=0x1f, offset1=4),
+      s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+      s_mov_b32(EXEC_LO, 0xffffffff),
+    ], n_lanes=32)
+    self.assertEqual([st.vgpr[i][1] for i in range(32)], [0, 99]*16)
+
 class TestDS2Addr(unittest.TestCase):
   """Tests for DS_*_2ADDR instructions."""
 

test/amd/hw/test_vop3p.py

@@ -457,6 +457,21 @@ class TestWMMAF16(unittest.TestCase):
         self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
         self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
 
+  def test_v_wmma_f16_16x16x16_f16_inline_zero_accumulator(self):
+    """V_WMMA_F16_16X16X16_F16 with the inline constant 0 as C: D = A @ B, whatever v[128:135] holds."""
+    instructions: list[Inst] = []
+    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
+    for i in range(16, 32):
+      instructions.append(v_mov_b32_e32(v[i], s[0]))
+    instructions.append(s_mov_b32(s[1], 0x57b057b0))  # packed f16 123.0, poison where a VGPR read of "128" would land
+    for i in range(128, 136):
+      instructions.append(v_mov_b32_e32(v[i], s[1]))
+    instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], 0))
+    st = run_program(instructions, n_lanes=32)
+    for lane in range(32):
+      for reg in range(8):
+        self.assertEqual(st.vgpr[lane][reg], 0x4c00, msg=f"v[{reg}] lane {lane}")
+
   def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
     """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
     instructions: list[Inst] = []

test/amd/test_emu2_pcode.py

@@ -17,6 +17,34 @@ def _srcs():
 class TestBasicParsing(unittest.TestCase):
   """Test basic pcode parsing for common instruction patterns."""
 
+  def test_c_style_blocks_and_array_access(self):
+    code = """
+      for (i = 0; i < 4; i+=2) {
+        if (mode == 0) {
+          out[i+0] = input[i+1];
+          out[i+1] = input[i+0];
+        } elsif (mode == 1) {
+          out[i+0] = 7;
+          out[i+1] = 8;
+        } else { // identity
+          out[i+0] = input[i+0];
+          out[i+1] = input[i+1];
+        }
+      }
+    """
+    for mode, expected in enumerate(([11, 10, 13, 12], [7, 8, 7, 8], [10, 11, 12, 13])):
+      with self.subTest(mode=mode):
+        result, _ = parse_pcode(code, {'mode': UOp.const(mode, dtypes.uint32)}, {'input': lambda i: i + 10})
+        self.assertEqual([result[f'out@{i}'].simplify().val for i in range(4)], expected)
+
+  def test_colon_concatenation(self):
+    result, _ = parse_pcode('offset = hi:lo;', {'hi': UOp.const(0x12, dtypes.uint8), 'lo': UOp.const(0x34, dtypes.uint8)})
+    self.assertEqual(result['offset'].simplify().val, 0x1234)
+
+  def test_unclosed_c_block(self):
+    with self.assertRaisesRegex(AssertionError, 'unclosed pcode block'):
+      parse_pcode('if (1) {\nvalue = 2;')
+
   def test_v_add_f32(self):
     """Test parsing V_ADD_F32 pcode."""
     _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())

test/amd/test_pdf.py

@@ -43,6 +43,16 @@ class TestPcodePDF(unittest.TestCase):
     self.assertEqual(pcode[('S_CMOVK_I32', 2)],
       "if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
 
+  def test_swizzle_spans_blocks_and_pages(self):
+    for arch in ('rdna3', 'rdna4'):
+      with self.subTest(arch=arch):
+        code = self.pcode[arch][('DS_SWIZZLE_B32', 53)]
+        self.assertIn('} elsif (offset >= 0xc000) {', code)
+        self.assertIn('thread_out[i+3]', code)
+        self.assertIn('xor_mask = offset[14:10];', code)
+        self.assertEqual(code.count('{'), code.count('}'))
+        self.assertTrue(code.endswith('\n}'))
+
   def test_pcode_no_examples(self):
     """Pseudocode should not contain example lines with '=>'."""
     for name in ARCHS:

test/amd/test_sqtt_profiler.py

@@ -1,43 +1,82 @@
 import unittest, contextlib
-from tinygrad import Device, Tensor, Context, TinyJit
+from tinygrad import Device, Tensor, Context, TinyJit, dtypes
+from test.helpers import is_hcq2_device
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Compiled, ProfileProgramEvent
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
+from tinygrad.renderer.amd.sqtt import decode, print_packets
+from tinygrad.renderer.amd.dsl import s
 
 @contextlib.contextmanager
 def save_sqtt():
   Device[Device.DEFAULT].synchronize()
   profile_start = len(Compiled.profile_events)
-  data = VizData()
-  yield data.ctxs
+  data = []
+  yield data
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)] +
-                          Compiled.profile_events[profile_start:])
-  data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
+  data[:] = [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)]+Compiled.profile_events[profile_start:]
+
+def map_sqtt(profile:list) -> list[dict]:
+  load_amd_counters(data:=VizData(), profile)
+  return [r for r in data.ctxs if r["name"].startswith("SQTT")]
+
+def custom_asm_cdna(A:UOp):
+  import tinygrad.runtime.autogen.amd.cdna.ins as cdna
+  WAVE_SIZE = 64
+  insts = [cdna.s_nop(0), cdna.s_mov_b32(s[0], 10)]
+  return custom_asm(A, insts+[cdna.s_endpgm()], WAVE_SIZE*2)
+
+def custom_asm_rdna(A:UOp):
+  import tinygrad.runtime.autogen.amd.rdna3.ins as rdna3
+  WAVE_SIZE = 32
+  insts = [rdna3.s_nop(0), rdna3.s_mov_b32(s[0], 10)]
+  return custom_asm(A, insts+[rdna3.s_endpgm()], WAVE_SIZE*2)
+
+def custom_asm(A, insts, num_threads) -> UOp:
+  return UOp(Ops.PROGRAM, src=(UOp.sink(A, UOp.special(num_threads, "lidx0"), arg=KernelInfo("asm")), \
+      UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS,arg=(x,dtypes.void)) for x in insts]))))
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
+    cls.arch = Device[Device.DEFAULT].arch
 
   def test_simple(self):
     t = Tensor.empty(1) + 1
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       linear = t.schedule_linear()
       run_linear(linear)
     fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), 1)
     self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}")
 
+  def test_asm(self):
+    t = Tensor.empty(1)
+    with save_sqtt() as data:
+      t.custom_kernel(fxn=custom_asm_cdna if self.arch == "gfx950" else custom_asm_rdna)[0].realize()
+    for event in data:
+      if not isinstance(event, ProfileSQTTEvent) or not event.itrace: continue
+      print(f"\n=== SE {event.se} ===")
+      print_packets(decode(event.blob))
+    from test.null.test_viz import write_files, run_cli
+    with write_files(profile=data) as files:
+      out = run_cli(*files, "-s", "asm SQTT SE:0 PKTS", json_fmt=False)[0]["out"]
+    print(out)
+
   def test_multiple_runs(self):
     t = Tensor.empty(1) + 1
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       linear = t.schedule_linear()
       for _ in range(N:=3): run_linear(linear)
     fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), N)
     for i in range(1, N):
       self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}")
@@ -45,8 +84,9 @@ class TestSQTTProfiler(unittest.TestCase):
   def test_multiple_kernels(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
     linear = t.schedule_linear()
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       run_linear(linear)
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), len(linear.src))
     for i,call in enumerate(linear.src):
       fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
@@ -55,8 +95,9 @@ class TestSQTTProfiler(unittest.TestCase):
   def test_multiple_kernels_lower(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
     linear = t.schedule_linear()
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       run_linear(linear)
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), len(linear.src))
     for i,call in enumerate(linear.src):
       fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
@@ -66,21 +107,22 @@ class TestSQTTProfiler(unittest.TestCase):
     @TinyJit
     def f(a): return a + 1
     t = Tensor.empty(1)
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       for _ in range(N:=5):
         f(t).realize()
+    sqtt = map_sqtt(data)
     self.assertEqual(len(sqtt), N)
     kernel_name = sqtt[0]["name"]
-    for i,s in enumerate(sqtt[1:], start=1): self.assertEqual(s["name"], f"{kernel_name} n{i+1}")
+    for i,e in enumerate(sqtt[1:], start=1): self.assertEqual(e["name"], f"{kernel_name} n{i+1}")
 
-  # TODO: can we trace SQTT for graphed kernels?
-  def test_jit_graph(self, kernel_count=3*1):
+  def test_jit_graph(self, kernel_count=3*(5 if is_hcq2_device() else 1)): # hcq2 traces the graphed kernels too
     @TinyJit
     def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
     t = Tensor.empty(32)
-    with save_sqtt() as sqtt:
+    with save_sqtt() as data:
       for _ in range(5):
         f(t).realize()
+    sqtt = map_sqtt(data)
     names = [s["name"] for s in sqtt]
     k0, k1, k2 = names[:3]
     for i in range(3, len(sqtt), 3):

test/backend/test_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException, is_hcq2_device
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -153,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(2)
+      assert_kernel_count(3 if is_hcq2_device() else 2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu
+from test.helpers import assert_kernel_count, needs_second_gpu, is_hcq2_device
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -22,7 +22,7 @@ class TestAssign(unittest.TestCase):
     assert ba1 == ba2 and ba1 != bb1
     np.testing.assert_allclose(a.numpy(), (np.arange(N*N)*2).reshape((N,N)))
 
-  def test_assign_zeros_good(self):
+  def test_assign_keeps_identical_tensor(self):
     a = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
     b = Tensor.zeros(10,10).contiguous()
@@ -30,7 +30,7 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), 0)
 
   @unittest.skip("TODO: this often crashes in CI")
-  def test_assign_zeros(self):
+  def test_assign_keeps_earlier_identical_tensor(self):
     a = Tensor.zeros(10,10).contiguous()
     b = Tensor.zeros(10,10).contiguous()
     a.assign(Tensor.ones(10,10))
@@ -40,10 +40,10 @@ class TestAssign(unittest.TestCase):
   def test_assign_copy(self):
     a = Tensor([1.,2,3], device="PYTHON")
     c = Tensor.empty(3).assign(a.to(None))
-    # it should copy into the empty buffer
+    # The creation copy has its own storage, independent of the assignment destination.
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(3 if is_hcq2_device() else 2)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -114,15 +114,6 @@ class TestAssign(unittest.TestCase):
     x.assign(x + 1)
     assert [y0.item(), y1.item(), y2.item(), x.item()] == [0.0, 1.0, 2.0, 3.0]
 
-  def test_assign_add_jit(self):
-    @TinyJit
-    def f(x):
-      x += 1
-      x.realize()
-    x = Tensor([0])
-    for _ in range(5): f(x)
-    assert x.item() == 5
-
   def test_assign_add_jit_other(self):
     @TinyJit
     def f(x):
@@ -180,21 +171,20 @@ class TestAssign(unittest.TestCase):
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
     self.assertEqual((a + b).item(), 3)
 
-  def test_assign_diamond_cycle(self):
-    # NOTE: should *not* raise AssertionError from numpy
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.ones(4).contiguous().realize()
-      times_a = a*3
-      a.assign(Tensor.full((4,), 2.).contiguous())
-      new = a + (times_a-1)
+  def test_assign_diamond(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.).contiguous())
+    new = a + (times_a-1)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       np.testing.assert_allclose(new.numpy(), 4)
 
-  def test_assign_diamond_contiguous_cycle(self):
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.ones(4).contiguous().realize()
-      times_a = a*3
-      a.assign(Tensor.full((4,), 2.))
-      new = a.contiguous() + times_a-1
+  def test_assign_diamond_contiguous(self):
+    a = Tensor.ones(4).contiguous().realize()
+    times_a = a*3
+    a.assign(Tensor.full((4,), 2.))
+    new = a.contiguous() + times_a-1
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       np.testing.assert_allclose(new.numpy(), 4)
 
   def test_assign_diamond_possible(self):
@@ -267,13 +257,12 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(b1.numpy(), 608)
 
   def test_crossunder_assign(self):
-    # NOTE: should *not* raise AssertionError from numpy
-    with self.assertRaisesRegex(RuntimeError, "cycle"):
-      a = Tensor.full((4,), 2).contiguous().realize()
-      b = Tensor.full((4,), 3).contiguous().realize()
-      c = a+9
-      a += b
-      b += c
+    a = Tensor.full((4,), 2).contiguous().realize()
+    b = Tensor.full((4,), 3).contiguous().realize()
+    c = a+9
+    a += b
+    b += c
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
       Tensor.realize(a,b)
       np.testing.assert_allclose(a.numpy(), 2+3)
       np.testing.assert_allclose(b.numpy(), 3+2+9)
@@ -356,49 +345,17 @@ class TestAssign(unittest.TestCase):
     # permute and base are the same buffer
     assert ba1 == ba2 and ba1 != bb1
 
-  def test_post_permuted_assignment(self):
-    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
-    b = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
-    a.realize()
-    b.realize()
-    #GlobalCounters.cache = []
-    ba1 = a.uop.base.realized # noqa: F841
-    bb1 = b.uop.base.realized # noqa: F841
-    a.assign(a.permute(1,0) + b)   # this should not work!
-    a.realize()
-    ba2 = a.uop.base.realized # noqa: F841
-    # NOTE: don't test that it's assigned
-    #assert ba1 == ba2 and ba1 != bb1
-    np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
-
-  def test_post_permuted_assignment_alt(self):
+  def _assign_view_of_self(self, view):
     a = Tensor.arange(N*N).reshape(N,N).clone().realize()
     b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.T+b).numpy()
-    a.assign(a.T+b)
+    new_a = (view(a)+b).numpy()
+    a.assign(view(a)+b)
     np.testing.assert_allclose(a.numpy(), new_a)
 
-  def test_post_flipped_assignment(self):
-    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.flip(0)+b).numpy()
-    a.assign(a.flip(0)+b)
-    np.testing.assert_allclose(a.numpy(), new_a)
-
-  def test_post_flipped_assignment_axis1(self):
-    a = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N,N).clone().realize()
-    new_a = (a.flip(1)+b).numpy()
-    a.assign(a.flip(1)+b)
-    np.testing.assert_allclose(a.numpy(), new_a)
-
-  def test_post_reshape_assignment_fine(self):
-    a = Tensor.arange(N*N).reshape(N, N).clone().realize()
-    b = Tensor.arange(N*N).reshape(N, N).clone().realize()
-    rhs = a.reshape(-1).reshape(N, N)
-    new_a = (rhs+b).numpy()
-    a.assign(rhs+b)  # self-assign with reshape view is fine
-    np.testing.assert_allclose(a.numpy(), new_a)
+  def test_post_permuted_assignment(self): self._assign_view_of_self(lambda a: a.T)
+  def test_post_flipped_assignment(self): self._assign_view_of_self(lambda a: a.flip(0))
+  def test_post_flipped_assignment_axis1(self): self._assign_view_of_self(lambda a: a.flip(1))
+  def test_post_reshape_assignment(self): self._assign_view_of_self(lambda a: a.reshape(-1).reshape(N,N))
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
@@ -421,14 +378,6 @@ class TestAssign(unittest.TestCase):
 
   # NOTE: if the assign target is read/write in a single kernel, it should be contiguous
 
-  def test_permuted_assignment_correct(self):
-    a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
-    b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
-    a = a.permute(1, 0)
-    new_val = a + b
-    a.assign(new_val)
-    np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
-
   def test_permuted_reduceop_child_dual_use(self):
     a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
     b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
@@ -526,34 +475,34 @@ class TestAssign(unittest.TestCase):
     a[2:5] = [1, 2, 3]
     np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.])
 
+  # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
+  REVERSED = [0x40800000, 0x40400000, 0x40000000, 0x3f800000]
+
   def test_assign_bitcast(self):
-    # assign to a bitcast view should modify the underlying buffer
     a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    # IEEE 754: 1.0f = 0x3f800000, 2.0f = 0x40000000, 3.0f = 0x40400000, 4.0f = 0x40800000
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
-    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
-    # double bitcast
-    b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.int32)).realize()
-    np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
-    # shrink then bitcast
-    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
-    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
-    # without .realize()
-    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
+    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
 
+  def test_assign_bitcast_unrealized(self):
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor(self.REVERSED, dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+
+  def test_assign_double_bitcast(self):
+    b = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    b.bitcast(dtypes.uint32).bitcast(dtypes.int32).assign(Tensor(self.REVERSED, dtype=dtypes.int32)).realize()
+    np.testing.assert_allclose(b.numpy(), [4.0, 3.0, 2.0, 1.0])
+
+  def test_assign_shrink_then_bitcast(self):
+    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    c[0:2].bitcast(dtypes.uint32).assign(Tensor(self.REVERSED[:2], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+
   def test_assign_bitcast_different_size(self):
-    # assign to a shape-changing bitcast view (only works on DISK currently)
+    # assign to a shape-changing bitcast view
     a = Tensor([0]*8, dtype=dtypes.uint8).realize()
     a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    try:
-      np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
-    except AssertionError:
-      # TODO: broken now
-      np.testing.assert_equal(a.numpy(), [0]*8)
+    np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
 
   def test_assign_dtype_mismatch(self):
     # assign should not implicitly cast dtypes - this can lose precision
@@ -562,13 +511,6 @@ class TestAssign(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
       a.assign(b)
 
-  def test_assign_dtype_mismatch_int64_to_float32(self):
-    # int64 -> float32 loses precision for large values, should not be implicit
-    a = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize()
-    b = Tensor([16777217], dtype=dtypes.int64)  # 2^24 + 1, not exactly representable in float32
-    with self.assertRaisesRegex(RuntimeError, "assign dtype mismatch"):
-      a.assign(b)
-
   def test_assign_shape_broadcast(self):
     # shape broadcasting should work when dtypes match
     a = Tensor.zeros(3, 5, dtype=dtypes.float32).contiguous().realize()
@@ -677,7 +619,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(2)  # TODO: first copy is dead, could be 1
+    assert_kernel_count(6 if is_hcq2_device() else 4)  # TODO: first copy is dead
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -687,9 +629,17 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(1)
+    assert_kernel_count(3 if is_hcq2_device() else 2)
     self.assertEqual(base.tolist(), [1,4,3])
 
+  def test_assign_temporary_copy_reshape(self):
+    a = Tensor([[1., 2], [3, 4]], device="PYTHON")
+    c = Tensor.empty(2, 2).assign(a.to(None))
+    GlobalCounters.reset()
+    c.realize()
+    assert_kernel_count(3 if is_hcq2_device() else 2)
+    self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
+
 class TestAssignOrdering(unittest.TestCase):
   """Tests for complex assign orderings that could differ between lazy and eager execution.
 
@@ -878,17 +828,22 @@ class TestAssignOrdering(unittest.TestCase):
       b_np *= 0.9
     np.testing.assert_allclose(param.item(), p_np, atol=1e-5)
 
+  def test_after_store_to_other_buffer(self):
+    x, state = Tensor([2.]).realize(), Tensor([0.]).realize()
+    ordered = Tensor(x.uop.after(state.uop.store(x.uop * 3)))
+    self.assertEqual((ordered + x).tolist(), [4.])
+    self.assertEqual(state.tolist(), [6.])
+    self.assertEqual(x.tolist(), [2.])
+
   def test_war_reader_already_depends_on_write(self):
     x = Tensor([1.0]).contiguous().realize()
     y = Tensor([2.0]).contiguous().realize()
-    x_expr = x + 10
+    x_expr = x + 10          # 11, x is read here, before the assign
     x.assign(x * 2)
     y.assign(y + x)
     z = y + x_expr
-    Tensor.realize(x, y, z)
-    # TODO: z should be 15: x_expr means 11 (x captured at build time), but the read is fused past the assign and
-    # sees the new bytes. once stale readers are scheduled before the overwrite, update this to 15
-    np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      Tensor.realize(x, y, z)
 
   def test_war_multi_read_then_assign(self):
     devices = ("CPU:0", "CPU:1")
@@ -909,6 +864,139 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(buf.sum().realize().item(), 6.0)
 
 # TODO: assigns into views of unrealized non-BUFFER bases are silently dropped
+  def test_read_before_two_assigns(self):
+    g = Tensor.full((2,), 4.0).realize()
+    before = g + 1                                     # 5
+    g.assign(0.0)
+    g.assign(g + 4)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+      np.testing.assert_allclose((before + g).numpy(), 9)
+
+  def test_read_between_two_assigns(self):
+    a = Tensor.ones(4).realize()
+    b = Tensor.full((4,), 10.).realize()
+    a.assign(b + 1)                                    # a == 11
+    v1 = a * 3                                         # reads 11 -> 33
+    a.assign(b + 100)                                  # a == 110
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      (a + v1).numpy()
+
+  def test_two_reads_between_three_assigns(self):
+    a = Tensor.zeros(4).realize()
+    first = a + 100
+    a.assign(Tensor([1., 2., 0., 0.]))
+    second = a + 0
+    a.assign(a + 10)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, raises
+      np.testing.assert_allclose((first + second + a).numpy(), [112, 114, 110, 110])
+
+  def test_read_before_slice_assign(self):
+    a = Tensor.ones(4).realize()
+    before = a * 3
+    a[0:2].assign(Tensor.full((2,), 2.))
+    out = (a + (before - 1)).numpy()
+    try:
+      np.testing.assert_allclose(out, [4, 4, 3, 3])
+    except AssertionError:
+      # TODO: broken now, before reads the two assigned elements after the assign
+      np.testing.assert_allclose(out, [7, 7, 3, 3])
+
+  def test_read_before_assign_survives_a_realize(self):
+    a = Tensor.ones(4).realize()
+    before = a * 3
+    a.assign(Tensor.full((4,), 5.))
+    a.realize()
+    out = before.numpy()
+    try:
+      np.testing.assert_allclose(out, 3)
+    except AssertionError:
+      # TODO: broken now, before is computed again from the assigned value
+      np.testing.assert_allclose(out, 15)
+
+  def test_loss_read_after_step_is_the_pre_step_loss(self):
+    from tinygrad import nn
+    w = Tensor([2.]).contiguous().realize()
+    x = Tensor([3.]).realize()
+    opt = nn.optim.SGD([w], lr=0.1)
+    with Context(TRAINING=1):
+      loss = (w*x).sum()                               # 6.0
+      loss.backward()
+      opt.step()                                       # w becomes 1.7
+    out = loss.item()
+    try:
+      self.assertAlmostEqual(out, 6.0, places=5)
+    except AssertionError:
+      # TODO: broken now, loss is computed again from the updated weight
+      self.assertAlmostEqual(out, 5.1, places=5)
+
+  def test_rand_realized_out_of_order(self):
+    Tensor.manual_seed(1)
+    r = [Tensor.rand(4) for _ in range(4)]
+    r[3].realize()
+    out_of_order = r[0].numpy()
+    Tensor.manual_seed(1)
+    in_order = [Tensor.rand(4).numpy() for _ in range(4)]
+    try:
+      np.testing.assert_equal(out_of_order, in_order[0])
+    except AssertionError:
+      # TODO: broken now, r[0] returns the fourth set of numbers
+      np.testing.assert_equal(out_of_order, in_order[3])
+
+  def test_batchnorm_stats_are_realized(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).realize()
+    with Context(TRAINING=1): bn(x).realize()
+    try:
+      self.assertTrue(bn.running_mean.uop.base.is_realized)
+    except AssertionError:
+      # TODO: broken now, the stat update is never run because nothing reads it
+      self.assertFalse(bn.running_mean.uop.base.is_realized)
+
+  def test_batchnorm_under_jit_counts_every_call(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(8, 4, 2, 2).realize()
+    @TinyJit
+    def step(t):
+      with Context(TRAINING=1): return bn(t).sum().realize()
+    for _ in range(4): step(x)
+    out = bn.num_batches_tracked.item()
+    try:
+      self.assertEqual(out, 4)
+    except AssertionError:
+      # TODO: broken now, only the calls whose stat update happened to be captured are counted
+      self.assertEqual(out, 2)
+
+  def test_assign_from_unrealized_tensor_does_not_alias(self):
+    a = Tensor.full((4,), 7.).realize()
+    b = Tensor.ones(4) * 1
+    b.assign(a)
+    b.assign(Tensor.zeros(4))
+    b.realize()
+    self.assertListEqual(a.tolist(), [7., 7., 7., 7.])
+
+  def test_assign_to_function_output(self):
+    from tinygrad import function
+    @function
+    def f(x:Tensor) -> Tensor: return x*2
+    out = f(Tensor.ones(4).realize())
+    out.assign(Tensor.full((4,), 9.).realize())
+    self.assertListEqual(out.tolist(), [9., 9., 9., 9.])
+
+  def test_nested_function_assign(self):
+    from tinygrad import function
+    @function
+    def inner(x:Tensor) -> Tensor:
+      x.assign(x+1)
+      return x*2
+    @function
+    def outer(x:Tensor) -> Tensor:
+      y = inner(x)
+      x.assign(x+1)
+      return y+x
+    a = Tensor([1.]).realize()
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      outer(a).item()
+
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
@@ -926,11 +1014,7 @@ class TestAssignToUnrealizedView(unittest.TestCase):
     c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
     self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(c.tolist(), [[1,1],[2,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(c.tolist(), [[1,3],[2,4]])
+    self.assertEqual(c.tolist(), [[1,1],[2,1]])
 
   def test_contiguous_backward(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
@@ -959,11 +1043,7 @@ class TestAssignToUnrealizedView(unittest.TestCase):
     d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
     self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
     d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(d.tolist(), [[1,1],[2,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(d.tolist(), [[1,3],[2,4]])
+    self.assertEqual(d.tolist(), [[1,1],[2,1]])
 
   def test_alu(self):
     a = Tensor([1,2,3,4]).contiguous().realize()
@@ -1009,6 +1089,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
+  def test_detach_assignment_preserves_earlier_update(self):
+    x = Tensor([1., 2.]).detach()
+    state = Tensor([0., 0.]).detach()
+    state.assign(state + x * 2)
+    result = state + 1
+    x.assign(x + 1).realize(state, result)
+    self.assertEqual(x.tolist(), [2., 3.])
+    self.assertEqual(state.tolist(), [2., 4.])
+    self.assertEqual(result.tolist(), [3., 5.])
+
 class TestPartialAssignToSharedBuffer(unittest.TestCase):
   def test_five_slices(self):
     big = Tensor.zeros(50).contiguous().realize()
@@ -1040,7 +1130,6 @@ class TestPartialAssignToSharedBuffer(unittest.TestCase):
     for v, s in zip(views, shapes):
       np.testing.assert_allclose(v.numpy(), np.ones(s))
 
-
 class TestAfterCachePatterns(unittest.TestCase):
   def test_double_store_after(self):
     a = Tensor.zeros(10).contiguous()
@@ -1071,14 +1160,6 @@ class TestAfterCachePatterns(unittest.TestCase):
     np.testing.assert_array_equal(head.numpy(), [3])
     np.testing.assert_array_equal(full.numpy(), [1, 2])
 
-class TestBatchNormRunningStats(unittest.TestCase):
-  @unittest.expectedFailure  # TODO: nothing reads the stat update so it is never scheduled, and the chain grows every step
-  def test_running_stats_are_realized(self):
-    from tinygrad import nn
-    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).contiguous().realize()
-    with Context(TRAINING=1): bn(x).realize()
-    self.assertTrue(bn.running_mean.uop.base.is_realized)
-
 class TestMultiAssign(unittest.TestCase):
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
 
@@ -1125,12 +1206,15 @@ class TestMultiAssign(unittest.TestCase):
     out[:, 2:3].assign(ones).realize()
     self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
 
-  @unittest.expectedFailure
   def test_multi_assign_piece_unrealized(self):
     out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
     ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
     out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+    try:
+      self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+    except AssertionError:
+      # TODO: broken now, the write is dropped
+      self.assertListEqual(out.tolist(), [[0,0,0,0], [0,0,0,0], [0,0,0,0], [0,0,0,0]])
 
   def test_multi_assign_var_offset(self):
     out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
@@ -1154,5 +1238,6 @@ class TestMultiAssign(unittest.TestCase):
       GlobalCounters.reset()
       f(out, vi.bind(i))
     self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
+
 if __name__ == "__main__":
   unittest.main()

test/backend/test_custom_kernel.py

@@ -105,6 +105,47 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 # **** tests ****
 
 class TestCustomKernel(unittest.TestCase):
+  def test_readonly_after_args(self):
+    for chained in (False, True):
+      for corealize in (False, True):
+        with self.subTest(chained=chained, corealize=corealize):
+          x = Tensor([2.]).realize()
+          a, x1 = Tensor.empty(1).custom_kernel(x, fxn=custom_add_one_kernel)
+          b, x2 = Tensor.empty(1).custom_kernel(x1 if chained else x, fxn=custom_add_one_kernel)
+          if corealize:
+            y = x1 + b
+            Tensor.realize(y, x2)
+            self.assertEqual(y.tolist(), [5.])
+          else:
+            self.assertEqual((x1 + x2).tolist(), [4.])
+          self.assertEqual(a.tolist(), [3.])
+          self.assertEqual(b.tolist(), [3.])
+          self.assertEqual(x.tolist(), [2.])
+
+  def test_aliased_args_different_sizes(self):
+    def kernel(out:UOp, a:UOp, b:UOp):
+      i = UOp.range(4, 0)
+      return out[i].store(a[i] + b[0]).end(i).sink(arg=KernelInfo(name="aliased_sizes"))
+    x = Tensor([1., 2., 3., 4.]).realize()
+    out = Tensor.empty(4).custom_kernel(x, x[:1], fxn=kernel)[0]
+    self.assertEqual(out.tolist(), [2., 3., 4., 5.])
+
+  def test_unindexed_access_before_assign(self):
+    def kernel(out:UOp, x:UOp): return out.store(x + 1).sink(arg=KernelInfo(name="unindexed"))
+    x = Tensor([2.]).realize()
+    y = Tensor.empty(1).custom_kernel(x, fxn=kernel)[0]
+    x.assign(x * 2)
+    Tensor.realize(x, y)
+    self.assertEqual(x.tolist(), [4.])
+    self.assertEqual(y.tolist(), [3.])
+
+  def test_readonly_after_does_not_hide_write(self):
+    x = Tensor([2.]).realize()
+    _, before = Tensor.empty(1).custom_kernel(x, fxn=custom_add_one_kernel)
+    x.assign(x * 2)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      (before + x).realize()
+
   def test_empty(self):
     a = Tensor.empty(1)
     a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]

test/backend/test_dtype.py

@@ -9,7 +9,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype, min_normal
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -19,7 +19,8 @@ settings.load_profile("my_profile")
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
 
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
+  emulatable = dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)
+  dts = [v for v in dict.fromkeys(DTYPES_DICT.values()) if v != dtype and (v in supported_dtypes or v in emulatable)]
   if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
     return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
   if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
@@ -71,14 +72,14 @@ class TestDType(unittest.TestCase):
     self.assertEqual(a.dtype, self.DTYPE)
     _test_to_np(a, _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
 
-  def test_casts_to(self):
-    for dtype in get_available_cast_dtypes(self.DTYPE):
-      _test_cast(Tensor(self.DATA, dtype=dtype), self.DTYPE)
-
   def test_casts_from(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       _test_cast(Tensor(self.DATA, dtype=self.DTYPE), dtype)
 
+  def test_const_kernel(self):
+    if not get_available_cast_dtypes(self.DTYPE): raise unittest.SkipTest("dtype does not run here")
+    _assert_eq(Tensor.ones((4,4), dtype=self.DTYPE).clone(), self.DTYPE, np.ones((4,4)))
+
   def test_same_size_ops(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       if dtype.itemsize == self.DTYPE.itemsize:
@@ -89,10 +90,10 @@ class TestDType(unittest.TestCase):
       if dtype.itemsize > self.DTYPE.itemsize:
         _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
 
-  def test_upcast_to_ops(self):
+  def test_downcast_ops(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       if dtype.itemsize < self.DTYPE.itemsize:
-        _test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
+        _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
 
   def test_bitcast(self):
     if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
@@ -112,12 +113,7 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
   target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
   if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
   _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
-  _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
-  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
   _assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
-  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
-  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
-  _assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
 
 class TestFp8sConversions(unittest.TestCase):
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -288,14 +284,10 @@ class TestUint8DType(TestDType):
     _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
 
 class TestBitCast(unittest.TestCase):
-  @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
-  def test_shape_change_bitcast(self, dt1, dt2):
-    data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
-    a = Tensor(data, dtype=dt1)
-    expected = _to_torch_storage(a).view(_to_torch_dtype(dt2))
-    if dt2 in dtypes.fp8s:
-      expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
-    _test_op(lambda: a.bitcast(dt2), dt2, expected.tolist())
+  def test_shape_change_bitcast(self):
+    for dt1, dt2 in [(dtypes.uint8, dtypes.int64), (dtypes.int64, dtypes.uint8)]:
+      a = Tensor(rand_for_dtype(dt1, 32).reshape(2, 2, 8), dtype=dt1)
+      _test_op(lambda: a.bitcast(dt2), dt2, _to_torch_storage(a).view(_to_torch_dtype(dt2)).tolist())
 
   def test_shape_change_bitcast_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -401,6 +393,9 @@ class TestEmulatedFp8e5m2(TestFp8e5m2):
   @classmethod
   def tearDownClass(cls): cls.stack.close()
 
+class TestFp8e4m3fnuz(TestDType): DTYPE = dtypes.fp8e4m3fnuz
+class TestFp8e5m2fnuz(TestDType): DTYPE = dtypes.fp8e5m2fnuz
+
 class TestImplicitFunctionTypeChange(unittest.TestCase):
   def test_functions(self):
     result = []

test/backend/test_jit.py

@@ -2,11 +2,11 @@
 import unittest
 import numpy as np
 
-from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
+from test.helpers import is_hcq2_device, assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
-from tinygrad.helpers import JIT, DEV, GlobalCounters, HCQ2
+from tinygrad.helpers import JIT, DEV, GlobalCounters
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer
 
@@ -222,7 +222,7 @@ class TestJitPrune(unittest.TestCase):
     assert_jit_cache_len(w2_prune, 1)
 
 class TestJitFree(unittest.TestCase):
-  @unittest.skipIf(HCQ2, "hcq2 keeps refs to intermediate buffers")
+  @unittest.skipIf(is_hcq2_device(), "hcq2 keeps refs to intermediate buffers")
   def test_free_intermediates(self):
     ext_tensor = Tensor([1,24,23,45,1])
     @TinyJit

test/backend/test_multitensor.py

@@ -1,16 +1,12 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
-from tinygrad.helpers import getenv, prod, Context
+from tinygrad.helpers import prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, lower_and_compile, pm_beam
 import numpy as np
-from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
 
-settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
-settings.load_profile("my_profile")
-
 d0 = f"{Device.DEFAULT}:0"
 d1 = f"{Device.DEFAULT}:1"
 d2 = f"{Device.DEFAULT}:2"
@@ -76,7 +72,7 @@ class TestMultiTensor(unittest.TestCase):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
     out = (X + X)
-    linear = compile_linear(out.schedule_linear())
+    linear = lower_and_compile(out.schedule_linear())
     uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
     run_linear(linear)
     self.assertEqual(len(set(uops)), 1, "function was relinearized")
@@ -129,17 +125,21 @@ class TestMultiTensor(unittest.TestCase):
       run_linear(linear, var_vals)
       np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])
 
-  @given(strat.sampled_from((devices_2, devices_3)),
-         strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)),
-         strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)))
-  def test_simple_reduce(self, devices, rop, shard_axis, reduce_axis):
-    N = 4 * len(devices)
-    X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis)
-    n = X.numpy()
-    f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop]
-    fX = f(X)
-    fn = f(n)
-    np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
+  def test_simple_reduce(self):
+    for devices, rop, shard_axis, reduce_axis in [
+      (devices_2, Ops.ADD, None, None), (devices_2, Ops.ADD, 0, 0), (devices_2, Ops.ADD, 0, 1),
+      (devices_2, Ops.ADD, 1, 0), (devices_2, Ops.ADD, 1, 1),
+      (devices_3, Ops.ADD, 0, 0), (devices_3, Ops.ADD, 1, 0),
+      (devices_2, Ops.MUL, 0, 1), (devices_2, Ops.MUL, 1, 1), (devices_3, Ops.MUL, 0, 0),
+      (devices_2, Ops.MAX, 0, 1), (devices_3, Ops.MAX, 1, 0)]:
+      with self.subTest(devices=len(devices), op=rop.name, shard_axis=shard_axis, reduce_axis=reduce_axis):
+        N = 4 * len(devices)
+        X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis)
+        n = X.numpy()
+        f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop]
+        fX = f(X)
+        fn = f(n)
+        np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
 
   def test_stack(self):
     X = Tensor.rand(4, 4).shard_(devices_2, 0)
@@ -176,21 +176,21 @@ class TestMultiTensor(unittest.TestCase):
   def test_allreduce_naive_jit(self):
     with Context(RING=0):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
   def test_allreduce_ring_jit(self):
     with Context(RING=2):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
   def test_allreduce_all2all_jit(self):
     with Context(ALL2ALL=2):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
@@ -212,7 +212,7 @@ class TestMultiTensor(unittest.TestCase):
 
   def test_fuzz_allreduce(self):
     random.seed(41)
-    for it in range(2):
+    for it in range(1):
       for n in range(2, 4+1):
         shape = tuple([(n if i == 0 else 1) * random.randint(1, 10) for i in range(random.randint(1, 4))])
         t = Tensor.rand(shape).shard_(tuple([d0, d1, d2, d3][:n]), 0)
@@ -445,6 +445,7 @@ class TestMultiBufferView(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class Test2DShard(unittest.TestCase):
+  @needs_second_gpu
   def setUp(self):
     self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
     self.rng = UOp.range(4, -1, AxisType.DEVICE)
@@ -460,6 +461,15 @@ class Test2DShard(unittest.TestCase):
     out = t.contiguous().realize()
     np.testing.assert_equal(out.numpy(), ref.numpy())
 
+  def test_2d_shard_clone(self):
+    ref = Tensor.arange(16).reshape(4, 4).realize()
+    t = self._shard_2d(ref)
+    out = t.clone().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy())
+    out.assign(out + 1).realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
+    np.testing.assert_equal(t.numpy(), ref.numpy())
+
   def test_2d_shard_elementwise(self):
     ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
     t = self._shard_2d(ref)
@@ -513,7 +523,8 @@ class TestMultiTransformer(unittest.TestCase):
       else: v.shard_(device, axis=None)
 
     last_tok = 0
-    for i in range(5):
+    # i=0: bypasses jit, i=1: jit warmup, i=2: capture and run, i>=3: re-execute jit with new start_pos (catches stale bindings)
+    for i in range(4):
       real_tok = real_model(Tensor([[last_tok]], device=Device.DEFAULT), i).item()
       shard_tok = shard_model(Tensor([[last_tok]], device=device), i).item()
 

test/backend/test_ops.py

@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -816,8 +815,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
-  # TODO: x86 PARAM dtype fails SPEC=2
-  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -825,6 +822,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor&tor, lambda: ten&ten, forward_only=True)
     helper_test_op([], lambda: tor&0x1337, lambda: ten&0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337&tor, lambda: 0x1337&ten, forward_only=True)
+    helper_test_op([], lambda: (tor&12)&tor, lambda: (ten&12)&ten, forward_only=True)
 
     data = [[True, True, False, False], [True, False, True, False]]
     tor0, tor1 = torch.tensor(data[0], dtype=torch.bool),  torch.tensor(data[1], dtype=torch.bool)
@@ -3109,6 +3107,13 @@ class TestOps(unittest.TestCase):
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])
 
+  def test_gather_bool_index(self):
+    helper_test_op(None, lambda x,y: x.gather(dim=0, index=y.bool().long()),
+                         lambda x,y: x.gather(dim=0, index=y.cast(dtypes.bool).cast(dtypes.int)),
+                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
+    helper_test_op(None, lambda x,y: x[y.bool().long()], lambda x,y: x[y.cast(dtypes.bool).cast(dtypes.int)],
+                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
+
   def test_scatter(self):
     b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
     a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)

test/backend/test_profiler.py

@@ -3,6 +3,7 @@ from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 
@@ -34,7 +35,18 @@ def helper_profile_filter_device(profile, device:str):
   assert len(dev_events) == 1, "only one device registration event is expected"
   return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
 
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT], (HCQCompiled, HCQ2Compiled)) or Device.DEFAULT == "METAL", "Dev not supported")
+class TestSimpleProfiler(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT == "CPU", "fails in CPU")
+  def test_profiler(self):
+    start = len(Compiled.profile_events)
+    with Context(PROFILE=1):
+      Tensor.empty(32).add(1).realize()
+      Device[Device.DEFAULT].synchronize()
+    self.assertTrue(any(isinstance(e, (ProfileRangeEvent, ProfileGraphEvent)) for e in Compiled.profile_events[start:]))
+
 # TODO: support in HCQCompiled
+# TODO: support these tests in HCQ2
 is_cpu_hcq = Device.DEFAULT in {"CPU"}
 
 @unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")

test/backend/test_setitem.py

@@ -1,4 +1,4 @@
-import unittest
+import unittest, operator
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
 from tinygrad.helpers import Context
 import numpy as np
@@ -378,6 +378,32 @@ class TestWithGrad(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       y[0] = 99.0
 
+  def test_unrealized_inplace_keeps_storage(self):
+    x = Tensor([1., 2.]).clone()
+    view = x[:1]
+    x += 3
+    x.realize()
+    self.assertEqual(x.tolist(), [4., 5.])
+    self.assertEqual(view.tolist(), [4.])
+
+  def test_unrealized_view_inplace_keeps_storage(self):
+    x = Tensor([1., 2.]).clone()
+    view = x[:1]
+    view += 3
+    view.realize()
+    self.assertEqual(x.tolist(), [4., 2.])
+    self.assertEqual(view.tolist(), [4.])
+
+  def test_set_augmented_backward(self):
+    for op, expected in ((operator.isub, [-1., -1.]), (operator.imul, [1., 2.]), (operator.itruediv, [-0.01, -0.005])):
+      with self.subTest(op=op.__name__):
+        z = Tensor([1.0, 2.0, 3.0, 4.0])
+        x = Tensor([10.0, 20.0])
+        z[:2] = op(z[:2], x)
+        z.sum().backward()
+        np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
+        np.testing.assert_allclose(x.grad.numpy(), expected)
+
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):
     N = 10

test/backend/test_tensor.py

@@ -737,6 +737,25 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_allclose(a.numpy(), b.numpy())
     self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
 
+  def test_clone_unrealized_copy_does_not_alias(self):
+    for realize_clone in (False, True):
+      with self.subTest(realize_clone=realize_clone):
+        a = Tensor([2.])
+        b = a.clone()
+        if realize_clone: b.realize()
+        b.assign(7.).realize()
+        self.assertEqual(a.tolist(), [2.])
+        self.assertEqual(b.tolist(), [7.])
+        self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+  def test_clone_preserves_creation_copy(self):
+    source = Tensor([2.], device="PYTHON")
+    copied = source.to("CPU")
+    cloned = copied.clone().realize()
+    source.assign(7.).realize()
+    self.assertEqual(copied.tolist(), [2.])
+    self.assertEqual(cloned.tolist(), [2.])
+
   def test_clone_deviceless_const(self):
     t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
     np.testing.assert_equal(t.numpy(), 2.0)

test/backend/test_uops.py

@@ -11,6 +11,7 @@ from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.runtime.ops_python import PythonRenderer
 from test.helpers import to_uops_list
 
 def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
@@ -56,8 +57,8 @@ def _test_uops_result(output_dtype, uops, res):
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, PythonRenderer)) and
+                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires buffer bitcast and 64-bit ints")
 class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_render(self):
@@ -85,6 +86,16 @@ class TestBitcastBufferView(unittest.TestCase):
     run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
     self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
 
+  def test_vector_load_store(self):
+    for src_dt, dst_dt in [(dtypes.uint8, dtypes.uint32), (dtypes.uint32, dtypes.uint8)]:
+      with self.subTest(src=src_dt, dst=dst_dt):
+        src, dst = [UOp.param(i, dt, 16 // dt.itemsize) for i, dt in enumerate((src_dt, dst_dt))]
+        src, dst = [b.bitcast(dtypes.uint32).index(UOp.stack(*[UOp.const(i) for i in range(4)])) for b in (src, dst)]
+        bufs = [Buffer(Device.DEFAULT, 16 // dt.itemsize, dt, initial_value=bytes(range(16)) if i == 0 else bytes(16))
+                for i, dt in enumerate((src_dt, dst_dt))]
+        run_uops([dst.store(src.load())], bufs)
+        self.assertEqual(bytes(bufs[1].as_memoryview()), bytes(range(16)))
+
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):
     assert isinstance(v2, (float, int, bool))

test/backend/test_wait_loop.py

@@ -1,13 +1,12 @@
-import unittest, threading
-from tinygrad import Tensor, UOp
+import unittest, threading, functools
+from tinygrad import Tensor, UOp, Context
 from tinygrad.device import Device, Buffer, BufferSpec
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.engine.realize import run_linear
 from tinygrad.uop.ops import Ops, KernelInfo
+from tinygrad.renderer.isa.x86 import X86Renderer
 
-def wait_loop_kernel(C:UOp) -> UOp:
-  N = 10
-
+def wait_loop_kernel(C:UOp, N=10) -> UOp:
   # a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
   # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
   l = UOp.loop(0)
@@ -42,6 +41,19 @@ def nested_loop_kernel(C:UOp) -> UOp:
 
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
 
+def pressure_loop_kernel(C:UOp, n=13) -> UOp:
+  vs = [C[j+1].load() for j in range(n)]
+  l = UOp.loop(0)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l)[0].load() + 1
+  st = i[0].store(inc)
+  i = i.after(st.end(l, inc < sum(v & inc for v in vs)))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="pressure_loop", opts_to_apply=()))
+
 def wait_ext_kernel() -> UOp:
   sig = UOp.param(0, dtypes.int, 1, volatile=True)
   l = UOp.loop(0)
@@ -100,6 +112,25 @@ class TestWaitLoop(unittest.TestCase):
     c.realize()
     self.assertEqual(c.item(), 25)
 
+  # TODO: x86's lower_loop builds an Ops.IF node after regalloc, which fails spec_full
+  @(unittest.expectedFailure if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else lambda f: f)
+  def test_wait_loop_spec(self):
+    c = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), fxn=functools.partial(wait_loop_kernel, N=7))[0]
+    with Context(SPEC=2): c.realize()
+    self.assertEqual(c.item(), 7)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "TODO: do-while loop under register pressure segfaults on x86")
+  def test_loop_carried_registers(self):
+    # more loads live across the backedge than any register file (x86 15 gprs, arm64 31, sass 255, rdna3 256 vgprs)
+    c = Tensor.custom_kernel(Tensor.ones(301, dtype=dtypes.int), fxn=functools.partial(pressure_loop_kernel, n=300))[0]
+    self.assertEqual(c[0].item(), 2)
+
+  def test_register_pressure_loop(self):
+    c = Tensor.zeros(16, dtype=dtypes.int).contiguous()
+    c = Tensor.custom_kernel(c, fxn=pressure_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c[0].item(), 1)
+
   def test_loop_in_loop(self):
     c = Tensor.empty(1, dtype=dtypes.int)
     c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]

test/device/test_hcq2.py

@@ -1,72 +1,282 @@
-import unittest, numpy as np
+import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
-from tinygrad import Device, Tensor
-from tinygrad.device import Buffer
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import HCQ2
-from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in, hcq_compile_cache
+from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import Context, dedup, partition
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
+from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
+from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.runtime.autogen import libc
+from tinygrad.runtime.support.c import init_c_struct_t
+import tinygrad.runtime.support.hcq2 as hcq2
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, HCQ2Compiled, all_devices_in, hcq_compile_cache, link_linear_cache
+from test.helpers import call_is_hcq
 
-@unittest.skipUnless(HCQ2 and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
-class TestHCQ2(unittest.TestCase):
-  def test_copy_without_copy_queue(self):
-    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
-      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
+@contextlib.contextmanager
+def rt_views():
+  calls, orig = [], HCQ2Compiled.rt_view
+  def track(dev, *args, **kwargs):
+    calls.append(dev)
+    return orig(dev, *args, **kwargs)
+  with patch.object(HCQ2Compiled, "rt_view", track): yield calls
 
-  @unittest.skipIf(Device.DEFAULT == "CPU", "ping-pong needs a non-CPU hcq2 device")
-  def test_cpu_device_ping_pong(self):
-    # CPU submits run inline, so alternating dependencies must be submitted in schedule order to avoid blocking the host submitter.
-    x = Tensor.ones(16, device="CPU").contiguous().realize()
-    a = (x + 1).contiguous()
-    b = (a.to(Device.DEFAULT).contiguous() + 1).contiguous()
-    c = (b.to("CPU").contiguous() + 1).contiguous()
-    out = (c.to(Device.DEFAULT).contiguous() + 1).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), np.full(16, 5))
+def chain(x:Tensor, n:int) -> Tensor:
+  for _ in range(n): x = (x + 1).contiguous()
+  return x
 
-  @unittest.skipIf(Device.DEFAULT == "CPU", "staged copies need a non-CPU hcq2 device")
-  def test_staged_copy_slot_reuse(self):
-    # chunks of a staged copy rotate through the staging buffer slots, many rotations must stay bit-exact in both directions
-    import tinygrad.runtime.support.hcq2 as hcq2
-    buf = Buffer("CPU", 1 << 20, dtypes.uint8, preallocate=True)
-    data = np.random.default_rng(42).integers(0, 256, (5 << 20) + 123, dtype=np.uint8)
-    with patch.object(hcq2, "STAGING_SIZE", 1 << 20), patch.object(hcq2, "STAGING_SLOTS", 4), patch.object(hcq2, "_staging", lambda: buf):
-      np.testing.assert_equal(Tensor(data).to(Device.DEFAULT).realize().numpy(), data)
+@contextlib.contextmanager
+def encoded_batches():
+  batches, orig = [], hcq2.lower_and_compile
+  with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
+    yield batches
 
-  def test_overlapping_device_tuples(self):
-    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
-    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-    try: Device[d4[-1]]
-    except Exception: self.skipTest("needs four devices")
-    ref = Tensor.arange(16).contiguous().realize()
-    Tensor(ref.uop.copy_to_device(d4)).realize()
-    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), np.ones(8))
+def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
+  for _ in range(n): x = (x + 1).contiguous()
+  return x.realize()
+
+def patch_words(batch:UOp) -> list[UOp]:
+  return [w for s in batch.src[0].toposort() if s.op is Ops.STORE and s.src[0].op is Ops.INDEX and s.src[0].src[1].op is Ops.STACK
+          and s.src[1].op is Ops.STACK for w in s.src[1].src]
+
+def rt_params(batch:UOp) -> list[str]:
+  return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
+
+class TestHCQ2Deps(unittest.TestCase):
+  def test_disjoint_write_preserves_dependencies(self):
+    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
+    for write in ([], [0]):
+      tracker = hcq2.HCQDepsTracker()
+      tracker.access_resources([b.shrink(((0, 4),))], write, 0)
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 8),))], [0], 1), [])
+      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
+
+  def test_partial_write_preserves_dependencies(self):
+    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
+    for write in ([], [0]):
+      tracker = hcq2.HCQDepsTracker()
+      tracker.access_resources([b], write, 0)
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 1), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
+
+@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
+class TestHCQ2Core(unittest.TestCase):
+  @staticmethod
+  def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
+
+  def compiled(self, n:int, jit=False):
+    x, inputs = self.input(), []
+    if jit:
+      f = TinyJit(lambda a: chain(a, n).realize())
+      f(x)
+      return f(x), f.captured._linear, [x.uop.base]
+    out = chain(x, n)
+    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
+
+  def test_jit_has_no_rt_buffers(self):
+    dev = Device[Device.DEFAULT]
+    rings = [dev.rt_buffer(True, host) for host in (False, True)]
+    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
+    for n in (1, 65):
+      with self.subTest(kernels=n):
+        x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
+        for _ in range(2): f(x)
+        for u in f.captured.linear.toposort():
+          if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
+            addr = buf._buf.va_addr
+            self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
+
+  def test_small_eager_cached(self):
+    _, compiled, inputs = self.compiled(1)
+    linked = link_linear(compiled, input_uops=inputs)
+    self.assertIs(link_linear(compiled, input_uops=inputs), linked)
+
+  def test_large_eager_not_cached(self):
+    _, compiled, inputs = self.compiled(65)
+    linked = link_linear(compiled, input_uops=inputs)
+    self.assertIsNot(link_linear(compiled, input_uops=inputs), linked)
+    self.assertNotIn(compiled, link_linear_cache)
+
+  def test_double_compile(self):
+    for n in (1, 65):
+      for jit in (False, True):
+        with self.subTest(kernels=n, jit=jit):
+          out, compiled, inputs = self.compiled(n, jit=jit)
+          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
+          before = tuple(inputs)
+          with rt_views() as borrowed:
+            for linear in (compiled, linked):
+              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
+          self.assertEqual(tuple(inputs), before)
+          self.assertFalse(borrowed)
+          run_linear(linked, input_uops=inputs, jit=True, wait=True)
+          self.assertEqual(out.tolist(), [2 + n] * 4)
+
+  def test_double_link(self):
+    for n in (1, 65):
+      for jit in (False, True):
+        with self.subTest(kernels=n, jit=jit):
+          out, compiled, inputs = self.compiled(n, jit=jit)
+          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
+          with rt_views() as borrowed:
+            again = link_linear(linked, input_uops=inputs, allow_cache=not jit)
+          self.assertIs(again, linked)
+          self.assertFalse(borrowed)
+          run_linear(again, input_uops=inputs, jit=True, wait=True)
+          self.assertEqual(out.tolist(), [2 + n] * 4)
+
+  def test_jit_new_inputs_each_call(self):
+    @TinyJit
+    def f(a, b): return (a * b + a).contiguous().realize()
+    ins = [(Tensor.full((23,), float(i)).contiguous().realize(), Tensor.full((23,), 2.0).contiguous().realize()) for i in range(6)]
+    for a, b in ins[:3]: f(a, b).tolist() # warm the jit and the copyout
 
-  def relowers(self, t:Tensor) -> int: # a compile miss relowers the whole submit, a hit only links it
     before = len(hcq_compile_cache)
-    t.realize()
-    return len(hcq_compile_cache) - before
+    self.assertEqual([f(a, b).tolist() for a, b in ins[3:]], [[i * 3.0] * 23 for i in range(3, 6)])
+    self.assertEqual(len(hcq_compile_cache), before)
 
-  def test_relower_only_on_new_kernel(self):
-    a, b = (Tensor.empty(64, 64).contiguous().realize() for _ in range(2))
-    self.relowers(a.sin())
-    self.assertEqual(self.relowers(a.sin()), 0)  # nothing changed
-    self.assertEqual(self.relowers(b.sin()), 0)  # new buffers, patched in at link time
-    self.assertEqual(self.relowers(a.cos()), 1)  # new kernel, though only the code address moved
-    self.assertEqual(self.relowers(a.cos()), 0)
-    self.assertEqual(self.relowers(Tensor.empty(32, 32).contiguous().realize().sin()), 1)  # new shape
+  def test_jit_symbolic(self):
+    @TinyJit
+    def f(a): return (a + 1).sum().contiguous().realize()
+    a = Tensor.rand(3, 10).contiguous().realize()
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      np.testing.assert_allclose(f(a[:, :vi]).item(), (a[:, :i] + 1).sum().item(), atol=1e-5, rtol=1e-5)
 
-  def test_dtype_sweep_relowers_every_dtype(self):
-    # test_dtype sweeps dtypes at one shape, so nearly every kernel is new: this is where hcq2 ci time goes
-    src = Tensor.empty(64, 64).contiguous().realize()
-    dts = (dtypes.int8, dtypes.uint8, dtypes.int16, dtypes.uint16, dtypes.int32)
-    self.assertEqual([self.relowers(src.cast(dt).contiguous()) for dt in dts], [1] * len(dts))
+  def test_map_cpu_buffer_preserves_contents(self):
+    src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
+    data = bytes(range(16))
+    src.as_memoryview(force_zero_copy=True)[:] = data
+    src.get_buf(Device.DEFAULT)
+    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
+
+  def test_staged_copy_roundtrip(self):
+    # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
+    stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
+    for npdt in (np.uint8, np.float32):
+      with self.subTest(dtype=npdt.__name__):
+        n = (size // 2 // np.dtype(npdt).itemsize) * 9 + 7 # nine rotations of a two slot ring, plus a short tail
+        data = np.arange(n, dtype=np.int64).astype(npdt)
+        with patch.object(hcq2, "STAGING_SIZE", size), patch.object(hcq2, "STAGING_SLOTS", 2), patch.object(hcq2, "_staging", lambda: stage):
+          out = Tensor(data).to(Device.DEFAULT).contiguous().realize()
+          np.testing.assert_equal(out.numpy(), data)
+
+  def test_rt_patches_are_inputs_and_vars_only(self):
+    x = Tensor.rand(17, 33).contiguous().realize()
+    with encoded_batches() as batches:
+      @TinyJit
+      def f(a): return (a.sin() * 3).contiguous().realize()
+      for _ in range(3): f(x)
+      eager_chain(x)
+
+    jit, eager = partition(batches, lambda c: c.arg.aux.table >= 0)
+    self.assertTrue(jit and eager, f"want both kinds of batch, got {len(jit)} jit and {len(eager)} eager")
+    for c in batches:
+      self.assertTrue(all(n.startswith(("inputs_", "timeline_")) for n in rt_params(c)), f"runtime patch reads {rt_params(c)}")
+      self.assertFalse([u for w in patch_words(c) for u in w.toposort() if u.op is Ops.GETADDR], "addresses bake at link time")
+    self.assertTrue(any(n.startswith("inputs_") for c in jit for n in rt_params(c)), "the jit patches its input addresses in")
+    self.assertFalse(any(n.startswith("inputs_") for c in eager for n in rt_params(c)), "eager bakes its input addresses")
+
+  def test_programs_are_not_call_args(self):
+    # a program is a link-time patch a cmdbuf word addresses: it rides inside that word, no arg or param of its own
+    def nargs(n):
+      x = Tensor.ones(16).contiguous().realize()
+      with encoded_batches() as batches:
+        @TinyJit
+        def f(a):
+          for i in range(n): a = (a * (i + 1.5)).contiguous()
+          return a.realize()
+        for _ in range(3): f(x)
+      return max(c.arg.aux.nargs for c in batches)
+    self.assertEqual(nargs(2), nargs(12))
+
+  def test_caches_hold_no_buffers(self):
+    # an eager template caches without its buffers and the jit's linear compiles once uncached: freeing the tensors frees the device memory
+    def step(i):
+      x = Tensor(np.full(1024, i, np.float32)).to(Device.DEFAULT).realize()
+      @TinyJit
+      def f(a): return (a * 2 + 1).contiguous().realize()
+      for _ in range(3): out = f(x)
+      self.assertEqual(out.tolist(), [2.0 * i + 1] * 1024)
+    step(1) # warms the programs, templates and rings
+    gc.collect()
+    used = GlobalCounters.mem_used
+    for i in range(2, 5): step(i)
+    gc.collect()
+    self.assertEqual(GlobalCounters.mem_used, used)
+
+  def test_device_state_survives_as_link_refs(self):
+    # a buffer the commands only address, never a param of the body, is kept by the linked call as a ref of what its getaddr resolved into
+    dev = Device[Device.DEFAULT]
+    names = {"AMD": () if getattr(dev, "is_aql", False) else ("scratch",), # the aql descriptor holds the scratch, nothing addresses it
+             "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
+    @TinyJit
+    def f(a): return (a * 2 + 1).contiguous().realize()
+    x = Tensor.ones(16).contiguous().realize()
+    for _ in range(3): f(x)
+    call = f.captured.linear.src[0]
+    self.assertIs(call.op, Ops.AFTER, "the linked call sits after its refs")
+    refs = [u.buffer for u in call.src[1:] if u.op is Ops.BUFFER]
+    for n in names: self.assertTrue(any(r is getattr(dev, n) for r in refs), f"{n} is not a ref of the call")
+
+@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
+class TestHCQ2FFI(unittest.TestCase):
+  @staticmethod
+  def _run(body:UOp) -> list[Buffer]:
+    call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
+    assert call is not None
+    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), allow_cache=False)
+    run_linear(linear, jit=True)
+    return [u.buffer for u in linear.src[0].without_after.src[1:] if u.op is Ops.BUFFER]
+
+  def test_ffi_ccall(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      out = UOp.placeholder((1,), dtypes.int32, slot=1, device="CPU", volatile=True, tag="ffi_result")
+      bufs = self._run(out.index(0).store(hcq2.ccall(libc.dll.ffs, 0x10)))
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 5)
+
+  def test_ffi_cstruct(self):
+    struct_t = init_c_struct_t(16, (("u8", ctypes.c_uint8, 0), ("u16", ctypes.c_uint16, 2),
+                                  ("u32", ctypes.c_uint32, 4), ("u64", ctypes.c_uint64, 8)))
+    UOp.placeholder((1,), dtypes.uint8, device="CPU") # reserve slot zero for device-owned placeholders
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      s = hcq2.cstruct(struct_t, u8=0x12, u16=UOp.const(0x3456, dtypes.uint16), u32=0x789ABCDE, u64=0xFEDCBA9876543210)
+      bufs = self._run(s.index(0).load())
+    got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
+    self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
+
+  def test_device_lower_after_encode(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      out = UOp.placeholder((1,), dtypes.int32, device="CPU", tag="result")
+      encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_encode"), lambda: UOp.custom_function("test_lower"))])
+      lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_lower"), lambda out=out: out.index(0).store(42))])
+      with patch.object(Device["CPU"], "pm_encode", encode), patch.object(Device["CPU"], "pm_lower", lower):
+        bufs = self._run(UOp.custom_function("test_encode"))
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 42)
+
+  def test_nested_cstruct_patches(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      inner = hcq2.cstruct(init_c_struct_t(4, (("value", ctypes.c_uint32, 0),)), value=42)
+      outer = hcq2.cstruct(init_c_struct_t(8, (("ptr", ctypes.c_uint64, 0),)), ptr=inner.getaddr("CPU"))
+      out = UOp.placeholder((1,), dtypes.uint32, device="CPU", tag="result")
+      copied = hcq2.ccall(libc.memcpy, out.index(0), outer.bitcast(dtypes.uint64).index(0).load(), 4)
+      bufs = self._run(out.after(copied).index(0).load())
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32)._buf.cpu_view().view(fmt='I')[0], 42)
+
+
+class TestHCQ2Timeline(unittest.TestCase):
+  def test_reused_timeline_is_zeroed(self):
+    buf = Buffer("CPU", 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    addr = buf._buf.va_addr
+    buf._buf.cpu_view().view(fmt='B')[:] = b'\xff' * 16
+    buf.deallocate()
+    dev = HCQ2Compiled.__new__(HCQ2Compiled)
+    dev.device = "CPU"
+    self.assertEqual(dev.timeline._buf.va_addr, addr)
+    self.assertEqual(bytes(dev.timeline._buf.cpu_view()), bytes(16))
 
-  @unittest.skipIf(Device.DEFAULT == "CPU", "sharding needs a non-CPU hcq2 device")
-  def test_shard_from_host(self): # the host copy, the p2p copy of its second half and the lane kernels are one batch: the deps must chain
-    try: Device[d1:=f"{Device.DEFAULT}:1"]
-    except Exception: self.skipTest("needs a second device")
-    a = np.arange(64*64, dtype=np.float32).reshape(64, 64)
-    np.testing.assert_equal(Tensor(a).shard((Device.DEFAULT, d1), axis=0).realize().numpy(), a)
 
 if __name__ == "__main__":
   unittest.main()

test/external/external_test_usb_asm24.py

@@ -3,13 +3,14 @@ from tinygrad.helpers import Timing, getenv
 from tinygrad import Tensor, Device
 import numpy as np
 
-class TestDevCopySpeeds(unittest.TestCase):
+class USBTestCase(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
+class TestDevCopySpeeds(USBTestCase):
   def testCopyCPUtoDefault(self):
     for _ in range(10):
       t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
@@ -24,6 +25,7 @@ class TestDevCopySpeeds(unittest.TestCase):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
+class TestUSBIntegrity(USBTestCase):
   def testValidateCopies(self):
     t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
@@ -34,5 +36,67 @@ class TestDevCopySpeeds(unittest.TestCase):
     np.testing.assert_equal(t.numpy(), y.numpy())
     del x, y, t
 
+  def testCopyinBoundaries(self):
+    rng, chunk = np.random.default_rng(0), 0x40000 - 4
+    for size in (1, 3, 508, 509, 0x3ffc, 0x3ffd, chunk, chunk+1, 2*chunk+31):
+      with self.subTest(size=size):
+        a = rng.integers(0, 256, size, dtype=np.uint8)
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+
+  def testCopyinFenceWrap(self):
+    a = np.arange(2*(0x40000-4)+31, dtype=np.uint8)
+    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
+    self.dev.synchronize()
+    alloc, usb = self.dev.allocator, self.dev.iface.pci_dev.usb
+    clear = usb.read(0xA808, 1)
+    # Model a completed 256-chunk copy instead of the one-chunk warmup. The next clear tag must still change.
+    alloc._usb_seq += 255
+    usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
+    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
+    self.assertNotEqual(clear, usb.read(0xA808, 1))
+    for bits in (8, 24):
+      with self.subTest(bits=bits):
+        alloc._usb_seq = ((alloc._usb_seq >> bits)+2)*(1 << bits)-2
+        usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+
+  def testCopyinRingWrap(self):
+    rng = np.random.default_rng(0)
+    a = rng.integers(0, 256, 1 << 20, dtype=np.uint8)
+    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+    ring = self.dev.sdma_queue(0)
+    # A 16 MiB copyin needs more than 4 KiB of SDMA packets, forcing the submission to wrap.
+    target = ring.ring.nbytes - 0x1000
+    padding = target - ring.put_value % ring.ring.nbytes - 16  # four-dword timeline fence
+    self.assertGreaterEqual(padding, 0)
+    q = self.dev.hw_copy_queue_t()
+    q.q(*([0] * (padding // 4)))
+    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+    self.dev.synchronize()
+    before = ring.put_value // ring.ring.nbytes
+    a = rng.integers(0, 256, 16 << 20, dtype=np.uint8)
+    t = Tensor(a, device="AMD").realize()
+    self.assertGreater(ring.put_value // ring.ring.nbytes, before)
+    np.testing.assert_array_equal(a, t.numpy())
+
+  def testCopyinStaleSentinel(self):
+    a = np.arange(16, dtype=np.uint8)
+    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+    chunk = 0x40000 - 4
+    for case in ("copyout", "reuse"):
+      with self.subTest(case=case):
+        if case == "copyout":
+          # A 512 KiB copyin takes three chunks. Copyout then fills both SRAM windows with the next expected tag.
+          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 3) & 0xFFFFFF)
+          a = np.full(0x80000 // 4, tag, dtype=np.uint32)
+          np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+          a = np.arange(31, dtype=np.uint8)
+        else:
+          # The first full chunk contains the tag expected by the short third chunk in the same window.
+          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 2) & 0xFFFFFF)
+          a = np.arange(2 * chunk + 31, dtype=np.uint8)
+          a[:chunk].view(np.uint32)[:] = tag
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+
 if __name__ == "__main__":
   unittest.main()

test/external/speed_v_theoretical.py

@@ -1,6 +1,10 @@
 import unittest
-from tinygrad import Tensor, TinyJit, Device
-from tinygrad.helpers import Context, DEBUG, GlobalCounters
+from dataclasses import replace
+from itertools import islice
+from tinygrad import Tensor, Device
+from tinygrad.codegen import to_program
+from tinygrad.engine.realize import time_call
+from tinygrad.helpers import Context, DEBUG
 from tinygrad.nn import Conv2d
 from tinygrad.nn.state import get_parameters
 
@@ -10,6 +14,13 @@ class TestKernelSpeed(unittest.TestCase):
       # TODO: randn is 20% faster than rand for gemv
       return Tensor.randn(shape, dtype="half").realize()
 
+  def _time_kernel(self, out:Tensor, beam:int):
+    linear = out.schedule_linear()
+    self.assertEqual(len(linear.src), 1, "expected a single kernel")
+    call = linear.src[0]
+    prg = to_program(call.src[0].replace(arg=replace(call.src[0].arg, beam=beam)), Device[out.device].renderer)
+    return min(islice(time_call(call.replace(src=(prg, *call.src[1:])), clear_l2=True), 3, 10))
+
   def _compare(self, tm, tflops, gbs, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
       if DEBUG >= 1:
         print(f"{tm=:.6f}")
@@ -34,53 +45,30 @@ class TestKernelSpeed(unittest.TestCase):
 
   def _test_matmul(self, M, K=None, N=None, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
     # (MxK) @ (KxN)
-    @TinyJit
-    def f(a, b) -> Tensor: return (a @ b).realize()
-
     if N is None: N = M
     if K is None: K = M
-    tms = []
-    with Context(BEAM=3):
-      for i in range(10):
-        a = self._get_tensor(M, K)
-        b = self._get_tensor(K, N)
-        if i >= 3:
-          GlobalCounters.time_sum_s = 0
-          with Context(DEBUG=max(DEBUG.value, 2)): c = f(a, b)
-          tms.append(GlobalCounters.time_sum_s)
-        else:
-          c = f(a, b)
+    a = self._get_tensor(M, K)
+    b = self._get_tensor(K, N)
+    tm = self._time_kernel(c:=a @ b, beam=3)
 
     ops = 2 * M * N * K
     mems = a.dtype.itemsize * M * K + b.dtype.itemsize * K * N + c.dtype.itemsize * M * N
-    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
   def _test_conv_3x3(self, BS, CIN, COUT, H, W, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
-    @TinyJit
-    def f(conv, x) -> Tensor: return conv(x).realize()
-    tms = []
     K = 3
     with Context(BEAM=0, DEBUG=0):
       conv = Conv2d(CIN, COUT, K, padding=1)
       Tensor.realize(*get_parameters(conv))
 
-    with Context(BEAM=2):
-      for i in range(10):
-        x = self._get_tensor(BS, CIN, H, W)
-        if i >= 3:
-          GlobalCounters.time_sum_s = 0
-          with Context(DEBUG=max(DEBUG.value, 2)): _c = f(conv, x)
-          tms.append(GlobalCounters.time_sum_s)
-        else:
-          _c = f(conv, x)
+    x = self._get_tensor(BS, CIN, H, W)
+    tm = self._time_kernel(_c:=conv(x), beam=2)
 
     # naive algo
     ops = 2 * BS * CIN * COUT * K * K * H * W
     mems = x.nbytes() + conv.weight.nbytes() + conv.bias.nbytes() + _c.nbytes()
-    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)

test/helpers.py

@@ -65,13 +65,17 @@ def assert_kernel_count(expected:int):
   got = GlobalCounters.kernel_count
   if got != expected: raise KernelCountException(expected, got)
 
+def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
+  from tinygrad.runtime.support.hcq2 import HCQ_DEVS
+  return Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}
+
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
 
 def call_is_hcq(call:UOp) -> bool: # an hcq2 batch: a compiled body whose aux lists the kernels it submits
   from tinygrad.runtime.support.hcq2 import HCQInfo
-  return isinstance(getattr(call.arg, "aux", None), HCQInfo)
+  return isinstance(getattr(call.without_after.arg, "aux", None), HCQInfo)
 
 def jit_cache_count(linear:UOp) -> int:
   n = 0
@@ -87,7 +91,7 @@ def assert_jit_cache_len(fxn, expected_len):
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
   if expected_len and any(call_is_hcq(call) for call in linear.src): # HCQ2: kernels batch into submits, the finalizers carry the batch's kernels
-    count = sum(len(call.arg.aux.kernels) if call_is_hcq(call) else 1 for call in linear.src)
+    count = sum(len(call.without_after.arg.aux.kernels) if call_is_hcq(call) else 1 for call in linear.src)
     if count != expected_len: raise KernelCountException(expected_len, count)
     return
   if call_is_graph(linear.src[0]):

test/mockgpu/amd/amdgpu.py

@@ -327,8 +327,11 @@ class SDMAExecutor(AMDQueue):
 
   def _execute_copy(self):
     struct = sdma_pkts.copy_linear.from_address(self.base + self.rptr[0] % self.size)
-    count_cnt = to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF
-    ctypes.memmove(self.gpu.translate_addr(struct.dst_addr), self.gpu.translate_addr(struct.src_addr), count_cnt + 1)
+    count, off = (to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF) + 1, 0
+    while off < count: # a page at a time: the physical pages of a range needn't be contiguous
+      n = min(count - off, 0x1000 - ((struct.src_addr + off) & 0xfff), 0x1000 - ((struct.dst_addr + off) & 0xfff))
+      ctypes.memmove(self.gpu.translate_addr(struct.dst_addr + off), self.gpu.translate_addr(struct.src_addr + off), n)
+      off += n
     self.rptr[0] += ctypes.sizeof(struct)
 
 class AMDGPURegisters:

test/mockgpu/amd/emu.py

@@ -69,7 +69,7 @@ from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp, Inst
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
 from test.amd.helpers import decode_dpp16
-from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32
+from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32, _bitreverse, _countbits
 
 MASK32 = 0xFFFFFFFF
 
@@ -537,15 +537,20 @@ class _Ctx:
           stores.extend([self.wsgpr_dyn(_c(EXEC_LO.offset), lo), self.wsgpr_dyn(_c(EXEC_LO.offset + 1), hi)])
         else: stores.append(self.wsgpr_dyn(_c(EXEC_LO.offset), _to_u32(val)))
       elif dest.startswith('VCC'): stores.extend(self.wmask(_c(VCC_LO.offset), val))
+      elif dest.startswith('PC'):  # S_SETPC/S_SWAPPC jump: write PC directly (caller skips inc_pc)
+        lo, hi = _split64(val.cast(dtypes.uint64))
+        stores.extend([self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)])
     return stores
 
   def compile_sop_pcode(self, op, srcs: dict[str, UOp | int], sdst_reg: UOp, sdst_size: int) -> UOp:
     """Compile a scalar instruction with dynamic destination register."""
     pcode = get_pcode(op)
-    srcs.update(self.base_srcs(self.rexec()), VCC=self.rmask(_c(VCC_LO.offset)))
+    srcs.update(self.base_srcs(self.rexec()), VCC=self.rmask(_c(VCC_LO.offset)), PC=self.rpc().cast(dtypes.int64))
     if 'D0' not in srcs: srcs['D0'] = self.rsgpr_dyn(sdst_reg)  # D0 is current dest value for read-modify-write ops
     _, assigns = parse_pcode(pcode, srcs)
-    return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *self.inc_pc())
+    # PC-writing ops (S_SETPC/S_SWAPPC) jump instead of advancing to the next instruction
+    inc = [] if any(dest.startswith('PC') for dest, _ in assigns) else self.inc_pc()
+    return UOp.sink(*self.scalar_stores(assigns, sdst_reg, sdst_size), *inc)
 
   def compile_lane_pcode(self, op, inst) -> UOp:
     """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
@@ -678,7 +683,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
             'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
             'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
     for dest, val in parse_pcode(pcode, srcs)[1]:
-      if dest == 'PC' or dest.startswith('PC.'):
+      if dest.startswith('PC'):
         lo, hi = _split64(val.cast(dtypes.uint64))
         return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), lo), ctx.wsgpr_dyn(_c(PC_HI_IDX), hi))
   return UOp.sink(*ctx.inc_pc())
@@ -1323,7 +1328,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   vdst_reg = ctx.inst_field(type(inst).vdst)
   src0_r, src1_r = ctx.inst_field(type(inst).src0) - _c(256), ctx.inst_field(type(inst).src1) - _c(256)
   src2_r = ctx.inst_field(type(inst).src2)
-  src2_r = (src2_r >= 256).where(src2_r - _c(256), src2_r)
+  is_c_vgpr = src2_r >= _c(256)
+  src2_r = is_c_vgpr.where(src2_r - _c(256), src2_r)  # also keeps the unused VGPR-side index in bounds when src2 is a constant
   output_type = op_name.split("WMMA_", 1)[1].split("_", 1)[0]
   is_bf16, is_rdna4 = 'BF16' in op_name, isinstance(inst, ir4.VOP3P)
   cvt = _FUNCS['bf16_to_f32' if is_bf16 else 'f16_to_f32']
@@ -1353,12 +1359,15 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
     return n + lane_bit * 16, vgpr
 
   # Accumulator C. RDNA4 f16/bf16 packs two f32 accumulator VGPRs into one f16 VGPR; RDNA3 uses the lo half of each.
+  # src2 may be a VGPR or an inline/scalar constant (128 = int 0, the usual ", 0" C form); the runner must handle both dynamically
+  out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32
+  cbits = ctx.rsrc_dyn(src2_r, None, 32)
+  cval_const = cvt(cbits & UOp.const(0xFFFF, dtypes.uint32)) if output_type in ("F16", "BF16") else cbits.bitcast(out_dt)
   if output_type in ("F16", "BF16"):
-    mat_c = [gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0)))
+    mat_c = [is_c_vgpr.where(gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0))), cval_const)
              for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
   else:
-    out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32
-    mat_c = [ctx.rvgpr_dyn(src2_r + _c(vgpr), UOp.const(lane, dtypes.int)).bitcast(out_dt)
+    mat_c = [is_c_vgpr.where(ctx.rvgpr_dyn(src2_r + _c(vgpr), UOp.const(lane, dtypes.int)).bitcast(out_dt), cval_const)
              for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
   mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
 
@@ -1557,9 +1566,19 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)  # type: ignore[union-attr]
 
+  if is_lds and op_name == 'DS_SWIZZLE_B32':
+    # The manual's reverse_bits operates on five-bit lane indices; thread indices wrap within the wave.
+    funcs = {'reverse_bits': lambda x: _bitreverse(x, 32) >> _c(27), 'count_ones': _countbits,
+             'thread_in': lambda x: ctx.rvgpr_dyn(addr_reg, x & _c(ctx.wave_size - 1)),
+             'thread_valid': lambda x: _lane_active(exec_mask, x & _c(ctx.wave_size - 1))}
+    result, _ = parse_pcode(pcode, {'offset0': offset0.cast(dtypes.uint8), 'offset1': offset1.cast(dtypes.uint8)}, funcs)
+    values = [result[f'thread_out@{i}'] for i in range(ctx.wave_size)]
+    # Snapshot every source before writing: destination and source registers may be identical.
+    reads = UOp(Ops.STACK, src=tuple(values))
+    return UOp.sink(*(ctx.wvgpr_dyn(vdst_reg, _c(i), val, exec_mask, after=reads) for i, val in enumerate(values)), *ctx.inc_pc())
+
   # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
   if is_lds and 'PERMUTE' in op_name:
-    pcode = get_pcode(inst.op)
     srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
             'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
     _, assigns = parse_pcode(pcode, srcs)
@@ -1947,7 +1966,9 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
   # Use Buffer objects with external_ptr=0 for vmem
   vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
   lds_buf = Buffer('CPU', max(lds_size // 4, 1), dtypes.uint32).ensure_allocated()
-  scratch_buf = Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated() if scratch_size else None
+  # Scratch is per-lane private memory: each wave needs its own region so data spilled before s_barrier survives other waves' execution.
+  n_waves = -(-total_threads // wave_size)
+  scratch_buf = Buffer('CPU', scratch_size * wave_size * n_waves, dtypes.uint8).ensure_allocated() if scratch_size else None
 
   # Initialize SQTT encoder — emits packets inline as instructions execute (only when profiling)
   if PROFILE:
@@ -1971,9 +1992,10 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
     waves: list[tuple[WaveState, list]] = []
     for wave_start in range(0, total_threads, wave_size):
       st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data, wave_size)
+      scratch_base = scratch_buf._buf.va_addr + (wave_start // wave_size) * scratch_size * wave_size if scratch_buf else 0
       waves.append((st, [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
                          ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
-                         ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0),
+                         ctypes.c_uint64(scratch_base if scratch_buf else 0),
                          ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]))
     done = [False] * len(waves)
     for _ in range(10_000_000):

test/mockgpu/amd/pcode.py

@@ -630,6 +630,10 @@ class Parser:
         self.eat('DOT')
         dt_name = self.eat('IDENT').val
         return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
+      if name in self.funcs and self.try_eat('LBRACKET'):
+        index = self.parse()
+        self.eat('RBRACKET')
+        return self.funcs[name](index)
       if name == 'VGPR' and self.at('LBRACKET'):
         self.eat('LBRACKET')
         lane = self.parse()
@@ -904,20 +908,16 @@ class Parser:
         idx2 = (addr + _const(adt, 4)) >> _const(adt, 2)
         val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
-      elif dt in (dtypes.uint16, dtypes.int16):
-        val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
       else:
-        # Handle unaligned 32-bit loads: combine two consecutive dwords and shift.
-        # To avoid OOB at buffer boundaries for aligned loads, clamp idx_hi to idx (safe).
+        # Handle unaligned 16/32-bit loads: combine two consecutive dwords and shift.
+        # The next dword is only read when the value straddles into it, so a load at the end of a buffer stays in bounds.
         # Use int64 for the WHERE to avoid 32-bit int overflow in C pointer arithmetic (addr can be >8GB).
         byte_off = (addr & _const(adt, 3)).cast(dtypes.uint32)
-        is_unaligned = byte_off.ne(_u32(0))
         idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
         idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
-        safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
-        hi = mindex(safe_idx_hi)
+        hi = mindex((byte_off > _u32(4 - dt.itemsize)).where(idx_hi_native, idx_native))
         combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
-        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val)
+        val = (combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32)
     return _cast_to(val, dt)
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
@@ -1010,20 +1010,24 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
     # for loop
     if first == 'for':
-      # Parse: for VAR in [SIZE']START : [SIZE']END do
-      p = Parser(toks, env, funcs)
-      p.eat_val('for', 'IDENT')
-      loop_var = p.eat('IDENT').val
-      p.eat_val('in', 'IDENT')
-      def parse_bound():
-        if p.at('NUM') and p.peek(1).type == 'QUOTE':
-          p.eat('NUM')
-          p.eat('QUOTE')
-        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        return int(p.parse())
-      start_val = parse_bound()
-      p.eat('COLON')
-      end_val = parse_bound()
+      # C-style loops use an exclusive bound; for/in loops use an inclusive bound.
+      if m := re.fullmatch(r'for\s*\(\s*(\w+)\s*=\s*(\d+);\s*\1\s*<\s*(\d+);\s*\1\s*(\+\+|\+=\s*\d+)\s*\)', line):
+        loop_var, start_val, end_val = m[1], int(m[2]), int(m[3]) - 1
+        step = 1 if m[4] == '++' else int(m[4][2:])
+      else:
+        p = Parser(toks, env, funcs)
+        p.eat_val('for', 'IDENT')
+        loop_var = p.eat('IDENT').val
+        p.eat_val('in', 'IDENT')
+        def parse_bound():
+          if p.at('NUM') and p.peek(1).type == 'QUOTE':
+            p.eat('NUM')
+            p.eat('QUOTE')
+          if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
+          return int(p.parse())
+        start_val = parse_bound()
+        p.eat('COLON')
+        end_val, step = parse_bound(), 1
       # Collect body
       i += 1
       body_lines: list[str] = []
@@ -1039,7 +1043,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       has_break = any('break' in bl.lower() for bl in body_lines)
       found_var = f'_found_{next(_break_var_ids)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
-      for loop_i in range(start_val, end_val + 1):
+      for loop_i in range(start_val, end_val + 1, step):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
         _, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
         if has_break:
@@ -1228,9 +1232,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       var = toks[0].val
       j, idx_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'EQUALS':
+        idx_expr = parse_tokens(idx_toks, env, funcs)
         # Static index: var[NUM] = value
-        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
-          idx = int(idx_toks[0].val.rstrip('UuLl'))
+        if isinstance(idx := _single_value(idx_expr), int):
           val = parse_tokens(toks[j+1:], env, funcs)
           existing = block_assigns.get(var, env.get(var))
           if existing is not None and isinstance(existing, UOp):
@@ -1242,7 +1246,6 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         # Dynamic index: var[expr] = value where var has @-elements
         elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
         if elems:
-          idx_expr = parse_tokens(idx_toks, env, funcs)
           val = parse_tokens(toks[j+1:], env, funcs)
           for elem_idx_str, old_elem in elems:
             elem_idx = int(elem_idx_str)
@@ -1411,16 +1414,29 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
 
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dict | None = None) -> tuple[dict, list]:
   env: dict = srcs.copy() if srcs else {}
   assigns: list[tuple[str, UOp]] = []
-  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
-  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
   lines: list[str] = []
-  for l in raw_lines:
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
-    else: lines.append(l)
-  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
+  blocks: list[str] = []
+  for raw in pcode.splitlines():
+    line = raw.split('//')[0].strip().rstrip(';')
+    if not line: continue
+    # Both block syntaxes share the same parser; braces supply the implicit end markers.
+    if line.startswith('}') and blocks:
+      end = blocks.pop()
+      line = line[1:].strip()
+      if not line.startswith(('elsif', 'else')): lines.append(end)
+    if m := re.match(r'(if|elsif|else|for)\b.*\{$', line):
+      blocks.append('endfor' if m[1] == 'for' else 'endif')
+      line = line[:-1].rstrip()
+      if m[1] in ('if', 'elsif'): line += ' then'
+    if not line: continue
+    line = re.sub(r'=\s*(\w+):(\w+)$', r'= {\1, \2}', line)
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] += ' ' + line
+    else: lines.append(line)
+  assert not blocks, "unclosed pcode block"
+  _, final, _ = parse_block(lines, 0, env, {**_FUNCS, **funcs} if funcs else None, assigns=assigns)
   sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
   for var, val in final.items():
     if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):

test/mockgpu/nv/nvdriver.py

@@ -53,6 +53,7 @@ class NVDriver(VirtDriver):
                            VirtFile('/dev/nvidia-uvm', functools.partial(NVUVMFileDesc, driver=self))]
 
     self.root_handle = None
+    self.host_ranges: set[int] = set()
 
     self.gpus = {}
     self.next_fd = (1 << 29)
@@ -251,7 +252,9 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_ENABLE_PEER_ACCESS: pass # uvm and shared spaced are setup already, no emulation for now
     elif nr == nv_gpu.UVM_CREATE_EXTERNAL_RANGE:
       st = nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS.from_address(argp)
-      libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
+      # Registered host memory already has a CPU mapping; MAP_FIXED would discard its contents.
+      if st.base not in self.host_ranges:
+        libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
     elif nr == nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION:
       st = nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS.from_address(argp)
       for gpu_attr_id in range(st.gpuAttributesCount):
@@ -265,6 +268,7 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_REGISTER_CHANNEL: pass
     elif nr == nv_gpu.UVM_FREE:
       st = nv_gpu.UVM_FREE_PARAMS.from_address(argp)
+      self.host_ranges.discard(st.base)
       libc.munmap(st.base, st.length)
     else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
     return 0
@@ -276,6 +280,7 @@ class NVDriver(VirtDriver):
       st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
       # Track host memory (signal memory) - progress queues when written to
       if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
+        self.host_ranges.add(st.params.pMemory)
         self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
                            lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
     return 0

test/mockgpu/nv/nvgpu.py

@@ -100,11 +100,11 @@ class GPFIFO:
     if qmd.release0_enable:
       rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
       rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)
-      rel0[1] = int(time.perf_counter() * 1e9)
+      if qmd.release0_structure_size == 0: rel0[1] = int(time.perf_counter() * 1e9) # four words: the timestamp after the payload
     if qmd.release1_enable:
       rel1 = to_mv(qmd.release1_address_lower + (qmd.release1_address_upper << 32), 0x10).cast('Q')
       rel1[0] = qmd.release1_payload_lower + (qmd.release1_payload_upper << 32)
-      rel1[1] = int(time.perf_counter() * 1e9)
+      if qmd.release1_structure_size == 0: rel1[1] = int(time.perf_counter() * 1e9)
     if qmd.dependent_qmd0_enable:
       if qmd.dependent_qmd0_action == 1: self.execute_qmd(qmd.dependent_qmd0_pointer << 8)
       else: raise RuntimeError("unsupported dependent qmd action")
@@ -192,11 +192,10 @@ class GPFIFO:
       sz = self._state(nv_gpu.NVC6B5_LINE_LENGTH_IN)
       assert flags == 0x182, f"unsupported flags in _exec_nvc6b5_dma: {flags}"
       ctypes.memmove(dst, src, sz)
-    elif ((flags >> 3) & 0b11) != 0:
-      src = to_mv(self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 0x10).cast('Q')
-      val = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
-      src[0] = val
-      src[1] = int(time.perf_counter() * 1e9)
+    elif (semaphore_type:=((flags >> 3) & 0b11)) != 0:
+      to_mv(addr:=self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 4).cast('I')[0] = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
+      if semaphore_type == nv_gpu.NVC6B5_LAUNCH_DMA_SEMAPHORE_TYPE_RELEASE_FOUR_WORD_SEMAPHORE:
+        to_mv(addr + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
     else: raise RuntimeError("unknown nvc6b5_dma flags")
 
   def _exec_pcas2(self):

test/null/test_device.py

@@ -89,7 +89,7 @@ class TestDevice(unittest.TestCase):
       except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
 
       imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler; "
-                 "from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler")
+                 "from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler")
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, AMDLLVMCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "AMD:LLVM"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],

test/null/test_multitensor.py

@@ -185,7 +185,7 @@ class TestMultiScalarALU(unittest.TestCase):
       return (inner.sum(),)
     param = x.as_param(0)
     fxn = _fxn(param.uop, x.device)
-    per_dev_scalar = Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+    per_dev_scalar = Tensor(fxn[0].uop.call_with_output(x.uop))
     result = x * per_dev_scalar
     self.assertEqual(result.shape, (4, 4))
     self.assertEqual(result.uop.axis, 0)

test/null/test_simplify_valid_idx.py

@@ -60,6 +60,10 @@ class TestValidIdxSimplification(unittest.TestCase):
     valid = (alu0 < 57) & (alu0 >= 1)
     self.assertIsNone(simplify_valid(valid))
 
+  def test_bitwise_and_is_not_a_valid(self):
+    ridx0 = Range(0, 16)
+    self.assertEqual(simplify_valid_idx(UOp.sink((ridx0 & UOp.const(12, dtypes.int)) & ridx0)).src[0].render(), "((int)(r0)&12&(int)(r0))")
+
   def test_valid_order_matters1(self):
     ridx0 = Range(0, 2)
     v0 = ridx0<1

test/null/test_uop_graph.py

@@ -406,6 +406,18 @@ class TestUOpGraph(unittest.TestCase):
     a = c.after(e)
     self.assertNotIn(r, a.ranges)
 
+  def test_external_call_preserves_ranges(self):
+    r = UOp.range(4, 0, dtype=dtypes.int)
+    fn = UOp.custom_function("external", UOp.const(0, dtypes.uint64))
+    call = fn.call(r + 1, ret_dtype=dtypes.int)
+    self.assertEqual(set(call.ranges), {r})
+
+  def test_conditional_end_preserves_outer_range(self):
+    outer, inner = UOp.range(4, 0), UOp.loop(1)
+    end = UOp.const(1).end(inner, outer < 2)
+    self.assertEqual(set(end.ranges), {outer})
+    self.assertEqual(set((outer + 1).after(end).ranges), {outer})
+
 class TestReduceCollapse(unittest.TestCase):
   def test_multi_range_reduce_add(self):
     """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""

test/null/test_uop_vmin_vmax.py

@@ -167,6 +167,17 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(UOp.const(4.5).cast(dtypes.float).cast(dtypes.int)._min_max, (4, 4))
     x = UOp.const(4.5).cast(dtypes.float)
     self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True)
+    # a source reaching past the destination clamps to its edge
+    self.assertEqual(UOp.variable('x', 2e9, 3e9, dtypes.float).cast(dtypes.int)._min_max, (2000000000, dtypes.int.max))
+    # a source entirely past the destination has no value in it
+    self.assertEqual(UOp.variable('x', 3e9, 4e9, dtypes.float).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.variable('x', -4e9, -3e9, dtypes.float).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.variable('x', 200, 300, dtypes.int).cast(dtypes.char)._min_max, (dtypes.char.min, dtypes.char.max))
+    self.assertEqual(UOp.const(300, dtypes.char)._min_max, (dtypes.char.min, dtypes.char.max))
+    self.assertEqual(UOp.const(math.inf).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.const(math.nan, dtypes.float)._min_max, (-math.inf, math.inf))
+    # a weak destination has no width to clamp to
+    self.assertEqual(UOp.variable('x', 5, 7, dtypes.int).cast(dtypes.weakfloat)._min_max, (5, 7))
 
   def test_vmin_vmax_cast_int_to_float_grid(self):
     # a cast to float only takes values on the float grid, so its bounds are the source bounds rounded at the destination

test/null/test_validate_oob.py

@@ -122,6 +122,35 @@ class TestValidateOOB(unittest.TestCase):
       r = UOp.range(20, 0)
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
       to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
+      # a float entirely out of the int range has no value, not an empty one
+      f = UOp.variable("f", 3e9, 4e9, dtypes.float32, param=True).cast(dtypes.int)
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(f).load()])
+
+  def test_float_cast_in_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int, 1)
+      r = UOp.range(20, 0)
+      unknown = r.cast(dtypes.float).cast(dtypes.bool)  # a bool from a float is unconstrained
+      to_uops_list([buf.index(r.valid((r < 1) & unknown)).load()])
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid(unknown)).load()])
+
+  def test_bitcast_in_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int, 16)
+      r = UOp.range(16, 0)
+      # the WEBGPU shift: int -> uint, shift, back to int
+      i = (r.cast(dtypes.int).bitcast(dtypes.uint) << UOp.const(1).cast(dtypes.uint)).bitcast(dtypes.int)
+      to_uops_list([buf.index(i.valid(i < 16)).load()])
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(i).load()])  # 0..30 oob
+      # a negative char reads as a large uchar
+      c = Variable("c", -128, -113).cast(dtypes.char)
+      to_uops_list([UOp.param(1, dtypes.int, 144).index(c.bitcast(dtypes.uchar).cast(dtypes.int)).load()])  # 128..143 valid
+      # the bits of a float are any int
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.cast(dtypes.float).bitcast(dtypes.int)).load()])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):

test/null/test_viz.py

@@ -509,10 +509,11 @@ class TestVizIntegration(unittest.TestCase):
     with save_viz() as viz:
       x.realize()
     lst = viz.list_items()
-    codegen_idx = len(lst)-1
+    # the codegen item is not the last one: the hcq compile and link groups come after it
+    codegen_idx = next((i for i,it in enumerate(lst) if any(s["name"] == "View Source" for s in it["steps"])), None)
+    assert codegen_idx is not None, "must have source rendering in list"
     steps = lst[codegen_idx]["steps"]
-    src_idx = next((i for i,s in enumerate(steps) if s["name"] == "View Source"), None)
-    assert src_idx is not None, "must have source rendering in list"
+    src_idx = next(i for i,s in enumerate(steps) if s["name"] == "View Source")
     src_render = get_render(viz.data, steps[src_idx]["query"])["src"]
     self.assertEqual(src, src_render)
 
@@ -1015,18 +1016,19 @@ class TestCfg(unittest.TestCase):
     self.get_cfg("jump_back_to_end", k)
 
 # launch viz cli without subprocess
-def run_cli(*cli_args) -> list[dict]:
+def run_cli(*cli_args, json_fmt=True) -> list[dict]:
   from tinygrad.viz.cli import main, get_arg_parser
-  args = get_arg_parser().parse_args(cli_args+("--json",))
+  args = get_arg_parser().parse_args(cli_args+(("--json",) if json_fmt else ()))
   with contextlib.redirect_stdout(buf:=io.StringIO()):
     main(args)
-  return [json.loads(line) for line in buf.getvalue().strip().splitlines()]
+  stdout = buf.getvalue().strip()
+  return [json.loads(line) for line in stdout.splitlines()] if json_fmt else [{"out":stdout}]
 
 @contextlib.contextmanager
-def write_files(viz) -> list[str]:
+def write_files(rewrites=None, profile=cpu_events) -> list[str]:
   with tempfile.TemporaryDirectory() as tmpdir:
-    (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
-    (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+    (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps((rewrites.data if rewrites is not None else VizData()).trace))
+    (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(profile))
     yield ["--rewrites-path", str(r), "--profile-path", str(p)]
 
 class TestCLI(unittest.TestCase):
@@ -1071,10 +1073,11 @@ class TestCLI(unittest.TestCase):
       out = run_cli(*files, "-s", "NULL")
       aggregate = run_cli(*files, "-s", "NULL", "-t")
     self.assertEqual(len(out), 3*2)
-    # flops increases as N gets larger
+    # Operation count increases with N; FLOPS is a rate and also depends on the measured duration.
     gflops = [row["fmt"]["FLOPS"] for row in out]
-    self.assertGreater(gflops[4], gflops[2])
-    self.assertGreater(gflops[5], gflops[3])
+    flops = [rate * row["dur_ms"] * 1e-3 for rate, row in zip(gflops, out)]
+    self.assertGreater(flops[4], flops[2])
+    self.assertGreater(flops[5], flops[3])
     # aggregate flops
     self.assertEqual(len(aggregate), 2)
     agg_gflops = [row["fmt"]["FLOPS"] for row in aggregate]

test/opt/test_gen_float4.py

@@ -12,10 +12,6 @@ class TestFloat4(unittest.TestCase):
   def count_float4(uops: list[UOp], n=4):
     return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float and uop.shape == (4,)]),
             len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float and uop.shape == (4,)]))
-  @staticmethod
-  def count_half4(uops: list[UOp]):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half and uop.shape == (4,)]))
 
   def test_float4_basic(self):
     a = Tensor.empty(2, 8).realize()

test/opt/test_kernel_opts.py

@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import Context
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.uop.ops import AxisType
 
@@ -19,13 +18,10 @@ class TestKernelOpts(unittest.TestCase):
     r = (b.sqrt() + ((a+1).sum(axis=3).exp()))
     helper_linearizer_opt(r, [
       [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))],
-      [Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL))],
       [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL))], # Checking how it works with locals
       [Opt(OptOps.SPLIT, 1, (2, AxisType.GROUP_REDUCE, True))],
-      [Opt(OptOps.SPLIT, 1, (32, AxisType.GROUP_REDUCE, True))],
       [Opt(OptOps.SPLIT, 1, (64, AxisType.GROUP_REDUCE, True))], # Checking how it works with grouped reduce
       [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))],
-      [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (16, AxisType.GROUP_REDUCE, True))],
       [Opt(OptOps.SPLIT, 0, (32, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))],
       # Checking how it works with locals + grouped reduce
       [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (64, AxisType.GROUP_REDUCE, True))],
@@ -42,6 +38,32 @@ class TestKernelOpts(unittest.TestCase):
        Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 8, (2, AxisType.GROUP_REDUCE))],
     ])
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_grouped_reduce_with_local_upcast_padto(self):
+    Tensor.manual_seed(7)
+    a = Tensor.rand(7, 11, 13)
+    helper_linearizer_opt(a.sum((1, 2)) + a.max((1, 2)), [
+      [Opt(OptOps.SPLIT, 0, (0, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (11, AxisType.UNROLL)),
+       Opt(OptOps.SPLIT, 2, (0, AxisType.GROUP_REDUCE, True)), Opt(OptOps.PADTO, 2, 32)],
+    ])
+    b = Tensor.rand(17, 19)
+    helper_linearizer_opt(b.flip(0).pad(((2, 3), (0, 0))).sum(0), [
+      [Opt(OptOps.SPLIT, 1, (0, AxisType.GROUP_REDUCE, True)), Opt(OptOps.PADTO, 0, 8),
+       Opt(OptOps.SPLIT, 0, (12, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (0, AxisType.LOCAL))],
+    ])
+    x, w = Tensor.rand(1, 3, 15, 15), Tensor.rand(4, 3, 3, 3)
+    helper_linearizer_opt(x.conv2d(w, padding=1, stride=2), [
+      [Opt(OptOps.SPLIT, 5, (0, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 1, (0, AxisType.LOCAL))],
+    ])
+
+  def test_unrolled_padded_cumsum(self):
+    Tensor.manual_seed(7)
+    a = Tensor.rand(13, 17)
+    helper_linearizer_opt(a.cumsum(1), [
+      [Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 0, 4)],
+    ])
+
   def test_upcasts(self):
     N = 16
     Tensor.manual_seed(1772)
@@ -65,7 +87,6 @@ class TestKernelOpts(unittest.TestCase):
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
-  @unittest.skipIf(Device.DEFAULT == "AMD", "TODO: too slow on MOCKKFD, hits the test timeout in CI")
   def test_matmul(self):
     N = 128
     Tensor.manual_seed(1552)
@@ -73,19 +94,13 @@ class TestKernelOpts(unittest.TestCase):
     b = Tensor.rand(N, N)
     r = a@b
     helper_linearizer_opt(r, [
-      [Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], # Checking how it works with upcasts
-      [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))],
       [Opt(OptOps.SPLIT, 1, (32, AxisType.LOCAL))],
       [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL))],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (32, AxisType.LOCAL))],
       [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (8, AxisType.LOCAL))], # Checking how it works with locals
-      [Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))],
       [Opt(OptOps.SPLIT, 2, (32, AxisType.GROUP_REDUCE, True))],
       [Opt(OptOps.SPLIT, 2, (32, AxisType.GROUP_REDUCE, True)),
        Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL))], # Checking how it works with grouped_reduce
       [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (32, AxisType.GROUP_REDUCE, True))],
-      [Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL)), Opt(OptOps.SPLIT, 3, (32, AxisType.GROUP_REDUCE, True))],
       [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL)),
        Opt(OptOps.SPLIT, 4, (4, AxisType.GROUP_REDUCE, True))], # Checking how it works with local+grouped_reduce
       # Checking all together
@@ -136,52 +151,6 @@ class TestKernelOpts(unittest.TestCase):
        Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], # No globals
     ])
 
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
-                      "test requires tensor cores with accumulation in half") # testing with half suffices.
-  @unittest.skipIf(Device.DEFAULT == "AMD", "TODO: the UNROLL axis is hardcoded for the METAL tensor core shape")
-  def test_tensor_core_opts(self):
-    N = 128
-    Tensor.manual_seed(1552)
-    a, b = Tensor.rand(N, N, dtype=dtypes.half), Tensor.rand(N, N, dtype=dtypes.half)
-    r = a.matmul(b, dtype=dtypes.half)
-    atol, rtol = 0.25, 0.01
-    helper_linearizer_opt(r, [
-      [],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))],
-      [Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], # check upcasts
-      [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))], # check unroll
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 5, (2, AxisType.UNROLL))], # check combo of unroll and upcast
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (2, AxisType.UNROLL))],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL))],
-      [Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], # check permutations
-      [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))],
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 5, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))],
-      [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)),
-       Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL))],
-    ], apply_tc=True, atol=atol, rtol=rtol)
-
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
-                      "test requires tensor cores with accumulation in half") # testing with half suffices.
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
-  @unittest.skipIf(Device.DEFAULT == "AMD", "TODO: the UNROLL axis is hardcoded for the METAL tensor core shape")
-  def test_tensor_core_opts_locals(self):
-    N = 128
-    Tensor.manual_seed(1552)
-    a, b = Tensor.rand(N, N, dtype=dtypes.half), Tensor.rand(N, N, dtype=dtypes.half)
-    r = a.matmul(b, dtype=dtypes.half)
-    atol, rtol = 0.25, 0.01
-    helper_linearizer_opt(r, [
-      [Opt(OptOps.SPLIT, 4, (0, AxisType.UNROLL))], # check full unroll of reduce with locals
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL))], # check local
-      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL)),
-       Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))],
-      [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (2, AxisType.UNROLL)),
-       Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))],
-    ], apply_tc=True, atol=atol, rtol=rtol)
-
   def test_padto_matmul(self):
     N = 17
     Tensor.manual_seed(289)
@@ -216,7 +185,6 @@ class TestKernelOpts(unittest.TestCase):
     with self.assertRaises(KernelOptError):
       helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL)), Opt(OptOps.PADTO, 2, 8)]])
 
-  @unittest.skipIf(Device.DEFAULT == "AMD", "TODO: off by one on MOCKKFD in CI, passes locally")
   def test_padto_sum_ok(self):
     N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
@@ -271,29 +239,6 @@ class TestKernelOpts(unittest.TestCase):
     helper_linearizer_opt(a.sum(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (0, AxisType.UNROLL)),
                                       Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))]])
 
-  @unittest.skipUnless(any(tc.dtype_in in (dtypes.half, dtypes.float) for tc in Device[Device.DEFAULT].renderer.tensor_cores),
-                       "test requires half or float tensor cores")
-  def test_tc_shape_padded(self):
-    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in in (dtypes.half, dtypes.float))
-    Tensor.manual_seed(3)
-    a, b = Tensor.rand(17, 23, dtype=tc.dtype_in).realize(), Tensor.rand(23, 29, dtype=tc.dtype_in).realize()
-    with Context(ALLOW_TF32=1):
-      helper_linearizer_opt(a.matmul(b, dtype=tc.dtype_out), [[Opt(OptOps.TC, 0, (-1, 2, 2))]], check_default_opt=False, atol=3e-2, rtol=1e-3)
-
-  @unittest.skipUnless(any(tc.dtype_in in (dtypes.half, dtypes.float) for tc in Device[Device.DEFAULT].renderer.tensor_cores),
-                       "test requires half or float tensor cores")
-  @unittest.skipIf(Device.DEFAULT == "AMD" and Device[Device.DEFAULT].renderer.target.arch.startswith(("gfx11", "gfx12")),
-                   "TODO: LLVM AMDGPU miscompiles RDNA WMMA with masked operands, passes on PYTHON::gfx1100")
-  def test_tc_padto_full_upcast(self):
-    # a fully upcast pad lane makes a WMMA operand entirely Invalid
-    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in in (dtypes.half, dtypes.float))
-    Tensor.manual_seed(3)
-    a, b = Tensor.rand(17, 23, dtype=tc.dtype_in).realize(), Tensor.rand(23, 29, dtype=tc.dtype_in).realize()
-    with Context(ALLOW_TF32=1):
-      helper_linearizer_opt(a.matmul(b, dtype=tc.dtype_out),
-                            [[Opt(OptOps.TC, 0, (-1, 2, 1)), Opt(OptOps.PADTO, 0, 4), Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST))]],
-                            check_default_opt=False, atol=3e-2, rtol=1e-3)
-
   def test_padto_nested_reduce(self):
     a = (Tensor.arange(2*3, dtype=dtypes.float).reshape(2, 3) + 1).clone().realize()  # [[1, 2, 3], [4, 5, 6]]
     # the pad gate has the outer reduce's range, the inner reduce must not resolve it with its own identity
@@ -377,14 +322,9 @@ class TestKernelOpts(unittest.TestCase):
        [("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]),
       # check to ensure local_dims are stable for full UNROLL of the first reduce
       ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.UNROLL))], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
-      ([Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL)),Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
       # check behavior for full UNROLL on an existing GROUP
       ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 3, (2, AxisType.UNROLL))],
        [("blue",16),("blue",32),("cyan",2),("green",16),("magenta",2)]),
-      ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 3, (0, AxisType.UNROLL))],
-       [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
-      ([Opt(OptOps.SPLIT, 2, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL))],
-       [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
       ([Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL))],
        [("blue",32),("blue",32),("red",16),("magenta",2)]),
     ]

test/opt/test_tensor_cores.py

@@ -6,7 +6,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.uop.ops import Ops, UOp, AxisType
 from tinygrad.dtype import DType
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, TC_SELECT, TC_OPT
 from test.helpers import slow, replace_opts
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
@@ -32,6 +32,11 @@ def _skip_unsupported_tc_dtypes(dtype_in:DType, dtype_out:DType):
   if unsupported := [f"{name}={dtype}" for name,dtype in (("dtype_in", dtype_in), ("dtype_out", dtype_out)) if dtype not in supported_dtypes]:
     raise unittest.SkipTest(f"tensor core requires unsupported renderer dtype: {', '.join(unsupported)}")
 
+def tc_reduce_axis(r:Tensor) -> int:
+  sche = Scheduler(r.schedule_linear().src[-1].src[0], Device[Device.DEFAULT].renderer)
+  sche.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1)))
+  return sche.axis_types.index(AxisType.REDUCE)
+
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                          ensure_triggered:bool=True):
   _skip_unsupported_tc_dtypes(dtype_in, dtype_out)
@@ -237,15 +242,13 @@ class TestTensorCores(unittest.TestCase):
   @Context(ALLOW_TF32=1)
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
-  @unittest.skipIf(Device.DEFAULT == "AMD" and Device[Device.DEFAULT].renderer.target.arch.startswith("gfx9"),
-                   "TODO: the UNROLL axis is hardcoded for the METAL tensor core shape")
   def test_tensor_cores_unroll_phi(self):
     # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
     # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
     tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
-    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in).realize(), Tensor.rand(64, 16, dtype=tc.dtype_in).realize()
+    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, tc_reduce_axis(x.matmul(y, dtype=tc.dtype_out)), (2, AxisType.UNROLL))]
     r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))]
     ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA]
     self.assertGreater(len(wmmas), 0)
@@ -255,13 +258,11 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
-  @unittest.skipIf(Device.DEFAULT == "AMD" and Device[Device.DEFAULT].renderer.target.arch.startswith("gfx9"),
-                   "TODO: the UNROLL axis is hardcoded for the METAL tensor core shape")
   def test_tensor_cores_unroll_casted_phi(self):
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in).realize(), Tensor.rand(64, 16, dtype=tc.dtype_in).realize()
+    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, tc_reduce_axis(x.matmul(y, dtype=tc.dtype_out)), (2, AxisType.UNROLL))]
     r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))]
     ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA]
     self.assertGreater(len(wmmas), 0)
@@ -271,18 +272,61 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
-  @unittest.skipIf(Device.DEFAULT == "AMD" and Device[Device.DEFAULT].renderer.target.arch.startswith("gfx9"),
-                   "TODO: the UNROLL axis is hardcoded for the METAL tensor core shape")
   def test_tensor_cores_unroll_casted_phi_with_children(self):
     # all STORE children are outside the loop
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in).realize(), Tensor.rand(64, 16, dtype=tc.dtype_in).realize()
+    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, tc_reduce_axis(x.matmul(y, dtype=tc.dtype_out).relu()), (2, AxisType.UNROLL))]
     r = x.matmul(y, dtype=tc.dtype_out).relu()
-    opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))]
     ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA]
     self.assertGreater(len(wmmas), 0)
     for u in wmmas: assert u.src[-1].src[0].op != Ops.STORE
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
+  @unittest.skipUnless(any(tc.dtype_in == tc.dtype_out == dtypes.half for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                      "test requires tensor cores with accumulation in half") # testing with half suffices.
+  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
+  def test_tensor_core_opts(self):
+    N = 128
+    Tensor.manual_seed(1552)
+    a, b = Tensor.rand(N, N, dtype=dtypes.half).realize(), Tensor.rand(N, N, dtype=dtypes.half).realize()
+    R = tc_reduce_axis(a.matmul(b, dtype=dtypes.half))
+    r = a.matmul(b, dtype=dtypes.half)
+    atol, rtol = 0.25, 0.01
+    helper_linearizer_opt(r, [
+      [],
+      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))],
+      [Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))],
+      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], # check upcasts
+      [Opt(OptOps.SPLIT, R, (2, AxisType.UNROLL))], # check unroll
+      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, R+1, (2, AxisType.UNROLL))], # check combo of unroll and upcast
+      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, R+2, (2, AxisType.UNROLL))],
+      [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, R+2, (4, AxisType.UNROLL))],
+    ], apply_tc=True, atol=atol, rtol=rtol)
+
+  @unittest.skipUnless(any(tc.dtype_in in (dtypes.half, dtypes.float) for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                       "test requires half or float tensor cores")
+  def test_tc_shape_padded(self):
+    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in in (dtypes.half, dtypes.float))
+    Tensor.manual_seed(3)
+    a, b = Tensor.rand(17, 23, dtype=tc.dtype_in).realize(), Tensor.rand(23, 29, dtype=tc.dtype_in).realize()
+    with Context(ALLOW_TF32=1):
+      helper_linearizer_opt(a.matmul(b, dtype=tc.dtype_out), [[Opt(OptOps.TC, 0, (-1, 2, 2))]], check_default_opt=False, atol=3e-2, rtol=1e-3)
+
+  @unittest.skipUnless(any(tc.dtype_in in (dtypes.half, dtypes.float) for tc in Device[Device.DEFAULT].renderer.tensor_cores),
+                       "test requires half or float tensor cores")
+  @unittest.skipIf(Device.DEFAULT == "AMD" and Device[Device.DEFAULT].renderer.target.arch.startswith(("gfx11", "gfx12")),
+                   "TODO: LLVM AMDGPU miscompiles RDNA WMMA with masked operands, passes on PYTHON::gfx1100")
+  def test_tc_padto_full_upcast(self):
+    # a fully upcast pad lane makes a WMMA operand entirely Invalid
+    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in in (dtypes.half, dtypes.float))
+    Tensor.manual_seed(3)
+    a, b = Tensor.rand(17, 23, dtype=tc.dtype_in).realize(), Tensor.rand(23, 29, dtype=tc.dtype_in).realize()
+    with Context(ALLOW_TF32=1):
+      helper_linearizer_opt(a.matmul(b, dtype=tc.dtype_out),
+                            [[Opt(OptOps.TC, 0, (-1, 2, 1)), Opt(OptOps.PADTO, 0, 4), Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST))]],
+                            check_default_opt=False, atol=3e-2, rtol=1e-3)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_call.py

@@ -144,6 +144,33 @@ class TestCallShape(unittest.TestCase):
     self.assertEqual(shape[0], sz.bind(5))
 
 class TestCallSchedule(unittest.TestCase):
+  def test_precompile_slice_assign(self):
+    @function(precompile=True)
+    def f(x:Tensor) -> Tensor: return x * 2 + 1
+    a = Tensor.arange(8).float().realize()
+    cache = Tensor.zeros(16)
+    # the output must land at the slice offset, not at the start of the base buffer
+    cache[4:12].assign(f(a)).realize()
+    np.testing.assert_equal(cache.numpy(), np.concatenate([np.zeros(4), np.arange(8)*2+1, np.zeros(4)]).astype(np.float32))
+
+  def test_precompile_slice_assign_2d(self):
+    @function(precompile=True)
+    def f(x:Tensor) -> Tensor: return x + 1
+    a = Tensor.arange(8).reshape(2, 4).float().realize()
+    big = Tensor.zeros(4, 8)
+    big[1:3, 2:6].assign(f(a)).realize()
+    ref = np.zeros((4, 8), dtype=np.float32)
+    ref[1:3, 2:6] = np.arange(8).reshape(2, 4) + 1
+    np.testing.assert_equal(big.numpy(), ref)
+
+  def test_precompile_full_buffer_assign(self):
+    @function(precompile=True)
+    def f(x:Tensor) -> Tensor: return x * 2 + 1
+    a = Tensor.arange(8).float().realize()
+    cache = Tensor.zeros(8).realize()
+    cache.assign(f(a)).realize()
+    np.testing.assert_equal(cache.numpy(), np.arange(8)*2+1)
+
   def test_reshape_precompile(self):
     a = Tensor.empty(4, 8).realize()
     a = a.reshape(4,4,2).assign(Tensor.empty(4,4,2)).reshape(8,4)
@@ -252,8 +279,8 @@ class TestCallSchedule(unittest.TestCase):
     a = Tensor.empty(4, 8)
     b = Tensor.empty(4, 8)
     r0, r1 = f(a), f(b)
-    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.CALL and u.num_returned)
-    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.CALL and u.num_returned)
+    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.CALL and u.has_unbound_outputs)
+    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.CALL and u.has_unbound_outputs)
     # output identities stay unique per call; they canonicalize only when combined into a scheduling scope
     self.assertIsNot(c0.src[-1], c1.src[-1])
     self.assertEqual(sched_key(r0), sched_key(r1))
@@ -287,26 +314,61 @@ class TestCallSchedule(unittest.TestCase):
     np.testing.assert_allclose(out.numpy(), np.arange(8, dtype=np.float32).reshape(4, 2) + 3)
 
 class TestArgOrder(unittest.TestCase):
-  """RETURNED placeholders can appear anywhere in a call's srcs: slots are src positions, nothing reorders"""
+  """outputs can appear anywhere in a call's srcs (output_pos): slots are src positions, nothing reorders"""
+  def _dev(self, x): return x.device if isinstance(x.device, str) else (x.device or (Device.DEFAULT,))[0]
   def make_intersperse_call(self, x, precompile=False):
-    # call with sources (body, returned, input(slot=1)): the input is the input, the output binds the RETURNED
-    dev = x.device if isinstance(x.device, str) else (x.device or (Device.DEFAULT,))[0]
-    r0 = UOp.returned(x.dtype, x.shape, device=dev)
-    o0 = UOp.param(0, x.dtype, x.shape, dev)
-    p1 = UOp.param(1, x.dtype, x.shape, dev)
-    from tinygrad.uop.ops import CallInfo
-    return UOp(Ops.CALL, src=(UOp.sink(o0.store(p1.reshape(x.shape) * 2)), r0, x.uop),
-               arg=CallInfo(None, 't', precompile, False, None))
+    # the output is at position 0, the input (param slot 1) at position 1 in the call's args
+    val = UOp.param(1, x.dtype, x.shape, self._dev(x)).reshape(x.shape) * 2
+    return UOp.call_with_outputs((val,), x.uop, name='t', output_pos=(0,), precompile=precompile)
 
   def test_intersperse_returned(self):
     x = Tensor.arange(3, dtype=dtypes.int).realize()
-    call = self.make_intersperse_call(x)
-    out = Tensor(call.returned_outputs[0], device=x.device) + 1
+    outs = self.make_intersperse_call(x)
+    out = Tensor(outs[0], device=x.device) + 1
     np.testing.assert_equal(out.numpy(), [1, 3, 5])
 
+  def test_outputs_arbitrary_order(self):
+    x = Tensor([1.0, 2.0, 3.0])
+    y = Tensor([4.0, 5.0, 6.0])
+    x.requires_grad = True
+    y.requires_grad = True
+    x, y = x.realize(), y.realize()
+    dev = self._dev(x)
+    # args (out0, in0, out1, in1): outputs at positions 0 and 2, input params slotted at their final positions 1 and 3
+    p1, p3 = UOp.param(1, x.dtype, x.shape, dev), UOp.param(3, y.dtype, y.shape, dev)
+    outs = UOp.call_with_outputs((p1.reshape(x.shape) * 2, p3.reshape(y.shape) + p1.reshape(y.shape)), x.uop, y.uop,
+                                 output_pos=(0, 2))
+    np.testing.assert_equal(Tensor(outs[0]).numpy(), [2, 4, 6])
+    np.testing.assert_equal(Tensor(outs[1]).numpy(), [5, 7, 9])
+    # the auto gradient path (no grad_fxn) resolves outputs and gradients positionally at any position
+    (Tensor(outs[0]).sum() + Tensor(outs[1]).sum()).backward()
+    np.testing.assert_equal(x.grad.numpy(), [3, 3, 3])
+    np.testing.assert_equal(y.grad.numpy(), [1, 1, 1])
+
+  def test_output_pos_symbolic_shape(self):
+    # symbolic output shapes resolve against the final arg slots, not the input order (PARAM(2) in the shape, output at 0)
+    x = Tensor.empty(8).realize()
+    sz = UOp.variable('sz', 1, 8)
+    dev = self._dev(x)
+    p1, p2 = UOp.param(1, x.dtype, x.shape, dev), sz.param_like(2)
+    value = p1.reshape(x.shape).shrink_to((p2,))
+    bound = sz.bind(5)
+    outs = UOp.call_with_outputs((value,), x.uop, bound, output_pos=(0,))
+    # the minted output's shape substituted PARAM(2) with the bind arg from position 2 in the arg list
+    shp = outs[0].shape[0]
+    self.assertIsInstance(shp, UOp)
+    self.assertNotEqual(shp.op, Ops.PARAM)
+    self.assertEqual(shp, bound)
+
+  def test_output_pos_must_be_ascending(self):
+    x = Tensor.arange(3, dtype=dtypes.int).realize()
+    p1 = UOp.param(1, x.dtype, x.shape, self._dev(x))
+    with self.assertRaises(AssertionError):
+      UOp.call_with_outputs((p1.reshape(x.shape) * 2, p1.reshape(x.shape) + 1), x.uop, output_pos=(1, 0))
+
   def test_intersperse_returned_precompile(self):
     x = Tensor.arange(3, dtype=dtypes.int).realize()
-    call = self.make_intersperse_call(x, precompile=True)
+    call = self.make_intersperse_call(x, precompile=True)[0].src[1]
     # the transform must preserve the RETURNED's src position: its placeholder is at src 1, the input stays at src 2
     from tinygrad.tensor import transform_precompiled_call
     new = transform_precompiled_call(call)
@@ -323,14 +385,10 @@ class TestArgOrder(unittest.TestCase):
   def test_intersperse_returned_gradient(self):
     x = Tensor([1.0, 2.0, 3.0]).realize()
     x.requires_grad = True
-    dev = x.device if isinstance(x.device, str) else (x.device or (Device.DEFAULT,))[0]
-    r0 = UOp.returned(dtypes.float, x.shape, device=dev)
-    o0 = UOp.param(0, dtypes.float, x.shape, dev)
-    p1 = UOp.param(1, dtypes.float, x.shape, dev)
-    from tinygrad.uop.ops import CallInfo
-    body = UOp.sink(o0.store(p1.reshape(x.shape) * p1.reshape(x.shape)))
-    call = UOp(Ops.CALL, src=(body, r0, x.uop), arg=CallInfo(None, 't', False, False, None))
-    y = Tensor(call.returned_outputs[0], device=x.device)
+    p1 = UOp.param(1, dtypes.float, x.shape, self._dev(x))
+    val = p1.reshape(x.shape) * p1.reshape(x.shape)
+    outs = UOp.call_with_outputs((val,), x.uop, name='t', output_pos=(0,))
+    y = Tensor(outs[0], device=x.device)
     y.sum().backward()
     np.testing.assert_equal(x.grad.numpy(), [2, 4, 6])
 

test/unit/test_call_access.py

@@ -0,0 +1,66 @@
+import unittest
+from tinygrad import Tensor, Context, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
+
+
+class TestCallAccess(unittest.TestCase):
+  def test_computed_reads_writes_and_unused_arguments(self):
+    out, x, unused = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(3))
+    body = out.store(x + 1).sink(arg=KernelInfo())
+    self.assertEqual(body.call(out, x, unused).call_access(), ((x,), (out,)))
+
+  def test_computed_read_modify_write(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    body = out.store(out + x).sink(arg=KernelInfo())
+    self.assertEqual(body.call(out, x).call_access(), ((out, x), (out,)))
+
+  def test_computed_empty_effects(self):
+    x = UOp.param(0, dtypes.float, (1,), "CPU")
+    self.assertEqual(UOp.sink(x, arg=KernelInfo()).call(x).call_access(), ((), ()))
+
+  def test_computed_program_accesses(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    sink = out.store(x.load()).sink(arg=KernelInfo())
+    program = UOp(Ops.PROGRAM, src=(sink,), arg=ProgramInfo.from_sink(sink))
+    self.assertEqual(program.call(out, x).call_access(), ((x,), (out,)))
+
+  def test_nested_linear_parameter_scopes(self):
+    a, b, c = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(3))
+    inner = a.store(b + 1).sink(arg=KernelInfo()).call(b, a)
+    body = UOp(Ops.LINEAR, src=(inner,))
+    self.assertEqual(body.call(a, b, c).call_access(), ((a,), (b,)))
+
+  def test_copy_accesses(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    self.assertEqual(UOp(Ops.COPY, src=(x,), arg=out.device).call(out, x).call_access(), ((x,), (out,)))
+
+  def test_unknown_opaque_accesses_reject(self):
+    x = UOp.param(0, dtypes.float, (1,), "CPU")
+    bodies = (UOp(Ops.PROGRAM, src=(UOp.sink(x),)),
+              UOp(Ops.CUSTOM, src=(x,), arg=("", dtypes.void)).sink(arg=KernelInfo()))
+    for body in bodies:
+      with self.assertRaisesRegex(RuntimeError, "cannot compute accesses"): body.call(x).call_access()
+
+  @Context(DEV="CPU")
+  def test_unknown_effects_do_not_replace_tensors_on_failure(self):
+    def kernel(x): return UOp(Ops.PROGRAM, src=(UOp.sink(x, arg=KernelInfo()),))
+    x = Tensor([2.]).realize().custom_kernel(fxn=kernel)[0]
+    before = x.uop
+    for _ in range(2):
+      with self.assertRaisesRegex(RuntimeError, "cannot compute accesses"): x.realize()
+      self.assertIs(x.uop, before)
+
+  def test_bad_access_slots(self):
+    arg = UOp.param(0, dtypes.float, (1,), "CPU")
+    for slot in (-1, 1):
+      p = UOp(Ops.PROGRAM, src=(UOp.sink(arg),), arg=ProgramInfo(globals=(0,), ins=(slot,), outs=()))
+      with self.assertRaisesRegex(RuntimeError, "invalid CALL access slot"): p.call(arg, arg).call_access()
+
+  def test_compiled_writable_alias_rejects(self):
+    a, b = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    sink = a.store(b.load()).sink(arg=KernelInfo())
+    program = UOp(Ops.PROGRAM, src=(sink,), arg=ProgramInfo.from_sink(sink))
+    with self.assertRaisesRegex(RuntimeError, "aliased opaque"): program.call(a, a).call_access()
+
+
+if __name__ == "__main__": unittest.main()

test/unit/test_callify.py

@@ -107,5 +107,20 @@ class TestCallify(unittest.TestCase):
     self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
     self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
 
+  def test_intermediate_clone_persists(self):
+    x = (Tensor([1, 2, 3]).realize() + 1).clone()
+    y = (x * 2).realize()
+    self.assertTrue(x.uop.has_buffer_identity())
+    self.assertEqual(x.tolist(), [2, 3, 4])
+    self.assertEqual(y.tolist(), [4, 6, 8])
+
+  def test_zero_size_cat_with_rng(self):
+    # Empty outputs must not replay a pending RNG counter update.
+    a = Tensor.rand(2, 2)
+    b = Tensor.rand(2, 0)
+    t = a.cat(b, dim=1).realize()
+    self.assertEqual(t.shape, (2, 2))
+    self.assertListEqual(t.tolist(), a.tolist())
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_weak.py

@@ -108,6 +108,13 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertIs(stacked.dtype, dtypes.weakfloat)
     self.assertEqual(stacked.tolist(), [2.0, -3.0])
 
+  def test_weakint_cast_truncates_for_every_consumer(self):
+    # a weakint cast of a float is a truncation whether a cast, a compare or an arithmetic op consumes it
+    x = Tensor([2.5, -3.5], dtype=dtypes.float32, device="CPU")
+    self.assertEqual(x.cast(dtypes.weakint).cast(dtypes.float32).tolist(), [2.0, -3.0])
+    self.assertEqual((x.cast(dtypes.weakint) * x).tolist(), [5.0, 10.5])
+    self.assertEqual(Tensor([0.5, -0.5], dtype=dtypes.float32, device="CPU").cast(dtypes.weakint).cast(dtypes.bool).tolist(), [False, False])
+
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),

test/unit/test_gradient.py

@@ -88,6 +88,72 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])     # gradient flows through clone
     np.testing.assert_allclose(base.grad.numpy(), [0.0, 0.0, 0.0, 0.0])  # ...but detach blocks it from base
 
+  def test_gradient_through_single_assign(self):
+    x = Tensor([2., 3.]).realize()
+    y = x.clone()
+    y.assign(y.square())
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [4., 6.])
+
+  def test_gradient_through_assign_requires_old_versions(self):
+    for count in (2, 3):
+      with self.subTest(count=count):
+        x = Tensor([2., 3.]).realize()
+        y = x.clone()
+        for _ in range(count): y.assign(y.square())
+        g = y.sum().gradient(x)[0]
+        before = (x.uop, y.uop, g.uop)
+        # Reject incompatible versions, including on retry: failed scheduling must not replace them with buffers.
+        for _ in range(2):
+          with self.assertRaisesRegex(RuntimeError, "cycle"): g.realize()
+          self.assertEqual((x.uop, y.uop, g.uop), before)
+
+  def test_gradient_through_assign_with_snapshots(self):
+    x = Tensor([2., 3.]).realize()
+    y = x.clone()
+    for _ in range(2): y.assign(y.clone().square())
+    g = y.sum().gradient(x)[0]
+    gg = g.sum().gradient(x)[0]
+    Tensor.realize(g, gg)
+    self.assertEqual(g.tolist(), [32., 108.])
+    self.assertEqual(gg.tolist(), [48., 108.])
+
+  def test_gradient_after_unrelated_store(self):
+    x, v, dst = Tensor([2.]).realize(), Tensor([3.]).realize(), Tensor.empty(1)
+    y = Tensor(x.uop.after(dst.uop.store(v.uop)))
+    self.assertEqual([g.tolist() for g in y.sum().gradient(x, v)], [[1.], [0.]])
+    self.assertEqual(y.tolist(), [2.])
+    self.assertEqual(dst.tolist(), [3.])
+
+  def test_gradient_after_multiple_unrelated_stores(self):
+    x, a, b = Tensor([2.]).realize(), Tensor.empty(1), Tensor.empty(1)
+    y = Tensor(x.uop.after(a.uop.store(x.uop * 3), b.uop.store(x.uop * 4)))
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [1.])
+
+  def test_gradient_after_readonly_call(self):
+    x = Tensor([2.]).realize()
+    def kernel(dst, src): return dst.store(src * 3).sink(arg=KernelInfo())
+    for grad_fxn in (None, lambda g, k: (None, g * 3)):
+      _, unchanged = Tensor.empty(1).custom_kernel(x, fxn=kernel, grad_fxn=grad_fxn)
+      self.assertEqual(unchanged.sum().gradient(x)[0].tolist(), [1.])
+
+  def test_gradient_after_unrelated_call(self):
+    x, v, dst = Tensor([2.]).realize(), Tensor([3.]).realize(), Tensor.empty(1)
+    p, q = dst.uop.param_like(0), v.uop.param_like(1)
+    call = p.store(q * 3).sink(arg=KernelInfo()).call(dst.uop, v.uop, grad_fxn=lambda g, k: (None, g * 3))
+    y = Tensor(x.uop.after(call))
+    self.assertEqual([g.tolist() for g in y.sum().gradient(x, v)], [[1.], [0.]])
+
+  def test_gradient_after_aliased_store_view_rejects(self):
+    x = Tensor([2., 3.]).realize()
+    y = Tensor(x.uop.after(x.uop.shrink(((0, 1),)).store(4.)))
+    with self.assertRaisesRegex(RuntimeError, "aliased write"): y.sum().gradient(x)
+
+  def test_gradient_after_duplicate_call_output_rejects(self):
+    x = Tensor([2.]).realize()
+    def kernel(a, b): return a.store(b * 2).sink(arg=KernelInfo())
+    y = x.custom_kernel(x, fxn=kernel, grad_fxn=lambda g, k: (g, g))[0]
+    with self.assertRaisesRegex(RuntimeError, "ambiguous CALL"): y.sum().gradient(x)
+
   def test_setitem_on_grad_used_tensor_raises(self):
     x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
     _ = (x * 2.0).sum()
@@ -136,6 +202,21 @@ class TestTensorGradient(unittest.TestCase):
     self.assertIsNone(w.grad)
 
 class TestMultiOutputGradient(unittest.TestCase):
+  def test_custom_kernel_inplace_gradient(self):
+    def double(x:UOp): return x[0].store(x[0]*2).sink(arg=KernelInfo(name="double_inplace"))
+    def backward(g:UOp, call:UOp): return (g*2,)
+    x = Tensor([2.]).realize()
+    y = x.custom_kernel(fxn=double, grad_fxn=backward)[0]
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [2.])
+    self.assertEqual(y.tolist(), [4.])
+
+  def test_custom_kernel_unchanged_output_gradient(self):
+    def noop(x:UOp): return x[0].store(x[0]).sink(arg=KernelInfo(name="identity"))
+    def backward(g:UOp, call:UOp): return (g,)
+    x = Tensor([2.]).realize()
+    y = x.custom_kernel(fxn=noop, grad_fxn=backward)[0]
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [1.])
+
   @staticmethod
   def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
     C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()

test/unit/test_llm_amd.py

@@ -1,7 +1,8 @@
 import unittest
+from unittest.mock import patch
 import numpy as np
 from tinygrad import Tensor, UOp, dtypes, nn, function
-from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention
+from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention, gated_delta_prefill
 from tinygrad.llm.gguf import ggml_data_to_tensor
 
 class TestQ8Quantize(unittest.TestCase):
@@ -28,6 +29,12 @@ class TestQ8Quantize(unittest.TestCase):
     # xsum holds the two per-16 sums per 32-wide group
     np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))
 
+  def test_quantize_rounding_ties(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    values = np.array([-127,127]+[i+0.5 for i in range(-15,15)],dtype=np.float32)
+    quant,_,_ = q8_quantize(Tensor(values),1,32)
+    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(32).numpy(),np.rint(values).astype(np.int8))
+
   def test_q6_linear_compiles_in_function(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
@@ -44,22 +51,125 @@ class TestQ8Quantize(unittest.TestCase):
     self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
     self.assertEqual(linear.weight.dtype, dtypes.uint32)
 
-  def test_q4_k_linear(self):
+  def test_q4_k_linear(self): self._test_quant_linear(12, 144)
+  def test_iq4_linear(self): self._test_quant_linear(23, 136)
+  def test_q5_linear(self): self._test_quant_linear(13, 176)
+
+  def test_quant_linear_partial_output_tile(self):
+    # Cover a sub-tile output, a trailing tile, and IQ4's larger-output tile selection.
+    for typ, size, outputs, tokens in ((12, 144, 16, 16), (12, 144, 48, 32), (13, 176, 48, 16), (23, 136, 4112, 32)):
+      with self.subTest(ggml_type=typ, out_features=outputs):
+        self._test_quant_linear(typ, size, in_features=256, out_features=outputs, token_counts=(tokens,))
+
+  def test_quant_linear_preserves_rope_permutation(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
-    in_features, blocks = 2048, 16*2048//256
-    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
-    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
+    for typ, size in ((12, 144), (13, 176), (14, 210), (23, 136)):
+      with self.subTest(ggml_type=typ):
+        packed = rng.integers(0, 256, (16, size), dtype=np.uint8)
+        packed[:, -2:] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        if typ != 14: packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        if typ in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
+        raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 16*256, typ).reshape(16, 256).half()
+        original = decoded.numpy()
+        x = rng.normal(size=(3, 256)).astype(np.float16)
+        for prefix in (None, 0, 4):
+          with self.subTest(prefix=prefix):
+            w = decoded.reshape(2, 8, 256)
+            if prefix is None:
+              weight = w.rearrange("n (h two) d -> n (two h) d", two=2)
+            else:
+              weight = w[:, :prefix].cat(w[:, prefix:].rearrange("n (h two) d -> n (two h) d", two=2), dim=1)
+            start = prefix or 0
+            rows = np.arange(16).reshape(2, 8)
+            order = np.concatenate((rows[:, :start], rows[:, start:].reshape(2, -1, 2).transpose(0, 2, 1).reshape(2, -1)), axis=1)
+            linear = Linear(256, 16, bias=False)
+            linear.weight = weight.reshape(16, 256)
+            np.testing.assert_allclose(linear(Tensor(x)).numpy(), x.astype(np.float32) @ original[order.flatten()].astype(np.float32).T,
+                                       rtol=3e-3, atol=2e-2)
+            self.assertIsNone(linear.ggml_type)
+
+  def test_quant_linear_rejects_unaligned_rows_and_integer_casts(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for width in (128, 256):
+      with self.subTest(width=width):
+        packed = np.zeros((2*width//256, 136), dtype=np.uint8)
+        packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        packed[:, 8:] = np.arange(128, dtype=np.uint8)
+        raw = Tensor(np.pad(packed.flatten(), (4, 0))).realize()[4:]
+        weight = ggml_data_to_tensor(raw, 2*width, 23).reshape(2, width)
+        if width == 256: weight = weight.int().float()
+        expected = weight.numpy().sum(-1)[None]
+        linear = Linear(width, 2, bias=False)
+        linear.weight = weight
+        np.testing.assert_allclose(linear(Tensor.ones(1, width)).numpy(), expected, rtol=1e-3, atol=1e-3)
+        self.assertIsNone(linear.ggml_type)
+
+  def test_dense_gemv_preserves_integer_casts(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    linear = Linear(128, 1)
+    linear.weight = Tensor.full((1, 128), 0.75).contiguous().realize().int().float()
+    linear.bias = Tensor.full((1,), 0.75).contiguous().realize().int().float()
+    np.testing.assert_array_equal(linear(Tensor.ones(1, 128)).numpy(), 0)
+
+  def test_dense_gemv_float32_range(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    linear = Linear(128, 1, bias=False)
+    linear.weight = Tensor.full((1, 128), 1/128, dtype=dtypes.float32).realize()
+    np.testing.assert_array_equal(linear(Tensor.full((1, 128), 65536, dtype=dtypes.float32)).numpy(), 65536)
+
+  def test_gated_delta_state_and_precision(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for case in ("view", "reset", "half"):
+      with self.subTest(case=case):
+        q = Tensor.full((1, 1, 1, 32), 256 if case == "half" else 1, dtype=dtypes.half if case == "half" else dtypes.float32)
+        state = Tensor.full((1, 1, 32, 4), int(case == "reset"), dtype=dtypes.float32).contiguous().realize().transpose(-1, -2)
+        if case != "view": state = state.contiguous().realize()
+        start = Tensor(UOp.variable("start_pos", 0, 10).bind(0)) if case == "reset" else None
+        beta = Tensor.full((1, 1, 1), 1/2097152 if case == "half" else 1, dtype=dtypes.float32)
+        if case != "reset":
+          message = "recurrent state must be contiguous" if case == "view" else "recurrent Q/K must be float32"
+          with self.assertRaisesRegex(AssertionError, message):
+            gated_delta_prefill(q, q, Tensor.ones(1, 1, 1, 4), beta, Tensor.ones(1, 1, 1), state, start)
+          continue
+        out = gated_delta_prefill(q, q, Tensor.ones(1, 1, 1, 4), beta, Tensor.ones(1, 1, 1), state, start)
+        np.testing.assert_array_equal(out.numpy(), 32)
+        np.testing.assert_array_equal(state.numpy(), 1)
+
+  def test_dense_gemv_bias(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    w, bias = rng.normal(size=(32, 128)).astype(np.float16), rng.normal(size=32).astype(np.float16)
+    linear = Linear(128, 32)
+    linear.weight, linear.bias = Tensor(w), Tensor(bias)
+    for tokens in (1, 3):
+      with self.subTest(tokens=tokens):
+        x = rng.normal(size=(tokens, 128)).astype(np.float16)
+        np.testing.assert_allclose(linear(Tensor(x)).numpy(), x.astype(np.float32) @ w.astype(np.float32).T + bias, rtol=2e-3, atol=2e-3)
+
+  def _test_quant_linear(self, ggml_type, block_bytes, in_features=2048, out_features=64, token_counts=(1, 3, 32, 64, 128)):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, (out_features*in_features//256, block_bytes), dtype=np.uint8)
+    packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+    if ggml_type in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, out_features*in_features, ggml_type).reshape(out_features, in_features)
     weight = decoded.numpy()
-    linear = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    x = rng.normal(size=(3, in_features)).astype(np.float32)
-    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
-    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
-    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, 12)
+    linear = Linear(in_features, out_features, bias=False)
+    linear.weight = decoded
+    for tokens in token_counts:
+      with self.subTest(tokens=tokens):
+        x = rng.normal(size=(tokens, in_features)).astype(np.float32 if tokens == 3 else np.float16)
+        reference_x = x.astype(np.float32)
+        if tokens < 16:
+          grouped = reference_x.reshape(tokens, -1, 32)
+          scale = np.maximum(np.abs(grouped).max(-1, keepdims=True) / 127, 1e-8)
+          reference_x = (np.clip(np.rint(grouped/scale), -127, 127)*scale).reshape(tokens, in_features)
+        reference_w = weight if tokens < 16 else weight.astype(np.float16).astype(np.float32)
+        np.testing.assert_allclose(linear(Tensor(x)).numpy(), reference_x @ reference_w.T, rtol=3e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, ggml_type)
 
   def test_q6_linear_multiple_tokens(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -86,6 +196,18 @@ class TestQ8Quantize(unittest.TestCase):
     self.assertTrue(generic.use_custom_quant)
     self.assertEqual(generic.ggml_type, 14)
 
+  def test_attention_fallback_shapes(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for tokens, capacity, dim in ((1, 65, 64), (32, 64, 32), (32, 64, 384), (32, 64, 512)):
+      with self.subTest(tokens=tokens, capacity=capacity, dim=dim):
+        valid = 33
+        cache = np.full((2, 1, 1, capacity, dim), np.nan, dtype=np.float16)
+        cache[0, :, :, :valid] = 0
+        cache[1, :, :, :valid] = np.arange(valid)[:, None]
+        q = Tensor.zeros(1, 2, tokens, dim, dtype=dtypes.half)
+        expected = np.broadcast_to(np.arange(valid-tokens, valid)[None, None, :, None]/2, q.shape)
+        np.testing.assert_allclose(flash_attention(q, Tensor(cache), valid).numpy(), expected, rtol=1e-3, atol=1e-3)
+
   def test_attention_uses_physical_cache_length(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
@@ -94,14 +216,71 @@ class TestQ8Quantize(unittest.TestCase):
     out = flash_attention(q, assigned, 1).realize()
     np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
 
-  def test_flash_attention_decode_gqa_output_layout(self):
+  def test_flash_attention_decode_symbolic_gqa(self):
+    with patch.object(Tensor, "scaled_dot_product_attention", side_effect=AssertionError("expected custom decode")):
+      self._test_flash_decode(8, 2, 256, 128, 37, symbolic=True)
+
+  def test_flash_attention_decode_gqa_tail(self): self._test_flash_decode(3, 1, 192, 64, 37)
+
+  def test_flash_attention_decode_gqa_output_layout(self): self._test_flash_decode(4, 1, 128, 256, 3)
+  def test_flash_attention_decode_large_gqa_group(self): self._test_flash_decode(8, 1, 256, 256, 73)
+
+  def _test_flash_decode(self, heads, kv_heads, dim, n, valid, symbolic=False):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    Tensor.manual_seed(42)
-    q = Tensor.randn(1, 4, 1, 128, dtype=dtypes.half).realize()
-    cache = Tensor.randn(2, 1, 1, 256, 128, dtype=dtypes.half).realize()
-    out = flash_attention(q, cache, 3).realize()
-    expected = q.scaled_dot_product_attention(cache[0, :, :, :3], cache[1, :, :, :3], enable_gqa=True)
-    np.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=2e-3, atol=2e-3)
+    rng = np.random.default_rng(42)
+    q = rng.normal(size=(1, heads, 1, dim)).astype(np.float16)
+    cache = rng.normal(size=(2, 1, kv_heads, n, dim)).astype(np.float16)
+    k, v = (np.repeat(c[0, :, :valid].astype(np.float32), heads//kv_heads, axis=0) for c in cache)
+    scores = q[0].astype(np.float32) @ k.transpose(0, 2, 1) / np.sqrt(dim)
+    probs = np.exp(scores - scores.max(-1, keepdims=True))
+    expected = (probs / probs.sum(-1, keepdims=True)) @ v
+    cache_tensor = Tensor(cache)
+    if symbolic:
+      start_pos = UOp.variable("start_pos", 0, n-1).bind(valid-1)
+      valid = start_pos + 1
+      cache_tensor = Tensor(cache_tensor.realize().uop.after(Tensor(start_pos).uop))
+    np.testing.assert_allclose(flash_attention(Tensor(q), cache_tensor, valid).numpy(), expected[None], rtol=2e-3, atol=2e-3)
+
+  def test_prefill_attention_nonfinite_cache_tail(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    q = Tensor.zeros(1, 2, 32, 128, dtype=dtypes.half)
+    values = rng.normal(size=(33, 128)).astype(np.float16)
+    expected = np.stack([values[:i+2].astype(np.float32).mean(0) for i in range(32)])[None, None].repeat(2, axis=1)
+    for tail in (np.nan, np.inf, -np.inf):
+      with self.subTest(tail=tail):
+        cache = np.full((2, 1, 1, 64, 128), tail, dtype=np.float16)
+        cache[0, :, :, :33] = 0
+        cache[1, :, :, :33] = values
+        valid = UOp.variable("valid_end", 32, 64).bind(33)
+        cache_tensor = Tensor(cache).realize()
+        assigned = Tensor(cache_tensor.uop.after(Tensor(valid).uop))
+        out = flash_attention(q, assigned, valid)
+        np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+  def test_flash_attention_decode_beyond_256_chunks(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    n = 257 * 64
+    q = Tensor.zeros(1, 1, 1, 32, dtype=dtypes.half).realize()
+    k = Tensor.zeros(1, 1, n, 32, dtype=dtypes.half)
+    v = Tensor.zeros(1, 1, n-64, 32, dtype=dtypes.half).cat(Tensor.ones(1, 1, 64, 32, dtype=dtypes.half), dim=2)
+    cache = Tensor.stack(k, v).contiguous().realize()
+    for valid, expected in ((1, 0), (n, 1/257)):
+      with self.subTest(valid=valid):
+        valid_kv_len = UOp.variable("valid_kv_len", 1, n).bind(valid)
+        assigned = Tensor(cache.uop.after(Tensor(valid_kv_len).uop))
+        np.testing.assert_allclose(flash_attention(q, assigned, valid_kv_len).numpy(), expected, rtol=2e-3, atol=2e-4)
+
+  def test_flash_attention_decode_long_context_random(self):
+    self._test_flash_decode(8, 2, 128, 257*64, 257*64-13)  # past 256 chunks, with a ragged tail
+
+  def test_flash_attention_decode_chunk_round_accumulator_range(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    valid_kv_len, max_kv_len = 6749, 6784  # three chunk rounds, with a ragged tail
+    q = Tensor.zeros(1, 8, 1, 32, dtype=dtypes.half).realize()
+    cache = Tensor.stack(Tensor.zeros(1, 1, max_kv_len, 32, dtype=dtypes.half),
+                         Tensor.full((1, 1, max_kv_len, 32), 5500, dtype=dtypes.half)).contiguous().realize()
+    np.testing.assert_allclose(flash_attention(q, cache, valid_kv_len).numpy(), 5500, rtol=2e-3, atol=2e-3)
 
   def test_prefill_attention_unaligned_start(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")

tinygrad/codegen/__init__.py

@@ -439,7 +439,7 @@ def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
   lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
   # isa renderers need to allocate registers
   if isinstance(ctx, ISARenderer):
-    if ctx.pre_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
+    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
     # register definitions (INS without srcs) move to the top so regalloc sees their live ranges span the whole program (callee saved regs)
     lst = sorted(lst, key=lambda u: u.op is not Ops.INS or bool(u.src))
     regalloc_ctx = LinearScanRegallocContext(lst, ctx)

tinygrad/codegen/late/regalloc.py

@@ -19,15 +19,15 @@ class LinearScanRegallocContext:
     # compute live ranges
     self.live_range: dict[Register, list[int]] = {}
     lr = self.live_range
-    ranges: list[Register] = []
-    for i,u in enumerate(reversed(uops)):
+    loops: dict[int, int] = {} # the interval of each loop, from its RANGE to the last uop that reads that RANGE
+    for idx,u in reversed(list(enumerate(uops))):
       if u.op in PSEUDO_OPS: continue
       defs = u.tag if isinstance(u.tag, tuple) else ()
       for v in defs + tuple(greg(s) for s in dedup(u.src)):
-        if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i)
+        if isinstance(v, Register): lr.setdefault(v, []).insert(0, idx)
       for v in defs:
-        if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n)
-      if u.op is Ops.RANGE: ranges.append(greg(u))
+        if v in lr and (n:=max((e for s,e in loops.items() if s <= lr[v][-1] < e), default=None)): lr[v].append(n)
+      if u.op is Ops.RANGE: loops[idx] = max(j for j,x in enumerate(uops) if u in x.src)
 
     # allocate registers
     self.stack_size: int = 0
@@ -90,7 +90,7 @@ class LinearScanRegallocContext:
       # loop prologue, avoid loading inside the loop
       if u.op is Ops.RANGE:
         # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
-        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < lr[greg(u)][-1] for l in lr[v])]
+        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < loops[i] for l in lr[v])]
         sorted_uses = sorted(used_in_loop, key=lambda k: (next(l-i for l in lr[k] if l >= i), lr[k][0], k.name, k.index))
         live_in: dict[Register, Register] = {}
         for v in sorted_uses:

tinygrad/device.py

@@ -187,7 +187,7 @@ class Buffer:
     # zero copy with as_memoryview (disabled by default due to use after free)
     if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'):
       if not no_sync: self.allocator.dev.synchronize()
-      return self.allocator._as_buffer(self._buf)
+      if (mv:=self.allocator._as_buffer(self._buf)) is not None: return mv
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv
@@ -280,9 +280,13 @@ class DepsTracker:
       if i in write:
         for dmap in [self.w_dependency_map, self.r_dependency_map]:
           kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
+          for entry in dmap[key]:
+            st, en, dep = entry
+            if st == en: continue
+            if en <= s or e <= st: kept.append(entry)
+            else:
+              if st < s: kept.append((st, s, dep))
+              if e < en: kept.append((e, en, dep))
           dmap[key] = kept
         self.w_dependency_map[key].append((s, e, new_dependency))
       else: self.r_dependency_map[key].append((s, e, new_dependency))
@@ -337,8 +341,9 @@ class Compiled:
 
   has_copy_queue:bool = True
 
-  pm_encode:Any = None # per queue kind: queue ops -> flat command words
-  pm_lower:Any = None # per queue kind: custom_function(submit, cmdbuf) -> the queue push
+  pm_batch:Any = None
+  pm_encode:Any = None
+  pm_lower:Any = None
   pm_bufferize:Any = None
 
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):

tinygrad/engine/jit.py

@@ -166,7 +166,7 @@ class CapturedJit(Generic[ReturnType]):
   expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]]  # (view, variables, dtype, device) per input
 
   @functools.cached_property
-  def linear(self) -> UOp: return link_linear(self._linear)
+  def linear(self) -> UOp: return link_linear(self._linear, allow_cache=False) # do not cache jit
 
   def __reduce__(self): return self.__class__, (self.ret, self._linear, self.expected_names, self.expected_input_info)
 

tinygrad/engine/realize.py

@@ -3,7 +3,7 @@ from typing import cast, Iterator, Any, Sequence
 import weakref, decimal, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, prod, flatten, Context, to_tuple, tqdm, dedup
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, HCQ2, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates, Renderer
@@ -130,7 +130,7 @@ class ExecContext:
   cache: bool = True
 
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.MSELECT, Ops.SHRINK): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK, Ops.BITCAST): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
   if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
   return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -194,10 +194,11 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   if (info:=call.arg.aux).inputs:
     addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr for u, dev in info.inputs]
     cast(Buffer, call.src[1 + info.table].buffer)._buf.cpu_view().view(fmt='Q')[:] = array.array('Q', addrs)
-  ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,)) # the body runs on the runtime device, it drives every device's queues
+  ctx = replace(ctx, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}})
+  ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,))
   if not (ctx.wait or PROFILE): return ets
 
-  slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots} # the batch's timestamps live in its slots
+  slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots}
   def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
     (d:=cast(Any, Device[device])).prof_ents[(slots[device], prof[0])] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
     if not ctx.wait: return None
@@ -279,24 +280,25 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = lower_and_compile(linear)
-  if HCQ2: linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
+  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
-def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if HCQ2 else linear
+def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
+  return hcq_link(linear, input_uops=input_uops, allow_cache=allow_cache)
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), cache=False) # a one-shot link
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), input_uops=inputs)
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
-  for call in linear.src: track_stats(ctx, call, perf_counter_us(), pm_exec.rewrite(call, ctx))
+  for call in linear.src: track_stats(ctx, call.without_after, perf_counter_us(), pm_exec.rewrite(call.without_after, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), allow_cache=ctx.cache)
   while True:
     if clear_l2:
       if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
       else:
         from tinygrad.tensor import Tensor
         with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
-    yield max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])
+    yield max(pm_exec.rewrite(linear.src[0].without_after, ctx) or [0.0])

tinygrad/function.py

@@ -26,8 +26,9 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
 def renumber_invalid_outputs(uret:UOp) -> UOp:
+  invalid = invalid_outputs(uret)
   return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
-                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid)})
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
@@ -79,13 +80,12 @@ class _function(Generic[ReturnType]):
         buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.max_numel()}, device={b.device}" for i,b in enumerate(implicit_buffers))
         raise RuntimeError(f"function {name} has {len(implicit_buffers)} implicit buffer(s), but allow_implicit=False\n  {buf_strs}")
 
-    fret = UOp.call_outputs(uret.src if isinstance(ret, tuple) else (uret,), *call_uops, grad_fxn=self.grad_fxn, name=name,
-                            precompile=self.precompile, precompile_backward=self.precompile_backward)
+    outs = UOp.call_with_outputs(uret.src if isinstance(ret, tuple) else (uret,), *call_uops, grad_fxn=self.grad_fxn, name=name,
+                                 precompile=self.precompile, precompile_backward=self.precompile_backward)
 
     if DEBUG >= 2:
       print("  "*_function.depth+f"function {uret.key.hex()[:8]} in {(time.perf_counter()-st)*1000:8.2f} ms: {name}")
 
-    outs = fret.returned_outputs
     if isinstance(ret, tuple):
       return cast(ReturnType, tuple(Tensor(o) for o in outs))
     else:

tinygrad/helpers.py

@@ -395,7 +395,7 @@ if getenv("DEBUG_GC"):
 cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches" if OSX else "~/.cache")), "tinygrad")
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
 
-VERSION = 22
+VERSION = 23
 _db_connection = threading.local()
 def db_connection():
   if (conn:=getattr(_db_connection, "conn", None)) is None:

tinygrad/llm/chat.html

@@ -1,24 +1,36 @@
-<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
+<!DOCTYPE html><html><head><meta charset="utf-8"><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 }
+  table { border-collapse: collapse; table-layout: fixed; width: 100%; overflow-wrap: anywhere }
+  th, td { border: 1px solid #555; padding: 6px 10px; text-align: left }
+  a { color: #8ab4f8 } hr { border: 0; border-top: 1px solid #555 }
+  .answer { white-space: normal; line-height: 1.65 } .answer > * { margin: 12px 0 }
+  pre, blockquote { background: #2f2f2f; padding: 12px 16px; border-radius: 8px } pre { white-space: pre-wrap }
   .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" autofocus></textarea>
+<script src="/assets/cdn.jsdelivr.net/npm/[email protected]/dist/browser/markdown-it.umd.min.js"></script>
 <script>
-  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
+  let generating = false;
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
+    e.preventDefault(); if (generating) return;
+    generating = true; send().finally(() => generating = false);
+  } };
   const msgs = [];
+  const md = markdownit();
   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);
+    d.innerHTML = '<span style="color:#888"></span><div class="answer"></div>'; const [thinking, answer] = d.children;
     const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
       body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
     let buf = '', txt = '', rsn = '';
@@ -29,12 +41,11 @@
       const lines = buf.split('\n');
       buf = lines.pop();
       for (const ln of lines)
-        if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
-          try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
-            if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888';
-              s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) }
-            if (dl?.content) { const s = document.createElement('span');
-              s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {}
+        if (ln.startsWith('data: ') && !ln.includes('[DONE]')) {
+          const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
+          if (dl?.reasoning_content) { rsn += dl.reasoning_content; thinking.textContent = rsn }
+          if (dl?.content) { txt += dl.content; answer.innerHTML = md.render(txt) }
+        }
       chat.scrollTop = chat.scrollHeight;
     }
     const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m);

tinygrad/llm/gguf.py

@@ -129,7 +129,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       return (dl * (grid + delta)).flatten(-3)
     if ggml_type == 20:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32)
-      return d * Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)[q_to_uint8(blocks[:, 2:], 4)]
+      return d * Tensor.const(tuple(_ggml.kvalues_iq4nl), dtypes.float32)[q_to_uint8(blocks[:, 2:], 4)]
     if ggml_type == 21:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
       scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
@@ -147,7 +147,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     if ggml_type == 23:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1))
       scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16)
-      iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)
+      iq4_xs_lut = Tensor.const(tuple(_ggml.kvalues_iq4nl), dtypes.float32)
       scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8))
       scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8)
       scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))

tinygrad/llm/kernels/amd.py

@@ -3,9 +3,11 @@ import functools, math
 from typing import Callable, cast
 from tinygrad import Tensor, UOp, nn, Device, Context
 from tinygrad.device import Buffer
+from tinygrad.llm.gguf import ggml_data_to_tensor
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
+from tinygrad.renderer.cstyle import HIPRenderer
 
 BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
@@ -31,7 +33,7 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
   if device is None or device.split(":")[0] != "AMD": return False
   # @function contexts set ALLOW_DEVICE_USAGE=0 (scheduling must not open devices); the device is always open here
   with Context(ALLOW_DEVICE_USAGE=1):
-    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11 and isinstance(Device[device].renderer, HIPRenderer)
 
 def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
   for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
@@ -54,15 +56,20 @@ class Linear(nn.Linear):
     super().__init__(in_features, out_features, bias)
     self.in_features, self.out_features = in_features, out_features
   def set_quantized(self, decoded:Tensor):
+    if self.in_features % GGML_BLOCK_SIZE: return
     packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()}
     graph = decoded.uop.toposort()
     raw = next((u for u in graph if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
     if raw is None: return
     ggml_type = packed_sizes[prod(raw.shape)]
-    # the packed byte rate alone can't distinguish same-rate formats (Q4_0 vs Q4_K, Q5_0 vs Q5_K, MXFP4 vs IQ4_XS).
-    # the supported formats are 256-wide superblocks: their decode views the packed bytes at the superblock width
-    # (ggml_data_to_tensor reshapes to (-1, QUANT_SIZES[type])), while same-rate 32-wide formats reshape to 17-22
-    if not any(u.op is Ops.RESHAPE and u.shape[-1:] == (QUANT_SIZES[ggml_type],) for u in graph): return
+    # Only unwrap storage/order-preserving views, then require the exact dequantization expression.
+    # This rejects subsequent arithmetic and permutations, including RoPE's concatenated query weights.
+    def unwrapped(u:UOp) -> UOp:
+      while u.op in (Ops.RESHAPE, Ops.CONTIGUOUS) or (u.op is Ops.CAST and dtypes.is_float(u.dtype) and dtypes.is_float(u.src[0].dtype)):
+        u = u.src[0]
+      return u
+    expected = ggml_data_to_tensor(Tensor(raw), self.in_features * self.out_features, ggml_type)
+    if unwrapped(decoded.uop).key != unwrapped(expected.uop).key: return
     raw_offset = raw.contiguous_view_offset()
     assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
     self.ggml_type = ggml_type
@@ -100,23 +107,18 @@ class Linear(nn.Linear):
     return super().__call__(x)
 
 def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  # int8 4-wide dot, widened to scalar multiply-adds (2% decode slower than the sudot4 builtin, but portable)
-  for i in range(4):
-    av = ((a >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    bv = ((b >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    c = c + av*bv
-  return c
+  return UOp(Ops.CUSTOMI, src=(a, b, c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32))
 
 def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
   return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32))
 
-def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
+def _amd_load(ptr:UOp, lanes:int|None=None, stream:bool=False) -> UOp:
   assert ptr.op is Ops.INDEX
   # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
   if lanes is None: return ptr.load(arg="nontemporal")
   buf, coords = ptr.src[0], ptr.src[1:]
   idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(arg="nontemporal" if stream else None)
 
 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
@@ -152,22 +154,19 @@ def iq4_half_lut(device:str) -> Tensor:
 @functools.cache
 def _q8_quantize_kernel(q:UOp, scale:UOp, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
   groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token_group, lane = UOp.range(tokens*groups, 0, AxisType.GLOBAL), UOp.range(32, -1, AxisType.WARP)
   token, group = token_group//groups, token_group%groups
-  x = x.reshape(tokens, groups, 32)
-  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
-  word_lane = lane.minimum(7)
-  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
-  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
-  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
-  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
-  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
-  store_half = (lane & 4) >> 2
-  stores = (q[token, group, lane.valid(lane < 8)].store(word),
-            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
-                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
-                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
+  value = x.reshape(tokens, groups, 32)[token, group, lane].float()
+  # Quantize each input once, then pack four neighboring lanes into one word.
+  d = (warp_reduce(value.abs(), maximum=True, full_wave=True)/127).maximum(1e-8)
+  rounded = UOp(Ops.CUSTOM, src=(value/d,), arg=("__builtin_nearbyintf({0})", dtypes.float))
+  quant = rounded.clip(-127, 127).cast(dtypes.int8)
+  word = quant.cast(dtypes.uint8).cast(dtypes.uint32) << ((lane%4)*8).cast(dtypes.uint32)
+  for offset in (1, 2):
+    word |= UOp(Ops.CUSTOM, src=(word,), arg=(f"__builtin_amdgcn_ds_swizzle({{0}}, {0x1f | offset<<10})", dtypes.uint32))
+  stores = (q[token, group, (lane//4).valid((lane%4).eq(0))].store(word),
+            scale[token, group.valid(lane.eq(0))].store(d),
+            xsum[token, group, (lane//16).valid((lane%16).eq(0))].store(warp_reduce(quant.float())))
   return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
 
 def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
@@ -221,8 +220,8 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
     # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
     base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
     # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
-    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
-    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
+    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4, stream=True) for half in range(2))
+    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4, stream=True) for half in range(2))
     dots = [UOp.const(0, dtypes.int32)] * 2
     for word_idx in range(8):
       within = (subgroup*32 + word_idx*4)%128
@@ -240,6 +239,7 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
   return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])
 
 def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
+  if out_features % (16*output_tiles): output_tiles = 1
   output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
   token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
   # lane is a hardware WARP range (like the flash kernel): the fragment math stays visible without being
@@ -318,15 +318,9 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:i
   def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
     d, scale = _iq4_scales(raw, base, subgroup)
     scale = scale * d
-    if out_features <= 6144:
-      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
-                    for word in range(4) for byte in range(4))
-      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    # a subgroup-half gathers the lo (half=0) or hi (half=1) nibbles of byte pairs of each packed word
-    lut_pairs = (lut[(((raw[base+2+subgroup*4+i] >> (8*j+4*half)) & 15) |
-                      (((raw[base+2+subgroup*4+i] >> (8*j+8+4*half)) & 15) << 4)).cast(dtypes.weakint)]
-                 for i in range(4) for j in (0, 2))
-    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
+    pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
+                  for word in range(4) for byte in range(4))
+    return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
   return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
 
 def q8_linear(layer:Linear, x:Tensor) -> Tensor:
@@ -369,21 +363,20 @@ def _amd_f16_gemv_kernel(out:UOp, w:UOp, x:UOp, *rest:UOp, in_features:int, out_
     for j in range(val_chunk):
       acc = acc + w[out_row, i, lane*val_chunk + j].load().float() * x[token, i, lane*val_chunk + j].load().float()
   total = warp_reduce(acc, full_wave=True)
-  if bias is not None: total = total + bias[token, out_row].load().float()
+  if bias is not None: total = total + bias[out_row].load().float()
   return out[token, out_row.valid(lane.eq(0))].store(total).end(token, out_row, lane).sink(arg=KernelInfo(name="linear_f16_gemv", opts_to_apply=()))
 
 def _view_back(t:Tensor) -> Tensor:
-  """strip top-of-chain CAST(s) from a lazy weight: reading the raw file bytes in the kernel instead of
-  materializing the cast into a fresh buffer every step"""
+  # Widening half to float is exact; preserve casts that round or change the values.
   uop = t.uop
-  while uop.op is Ops.CAST: uop = uop.src[0]
+  while uop.op is Ops.CAST and uop.dtype == dtypes.float32 and uop.src[0].dtype in (dtypes.half, dtypes.bfloat16): uop = uop.src[0]
   return Tensor(uop).reshape(t.shape)
 
 def f16_gemv(layer:Linear, x:Tensor) -> Tensor:
   tokens = prod(x.shape[:-1])
   assert isinstance(tokens, int)
   weight = _view_back(layer.weight)
-  x = x.contiguous() if x.dtype == dtypes.half else x.cast(dtypes.half).contiguous()
+  x = x.contiguous()
   out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
   fxn = functools.partial(_amd_f16_gemv_kernel, in_features=layer.in_features, out_features=layer.out_features, tokens=tokens)
   srcs = (out, weight.reshape(-1), x.reshape(tokens, layer.in_features)) + (() if layer.bias is None else (_view_back(layer.bias),))
@@ -402,22 +395,25 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
   _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
   assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
+  G, CHUNK, DPL, WAVES, PARTIALS = H // H_KV, block_n, D // WARP_SIZE, waves, out.shape[2]
   assert CHUNK % WAVES == 0
   SEC = CHUNK // WAVES  # keys each wave scans independently
-  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
+  total_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  live_chunks = min(total_chunks, PARTIALS) if isinstance(total_chunks, int) else total_chunks.minimum(PARTIALS)
   block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
   lane, wave = UOp.range(WARP_SIZE, -1, axis_type=AxisType.WARP), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
   b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
   # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
   qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
   zerof = UOp.const(0, dtypes.float)
+  # Each block scans every PARTIALS-th chunk, keeping an online softmax across rounds.
+  chunk_round = UOp.range((total_chunks-1-block_chunk)//PARTIALS+1, 4, AxisType.REDUCE)
+  chunk_id = block_chunk + chunk_round*PARTIALS
   valids: list[UOp] = []
   scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
   vfrags: list[tuple[UOp, ...]] = [()]*SEC
   for j in range(SEC):
-    key = block_chunk*CHUNK + wave*SEC + j
+    key = chunk_id*CHUNK + wave*SEC + j
     valid = key < valid_kv_len
     valids.append(valid)
     kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
@@ -425,23 +421,32 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
     vfrags[j] = tuple(valid.where(v, zerof) for v in _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL))
     for h in range(G):
       s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
-      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
-  ninf = UOp.const(-math.inf, dtypes.float)
-  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
-  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
-  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
+      scores[j][h] = valid.where(s, UOp.const(-1e30, dtypes.float))
+  # A finite initial max keeps fully masked waves from computing exp(-inf - -inf).
+  acc_reg, max_reg, sum_reg = _reg((G, DPL), 2, 0), _reg((G,), 3, -1e30), _reg((G,), 4, 0)
+  prev_acc, prev_max, prev_sum = acc_reg.after(chunk_round), max_reg.after(chunk_round), sum_reg.after(chunk_round)
+  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), prev_max[h].load()) for h in range(G)]
+  # Rescale the previous rounds to the new max, then accumulate this round's keys.
+  alpha = [((prev_max[h].load()-row_max[h])*LOG2E).exp2() for h in range(G)]
+  accs = [[alpha[h]*prev_acc[h, i].load() for i in range(DPL)] for h in range(G)]
+  row_sums = [alpha[h]*prev_sum[h].load() for h in range(G)]
   for j in range(SEC):
     for h in range(G):
-      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
+      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), zerof)
       accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
       row_sums[h] = row_sums[h] + beta
+  update = UOp.group(acc_reg.store(UOp.stack(*(x for acc in accs for x in acc)).reshape(G, DPL)),
+                     max_reg.store(UOp.stack(*row_max)), sum_reg.store(UOp.stack(*row_sums))).end(chunk_round)
+  acc_reg, max_reg, sum_reg = acc_reg.after(update), max_reg.after(update), sum_reg.after(update)
   # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
-  acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  # Matching cache/LDS strides can reuse a loop-local cache index outside the loop. Pad that layout.
+  acc_lds = UOp.placeholder((WAVES, G, D + (LDS_PAD if G == SEC else 0)), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)[:, :, :D]
   ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
-  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
+  # Normalize before fp16 to avoid overflow. Nonempty waves have sum >= 1; empty waves keep their zero accumulator.
+  stores = [lds_acc[wave, h, lane].store((acc_reg[h].load() / sum_reg[h].load().maximum(1)).cast(dtypes.half)) for h in range(G)]
   # NOTE: duplicate stores of the same value from every lane are harmless here
-  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
+  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((max_reg[h].load(), sum_reg[h].load()))]
   barrier = UOp.barrier(UOp.group(*stores))
   acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
   tid = wave*WARP_SIZE + lane
@@ -449,14 +454,16 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   for i in range(-(-G*D//(WAVES*WARP_SIZE))):
     flat = tid + i*WAVES*WARP_SIZE
     h, d = flat // D, flat % D
-    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
-    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
+    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)))
+    # LDS holds normalized values; restore each wave's sum before combining.
+    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * ml_lds[w, h, 1].load() * acc_lds[w, h, d].load().float()
+               for w in range(WAVES)), zerof)
     oidx = out[b, kv_head*G + h, block_chunk, d]
     if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
     final_stores.append(oidx.store(val))
   hstat = tid
-  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
-  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
+  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)))
+  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), zerof)
   q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
   final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
   return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
@@ -493,10 +500,13 @@ def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UO
 
 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
   B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(256, max_kv_len // 64)
+  chunks = min(48, max_kv_len // 64)
   partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
   stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
+  waves, group = 16, H // cache_kv.shape[2]
+  while waves * group * ((D+LDS_PAD)*2 + 8) > 65536: waves //= 2
+  assert waves > 0, "attention head group exceeds shared memory capacity"
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=waves)
   partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
   live = (valid_kv_len+63)//64
   live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
@@ -512,7 +522,7 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
   assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
   gqa_group = BH // k.shape[0]
-  if isinstance(M, int) and isinstance(valid_kv_len, int): assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
+  if isinstance(M, int): assert M % BLOCK_M == 0
   assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
   TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
   # query row 0 sits at sequence position q_base (the queries may be padded beyond valid_kv_len - q_base rows)
@@ -540,7 +550,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
   q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
   k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
-  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
+  # All waves must finish reading Q/K before their shared memory is reused for P/V.
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk).barrier()
   S_reg = S_reg.after(qk_done, S_reg.store(S_reg * SCALE))
   rm, rn = UOp.range(TM, 250), UOp.range(TN, 251)
   q_idx = q_base + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
@@ -567,7 +578,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
   V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
   V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390)
-  vval = v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float()
+  v_pos = n_tile*BLOCK_N + (tid*KV_ELEMS_PER_THREAD + load_v)//D
+  vval = (v_pos < valid_kv_len).where(v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float(), 0)
   V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
   pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
   P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
@@ -589,7 +601,16 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
 def flash_attention(q:Tensor, assigned_kv:Tensor, valid_end:int|UOp) -> Tensor:
   # cached flash attention on the half KV cache (already written through assigned_kv); valid_end stays bound at the graph level
   T_real, q_start = q.shape[2], None
-  if resolve(T_real == 1): return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, assigned_kv.shape[3]))
+  D, N, group = q.shape[3], assigned_kv.shape[3], q.shape[1] // assigned_kv.shape[2]
+  decode = resolve(T_real == 1, False)
+  # Non-power-of-two decode dimensions can lose tail-store masks. Q/P, K, and V use separate LDS allocations.
+  supported = D % 32 == 0 and (D & (D-1) == 0 and N % 64 == 0 and group*((D+LDS_PAD)*2+8) <= 65536 if decode else
+    D >= 64 and 2*(2*BLOCK_M*(D+LDS_PAD) + D*(BLOCK_N+LDS_PAD)) <= 65536 and N % BLOCK_N == 0 and q.max_shape[2] % BLOCK_M == 0)
+  if not supported:
+    k, v = (assigned_kv[i, :, :, :valid_end].float() for i in range(2))
+    mask = None if decode else Tensor.full((T_real, valid_end), -math.inf, dtype=dtypes.float32, device=q.device).triu(valid_end-T_real+1)
+    return q.float().scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)
+  if decode: return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, N))
   if isinstance(T_real, UOp):
     # symbolic chunk: pad the queries to the static tile; garbage rows are sliced off
     T_pad = q.max_shape[2]
@@ -643,12 +664,14 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor,
   assert q.shape == k.shape and v.shape[:3] == beta.shape == (batch, heads, tokens) and state.shape == (batch, heads, value_dim, key_dim)
   assert alpha.shape[:3] == (batch, heads, tokens) and (len(alpha.shape) == 3 or alpha.shape[-1] in (1, value_dim))
   assert key_dim % 32 == 0 and value_dim % 4 == 0
+  assert q.dtype == k.dtype == dtypes.float32, "recurrent Q/K must be float32"
+  assert state.uop.contiguous_view_offset() is not None, "recurrent state must be contiguous"
+  if start_pos is not None:
+    assert start_pos.uop.is_bound_var
+    state = Tensor(state.uop.after(start_pos.uop))
   core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
   srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
-  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=_gated_delta_prefill_kernel)[0]
   contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
   params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
-  assert start_pos.uop.is_bound_var
-  # the bound start_pos reaches the graph through the state AFTER chain, like the flash kernels' valid_end
-  call = _gated_delta_prefill_kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig)
+  call = _gated_delta_prefill_kernel(*params, None if start_pos is None else kernel_var(start_pos.uop.src[0])).call(*contig)
   return Tensor(contig[0].after(call))

tinygrad/llm/serve.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import json, pathlib, re, time, typing, uuid
 from typing import TYPE_CHECKING
 from tinygrad.helpers import DEBUG, colored, stderr_log
-from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+from tinygrad.viz.serve import TCPServerWithReuse, Handler as VizHandler
 if TYPE_CHECKING:
   from tinygrad.llm.cli import SimpleTokenizer
   from tinygrad.llm.model import Transformer
@@ -60,11 +60,12 @@ class StreamRouter:
     if emit: yield "content", emit
     if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
 
-class Handler(HTTPRequestHandler):
+class Handler(VizHandler):
   server: LLMServer
   def log_request(self, code='-', size='-'): pass
   def do_GET(self):
     if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
+    elif self.path.startswith("/assets/"): super().do_GET()
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
   def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
                 reasoning:bool=False):

tinygrad/mixin/gradient.py

@@ -3,6 +3,7 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
+from tinygrad.device import Buffer
 from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
@@ -35,7 +36,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
       return arg_grads(k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k))
     return arg_grads(k.arg.grad_fxn(on_dev(ctx, 0), k))
   # the RETURNED inputs are the call outputs: their positions in the args get the output gradients from the AFTER rule
-  assert fxn.op is Ops.SINK and k.num_returned, f"expected a CALL with RETURNED inputs or a grad_fxn, got {fxn.op}"
+  assert fxn.op is Ops.SINK and k.has_unbound_outputs, f"expected a CALL with unbound BUFFER outputs or a grad_fxn, got {fxn.op}"
   ret_pos = [i for i, a in enumerate(args) if a.unsharded_base.is_unbound]
   # the body stores the outputs into output PARAMs: the values are the stored values in slot order
   values = UOp.sink(*[st.src[1] for st in fxn.src if st.op is Ops.STORE])
@@ -50,20 +51,39 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   grads = compute_gradient(values, root_grad, set(params.values()))
   # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(values.src)} if k.arg.precompile else {}
-  fwd_outs = k.returned_outputs if k.arg.precompile else ()
+  fwd_outs = k.unbound_outputs if k.arg.precompile else ()
   # collect needed gradient bodies, compact unused params, create a single backward CALL
   grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.sink(*[gb for _, gb in grad_bodies]).substitute(fwd_subs, walk=True)
   bwd_body = renumber_invalid_outputs(bwd_body)
   # NOTE: args includes the RETURNED inputs so the param slots above line up; they are unused and compacted away
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
-  bwd_outs = UOp.call_outputs(bwd_body.src, *compact_args, name=(k.arg.name or "")+"_backward",
-                              precompile=k.arg.precompile_backward).returned_outputs
+  bwd_outs = UOp.call_with_outputs(bwd_body.src, *compact_args, name=(k.arg.name or "")+"_backward",
+                                   precompile=k.arg.precompile_backward)
   gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
   # align gradients with the original source positions: None at RETURNED positions, gradients elsewhere
   ret_set = set(ret_pos)
   return (None,) + tuple(None if i in ret_set else (bwd_outs[gb_map[i]] if i in gb_map else None) for i in range(len(args)))
 
+def after_gradient(ctx:UOp, ret:UOp):
+  value, *deps = ret.src
+  if len(deps) == 1:
+    dep = deps[0]
+    if dep.op is Ops.STORE and len(dep.src) == 2 and value is dep.src[0]: return (None, ctx)
+    if dep.op is Ops.CALL and (value.unsharded_base.is_unbound or value in dep.call_access()[1]):
+      if dep.src[1:].count(value) != 1: raise RuntimeError("ambiguous CALL output gradient")
+      return (None, UOp.sink(*(ctx if a is value else UOp(Ops.NOOP) for a in dep.src[1:])))
+  for dep in deps:
+    if dep.op is Ops.STORE: writes = dep.src[:1]
+    elif dep.op is Ops.CALL: _, writes = dep.call_access()
+    else: raise RuntimeError(f"gradient through {dep.op} ordering is unsupported")
+    for w in writes:
+      a, b = (u.storage_base.arg.buffer if u.storage_base.op is Ops.BUFFER else None for u in (value, w))
+      if not isinstance(a, Buffer) or not isinstance(b, Buffer) or a.base is b.base or \
+         any(buf.base.options is not None and buf.base.options.external_ptr is not None for buf in (a, b)):
+        raise RuntimeError("gradient through an aliased write is unsupported")
+  return (ctx,) + (None,)*len(deps)
+
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -94,10 +114,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device),)),
   (UPat(Ops.UNSHARD, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   (UPat(Ops.SINK), lambda ctx: ctx.src),
-  (UPat(Ops.AFTER, src=(UPat.var("d"), UPat(Ops.CALL, name="k"))), lambda ctx, d, k:
-    (ctx, UOp.sink(*([ctx if i == k.src.index(d)-1 else UOp(Ops.NOOP) for i in range(len(k.src)-1)])))),
-  # clone/assign gradient passes through to val
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE))), lambda ctx: (None, ctx)),
+  (UPat(Ops.AFTER, name="ret"), after_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/renderer/amd/generate.py

@@ -233,12 +233,15 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
       sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
       # Stop at large Y gaps (>30) - indicates section break
       filtered = [sorted_lines[0]]
+      depth = sorted_lines[0][2].count("{") - sorted_lines[0][2].count("}")
       for j in range(1, len(sorted_lines)):
         prev_page, prev_y, _ = sorted_lines[j-1]
         curr_page, curr_y, _ = sorted_lines[j]
-        if curr_page == prev_page and prev_y - curr_y > 30: break
-        if curr_page != prev_page and prev_y > 60 and curr_y < 730: break
+        if depth == 0 and curr_page == prev_page and prev_y - curr_y > 30: break
+        if depth == 0 and curr_page != prev_page and prev_y > 60 and curr_y < 730: break
         filtered.append(sorted_lines[j])
+        code = sorted_lines[j][2].split("//")[0]
+        depth += code.count("{") - code.count("}")
       pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
       if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
   return pcode

tinygrad/renderer/amd/sqtt.py

@@ -512,17 +512,7 @@ class CDNA_ISSUE(PacketType):
   """pkt_fmt=13: 32-bit (Issue)"""
   encoding = bits[3:0] == 13
   simd = bits[6:5]
-  _gap = bits[7:7]
-  inst0 = bits[9:8]
-  inst1 = bits[11:10]
-  inst2 = bits[13:12]
-  inst3 = bits[15:14]
-  inst4 = bits[17:16]
-  inst5 = bits[19:18]
-  inst6 = bits[21:20]
-  inst7 = bits[23:22]
-  inst8 = bits[25:24]
-  inst9 = bits[27:26]
+  inst = bits[27:8]
   _padding = bits[31:28]
 
 class CDNA_PERF(PacketType):
@@ -676,6 +666,12 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
           inst = pc_map[pc:=wave_pc[(simd, wave)]]
           wave_pc[(simd, wave)] += inst.size()
           yield (p, InstructionInfo(pc, wave, inst))
+    elif isinstance(p, CDNA_ISSUE):
+      for wave in range(10):
+        if (p.inst >> (wave * 2)) & 3 == 3:
+          inst = pc_map[pc:=wave_pc[(p.simd, wave)]]
+          wave_pc[(p.simd, wave)] += inst.size()
+          yield (p, InstructionInfo(pc, wave, inst))
     # map INST events on this SIMD to the program counter, we know the waves
     elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):
       inst = pc_map[pc:=wave_pc[(simd, p.wave)]]
@@ -725,7 +721,8 @@ def format_packet(p) -> str:
 
 def print_packets(packets) -> None:
   skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
-          "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
+          "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT",
+          "CDNA_MISC", "CDNA_REG_CS", "CDNA_REG_CS_PRIV", "CDNA_REG"} if not getenv("NOSKIP") else {"NOP"}
   for p in packets:
     if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p))
 

tinygrad/renderer/isa/__init__.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import itertools
 from dataclasses import dataclass, field
 from tinygrad.renderer import Renderer
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, consumer_map_from_toposort
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops
 
 @dataclass(frozen=True)
 class Register:
@@ -16,12 +16,9 @@ class Register:
 
 class IselContext:
   def __init__(self, sink:UOp):
-    self.uses = consumer_map_from_toposort(sink.toposort())
     self.reg_n = itertools.count()
-    def arg_key(u:UOp):
-      if u.op is Ops.SPECIAL: return (2, u.arg)
-      return (0, u.arg.slot) if u.arg.addrspace is not None else (1, u.expr)
-    self.func_args = sorted([u for u in self.uses if u.op in {Ops.PARAM, Ops.SPECIAL}], key=arg_key)
+    def arg_key(u:UOp): return (1, u.arg) if u.op is Ops.SPECIAL else (0, u.arg.slot)
+    self.func_args = sorted([u for u in sink.toposort() if u.op in {Ops.PARAM, Ops.SPECIAL}], key=arg_key)
 
   def vreg(self, cons:tuple[Register, ...]|Register):
     return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
@@ -34,17 +31,15 @@ def greg(u:UOp):
 @dataclass
 class PreRegAllocContext:
   lock: UOp|None = None
-  clobbered: set[UOp] = field(default_factory=set)
 
 class ISARenderer(Renderer):
   pre_isel_matcher: PatternMatcher
   isel_matcher: PatternMatcher
-  pre_regalloc_matcher: PatternMatcher|None = None
+  pre_regalloc_matcher: PatternMatcher
   post_regalloc_matcher: PatternMatcher
 
   def is_two_address(self, x:UOp) -> bool: return False
   def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific")
-  def copy(self, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
   def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
   def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
   def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")

tinygrad/renderer/isa/x86.py

@@ -23,18 +23,13 @@ class X86Ops(FastEnum):
   VMOVSSm = auto(); VMOVSDm = auto(); VMOVUPSm = auto()
   # casts
   MOVZX = auto(); MOVSX = auto(); MOVSXD = auto()
-  VPMOVZXBW = auto(); VPMOVZXBD = auto(); VPMOVZXBQ = auto()
-  VPMOVZXWD = auto(); VPMOVZXWQ = auto(); VPMOVZXDQ = auto()
-  VPMOVSXBW = auto(); VPMOVSXBD = auto(); VPMOVSXBQ = auto()
-  VPMOVSXWD = auto(); VPMOVSXWQ = auto(); VPMOVSXDQ = auto()
-  VCVTDQ2PS = auto(); VCVTDQ2PD = auto(); VCVTTPS2DQ = auto(); VCVTTPD2DQ = auto()
-  VCVTPH2PS = auto(); VCVTPS2PH = auto(); VCVTPS2PD = auto(); VCVTPD2PS = auto()
+  VCVTPH2PS = auto(); VCVTPS2PH = auto()
   VCVTSS2SD = auto(); VCVTSD2SS = auto(); VCVTSI2SS = auto(); VCVTSI2SD = auto()
   VCVTTSS2SI = auto(); VCVTTSD2SI = auto()
   # bitcasts
   VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
   # comparisons
-  VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
+  VCMPSS = auto(); VCMPSD = auto()
   SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
   # where
   CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
@@ -43,29 +38,17 @@ class X86Ops(FastEnum):
   JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
   # vectorize / gep
   VINSERTPS = auto(); VPSRLDQ = auto()
-  VPEXTRB = auto(); VPEXTRW = auto(); VPEXTRD = auto(); VPEXTRQ = auto()
-  VPINSRB = auto(); VPINSRW = auto(); VPINSRD = auto(); VPINSRQ = auto()
+  VPEXTRW = auto(); VPEXTRD = auto()
+  VPINSRW = auto(); VPINSRD = auto()
   # int binary
   IDIV = auto(); DIV = auto()
   ADD = auto(); ADDi = auto(); SUB = auto(); SUBi = auto(); IMUL = auto(); IMULi = auto()
   AND = auto(); ANDi = auto(); XOR = auto(); XORi = auto(); OR = auto(); ORi = auto()
   SHL = auto(); SHLi = auto(); SHR = auto(); SHRi = auto(); SAR = auto(); SARi = auto(); CMP = auto(); CMPi = auto()
   # float unary (sometimes not unary)
-  VROUNDSS = auto(); VROUNDSD = auto(); VROUNDPS = auto(); VROUNDPD = auto()
-  VSQRTSS = auto(); VSQRTSD = auto(); VSQRTPS = auto(); VSQRTPD = auto()
-  # float scalar / vector binary
-  VADDSS = auto(); VADDSD = auto(); VADDPS = auto(); VADDPD = auto()
-  VSUBSS = auto(); VSUBSD = auto(); VSUBPS = auto(); VSUBPD = auto()
-  VMULSS = auto(); VMULSD = auto(); VMULPS = auto(); VMULPD = auto()
-  VDIVSS = auto(); VDIVSD = auto(); VDIVPS = auto(); VDIVPD = auto()
-  # int vector binary
-  VPADDB = auto(); VPADDW = auto(); VPADDD = auto(); VPADDQ = auto()
-  VPSUBB = auto(); VPSUBW = auto(); VPSUBD = auto(); VPSUBQ = auto()
-  VPMULLW = auto(); VPMULLD = auto()
-  # packed bitwise
-  VPAND = auto(); VPOR = auto(); VPXOR = auto()
-  # packed variable shifts
-  VPSLLVD = auto(); VPSLLVQ = auto(); VPSRLVD = auto(); VPSRLVQ = auto(); VPSRAVD = auto()
+  VROUNDSS = auto(); VROUNDSD = auto(); VSQRTSS = auto(); VSQRTSD = auto()
+  # float binary
+  VADDSS = auto(); VADDSD = auto(); VSUBSS = auto(); VSUBSD = auto(); VMULSS = auto(); VMULSD = auto(); VDIVSS = auto(); VDIVSD = auto()
   # return
   RET = auto()
 
@@ -75,30 +58,18 @@ class X86GroupOp:
                 X86Ops.SUB, X86Ops.SUBi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi,
                 X86Ops.IDIV, X86Ops.DIV, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB}
 
-  # X86Ops whose first src can read from memory
-  ReadMem1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
-                X86Ops.VPMOVZXBW, X86Ops.VPMOVZXBD, X86Ops.VPMOVZXBQ, X86Ops.VPMOVZXWD, X86Ops.VPMOVZXWQ, X86Ops.VPMOVZXDQ,
-                X86Ops.VPMOVSXBW, X86Ops.VPMOVSXBD, X86Ops.VPMOVSXBQ, X86Ops.VPMOVSXWD, X86Ops.VPMOVSXWQ, X86Ops.VPMOVSXDQ,
-                X86Ops.VCVTDQ2PS, X86Ops.VCVTDQ2PD, X86Ops.VCVTTPS2DQ, X86Ops.VCVTTPD2DQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
-                X86Ops.VCVTPH2PS, X86Ops.VCVTPS2PD, X86Ops.VCVTPD2PS, X86Ops.VROUNDPS, X86Ops.VROUNDPD, X86Ops.VSQRTPS, X86Ops.VSQRTPD,
-                X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
-
-  # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
-  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
-                X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
-                X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
-                X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
-                X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
-                X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VINSERTPS,
-                X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
-                X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
-                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
+  # X86Ops whose second src is the rm field, so that src is what can be a memory operand
+  Rm2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
+           X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VDIVSS, X86Ops.VDIVSD,
+           X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD,
+           X86Ops.VINSERTPS, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
+           X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
 
   # X86Ops that can write to memory
   WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
               X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR,
               X86Ops.SARi, X86Ops.SETNE, X86Ops.SETE, X86Ops.SETL, X86Ops.SETB,
-              X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
+              X86Ops.VCVTPS2PH, X86Ops.VPEXTRW, X86Ops.VPEXTRD}
 
   # X86Ops that read flags
   ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
@@ -109,11 +80,9 @@ class X86GroupOp:
                 X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
                 X86Ops.OR, X86Ops.ORi}
 
-  # X86Ops whose first src is the rm field
-  Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
-
-  # X86Ops whose second src is the rm field
-  Rm2nd = ReadMem2nd
+  # X86Ops whose first src is the rm field. a TwoAddress op drops its first src post regalloc, so its Rm2nd src ends up first
+  Rm1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
+           X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI, X86Ops.VCVTPH2PS, X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA, X86Ops.VPSRLDQ} | (Rm2nd & TwoAddress)
 
 # ***** X86 legalization *****
 
@@ -136,15 +105,12 @@ extra_matcher = PatternMatcher([
   # no int8 mul or cmove, cast to int16
   (UPat.var("a", dtypes.int8s) * UPat.var("b"), lambda a,b: (a.cast(dtypes.int16) * b.cast(dtypes.int16)).cast(a.dtype)),
   (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
-   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.max_numel() == 1 else None),
+   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype)),
   # float16 alus are done in float32
   (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
    src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
   (UPat(GroupOp.Comparison, src=[UPat(dtype=dtypes.float16), UPat()], name="x"),
    lambda x: UOp(x.op, src=tuple(s.cast(dtypes.float32) for s in x.src)).cast(x.dtype)),
-  # no cmpne for packed ints, y != x => !(y==x)
-  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
-   lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
   # a float WHERE blends at the width of its value, so it needs a comparison at that width to make the mask
   (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
    lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and promo_dtype(m.src) is not w.dtype else None),
@@ -179,8 +145,8 @@ def flag_gate(m:UOp) -> UOp|None:
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
 pre_isel_matcher = PatternMatcher([
-  # widening a scalar uint32 is free, the 32bit write that produced it already zeroed the upper half
-  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.BITCAST) if y.max_numel() == 1 else None),
+  # widening a uint32 is free, the 32bit write that produced it already zeroed the upper half
+  (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.replace(op=Ops.BITCAST)),
   (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
    lambda y,x: x.replace(op=Ops.BITCAST) if x.dtype.itemsize == y.dtype.itemsize else None),
   # gated load/store become a conditional move on the address, the load/store are unconditional
@@ -231,8 +197,7 @@ def cmp(x:UOp) -> UOp:
 # comparisons that produce masks, the mask has the width of the operands
 def mask(x:UOp) -> UOp:
   dt, v = x.src[0].dtype, imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if dt is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, dtype=dt, src=x.src + (v,))
-  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, dtype=dt, src=x.src + (v,))
+  return x.ins(X86Ops.VCMPSS if dt is dtypes.float32 else X86Ops.VCMPSD, dtype=dt, src=x.src + (v,))
 
 # vinsertps xmm2, xmm0, xmm1, imm
 # inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
@@ -243,10 +208,10 @@ def vinsertps(x:UOp) -> UOp:
     return x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
   return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
 
-# vpinsq xmm2, xmm0, rax, imm
-# inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
+# vpinsrd xmm2, xmm0, eax, imm
+# inserts the element in eax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
-  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
+  op = {2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD}[x.dtype.itemsize]
   return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
@@ -306,8 +271,7 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None:
   # this move "cleanses" the abi register constraint
   return x.ins(X86Ops.MOV, dtype=dt, src=src)
 
-GPR_DEST_OPS = {X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
-                 X86Ops.VMOVDm, X86Ops.VMOVQm}
+GPR_DEST_OPS = {X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI, X86Ops.VMOVDm, X86Ops.VMOVQm}
 XMM_OPS = {op for op in X86Ops if op.name.startswith('V')} - GPR_DEST_OPS
 
 def _is_vec_xmm(y: UOp) -> bool:
@@ -351,7 +315,6 @@ isel_matcher = PatternMatcher([
   # **** Op -> Op ****
   # range is lowered to acc, cmp, jmp after regalloc
   (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])),
-  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
   # really all a backedge END is is an IF with a tag referencing the RANGE start label
   (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
     lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])),
@@ -399,50 +362,22 @@ isel_matcher = PatternMatcher([
   (UPat(Ops.CMPEQ, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
   (UPat(Ops.CMPNE, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
   # float unary
-  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPS)),
-  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPD)),
-  (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x:
-   x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))),
-  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
-   x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
+  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y))),
+  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y))),
+  (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3)))),
+  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3)))),
   # for float16 we route the srcs through gprs, this is suboptimal for values in xmms, in that case we want vpunpcklwd
   (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x: vpins(x, tuple(s.bitcast(dtypes.int16) for s in x.src))),
   (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
-  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), lambda x: vpins(x, x.src)),
+  (UPat(Ops.STACK, dtypes.int32s, name="x"), lambda x: vpins(x, x.src)),
   # INDEX on a vector register value extracts a single element
-  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.int32s).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.floats).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None),
-  # packed bitwise
-  ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None),
-  ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.max_numel() > 1 else None),
-  ((UPat() ^ UPat()).named("x"), lambda x: x.ins(X86Ops.VPXOR) if x.max_numel() > 1 else None),
-  # packed int binary
-  ((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVQ) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.uint32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.uint64) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVQ) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRAVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int8s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDB) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int16s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDW) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int32s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int64s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDQ) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int8s, name="x"), lambda x: x.ins(X86Ops.VPSUBB) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPSUBW) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPSUBD) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPSUBQ) if x.max_numel() > 1 else None),
-  (UPat(Ops.MUL, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMULLW) if x.max_numel() > 1 else None),
-  (UPat(Ops.MUL, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMULLD) if x.max_numel() > 1 else None),
-  # scalar int binary
+  # int binary
   ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv),
-  # scalar int binary with immediate
+  # int binary with immediate
   (UPat.var("a", dtypes.ints) << UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
   (UPat.var("a", dtypes.uints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
   (UPat.var("a", dtypes.sints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
@@ -456,7 +391,7 @@ isel_matcher = PatternMatcher([
    lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar().cast(name="c"))),
    lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  # scalar int binary with register
+  # int binary with register
   ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
   ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
   ((UPat(dtype=dtypes.sints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SAR)),
@@ -467,21 +402,15 @@ isel_matcher = PatternMatcher([
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.var("b"), lambda a,b: a.ins(X86Ops.XOR, src=(a, b))),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.var("b"))), lambda a,b: a.ins(X86Ops.SUB, src=(a, b))),
   # float binary
-  ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS if x.max_numel() == 1 else X86Ops.VADDPS)),
-  ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD if x.max_numel() == 1 else X86Ops.VADDPD)),
-  ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS if x.max_numel() == 1 else X86Ops.VMULPS)),
-  ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD if x.max_numel() == 1 else X86Ops.VMULPD)),
-  (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS if x.max_numel() == 1 else X86Ops.VSUBPS)),
-  (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.max_numel() == 1 else X86Ops.VSUBPD)),
-  (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.max_numel() == 1 else X86Ops.VDIVPS)),
-  (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.max_numel() == 1 else X86Ops.VDIVPD)),
+  ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS)),
+  ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD)),
+  ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS)),
+  ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD)),
+  (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS)),
+  (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD)),
+  (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS)),
+  (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD)),
   # casts
-  (UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float32).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPS2DQ) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.max_numel() > 1 else None),
   (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))),
   (UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)),
   (UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)),
@@ -491,23 +420,10 @@ isel_matcher = PatternMatcher([
   (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
   (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
   (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
-   x.ins(X86Ops.MOVZX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
-  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.max_numel() == 1 else None),
-  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x:
-   x.ins(X86Ops.MOVSX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
-  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else None),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
-  (UPat(dtype=dtypes.uint16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWD)),
-  (UPat(dtype=dtypes.uint16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWQ)),
-  (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXDQ)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBW)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBD)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBQ)),
-  (UPat(dtype=dtypes.int16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWD)),
-  (UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)),
-  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
+   x.ins(X86Ops.MOVZX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
+  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD)),
+  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
+  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV)),
   # bitcasts between scalar floats and ints
   (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
   (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), lambda x: vpins(x, x.src)),
@@ -518,40 +434,32 @@ isel_matcher = PatternMatcher([
   # index on a buffer (or the stack pointer) computes an address, addresses are 64bit values
   (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda x: lea(x) if not _is_vec_xmm(x.src[0]) else None),
   # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
-  # copy, load, store
-  # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
-  (UPat(Ops.COPY, dtypes.floats, name="x"), lambda x: x.ins(_xmm_sz(x))),
-  (UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else x.ins(_xmm_sz(x))),
+  # load, store
   (UPat(Ops.LOAD, dtypes.floats, src=(UPat(name="a"),), name="x"), lambda x,a:
    x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
    x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a:
-   x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else
-   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
-   x.ins(_xmm_sz(x), src=fold_address(a))),
+   x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat.var("a").store(UPat.var("b", dtypes.floats), name="x"), lambda a,b,x:
    x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() * b.dtype.itemsize == 2 else
    x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,))),
   (UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x:
-   x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() > 1 and b.max_numel() * b.dtype.itemsize == 2 else
    x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,)) if b.max_numel() > 1 else
    x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))),
   # allocate virtual registers
-  (UPat((Ops.INS, Ops.BUFFER), name="x"), alloc_vregs),
+  (UPat((Ops.INS, Ops.BUFFER, Ops.RANGE), name="x"), alloc_vregs),
 ])
 
 # ***** pre register allocation *****
-# this handles flag clobbers. Unfortunately x86 doesn't have a good way to store/restore the flag register (then regalloc would handle it)
-# so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
-# regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
+# the flags belong to the last instruction that wrote them. x86 has no good way to store/restore them (then regalloc would
+# handle it), so a consumer that no longer owns its compare re-emits it. Unlike a regalloc rematerialization this is not
+# optional, there is no fallback load from stack
 def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
-  flag_def = x if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags else x.src[-1] if x.arg[0] in X86GroupOp.ReadFlags else None
-  if flag_def is None: return None
-  if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
-  ctx.lock = flag_def
-  if flag_def not in ctx.clobbered: return None
-  ctx.clobbered.remove(flag_def)
-  return (x, [flag_def, x])
+  if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags: ctx.lock = x
+  elif x.arg[0] in X86GroupOp.ReadFlags and ctx.lock is not (flag_def:=x.src[-1]):
+    ctx.lock = flag_def
+    return (x, [flag_def, x])
+  return None
 
 pre_regalloc_matcher = PatternMatcher([
   (UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize),
@@ -707,20 +615,11 @@ encodings = {
   # casts
   X86Ops.MOVZX: lambda x: encode(x, 0x0FB7),
   X86Ops.MOVSX: lambda x: encode(x, 0x0FBF), X86Ops.MOVSXD: lambda x: encode(x, 0x63),
-  X86Ops.VPMOVZXBW: lambda x: encode(x, 0x30, pp=1, sel=2), X86Ops.VPMOVZXBD: lambda x: encode(x, 0x31, pp=1, sel=2),
-  X86Ops.VPMOVZXBQ: lambda x: encode(x, 0x32, pp=1, sel=2), X86Ops.VPMOVZXWD: lambda x: encode(x, 0x33, pp=1, sel=2),
-  X86Ops.VPMOVZXWQ: lambda x: encode(x, 0x34, pp=1, sel=2), X86Ops.VPMOVZXDQ: lambda x: encode(x, 0x35, pp=1, sel=2),
-  X86Ops.VPMOVSXBW: lambda x: encode(x, 0x20, pp=1, sel=2), X86Ops.VPMOVSXBD: lambda x: encode(x, 0x21, pp=1, sel=2),
-  X86Ops.VPMOVSXBQ: lambda x: encode(x, 0x22, pp=1, sel=2), X86Ops.VPMOVSXWD: lambda x: encode(x, 0x23, pp=1, sel=2),
-  X86Ops.VPMOVSXWQ: lambda x: encode(x, 0x24, pp=1, sel=2), X86Ops.VPMOVSXDQ: lambda x: encode(x, 0x25, pp=1, sel=2),
   X86Ops.VCVTSS2SD: lambda x: encode(x, 0x5A, pp=2, sel=1), X86Ops.VCVTSD2SS: lambda x: encode(x, 0x5A, pp=3, sel=1),
   X86Ops.VCVTPH2PS: lambda x: encode(x, 0x13, pp=1, sel=2), X86Ops.VCVTPS2PH: lambda x: encode(x, 0x1D, pp=1, sel=3),
-  X86Ops.VCVTDQ2PS: lambda x: encode(x, 0x5B, pp=0, sel=1), X86Ops.VCVTDQ2PD: lambda x: encode(x, 0xE6, pp=2, sel=1),
-  X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
-  X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
-  # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address
-  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
-  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  # the int src is the 2nd src (the rm field), its width picks the 32 or 64 bit form
+  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=x.src[1].dtype.itemsize == 8),
+  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=x.src[1].dtype.itemsize == 8),
   X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
   X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
   # int division
@@ -738,46 +637,25 @@ encodings = {
   X86Ops.IMUL: lambda x: encode(x, 0x0FAF), X86Ops.IMULi: lambda x: encode(x, 0x69),
   X86Ops.SETB: lambda x: encode(x, 0x0F92, reg=0), X86Ops.SETL: lambda x: encode(x, 0x0F9C, reg=0),
   X86Ops.SETE: lambda x: encode(x, 0x0F94, reg=0), X86Ops.SETNE: lambda x: encode(x, 0x0F95, reg=0),
-  # packed bitwise NOTE: only bitwise and packed
-  X86Ops.VPAND: lambda x: encode(x, 0xDB, pp=1, sel=1), X86Ops.VPXOR: lambda x: encode(x, 0xEF, pp=1, sel=1),
-  X86Ops.VPOR: lambda x: encode(x, 0xEB, pp=1, sel=1),
   # unary
-  X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTPS: lambda x: encode(x, 0x51, pp=0, sel=1),
-  X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1), X86Ops.VSQRTPD: lambda x: encode(x, 0x51, pp=1, sel=1),
-  X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDPS: lambda x: encode(x, 0x08, pp=1, sel=3),
-  X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3), X86Ops.VROUNDPD: lambda x: encode(x, 0x09, pp=1, sel=3),
-  # packed int binary
-  X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
-  X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
-  X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
-  X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
-  X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
-  X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
-  X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
-  X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
-  # scalar / packed float binary
-  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
-  X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
-  X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBPS: lambda x: encode(x, 0x5C, pp=0, sel=1),
-  X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1), X86Ops.VSUBPD: lambda x: encode(x, 0x5C, pp=1, sel=1),
-  X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULPS: lambda x: encode(x, 0x59, pp=0, sel=1),
-  X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1), X86Ops.VMULPD: lambda x: encode(x, 0x59, pp=1, sel=1),
-  X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVPS: lambda x: encode(x, 0x5E, pp=0, sel=1),
-  X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1), X86Ops.VDIVPD: lambda x: encode(x, 0x5E, pp=1, sel=1),
-  X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPPS: lambda x: encode(x, 0xC2, pp=0, sel=1),
-  X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1), X86Ops.VCMPPD: lambda x: encode(x, 0xC2, pp=1, sel=1),
+  X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1),
+  X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3),
+  # float binary
+  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1),
+  X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1),
+  X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1),
+  X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1),
+  X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1),
   # ternary
   X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
   X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
   X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
   # shuffles
   X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
-  X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
-  X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3), X86Ops.VPINSRQ: lambda x: encode(x, 0x22, pp=1, sel=3, we=1),
+  X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1), X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3),
   X86Ops.VINSERTPS: lambda x: encode(x, 0x21, pp=1, sel=3),
   # extract
-  X86Ops.VPEXTRB: lambda x: encode(x, 0x14, pp=1, sel=3), X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3),
-  X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3), X86Ops.VPEXTRQ: lambda x: encode(x, 0x16, pp=1, sel=3, we=1),
+  X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3), X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3),
   # jumps are encoded with a placeholder which gets patched later once the real offset is known
   X86Ops.JE: lambda x: bytes([0x0F, 0x84]) + int(0).to_bytes(4),
   X86Ops.JNE: lambda x: bytes([0x0F, 0x85]) + int(0).to_bytes(4),
@@ -805,19 +683,14 @@ class X86Renderer(ISARenderer):
     self.compiler = X86Compiler()
   def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
   def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
-  # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
-  def copy(self, x:UOp, reg:Register):
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
-    ret = isel_matcher.rewrite(UOp(Ops.COPY, src=(x,), tag=reg))
-    assert ret is not None, f"failed to copy {x}"
-    return ret
+  def copy(self, x:UOp, reg:Register) -> UOp: return x.ins(X86Ops.MOV, src=(x,), tag=reg)
 
   def spill(self, disp:UOp, x:UOp) -> UOp:
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
     return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
 
+  # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
   def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16
     dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype

tinygrad/runtime/autogen/am/__init__.py

@@ -22,7 +22,7 @@ reg_files = {
 reg_patterns = {
   "gc": ["GCVM", "GCMC_VM", "CP_(HQD|MQD|MEC|ME_CNTL|PERFMON|RB_WPTR_POLL_CNTL|INT_CNTL|STAT|PFP_PRGRM|ME_PRGRM|COHER_START)", "COMPUTE_",
          "(SQ|GL2C|TCC)_PERFCOUNTER", "SQ_THREAD_TRACE", "SPI_(CONFIG_CNTL|COMPUTE_QUEUE_RESET)", "GRBM", "SH_MEM", "RLC", "TCP", "GB_ADDR_CONFIG",
-         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[0-367]"],
+         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[0-35-7]"],
   "mmhub": ["MMVM", "MMMC_VM", "MM_ATC_L2_MISC_CG"],
   "nbio": (nbio:=["BIF_BX_PF[01]_GPU_HDP_FLUSH", "BIF_BX_PF0_RSMU", "BIF_BX0_(REMAP_HDP_MEM_FLUSH_CNTL|BIF_DOORBELL_INT_CNTL|PCIE_INDEX2|PCIE_DATA2)",
                   "BIFC_(DOORBELL_ACCESS_EN_PF|GFX_INT_MONITOR_MASK)", "XCC_DOORBELL_FENCE", "DOORBELL0_CTRL_ENTRY", "GDC_S2A0_S2A_DOORBELL_ENTRY",

tinygrad/runtime/autogen/am/regs.py

@@ -514,6 +514,7 @@ gc_9_4_3 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regCP_COHER_START_DELAY': (8315, 1, {'start_delay_count': (0, 5)}),
@@ -1799,6 +1800,7 @@ gc_11_0_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -3378,6 +3380,7 @@ gc_11_0_3 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -4807,6 +4810,7 @@ gc_11_5_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -6068,6 +6072,7 @@ gc_12_0_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),

tinygrad/runtime/autogen/amd/rdna3/str_pcode.py

@@ -47,7 +47,7 @@ PCODE = {
   DSOp.DS_MIN_RTN_F32: 'tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nMEM[ADDR].f32 = src < tmp ? src : tmp;\nRETURN_DATA.f32 = tmp',
   DSOp.DS_MAX_RTN_F32: 'tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nMEM[ADDR].f32 = src > tmp ? src : tmp;\nRETURN_DATA.f32 = tmp',
   DSOp.DS_WRAP_RTN_B32: 'tmp = MEM[ADDR].u32;\nMEM[ADDR].u32 = tmp >= DATA.u32 ? tmp - DATA.u32 : tmp + DATA2.u32;\nRETURN_DATA = tmp',
-  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}',
+  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset >= 0xc000) {\n// rotate\nrotate = offset[9:5];\nmask = offset[4:0];\nif (offset[10]) {\nrotate = -rotate;\n}\nfor (i = 0; i < 64; i++) {\nj = (i & mask) | ((i + rotate) & ~mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset[15]) {\n// full data sharing within 4 consecutive threads\nfor (i = 0; i < 64; i+=4) {\nthread_out[i+0] = thread_valid[i+offset[1:0]]?thread_in[i+offset[1:0]]:0;\nthread_out[i+1] = thread_valid[i+offset[3:2]]?thread_in[i+offset[3:2]]:0;\nthread_out[i+2] = thread_valid[i+offset[5:4]]?thread_in[i+offset[5:4]]:0;\nthread_out[i+3] = thread_valid[i+offset[7:6]]?thread_in[i+offset[7:6]]:0;\n}\n} else { // offset[15] == 0\n// limited data sharing within 32 consecutive threads\nxor_mask = offset[14:10];\nor_mask = offset[9:5];\nand_mask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = (((i & 0x1f) & and_mask) | or_mask) ^ xor_mask;\nj |= (i & 0x20); // which group of 32\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n}',
   DSOp.DS_LOAD_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET.u32].b32',
   DSOp.DS_LOAD_2ADDR_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET0.u32 * 4U].b32;\nRETURN_DATA[63 : 32] = MEM[ADDR + OFFSET1.u32 * 4U].b32',
   DSOp.DS_LOAD_2ADDR_STRIDE64_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET0.u32 * 256U].b32;\nRETURN_DATA[63 : 32] = MEM[ADDR + OFFSET1.u32 * 256U].b32',

tinygrad/runtime/autogen/amd/rdna4/str_pcode.py

@@ -44,7 +44,7 @@ PCODE = {
   DSOp.DS_CMPSTORE_RTN_B32: 'addr = CalcDsAddr(vgpr_a.b32, offset.b32);\ntmp = MEM[addr].b32;\nsrc = DATA.b32;\ncmp = DATA2.b32;\nMEM[addr].b32 = tmp == cmp ? src : tmp;\nRETURN_DATA.b32 = tmp',
   DSOp.DS_MIN_NUM_RTN_F32: "tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nif (isNAN(64'F(src.f32)) && isNAN(64'F(tmp.f32))) then\nMEM[ADDR].f32 = 32'F(cvtToQuietNAN(64'F(src.f32)))\nelsif isNAN(64'F(src.f32)) then\nMEM[ADDR].f32 = tmp.f32\nelsif isNAN(64'F(tmp.f32)) then\nMEM[ADDR].f32 = src.f32\nelsif ((src.f32 < tmp.f32) || ((abs(src.f32) == 0.0F) && (abs(tmp.f32) == 0.0F) && sign(src.f32) &&\n!sign(tmp.f32))) then\n// NOTE: -0<+0 is TRUE in this comparison\nMEM[ADDR].f32 = src.f32\nelse\nMEM[ADDR].f32 = tmp.f32\nendif;\nRETURN_DATA.f32 = tmp",
   DSOp.DS_MAX_NUM_RTN_F32: "tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nif (isNAN(64'F(src.f32)) && isNAN(64'F(tmp.f32))) then\nMEM[ADDR].f32 = 32'F(cvtToQuietNAN(64'F(src.f32)))\nelsif isNAN(64'F(src.f32)) then\nMEM[ADDR].f32 = tmp.f32\nelsif isNAN(64'F(tmp.f32)) then\nMEM[ADDR].f32 = src.f32\nelsif ((src.f32 > tmp.f32) || ((abs(src.f32) == 0.0F) && (abs(tmp.f32) == 0.0F) && !sign(src.f32) &&\nsign(tmp.f32))) then\n// NOTE: +0>-0 is TRUE in this comparison\nMEM[ADDR].f32 = src.f32\nelse\nMEM[ADDR].f32 = tmp.f32\nendif;\nRETURN_DATA.f32 = tmp",
-  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}',
+  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset >= 0xc000) {\n// rotate\nrotate = offset[9:5];\nmask = offset[4:0];\nif (offset[10]) {\nrotate = -rotate;\n}\nfor (i = 0; i < 64; i++) {\nj = (i & mask) | ((i + rotate) & ~mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset[15]) {\n// full data sharing within 4 consecutive threads\nfor (i = 0; i < 64; i+=4) {\nthread_out[i+0] = thread_valid[i+offset[1:0]]?thread_in[i+offset[1:0]]:0;\nthread_out[i+1] = thread_valid[i+offset[3:2]]?thread_in[i+offset[3:2]]:0;\nthread_out[i+2] = thread_valid[i+offset[5:4]]?thread_in[i+offset[5:4]]:0;\nthread_out[i+3] = thread_valid[i+offset[7:6]]?thread_in[i+offset[7:6]]:0;\n}\n} else { // offset[15] == 0\n// limited data sharing within 32 consecutive threads\nxor_mask = offset[14:10];\nor_mask = offset[9:5];\nand_mask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = (((i & 0x1f) & and_mask) | or_mask) ^ xor_mask;\nj |= (i & 0x20); // which group of 32\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n}',
   DSOp.DS_LOAD_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET.u32].b32',
   DSOp.DS_LOAD_2ADDR_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET0.u32 * 4U].b32;\naddr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[63 : 32] = MEM[addr + OFFSET1.u32 * 4U].b32',
   DSOp.DS_LOAD_2ADDR_STRIDE64_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET0.u32 * 256U].b32;\naddr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[63 : 32] = MEM[addr + OFFSET1.u32 * 256U].b32',

tinygrad/runtime/ops_amd.py

@@ -538,7 +538,7 @@ class AMDCopyQueue(HWQueue):
       tail_blit_dword += cmdsz
 
     # Force align of submits to hit our usb layer write cache.
-    if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0 and dev.is_usb(): tail_blit_dword = 0
+    if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0 and dev.is_usb(): tail_blit_dword, rem_packet_cnt = 0, len(cmds)
 
     # USB devices run in single-step mode, so they can't overrun the queue.
     total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -sdma_queue.put_value % sdma_queue.ring.nbytes) + rem_packet_cnt * 4
@@ -661,48 +661,67 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
     dev, usb, ts, sdma = self.dev, self.dev.iface.pci_dev.usb, self.dev.timeline_signal, self.dev.sdma
     CHUNK, src_mv = 0x40000 - 4, src.cast('B')  # payload per chunk: the 256KB window minus the 4B trailing sentinel
     nchunks = ceildiv(src.nbytes, CHUNK)
-    FENCE = 0xA800  # drain fence: the GPU writes it via sys_buf (PCIe 0x820800), the host reads it here (xdata)
-    if not hasattr(self, '_usb_seq'):  # one-time: clear the fence and zero both windows so garbage can't match a sentinel
-      self._usb_seq, self._usb_stage = 0, [alloc_cbuffer(0x40000) for _ in range(2)]  # (backing array, memoryview) pairs
+    if nchunks == 0: return
+    # Keep each submit below half of the 1 MiB SDMA ring, leaving room for wrap padding.
+    if src.nbytes > (batch := 1 << 30):
+      for off in range(0, src.nbytes, batch): self._copyin(dest.offset(off), src_mv[off:off+batch])
+      return
+    FENCE = 0xA800  # byte fences at FENCE and FENCE+8; GPU sys_buf offsets 0x800 and 0x808
+    if not hasattr(self, '_usb_seq'):
+      self._usb_seq, self._usb_clear, self._usb_stage = 0, 0, [alloc_cbuffer(0x40000) for _ in range(2)]
       self._usb_wins = (self.b[0].offset(0, 0x40000), self.b[0].offset(0x40000, 0x40000))  # two windows, engine slots 0/16
-      usb.write(FENCE, bytes(8))
-      for bi in range(2): usb.scsi_write(bytes(0x40000), slot_start=bi * 16)
+      usb.write(FENCE, bytes(16))
 
-    def wait_drain(count):  # spin until the drain fence reaches count, i.e. chunks 0..count-1 are fully in VRAM
+    def wait_fence(count, addr=FENCE, next_ok=False, current=None):
+      expected = (count & 0xff, (count + int(next_ok)) & 0xff)
       t0 = time.perf_counter()
-      while int.from_bytes(usb.read(FENCE, 8), 'little') < count:
-        if time.perf_counter() - t0 > 10: raise RuntimeError(f"GPU failed to drain USB copyin chunk {count - 1} (10s, hung GPU?)")
+      while (current if current is not None else usb.read(addr, 1)[0]) not in expected:
+        current = None
+        if time.perf_counter() - t0 > 10: raise RuntimeError(f"GPU failed to reach USB copyin fence {addr:#x} value {count} (10s, hung GPU?)")
 
     # build the whole ring upfront: per chunk, poll the sentinel, copy SRAM->VRAM, bump the fence; then one doorbell
     POLL_EQ = sdma.SDMA_OP_POLL_REGMEM | sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(3) | sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
     POLL_DW5 = sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)
-    q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
-    for c in range(nchunks):
-      seq, size = self._usb_seq + c, min(CHUNK, src.nbytes - c * CHUNK)
-      q.q(POLL_EQ, *data64_le(self._usb_wins[seq & 1].va_addr + round_up(size + 4, 512) - 4), 0x51000000 | (seq & 0xFFFFFF), 0xFFFFFFFF, POLL_DW5)
-      q.copy(dest.offset(c * CHUNK), self._usb_wins[seq & 1], size)
-      q.write(dev.iface.sys_buf.offset(0x800, 8), seq + 1, b64=True)
-    q.signal(ts, dev.next_timeline()).submit(dev)
+    def sentinel(c):
+      size = min(CHUNK, src.nbytes - c * CHUNK)
+      return self._usb_wins[(self._usb_seq + c) & 1].offset(round_up(size + 4, 512) - 4, 4)
 
-    # stream the chunks: stage the wire image [payload][sentinel], arm the window, send. A window is reusable once
-    # its previous occupant (seq-2) is both fully sent (tag reaped) and fully drained to VRAM (the fence).
-    inflight = [None, None]
+    q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
+    # A short chunk's sentinel can lie in old payload (including copyout data). Clear it before USB can fill the window.
+    for c in range(min(2, nchunks)): q.write(sentinel(c), 0)
+    # Toggle per submit, not per chunk: a 256-chunk transfer must not leave the next clear fence already satisfied.
+    self._usb_clear ^= 1
+    q.write(dev.iface.sys_buf.offset(0x808, 4), self._usb_clear)
     for c in range(nchunks):
       seq, size = self._usb_seq + c, min(CHUNK, src.nbytes - c * CHUNK)
-      if inflight[seq & 1] is not None: usb.usb.bulk_wait(inflight[seq & 1])
+      q.q(POLL_EQ, *data64_le(sentinel(c).va_addr), 0x51000000 | (seq & 0xFFFFFF), 0xFFFFFFFF, POLL_DW5)
+      q.copy(dest.offset(c * CHUNK), self._usb_wins[seq & 1], size)
+      # Clear the next occupant's sentinel only after this payload is copied, and before publishing the drain fence.
+      if c + 2 < nchunks: q.write(sentinel(c + 2), 0)
+      q.write(dev.iface.sys_buf.offset(0x800, 4), (seq + 1) & 0xff)
+    q.signal(ts, dev.next_timeline()).submit(dev)
+    wait_fence(self._usb_clear, FENCE + 8)
+
+    # Stage the next window while USB sends the previous one. F2 configures a single engine shared by both windows.
+    inflight = None
+    for c in range(nchunks):
+      seq, size = self._usb_seq + c, min(CHUNK, src.nbytes - c * CHUNK)
+      # Reading the drain fence does not reconfigure F2, so overlap it with USB sending the previous chunk and CPU staging.
+      rd_tag, rd_mv = usb.usb.control_read_async(0xE4, 1, value=FENCE)
       buf = self._usb_stage[seq & 1][1]
       buf[:size] = src_mv[c * CHUNK : c * CHUNK + size]
       wire = round_up(size + 4, 512)  # payload plus the sentinel, padded to 512B sectors (full window for max chunks)
       struct.pack_into('<I', buf, wire - 4, 0x51000000 | (seq & 0xFFFFFF))  # the sentinel is the last dword of the wire
-      arm_tag = usb.usb.control_write_async(0xF2, wire // 512, (seq & 1) * 16 | (ceildiv(wire, 0x4000) << 8))  # wValue=sectors, wIndex=slot|count
-      rd_tag, rd_mv = usb.usb.control_read_async(0xE4, 8, value=FENCE)  # arm and fence read fly in one round-trip window
-      usb.usb.bulk_wait(arm_tag)
+      # Rearming F2 recycles the slots immediately, before bulk data arrives. Drain seq-2 before the arm itself.
       usb.usb.bulk_wait(rd_tag)
-      if int.from_bytes(rd_mv, 'little') < seq - 1: wait_drain(seq - 1)  # rare: the drain lagged; spin on fresh reads
-      inflight[seq & 1] = usb.usb.bulk_write_async(buf[:wire])
-    for tag in inflight: usb.usb.bulk_wait(tag)
+      # One-byte reads cannot tear. At most two chunks are outstanding, so modulo-256 equality is unambiguous.
+      wait_fence(seq - 1, next_ok=True, current=rd_mv[0])
+      if inflight is not None: usb.usb.bulk_wait(inflight)
+      usb.usb.control_write(0xF2, wire // 512, (seq & 1) * 16 | (ceildiv(wire, 0x4000) << 8))  # wValue=sectors, wIndex=slot|count
+      inflight = usb.usb.bulk_write_async(buf[:wire])
+    if inflight is not None: usb.usb.bulk_wait(inflight)
     self._usb_seq += nchunks
-    wait_drain(self._usb_seq)  # copyin is synchronous: everything must be in VRAM before returning
+    wait_fence(self._usb_seq)  # copyin is synchronous: everything must be in VRAM before returning
 
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     if not self.dev.is_usb(): return super()._copyout(dest, src)

tinygrad/runtime/ops_cpu.py

@@ -87,7 +87,7 @@ class CPUAllocator(HCQAllocator['CPUDevice']):
     ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
-    ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+    dest[:] = self._as_buffer(src)[:len(dest)]
   def _do_map(self, buf:HCQBuffer):
     if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
     return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)

tinygrad/runtime/ops_nv.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
-import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref
+import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, itertools
 assert sys.platform != 'win32'
-from typing import cast
+from typing import Any
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
-from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, hcq_filter_visible_devices, hcq_profile
-from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec, TinyELF
-from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, hi32, lo32, PROFILE, ContextVar, VIZ, ProfileEvent
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, patch, to_name, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface, FileIOInterface, BumpAllocator, hcq_filter_visible_devices
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, TinyELF
+from tinygrad.dtype import dtypes, DType
+from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, PROFILE, ContextVar, VIZ
+from tinygrad.helpers import ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.runtime.autogen import nv_570, nv_580, nv_610, mesa
@@ -24,10 +27,7 @@ PMA = ContextVar("PMA", abs(VIZ.value)>=2)
 @dataclass(frozen=True)
 class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:int; profile_key:bytes|None=None # noqa: E702
 
-class NVSignal(HCQSignal):
-  def _sleep(self, time_spent_since_last_sleep_ms:int):
-    # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
-    if time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)
+def hilo(addr:UOp) -> tuple[UOp, UOp]: return (addr >> 32).cast(dtypes.uint32), addr.cast(dtypes.uint32)
 
 def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
 
@@ -41,10 +41,13 @@ def nv_iowr(fd:FileIOInterface, nr, args, cmd=None):
   ret = fd.ioctl(cmd or ((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF)), args)
   if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
 
+def nvm(subc:int, mthd:int, *vals, typ=2) -> list:
+  return [(typ << 28) | (sum(v.dtype.itemsize // 4 if isinstance(v, UOp) else 1 for v in vals) << 16) | (subc << 13) | (mthd >> 2), *vals]
+
 class QMD:
   fields: dict[str, dict[str, tuple[int, int]]] = {}
 
-  def __init__(self, dev:NVDevice, view:MMIOInterface|None=None, **kwargs):
+  def __init__(self, dev:NVDevice, blob:bytearray|None=None):
     self.ver, self.sz = (5, 0x60) if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else (3, 0x40)
 
     # Init fields from module
@@ -52,272 +55,240 @@ class QMD:
       QMD.fields[pref] = {**{name[len(pref)+1:]: dt for name,dt in nv_gpu.__dict__.items() if name.startswith(pref) and isinstance(dt, tuple)},
         **{name[len(pref)+1:]+f"_{i}": dt(i) for name,dt in nv_gpu.__dict__.items() for i in range(8) if name.startswith(pref) and callable(dt)}}
 
-    self.mv, self.pref = (memoryview(bytearray(self.sz * 4)) if view is None else view), pref
-    if kwargs: self.write(**kwargs)
+    self.mv, self.pref = (bytearray(self.sz * 4) if blob is None else blob), pref
+    self.patches:dict[int, UOp] = {}
 
-  def _rw_bits(self, hi:int, lo:int, value:int|None=None):
-    mask = ((1 << (width:=hi - lo + 1)) - 1) << (lo % 8)
-    num = int.from_bytes(self.mv[lo//8:hi//8+1], "little")
+  def read(self, k:str) -> int:
+    hi, lo = QMD.fields[self.pref][k.upper()]
+    return (int.from_bytes(self.mv[lo//8:hi//8+1], "little") >> (lo % 8)) & ((1 << (hi - lo + 1)) - 1)
 
-    if value is None: return (num & mask) >> (lo % 8)
+  def write(self, **kwargs:int|UOp):
+    for k, v in kwargs.items():
+      hi, lo = QMD.fields[self.pref][k.upper()]
+      if isinstance(v, UOp):
+        assert lo % 8 == 0, f"{k} is not byte aligned"
+        self.patches[lo // 8] = v.ccast(next(t for t in (dtypes.uint64, dtypes.uint32, dtypes.uint16, dtypes.uint8) if t.itemsize * 8 <= hi - lo + 1))
+      else:
+        if v >> (hi - lo + 1): raise ValueError(f"{k}={v:#x} does not fit")
+        mask, num = ((1 << (hi - lo + 1)) - 1) << (lo % 8), int.from_bytes(self.mv[lo//8:hi//8+1], "little")
+        self.mv[lo//8:hi//8+1] = ((num & ~mask) | (v << (lo % 8))).to_bytes(hi//8 - lo//8 + 1, "little")
 
-    if value >= (1 << width): raise ValueError(f"{value:#x} does not fit.")
-    self.mv[lo//8:hi//8+1] = int((num & ~mask) | ((value << (lo % 8)) & mask)).to_bytes((hi//8 - lo//8 + 1), "little")
+  def set_addr(self, name:str, addr:UOp, sfx:str=""): self.write(**{f"{name}_lower{sfx}": addr, f"{name}_upper{sfx}": addr >> 32})
+  def set_constant_buf_addr(self, i:int, addr:UOp):
+    self.set_addr("constant_buffer_addr", addr >> (6 if self.ver >= 4 else 0), f"_shifted6_{i}" if self.ver >= 4 else f"_{i}")
+  def set_program_addr(self, addr:UOp):
+    self.set_addr("program_address", addr >> (4 if self.ver >= 4 else 0), "_shifted4" if self.ver >= 4 else "")
+    self.set_addr("program_prefetch_addr", addr >> 8, "_shifted")
+  def set_release(self, addr:UOp, payload:UOp, timestamp:bool=False) -> bool:
+    if (i:=next((i for i in range(2) if not self.read(f"release{i}_enable")), None)) is None: return False
+    self.set_addr(f"release_semaphore{i}_addr" if self.ver >= 4 else f"release{i}_address", addr)
+    self.set_addr(f"release_semaphore{i}_payload" if self.ver >= 4 else f"release{i}_payload", payload)
+    self.write(**{f"release{i}_enable": 1, f"release_structure_size_{i}" if self.ver >= 4 else f"release{i}_structure_size": 0 if timestamp else 2},
+               **({} if self.ver >= 4 else {f"release{i}_payload64b": 1}))
+    return True
+  @property
+  def grid(self) -> tuple[str, ...]:
+    return ("grid_width", "grid_height", "grid_depth") if self.ver >= 4 else ("cta_raster_width", "cta_raster_height", "cta_raster_depth")
 
-  def write(self, **kwargs):
-    for k,val in kwargs.items(): self._rw_bits(*QMD.fields[self.pref][k.upper()], value=val) # type: ignore [misc]
+# *****************
+# queues
 
-  def read(self, k, val=0): return self._rw_bits(*QMD.fields[self.pref][k.upper()])
+class NVQueue(HWQueue):
+  dev:NVDevice
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
-  def field_offset(self, k): return QMD.fields[self.pref][k.upper()][1] // 8
+  def nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
-  def set_constant_buf_addr(self, i, addr):
-    if self.ver < 4: self.write(**{f'constant_buffer_addr_upper_{i}':hi32(addr), f'constant_buffer_addr_lower_{i}':lo32(addr)})
-    else: self.write(**{f'constant_buffer_addr_upper_shifted6_{i}':hi32(addr >> 6), f'constant_buffer_addr_lower_shifted6_{i}':lo32(addr >> 6)})
+  def sem(self, addr:UOp, value:UOp, **flags:str):
+    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, addr, value.ccast(dtypes.uint64), nv_flags("NVC56F_SEM_EXECUTE", payload_size="64bit", **flags))
 
-class NVCommandQueue(HWQueue[HCQSignal, 'NVDevice', 'NVProgram', 'NVArgsState']):
-  def __init__(self):
-    self.active_qmd = None
-    super().__init__()
+  def wait(self, signal:UOp, value:UOp): self.sem(signal.getaddr(self.devs), value, operation="acq_circ_geq")
+  def signal(self, signal:UOp, value:UOp): self.release(signal, value)
+  def timestamp(self, signal:UOp): self.release(signal, UOp.const(0, dtypes.uint64), timestamp=True)
+  def release(self, signal:UOp, value:UOp, timestamp:bool=False):
+    self.sem(signal.getaddr(self.devs), value, operation="release", release_wfi="en", release_timestamp="en" if timestamp else "dis")
+    if not timestamp: self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
 
-  def __del__(self):
-    if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))
+  def submit(self, cmdbuf:UOp) -> UOp:
+    fifo, ib, off = self.dev.fifos[self.queue], *unwrap_view(cmdbuf)
 
-  def nvm(self, subchannel, mthd, *args, typ=2): self.q((typ << 28) | (len(args) << 16) | (subchannel << 13) | (mthd >> 2), *args)
+    ring, gpput, doorbell, put, gpentry = [UOp.placeholder((sz,), dt, device=self.devs, volatile=True, tag=to_name(nm, self.queue))
+      for nm, dt, sz in (("ring", dtypes.uint64, fifo.entries), ("gpput", dtypes.uint32, 1), ("doorbell", dtypes.uint32, 1),
+                         ("put_value", dtypes.uint64, 1), ("gpentry", dtypes.uint64, 1))]
+    gpentry = patch(gpentry, [(0, ib.getaddr(self.devs) + UOp.const(off | (cmdbuf.max_numel() // 4 << 42) | (1 << 41), dtypes.uint64))])
 
-  def setup(self, compute_class=None, copy_class=None, local_mem_window=None, shared_mem_window=None, local_mem=None, local_mem_tpc_bytes=None):
-    if compute_class: self.nvm(1, nv_gpu.NVC6C0_SET_OBJECT, compute_class)
-    if copy_class: self.nvm(4, nv_gpu.NVC6C0_SET_OBJECT, copy_class)
-    if local_mem_window: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_WINDOW_A, *data64(local_mem_window))
-    if shared_mem_window: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_SHARED_MEMORY_WINDOW_A, *data64(shared_mem_window))
-    if local_mem: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(local_mem))
-    if local_mem_tpc_bytes: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(local_mem_tpc_bytes), 0xff)
-    return self
+    p = put.index(0).load()
+    written = UOp.barrier(ring.after(cmdbuf).index((p % fifo.entries).cast(dtypes.int)).store(gpentry.index(0).load()), put.index(0).store(p + 1))
+    queued = UOp.barrier(gpput.after(written).index(0).store(((p + 1) % fifo.entries).cast(dtypes.uint32)))
+    return doorbell.after(queued).index(0).store(UOp.const(fifo.token, dtypes.uint32))
 
-  def wait(self, signal:HCQSignal, value:sint=0):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit"))
-    self.active_qmd = None
-    return self
+class NVComputeQueue(NVQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+  ]) + NVQueue.q_rewrite
 
-  def timestamp(self, signal:HCQSignal): return self.signal(signal, 0)
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
 
-  def bind(self, dev:NVDevice):
-    self.binded_device = dev
-    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True))
-    hw_view = self.hw_page.cpu_view().view(fmt='I')
-    for i, value in enumerate(self._q): hw_view[i] = value
+    progs = [nv_build_program(self.dev, u.src[0], self.devs)[0] for u in self.lin.src if u.op is Ops.CALL]
+    self.qmd_sz = round_up(QMD(self.dev).sz * 4, 256)
+    self.stride = self.qmd_sz + max([p.kernargs_size for p in progs], default=0)
+    self.qmd_buf = UOp.placeholder((len(progs) * self.stride,), dtypes.uint8, device=self.devs, tag=to_name("qmd", self.queue))
+    self.qmds:list[QMD] = []
+    self.prev_qmd:QMD|None = None # the launch the next one chains onto
 
-    # From now on, the queue is on the device for faster submission.
-    self._q = hw_view
+  def wait(self, signal:UOp, value:UOp):
+    self.prev_qmd = None
+    super().wait(signal, value)
 
-  def _submit_to_gpfifo(self, dev:NVDevice, gpfifo:GPFifo):
-    if dev == self.binded_device: cmdq_addr = self.hw_page.va_addr
-    else:
-      cmdq_addr = dev.cmdq_allocator.alloc(len(self._q) * 4, 16)
-      cmdq_wptr = (cmdq_addr - dev.cmdq_page.va_addr) // 4
-      dev.cmdq[cmdq_wptr : cmdq_wptr + len(self._q)] = array.array('I', self._q)
+  def release(self, signal:UOp, value:UOp, timestamp:bool=False):
+    if self.prev_qmd is None or not self.prev_qmd.set_release(signal.getaddr(self.devs), value, timestamp):
+      self.prev_qmd = None
+      super().release(signal, value, timestamp)
 
-    gpfifo.ring[gpfifo.put_value % gpfifo.entries_count] = (cmdq_addr//4 << 2) | (len(self._q) << 42) | (1 << 41)
-    gpfifo.gpput[0] = (gpfifo.put_value + 1) % gpfifo.entries_count
+  def submit(self, cmdbuf:UOp) -> UOp:
+    if self.qmds:
+      patches = [(i * self.stride + off, w) for i, q in enumerate(self.qmds) for off, w in q.patches.items()]
+      cmdbuf = cmdbuf.after(patch(self.qmd_buf, patches, b"".join(q.mv for q in self.qmds)))
+    return super().submit(cmdbuf)
 
-    System.memory_barrier()
-    dev.gpu_mmio[0x90 // 4] = gpfifo.token
-    gpfifo.put_value += 1
-
-class NVComputeQueue(NVCommandQueue):
   def memory_barrier(self):
+    self.prev_qmd = None
     self.nvm(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI,
              nv_flags("NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI", instruction="true", global_data="true", constant="true"))
-    self.active_qmd:QMD|None = None
-    return self
 
-  def exec(self, prg:NVProgram, args_state:NVArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
-    self.bind_args_state(args_state)
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = nv_build_program(self.dev, prg, self.devs)
+    global_size, local_size = prg.arg.global_size, prg.arg.local_size
+    if prod(local_size) > 1024 or data.max_threads < prod(local_size):
+      raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {data.max_threads=}")
+    if any(g > mx for g,mx in zip(global_size, [2147483647, 65535, 65535]) if isinstance(g, int)) or \
+       any(l > mx for l,mx in zip(local_size, [1024, 1024, 64])):
+      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
 
-    qmd_buf = args_state.buf.offset(round_up(prg.constbufs[0][1], 1 << 8))
-    qmd_buf.cpu_view().view(size=prg.qmd.mv.nbytes, fmt='B')[:] = prg.qmd.mv
-    assert qmd_buf.va_addr < (1 << 40), f"large qmd addr {qmd_buf.va_addr:x}"
+    qmd_addr = self.qmd_buf.getaddr(self.devs) + UOp.const(len(self.qmds) * self.stride, dtypes.uint64)
+    qmd = QMD(self.dev, data.qmd.mv.ljust(self.stride, b"\0")) # the program's template, in a slot of its own
+    qmd.write(**dict(zip(qmd.grid, global_size)), **{f"cta_thread_dimension{j}": l for j, l in enumerate(local_size)})
+    qmd.set_program_addr(lib.getaddr(self.devs) + data.prog_off)
+    for j, (off, _) in data.constbufs.items():
+      qmd.set_constant_buf_addr(j, qmd_addr + UOp.const(self.qmd_sz, dtypes.uint64) if j == 0 else lib.getaddr(self.devs) + off)
+    bufs, vals = [get_call_arg_uops(call)[j] for j in prg.arg.globals], get_call_var_uops(call, prg)
+    qmd.mv[self.qmd_sz:(at:=self.qmd_sz + len(data.cbuf_0) * 4)] = array.array('I', data.cbuf_0).tobytes() # constant buffer 0: the driver params
+    qmd.patches |= {at + j * 8: b.getaddr(self.devs) for j, b in enumerate(bufs)} | {at + o: v.ccast(dt) for v, (o, dt) in zip(vals, data.vars)}
 
-    qmd = QMD(dev=prg.dev, view=qmd_buf.cpu_view()) # Save qmd for later update
+    if self.prev_qmd is None:
+      if self.dev.pma_enabled: self.nvm(1, nv_gpu.NVC6C0_PM_TRIGGER, 0)
+      self.nvm(1, nv_gpu.NVC6C0_SEND_PCAS_A, (qmd_addr >> 8).cast(dtypes.uint32))
+      self.nvm(1, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B_PCAS_ACTION_PREFETCH_SCHEDULE)
+    else: self.prev_qmd.write(dependent_qmd0_pointer=qmd_addr >> 8, dependent_qmd0_action=1, dependent_qmd0_prefetch=1, dependent_qmd0_enable=1)
+    self.qmds.append(qmd)
+    self.prev_qmd = qmd
 
-    self.bind_sints_to_mem(*global_size, mem=qmd_buf.cpu_view(), fmt='I', offset=qmd.field_offset('cta_raster_width' if qmd.ver<4 else 'grid_width'))
-    self.bind_sints_to_mem(*(local_size[:2]), mem=qmd_buf.cpu_view(), fmt='H', offset=qmd.field_offset('cta_thread_dimension0'))
-    self.bind_sints_to_mem(local_size[2], mem=qmd_buf.cpu_view(), fmt='B', offset=qmd.field_offset('cta_thread_dimension2'))
-    qmd.set_constant_buf_addr(0, args_state.buf.va_addr)
+class NVCopyQueue(NVQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ()),
+  ]) + NVQueue.q_rewrite
 
-    if self.active_qmd is None:
-      if prg.dev.pma_enabled: self.nvm(1, nv_gpu.NVC6C0_PM_TRIGGER, 0)
-      self.nvm(1, nv_gpu.NVC6C0_SEND_PCAS_A, qmd_buf.va_addr >> 8)
-      self.nvm(1, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B, 9)
-    else:
-      self.active_qmd.write(dependent_qmd0_pointer=qmd_buf.va_addr >> 8, dependent_qmd0_action=1, dependent_qmd0_prefetch=1, dependent_qmd0_enable=1)
-
-    self.active_qmd, self.active_qmd_buf = qmd, qmd_buf
-    return self
-
-  def signal(self, signal:HCQSignal, value:sint=0):
-    if self.active_qmd is not None:
-      for i in range(2):
-        if self.active_qmd.read(f'release{i}_enable') == 0:
-          self.active_qmd.write(**{f'release{i}_enable': 1})
-
-          addr_off = self.active_qmd.field_offset(f'release{i}_address_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_addr_lower')
-          self.bind_sints_to_mem(signal.value_addr & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=addr_off)
-          self.bind_sints_to_mem(signal.value_addr >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', mask=0xf, offset=addr_off+4)
-
-          val_off = self.active_qmd.field_offset(f'release{i}_payload_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_payload_lower')
-          self.bind_sints_to_mem(value & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off)
-          self.bind_sints_to_mem(value >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off+4)
-          return self
-
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit", release_timestamp="en"))
-    self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
-    self.active_qmd = None
-    return self
-
-  def write(self, b:HCQBuffer, val:sint, b64:bool=False):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(b.va_addr), *data64_le(val),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit" if b64 else "32bit"))
-    self.active_qmd = None
-    return self
-
-  def poll_bit(self, b:HCQBuffer, val:sint, mask:int):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(b.va_addr), *data64_le((~mask & 0xFFFFFFFF) if val == 0 else val),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_nor" if val == 0 else "acq_and", payload_size="32bit"))
-    self.active_qmd = None
-    return self
-
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.compute_gpfifo)
-
-class NVCopyQueue(NVCommandQueue):
-  def __init__(self, queue_idx=0):
-    self.queue_idx = queue_idx
-    super().__init__()
-
-  def copy(self, dest:HCQBuffer, src:HCQBuffer, copy_size:int):
-    for off in range(0, copy_size, step:=(1 << 31)):
-      self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src.va_addr+off), *data64(dest.va_addr+off))
-      self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(copy_size-off, step))
+  def copy(self, call:UOp):
+    dest, src = (a.getaddr(self.devs) for a in call.src[1:3])
+    for off in range(0, sz:=call.src[2].max_numel() * call.src[2].dtype.itemsize, step:=(1 << 31)):
+      self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *hilo(src + UOp.const(off, dtypes.uint64)), *hilo(dest + UOp.const(off, dtypes.uint64)))
+      self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(sz - off, step))
       self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA,
                nv_flags("NVC6B5_LAUNCH_DMA", data_transfer_type="non_pipelined", src_memory_layout="pitch", dst_memory_layout="pitch"))
-    return self
 
-  def signal(self, signal:HCQSignal, value:sint=0):
-    self.nvm(4, nv_gpu.NVC6B5_SET_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type="release_four_word_semaphore"))
-    return self
+  def semaphore(self, addr:UOp, value:UOp, typ:str): # a one word release writes just the payload, a four word one the timestamp after it
+    self.nvm(4, nv_gpu.NVC6B5_SET_SEMAPHORE_A, *hilo(addr), value.ccast(dtypes.uint32))
+    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type=f"release_{typ}_word_semaphore"))
+  def timestamp(self, signal:UOp): self.semaphore(signal.getaddr(self.devs), UOp.const(0, dtypes.uint32), "four")
+  def signal(self, signal:UOp, value:UOp): self.semaphore(signal.getaddr(self.devs), value, "one")
 
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.dma_gpfifo)
+# *****************
+# programs
 
-class NVVideoQueue(NVCommandQueue):
-  def decode_hevc_chunk(self, pic_desc:HCQBuffer, in_buf:HCQBuffer, out_buf:HCQBuffer, out_buf_pos:int, hist_bufs:list[HCQBuffer], hist_pos:list[int],
-                        chroma_off:int, coloc_buf:HCQBuffer, filter_buf:HCQBuffer, intra_top_off:int, intra_unk_off:int|None, status_buf:HCQBuffer):
-    self.nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
-    self.nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
-             err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
-    self.nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, pic_desc.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, in_buf.va_addr >> 8)
-    for pos, buf in zip(hist_pos + [out_buf_pos], hist_bufs + [out_buf]):
-      self.nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
-      self.nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, coloc_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, status_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, pic_desc.offset(0x200).va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, filter_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_buf.va_addr + intra_top_off) >> 8)
-    if intra_unk_off is not None: self.nvm(4, 0x4dc, (filter_buf.va_addr + intra_unk_off) >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
-    return self
-
-  def signal(self, signal:HCQSignal, value:sint=0):
-    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_D, nv_flags("NVC9B0_SEMAPHORE_D", structure_size="four", payload_size="64bit"))
-    return self
-
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.vid_gpfifo)
-
-class NVArgsState(CLikeArgsState):
-  def __init__(self, buf:HCQBuffer, prg:NVProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    if (is_mock:=isinstance(prg.dev.iface, MOCKIface)): prg.cbuf_0[80:82] = [len(bufs), len(vals)]
-    super().__init__(buf, prg, bufs, vals=() if is_mock else vals, prefix=prg.cbuf_0 or None)
-    # mock expects all vars to be 64 bit
-    if is_mock and vals: self.bind_sints_to_buf(*vals, buf=self.buf, fmt='q', offset=len(prg.cbuf_0)*4 + len(bufs)*8)
-
-class NVProgram(HCQProgram['NVDevice']):
+class NVProgramData:
   def __init__(self, dev:NVDevice, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
-    self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]]
+    name, signature, mock = obj.name, obj.signature, isinstance(dev.iface, MOCKIface)
+    self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[offset in the image, size]]
+    self.relocs: list[tuple[int, int, DType, int]] = [] # (byte offset in the image, symbol offset, width, shift) of the program's address
+    self.prog_off, self.cbuf_0, sections, relocs = 0, [], list[Any](), list[Any]()
+    image:bytes = obj.lib
 
     if (NAK:=isinstance(dev.renderer, NAKRenderer)):
-      image, self.cbuf_0 = memoryview(bytearray(obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):])), []
-      self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
-    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(obj.lib) + b'\x00' * (4 - len(obj.lib)%4)).cast("I"), [], [] # type: ignore
-    else: image, sections, relocs = elf_loader(self.lib, force_section_align=128)
-    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
-    self.lib_gpu = self.dev.allocator.alloc(round_up((prog_sz:=image.nbytes), 0x1000) + 0x1000, buf_spec:=BufferSpec(nolru=True))
-    prog_addr = self.lib_gpu.va_addr
+      image = obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):]
+      regs, shmem, lcmem = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
+    elif mock: image = obj.lib.ljust(round_up(len(obj.lib), 4), b'\x00') # for MOCKGPU the lib is PTX code, not an elf
+    else:
+      img, sections, relocs = elf_loader(obj.lib, force_section_align=128)
+      image = bytes(img)
+    prog_sz = len(image)
+
     if not NAK:
       # For MOCKGPU, the lib is PTX code, so some values are emulated.
-      self.regs_usage, self.shmem_usage, self.lcmem_usage, cbuf0_size = 0, 0x400, 0x240, 0x160 if isinstance(dev.iface, MOCKIface) else 0
-      for sh in sections: # pylint: disable=possibly-used-before-assignment
-        if sh.name == f".nv.shared.{self.name}": self.shmem_usage = round_up(0x400 + sh.header.sh_size, 128)
-        if sh.name == f".text.{self.name}": prog_addr, prog_sz = self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size
-        elif m:=re.match(r'\.nv\.constant(\d+)', sh.name):
-          self.constbufs[int(m.group(1))] = (self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size)
+      regs, shmem, lcmem, cbuf0_size = 0, 0x400, 0x240, 0x160 if mock else 0
+      for sh in sections:
+        if sh.name == f".nv.shared.{name}": shmem = round_up(0x400 + sh.header.sh_size, 128)
+        if sh.name == f".text.{name}": self.prog_off, prog_sz = sh.header.sh_addr, sh.header.sh_size
+        elif m:=re.match(r'\.nv\.constant(\d+)', sh.name): self.constbufs[int(m.group(1))] = (sh.header.sh_addr, sh.header.sh_size)
         elif sh.name.startswith(".nv.info"):
           for typ, param, data in self._parse_elf_info(sh):
             if sh.name == f".nv.info.{obj.name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
-            elif sh.name == ".nv.info" and param == 0x12: self.lcmem_usage = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
-            elif sh.name == ".nv.info" and param == 0x2f: self.regs_usage = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
+            elif sh.name == ".nv.info" and param == 0x12: lcmem = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
+            elif sh.name == ".nv.info" and param == 0x2f: regs = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
 
-      # Apply relocs
-      for apply_image_offset, rel_sym_offset, typ, _ in relocs: # pylint: disable=possibly-used-before-assignment
-        # These types are CUDA-specific, applying them here
-        if typ == 2: image[apply_image_offset:apply_image_offset+8] = struct.pack('<Q', self.lib_gpu.va_addr + rel_sym_offset) # R_CUDA_64
-        elif typ == 0x38: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) & 0xffffffff)
-        elif typ == 0x39: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) >> 32)
+      # These reloc types are CUDA-specific: they all want the program's own address, which is only known once the linear links.
+      for apply_image_offset, rel_sym_offset, typ, _ in relocs:
+        if typ == 2: self.relocs.append((apply_image_offset, rel_sym_offset, dtypes.uint64, 0)) # R_CUDA_64
+        elif typ == 0x38: self.relocs.append((apply_image_offset + 4, rel_sym_offset, dtypes.uint32, 0))
+        elif typ == 0x39: self.relocs.append((apply_image_offset + 4, rel_sym_offset, dtypes.uint32, 32))
         else: raise RuntimeError(f"unknown NV reloc {typ}")
 
       # Minimum cbuf_0 size for driver params: Blackwell needs index 223 (224 entries), older GPUs need index 11 (12 entries)
       min_cbuf0_entries = 224 if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 12
       self.cbuf_0 = [0] * max(cbuf0_size // 4, min_cbuf0_entries)
 
+    # the arguments follow the driver params in constant buffer 0: the buffers as 64 bit addresses, then the vars packed by their width
+    nbufs = sum(name is None for name, *_ in signature)
+    self.vars = list(TinyELF.iter_sig(signature[nbufs:], nbufs * 8))
+    if mock: # mockgpu reads the arg counts out of cbuf0 and wants every var 64 bit
+      self.cbuf_0[80:82], self.vars = [nbufs, len(self.vars)], [(nbufs * 8 + i * 8, dtypes.uint64) for i in range(len(self.vars))]
+
+    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
+    self.image = image.ljust(round_up(len(image), 0x1000) + 0x1000, b'\x00')
+    # constant buffer 0 holds the driver params and every argument after them, and starts 256 aligned like all constant buffers
+    self.kernargs_size = round_up(max(self.constbufs[0][1], len(self.cbuf_0) * 4 + len(signature) * 8), 256)
+
     # Ensure device has enough local memory to run the program
-    self.dev._ensure_has_local_memory(self.lcmem_usage)
-    self.dev.allocator._copyin(self.lib_gpu, image)
-    self.dev.synchronize()
+    dev._ensure_has_local_memory(lcmem)
 
     if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A:
-      if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0
-      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'program_address_upper_shifted4':hi32(prog_addr>>4),
-        'program_address_lower_shifted4':lo32(prog_addr>>4), 'register_count':self.regs_usage, 'shared_memory_size_shifted7':self.shmem_usage>>7,
-        f'shader_local_memory_{"low" if NAK else "high"}_size_shifted4': self.dev.slm_per_thread>>4}
+      if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(dev.shared_mem_window), *data64_le(dev.local_mem_window)], 0xfffdc0
+      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'register_count':regs,
+        'shared_memory_size_shifted7':shmem>>7, f'shader_local_memory_{"low" if NAK else "high"}_size_shifted4':dev.slm_per_thread>>4}
     else:
-      if not NAK: self.cbuf_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)]
-      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'program_address_upper':hi32(prog_addr), 'program_address_lower':lo32(prog_addr),
-        'shared_memory_size':self.shmem_usage, 'register_count_v':self.regs_usage,
-        f'shader_local_memory_{"low" if NAK else "high"}_size':self.dev.slm_per_thread}
+      if not NAK: self.cbuf_0[6:12] = [*data64_le(dev.shared_mem_window), *data64_le(dev.local_mem_window), *data64_le(0xfffdc0)]
+      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'shared_memory_size':shmem, 'register_count_v':regs,
+        f'shader_local_memory_{"low" if NAK else "high"}_size':dev.slm_per_thread}
 
-    smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= self.shmem_usage) // 4096 + 1
+    smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= shmem) // 4096 + 1
 
-    self.qmd:QMD = QMD(dev, **qmd, qmd_group_id=0x3f, invalidate_texture_header_cache=1, invalidate_texture_sampler_cache=1,
+    # the program and constant buffer addresses are patched into a copy of this at exec, everything else is the same for every launch
+    self.qmd = QMD(dev)
+    self.qmd.write(**qmd, qmd_group_id=0x3f, invalidate_texture_header_cache=1, invalidate_texture_sampler_cache=1,
       invalidate_texture_data_cache=1, invalidate_shader_data_cache=1, api_visible_call_limit=1, sampler_index=1, barrier_count=1,
       cwd_membar_type=nv_gpu.NVC6C0_QMDV03_00_CWD_MEMBAR_TYPE_L1_SYSMEMBAR, constant_buffer_invalidate_0=1, min_sm_config_shared_mem_size=smem_cfg,
       target_sm_config_shared_mem_size=smem_cfg, max_sm_config_shared_mem_size=0x1a, program_prefetch_size=min(prog_sz>>8, 0x1ff),
-      sass_version=dev.sass_version, program_prefetch_addr_upper_shifted=prog_addr>>40, program_prefetch_addr_lower_shifted=prog_addr>>8)
-
-    for i,(addr,sz) in self.constbufs.items():
-      self.qmd.set_constant_buf_addr(i, addr)
-      self.qmd.write(**{f'constant_buffer_size_shifted4_{i}': sz, f'constant_buffer_valid_{i}': 1})
+      sass_version=dev.sass_version)
+    for i,(_,sz) in self.constbufs.items(): self.qmd.write(**{f'constant_buffer_size_shifted4_{i}': sz, f'constant_buffer_valid_{i}': 1})
 
     # Registers allocation granularity per warp is 256, warp allocation granularity is 4. Register file size is 65536.
-    self.max_threads = ((65536 // round_up(max(1, self.regs_usage) * 32, 256)) // 4) * 4 * 32
-
-    # NV's kernargs is constbuffer, then arguments to the kernel follows. Kernargs also appends QMD at the end of the kernel.
-    super().__init__(NVArgsState, self.dev, obj, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8))
-    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
+    self.max_threads = ((65536 // round_up(max(1, regs) * 32, 256)) // 4) * 4 * 32
 
   def _parse_elf_info(self, sh, start_off=0):
     while start_off < sh.header.sh_size:
@@ -325,18 +296,14 @@ class NVProgram(HCQProgram['NVDevice']):
       yield typ, param, sh.content[start_off+4:start_off+sz+4] if typ == 0x4 else sz
       start_off += (sz if typ == 0x4 else 0) + 4
 
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
-               wait=False, timeout:int|None=None):
-    if prod(local_size) > 1024 or self.max_threads < prod(local_size) or self.lcmem_usage > self.dev.slm_per_thread:
-      raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {self.max_threads=}")
-    if any(cur > mx for cur,mx in zip(global_size, [2147483647, 65535, 65535])) or any(cur > mx for cur,mx in zip(local_size, [1024, 1024, 64])):
-      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
-    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
-    if self.dev.pma_enabled:
-      self.dev.synchronize()
-      if pma_blob:=self.dev._prof_readback():
-        Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob, self.dev.prof_exec_counter, self.profile_key)]
-    return res
+_nv_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[NVProgramData, UOp]] = {}
+def nv_build_program(dev:NVDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[NVProgramData, UOp]:
+  if (cached:=_nv_program_cache.get(key:=(prg.src[3].arg, devs))) is None:
+    data = NVProgramData(dev, prg.to_elf())
+    buf = UOp.placeholder((len(data.image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
+    rows = [(off, ((buf.getaddr(devs) + sym) >> sh).ccast(dt)) for off, sym, dt, sh in data.relocs]
+    cached = _nv_program_cache[key] = (data, patch(buf, rows, data.image))
+  return cached
 
 class NVAllocator(HCQAllocator['NVDevice']):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
@@ -350,22 +317,31 @@ class NVAllocator(HCQAllocator['NVDevice']):
     assert all(h.va_addr % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
 
     h, w = ((2 * shape[0]) // 3 if shape[0] % 3 == 0 else (2 * shape[0] - 1) // 3), shape[1]
-    self.dev._ensure_has_vid_hw(w, h)
+    dev, chroma_off = self.dev, round_up(w, 64) * round_up(h, 64)
+    dev._ensure_has_vid_hw(w, h)
 
-    q = NVVideoQueue().wait(self.dev.timeline_signal, self.dev.timeline_value - 1)
-    with hcq_profile(self.dev, queue=q, desc="HEVC Decode", enabled=PROFILE, dev_suff="NVDEC"):
-      q.decode_hevc_chunk(desc_buf, bufin, bufout, frame_pos, hist, [(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)],
-                          round_up(w, 64)*round_up(h, 64), self.dev.vid_coloc_buf, self.dev.vid_filter_buf, self.dev.intra_top_off,
-                          self.dev.intra_unk_off, self.dev.vid_stat_buf)
-    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+    cmds = nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
+                err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin.va_addr >> 8)
+    for pos, buf in zip([(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)] + [frame_pos], hist + [bufout]):
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, desc_buf.offset(0x200).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf.va_addr) >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_addr + dev.intra_top_off) >> 8)
+    if dev.intra_unk_off is not None: cmds += nvm(4, 0x4dc, (filter_addr + dev.intra_unk_off) >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
+    dev._submit_cmds(dev.fifos["NVDEC:0"], *cmds)
+
+# *****************
+# device
 
 @dataclass
-class GPFifo:
-  ring: MMIOInterface
-  gpput: MMIOInterface
-  entries_count: int
-  token: int
-  put_value: int = 0
+class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; entries: int; token: int # noqa: E702
 
 class NVKIface:
   root = None
@@ -454,7 +430,7 @@ class NVKIface:
     self.uvm(nv_gpu.UVM_REGISTER_GPU_VASPACE, nv_gpu.UVM_REGISTER_GPU_VASPACE_PARAMS(
       gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace))
 
-    for dev in [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, NVDevice) and not d.is_nvd()]:
+    for dev in [d for x in Device._opened_devices if isinstance(d:=Device[x], NVDevice) and not d.is_nvd()]:
       try: self.uvm(nv_gpu.UVM_ENABLE_PEER_ACCESS, nv_gpu.UVM_ENABLE_PEER_ACCESS_PARAMS(gpuUuidA=self.gpu_uuid, gpuUuidB=dev.iface.gpu_uuid))
       except RuntimeError as e: raise RuntimeError(f"{e}. Make sure GPUs #{self.gpu_minor} & #{dev.iface.gpu_minor} have P2P enabled.") from e
 
@@ -582,8 +558,13 @@ class PCIIface(PCIIfaceBase):
 
 class MOCKIface(NVKIface): count = 1
 
-class NVDevice(HCQCompiled[NVSignal]):
+class NVDevice(HCQ2Compiled):
   ifaces = [NVKIface, PCIIface, MOCKIface]
+  sleep_timeout_ms = 200
+  pm_encode = PatternMatcher([
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_compute", name="submit"), lambda ctx, submit: encode_submit(NVComputeQueue(ctx, submit))),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_copy", name="submit"), lambda ctx, submit: encode_submit(NVCopyQueue(ctx, submit))),
+  ])
 
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)
 
@@ -610,19 +591,11 @@ class NVDevice(HCQCompiled[NVSignal]):
     channel_params = nv_gpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS(engineType=nv_gpu.NV2080_ENGINE_TYPE_GRAPHICS)
     self.channel_group = self.iface.rm_alloc(self.nvdevice, nv_gpu.KEPLER_CHANNEL_GROUP_A, channel_params)
 
-    self.gpfifo_area = self.iface.alloc(0x300000, contiguous=True, cpu_access=True, force_devmem=True,
+    self.gpfifo_mem = self.iface.alloc(0x300000, contiguous=True, cpu_access=True, force_devmem=True,
       map_flags=(nv_gpu.NVOS33_FLAGS_CACHING_TYPE_WRITECOMBINED<<23))
 
-    ctxshare_params = nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC)
-    ctxshare = self.iface.rm_alloc(self.channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A, ctxshare_params)
-
-    self.compute_gpfifo = self._new_gpu_fifo(self.gpfifo_area, ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
-    self.dma_gpfifo = self._new_gpu_fifo(self.gpfifo_area, ctxshare, self.channel_group, offset=0x100000, entries=0x10000, compute=False)
-    self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
-
-    self.cmdq_page:HCQBuffer = self.iface.alloc(0x200000, cpu_access=True)
-    self.cmdq_allocator = BumpAllocator(size=self.cmdq_page.size, base=int(self.cmdq_page.va_addr), wrap=True)
-    self.cmdq = self.cmdq_page.cpu_view().view(fmt='I')
+    self.ctxshare = self.iface.rm_alloc(self.channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A,
+      nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC))
 
     self.num_gpcs, self.num_tpc_per_gpc, self.num_sm_per_tpc, self.max_warps_per_sm, self.sm_version = self._query_gpu_info('num_gpcs',
       'num_tpc_per_gpc', 'num_sm_per_tpc', 'max_warps_per_sm', 'sm_version')
@@ -631,19 +604,37 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], NVProgram, NVSignal, NVComputeQueue,
-                     NVCopyQueue, arch=self.arch)
+    self.slm_per_thread = 0
+    self.shader_local_mem:Buffer|None = None
+    # Set windows addresses to not collide with other allocated buffers.
+    self.shared_mem_window, self.local_mem_window = 0x729400000000, 0x729300000000
 
-    self.pma_enabled = PMA.value > 0 and PROFILE >= 1
-    if self.pma_enabled: self._prof_init()
+    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], None, arch=self.arch)
 
-    self._setup_gpfifos()
+    self.pma_enabled, self.pma_exec_counter = PMA.value > 0 and PROFILE >= 1, itertools.count(0)
 
-  def _new_gpu_fifo(self, gpfifo_area, ctxshare, channel_group, offset=0, entries=0x400, compute=False, video=False) -> GPFifo:
+  @functools.cached_property
+  def fifos(self) -> dict[str, GPFifo]:
+    self.gpfifo_buf = Buffer(self.device, self.gpfifo_mem.size, dtypes.uint8, options=BufferSpec(external_ptr=self.gpfifo_mem.va_addr, nolru=True)) \
+                        .allocate(opaque=self.gpfifo_mem)
+    compute = self._new_gpu_fifo("COMPUTE:0", self.ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
+    copy = self._new_gpu_fifo("COPY:0", self.ctxshare, self.channel_group, offset=0x100000, entries=0x10000)
+    self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
+
+    self._submit_cmds(compute, *nvm(1, nv_gpu.NVC6C0_SET_OBJECT, self.iface.compute_class),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_WINDOW_A, *data64(self.local_mem_window)),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_SHARED_MEMORY_WINDOW_A, *data64(self.shared_mem_window)))
+    self._submit_cmds(copy, *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.dma_class))
+
+    if self.pma_enabled: self._prof_init() # the sampler binds to the channel group, so it only comes up once the channels do
+    return {"COMPUTE:0": compute, "COPY:0": copy}
+
+  def _new_gpu_fifo(self, name:str, ctxshare, channel_group, offset=0, entries=0x400, compute=False, video=False) -> GPFifo:
     notifier = self.iface.alloc(48 << 20, uncached=True)
-    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=gpfifo_area.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
-      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else gpfifo_area.meta.hMemory,
-      hUserdMemory=(ctypes.c_uint32*8)(gpfifo_area.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset), engineType=19 if video else 0,
+    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_mem.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
+      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_mem.meta.hMemory,
+      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_mem.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
+      engineType=19 if video else 0,
       hVASpace=self.vaspace if video and self.is_nvd() else 0) # gsp has no default vaspace, rm maps the decoder ctx into its own
     gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params)
 
@@ -662,8 +653,14 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS(workSubmitToken=-1))
     if ctxshare != 0: self.iface.setup_gpfifo_vm(gpfifo)
 
-    return GPFifo(ring=gpfifo_area.cpu_view().view(offset, entries*8, fmt='Q'), entries_count=entries, token=ws_token_params.workSubmitToken,
-                  gpput=gpfifo_area.cpu_view().view(offset + entries*8 + getattr(nv_gpu.AmpereAControlGPFifo, 'GPPut').offset, fmt='I'))
+    gpput_off = offset + entries*8 + getattr(nv_gpu.AmpereAControlGPFifo, 'GPPut').offset
+    fifo = GPFifo(ring=self.gpfifo_buf.view(entries, dtypes.uint64, offset).ensure_allocated(),
+      gpput=self.gpfifo_buf.view(1, dtypes.uint32, gpput_off).ensure_allocated(),
+      doorbell=Buffer("CPU", 1, dtypes.uint32, options=BufferSpec(external_ptr=self.gpu_mmio.addr + 0x90), preallocate=True),
+      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), entries=entries, token=ws_token_params.workSubmitToken)
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag=to_name(n, name)), lambda ctx, b=getattr(fifo, n): b)
+                                        for n in ("ring", "gpput", "doorbell", "put_value")]) + self.pm_bufferize
+    return fifo
 
   def _query_gpu_info(self, *reqs):
     nvrs = [getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_'+r.upper(), getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_LITTER_'+r.upper(), None)) for r in reqs]
@@ -678,34 +675,35 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.NV2080_CTRL_GR_GET_INFO_PARAMS(grInfoListSize=len(infos), grInfoList=ctypes.addressof(infos)))
     return [x.data for x in infos]
 
-  def _setup_gpfifos(self):
-    self.slm_per_thread, self.shader_local_mem = 0, None
+  def _push(self, fifo:GPFifo, cmds:list[int]): # a pushbuffer built in python: channel setup and video decode
+    (buf:=self.rt_view(len(cmds) * 4))._buf.cpu_view().view(fmt='I')[:] = array.array('I', cmds)
 
-    # Set windows addresses to not collide with other allocated buffers.
-    self.shared_mem_window, self.local_mem_window = 0x729400000000, 0x729300000000
+    put = fifo.put_value._buf.view.view(fmt='Q')
+    fifo.ring._buf.cpu_view().view(fmt='Q')[put[0] % fifo.entries] = buf._buf.va_addr | (len(cmds) << 42) | (1 << 41)
+    fifo.gpput._buf.cpu_view().view(fmt='I')[0] = (put[0] + 1) % fifo.entries
 
-    NVComputeQueue().setup(compute_class=self.iface.compute_class, local_mem_window=self.local_mem_window, shared_mem_window=self.shared_mem_window) \
-                    .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    System.memory_barrier()
+    self.gpu_mmio[0x90 // 4] = fifo.token
+    put[0] += 1
 
-    NVCopyQueue().wait(self.timeline_signal, self.timeline_value - 1) \
-                 .setup(copy_class=self.iface.dma_class) \
-                 .signal(self.timeline_signal, self.next_timeline()).submit(self)
-
-    self.synchronize()
+  def _submit_cmds(self, fifo:GPFifo, *cmds:int): # runs cmds once everything already submitted is done, then bumps the timeline
+    tl, addr = self.timeline._buf.cpu_view().view(fmt='Q'), self.timeline._buf.va_addr
+    self._push(fifo, nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1]),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit")) + list(cmds) +
+                     nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1] + 1),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit")))
+    tl[1] += 1
 
   def _ensure_has_local_memory(self, required):
     if self.slm_per_thread >= required: return
 
-    self.slm_per_thread, old_slm_per_thread = round_up(required, 32), self.slm_per_thread
+    self.slm_per_thread = round_up(required, 32)
     bytes_per_tpc = round_up(round_up(self.slm_per_thread * 32, 0x200) * self.max_warps_per_sm * self.num_sm_per_tpc, 0x8000)
-    self.shader_local_mem, ok = self._realloc(self.shader_local_mem, round_up(bytes_per_tpc*self.num_tpc_per_gpc*self.num_gpcs, 0x20000))
+    self.shader_local_mem = Buffer(self.device, round_up(bytes_per_tpc*self.num_tpc_per_gpc*self.num_gpcs, 0x20000), dtypes.uint8,
+                                   options=BufferSpec(nolru=True), preallocate=True)
 
-    # Realloc failed, restore the old value.
-    if not ok: self.slm_per_thread = old_slm_per_thread
-
-    cast(NVComputeQueue, NVComputeQueue().wait(self.timeline_signal, self.timeline_value - 1)) \
-                                         .setup(local_mem=self.shader_local_mem.va_addr, local_mem_tpc_bytes=bytes_per_tpc) \
-                                         .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    self._submit_cmds(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf.va_addr)),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(bytes_per_tpc), 0xff))
 
   def _ensure_has_vid_hw(self, w, h):
     if self.iface.viddec_class is None: raise RuntimeError(f"{self.device} Video decoder class not available.")
@@ -716,18 +714,20 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.intra_unk_off = (round_up(self.intra_top_off, 0x10000) + (64 << 10)) if intra_unk_size > 0 else None
     filter_sz = round_up(round_up(self.intra_top_off, 0x10000) + (64 << 10) + intra_unk_size, 2 << 20)
 
-    if not hasattr(self, 'vid_gpfifo'):
-      self.vid_gpfifo = self._new_gpu_fifo(self.gpfifo_area, 0, self.nvdevice, offset=0x200000, entries=2048, compute=False, video=True)
-      self.vid_coloc_buf, self.vid_filter_buf = (self.allocator.alloc(sz, BufferSpec(zero=True)) for sz in [coloc_sz, filter_sz])
-      self.vid_stat_buf = self.allocator.alloc(0x1000, BufferSpec(zero=True))
-      NVVideoQueue().wait(self.timeline_signal, self.timeline_value - 1) \
-                    .setup(copy_class=self.iface.viddec_class) \
-                    .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    def _vid_buf(sz): return Buffer(self.device, sz, dtypes.uint8, options=BufferSpec(zero=True, nolru=True), preallocate=True)
+    if "NVDEC:0" not in self.fifos:
+      self.fifos["NVDEC:0"] = self._new_gpu_fifo("NVDEC:0", 0, self.nvdevice, offset=0x200000, entries=2048, video=True)
+      self.vid_coloc_buf, self.vid_filter_buf, self.vid_stat_buf = _vid_buf(coloc_sz), _vid_buf(filter_sz), _vid_buf(0x1000)
+      self._submit_cmds(self.fifos["NVDEC:0"], *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.viddec_class))
     else:
-      if coloc_sz > self.vid_coloc_buf.size: self.vid_coloc_buf,_= self._realloc(self.vid_coloc_buf, coloc_sz, BufferSpec(zero=True), force=True)
-      if filter_sz > self.vid_filter_buf.size: self.vid_filter_buf,_= self._realloc(self.vid_filter_buf, filter_sz, BufferSpec(zero=True), force=True)
+      if coloc_sz > self.vid_coloc_buf.nbytes: self.vid_coloc_buf = _vid_buf(coloc_sz)
+      if filter_sz > self.vid_filter_buf.nbytes: self.vid_filter_buf = _vid_buf(filter_sz)
 
-  def hw_copy_queues(self): return super().hw_copy_queues() + ([("NVDEC:0", NVVideoQueue)] if hasattr(self, 'vid_gpfifo') else [])
+  def collect_prof(self):
+    # the pc samples of a whole batch come back as one stream, so they are reported against the first kernel of it
+    if self.pma_enabled and (ents:=list(self.prof_ents.values())) and (blob:=self._prof_readback()) is not None:
+      Compiled.profile_events.append(ProfilePMAEvent(self.device, str(ents[0].name), blob, next(self.pma_exec_counter), ents[0].profile_key))
+    super().collect_prof()
 
   def invalidate_caches(self):
     if self.is_nvd(): self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_INTERNAL_BUS_FLUSH_WITH_SYSMEMBAR, None)
@@ -845,4 +845,4 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS(bytesConsumed=params.bytesAvailable))
     return pma_data
 
-  def device_props(self): return {'arch': self.arch, 'sm_version': self.sm_version}
+  def device_props(self) -> dict[str, Any]: return {'arch': self.arch, 'sm_version': self.sm_version}

tinygrad/runtime/ops_python.py

@@ -18,9 +18,11 @@ def _load(m, i, dtype: DType):
   if (w:=m.nbytes // len(m)) >= dtype.itemsize: return from_storage_scalar(m[i], dtype)
   return sum(m[i+k] << (8*w*k) for k in range(dtype.itemsize // w)) # a bitcast can read wider than the buffer, _store splits it the same way
 
+def _step(m, dtype: DType): return max(1, dtype.itemsize // (m.nbytes // len(m))) # storage elements per lane
+
 def load(inp, j, dtype: DType):
-  if len(inp) >= 3: return [_load(m, x+j if x is not None else None, dtype) if gate else default for (m,x),default,gate in zip(*inp[:3])]
-  return [_load(m, x+j if x is not None else None, dtype) for m,x in inp[0]]
+  if len(inp) >= 3: return [_load(m, x+j*_step(m, dtype) if x is not None else None, dtype) if gate else alt for (m,x),alt,gate in zip(*inp[:3])]
+  return [_load(m, x+j*_step(m, dtype) if x is not None else None, dtype) for m,x in inp[0]]
 
 def _store(m, i, v, dtype: DType):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
@@ -86,7 +88,7 @@ class PythonProgram(Program['PythonDevice']):
           store_gate = exec_masks[-1]
           for j,val in enumerate(src_values[1] if u.max_numel() > 1 else [src_values[1]]):
             for (m,o),v,g in zip(src_values[0], val, store_gate):
-              if g: _store(m, o+j, v, src_dtypes[1])
+              if g: _store(m, o+j*_step(m, src_dtypes[1]), v, src_dtypes[1])
           i += 1
           continue
         if u.op is Ops.AFTER or (u.op is Ops.BITCAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)): values[u] = src_values[0]

tinygrad/runtime/ops_qcom.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
-import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
+import os, ctypes, functools, mmap, struct, array, math, sys, contextlib
 assert sys.platform != 'win32'
-from typing import Any, cast
-from tinygrad.device import BufferSpec, Device, TinyELF
-from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
-from tinygrad.runtime.autogen import kgsl, mesa
+from typing import Any
+from tinygrad.device import BufferSpec, Buffer, Device, TinyELF
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, HCQ_RUNTIME_DEV, encode_submit, ccall, cstruct, patch, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface, MMIOInterface
+from tinygrad.runtime.autogen import kgsl, mesa, libc
 from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import IR3Renderer
-from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, cpu_profile, lo32, suppress_finalizing, is_image_shape
+from tinygrad.helpers import getenv, mv_address, round_up, ceildiv, prod, is_image_shape
 from tinygrad.helpers import next_power2, flatten, PROFILE, IMAGE
 from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
 
@@ -18,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 
 @functools.cache
 def dcache_flush():
-  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.uop.ops import KernelInfo
   from tinygrad.codegen import to_program
   buf, n = UOp.param(0, dtypes.uint8, 1), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
   i = UOp.range(n, 0, dtype=dtypes.int)
@@ -49,92 +51,107 @@ def pkt4_hdr(reg: int, cnt: int): return mesa.CP_TYPE4_PKT | cnt & 0x7F | parity
 
 def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 
-class QCOMSignal(HCQSignal):
-  def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 19.2})
-
-  def _sleep(self, time_spent_since_last_sleep_ms:int):
-    # Sleep only for timeline signals. Do it immediately to free cpu.
-    if self.is_timeline and self.owner is not None:
-      kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.owner.fd, context_id=self.owner.ctx, timestamp=self.owner.last_cmd, timeout=0xffffffff)
-
 class QCOMComputeQueue(HWQueue):
-  def __init__(self, dev:QCOMDevice):
-    self.dev = dev
-    super().__init__()
+  dev:QCOMDevice
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
-  @suppress_finalizing
-  def __del__(self):
-    if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))
+  def cmd(self, opcode:int, *vals): self.q(pkt7_hdr(opcode, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 
-  def cmd(self, opcode: int, *vals: int): self.q(pkt7_hdr(opcode, len(vals)), *vals)
-
-  def reg(self, reg: int, *vals: int): self.q(pkt4_hdr(reg, len(vals)), *vals)
+  def reg(self, reg:int, *vals): self.q(pkt4_hdr(reg, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 
   def _cache_flush(self, write_back=True, invalidate=False, sync=True, memsync=False):
     # TODO: 7xx support.
-    if write_back: self.cmd(mesa.CP_EVENT_WRITE, mesa.CACHE_FLUSH_TS, *data64_le(self.dev.dummy_addr), 0) # dirty cache write-back.
+    if write_back: # dirty cache write-back, into the device's dummy buffer
+      dummy = UOp.placeholder((0x1000,), dtypes.uint8, 0, device=self.devs, tag="dummy")
+      self.cmd(mesa.CP_EVENT_WRITE, mesa.CACHE_FLUSH_TS, dummy.getaddr(self.devs), 0)
     if invalidate: self.cmd(mesa.CP_EVENT_WRITE, mesa.CACHE_INVALIDATE) # invalidate cache lines (following reads from RAM).
     if memsync: self.cmd(mesa.CP_WAIT_MEM_WRITES)
     if sync: self.cmd(mesa.CP_WAIT_FOR_IDLE)
 
-  def memory_barrier(self):
-    self._cache_flush(write_back=True, invalidate=True, sync=True, memsync=True)
-    return self
+  def memory_barrier(self): self._cache_flush(write_back=True, invalidate=True, sync=True, memsync=True)
 
-  def signal(self, signal:QCOMSignal, value=0):
+  def signal(self, signal:UOp, value:UOp):
     self.cmd(mesa.CP_WAIT_FOR_IDLE)
     if self.dev.gpu_id[:2] < (7, 3):
-      self.cmd(mesa.CP_EVENT_WRITE, qreg.cp_event_write_0(event=mesa.CACHE_FLUSH_TS), *data64_le(signal.value_addr), lo32(value))
+      self.cmd(mesa.CP_EVENT_WRITE, qreg.cp_event_write_0(event=mesa.CACHE_FLUSH_TS), signal.getaddr(self.devs), value.cast(dtypes.uint32))
       self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
     else:
       # TODO: support devices starting with 8 Gen 1. Also, 700th series have convenient CP_GLOBAL_TIMESTAMP and CP_LOCAL_TIMESTAMP
       raise RuntimeError('CP_EVENT_WRITE7 is not supported')
-    return self
 
-  def timestamp(self, signal:QCOMSignal):
+  def timestamp(self, signal:UOp):
     self.cmd(mesa.CP_WAIT_FOR_IDLE)
-    self.cmd(mesa.CP_REG_TO_MEM, qreg.cp_reg_to_mem_0(reg=mesa.REG_A6XX_CP_ALWAYS_ON_COUNTER, cnt=2, _64b=True),*data64_le(signal.timestamp_addr))
-    return self
+    self.cmd(mesa.CP_REG_TO_MEM, qreg.cp_reg_to_mem_0(reg=mesa.REG_A6XX_CP_ALWAYS_ON_COUNTER, cnt=2, _64b=True), signal.getaddr(self.devs))
 
-  def wait(self, signal:QCOMSignal, value=0):
-    self.cmd(mesa.CP_WAIT_REG_MEM, qreg.cp_wait_reg_mem_0(function=mesa.WRITE_GE, poll=mesa.POLL_MEMORY),*data64_le(signal.value_addr),
-             qreg.cp_wait_reg_mem_3(ref=value&0xFFFFFFFF), qreg.cp_wait_reg_mem_4(mask=0xFFFFFFFF), qreg.cp_wait_reg_mem_5(delay_loop_cycles=32))
-    return self
+  def wait(self, signal:UOp, value:UOp):
+    self.cmd(mesa.CP_WAIT_REG_MEM, qreg.cp_wait_reg_mem_0(function=mesa.WRITE_GE, poll=mesa.POLL_MEMORY), signal.getaddr(self.devs),
+             value.cast(dtypes.uint32), qreg.cp_wait_reg_mem_4(mask=0xFFFFFFFF), qreg.cp_wait_reg_mem_5(delay_loop_cycles=32))
 
-  def _build_gpu_command(self, dev:QCOMDevice, hw_addr=None):
-    to_mv((hw_page_addr:=hw_addr or dev.cmd_buf_allocator.alloc(len(self._q) * 4)), len(self._q) * 4).cast('I')[:] = array.array('I', self._q)
-    obj = kgsl.struct_kgsl_command_object(gpuaddr=hw_page_addr, size=len(self._q) * 4, flags=kgsl.KGSL_CMDLIST_IB)
-    submit_req = kgsl.struct_kgsl_gpu_command(cmdlist=ctypes.addressof(obj), numcmds=1, context_id=dev.ctx,
-                                              cmdsize=ctypes.sizeof(kgsl.struct_kgsl_command_object))
-    return submit_req, obj
+  def kernargs(self, call:UOp, prg:UOp, data:QCOMProgramData) -> UOp:
+    bufs, vals = get_call_arg_uops(call), get_call_var_uops(call, prg)
+    ubos = [bufs[slot] for _,slot,_,shape in data.signature if slot < len(bufs) and not is_image_shape(shape)]
+    uavs = [(dt,shape,bufs[slot]) for _,slot,dt,shape in data.signature if slot < len(bufs) and is_image_shape(shape)]
+    # NIR can reorder images to different texture slots
+    ibos, texs = uavs[:data.ibo_cnt], [uavs[data.ibo_cnt + (data.tex_to_image[i] if data.NIR else i)] for i in range(data.tex_cnt)]
 
-  def bind(self, dev:QCOMDevice):
-    self.binded_device = dev
-    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True))
-    self.submit_req, self.obj = self._build_gpu_command(self.binded_device, self.hw_page.va_addr)
-    # From now on, the queue is on the device for faster submission.
-    self._q = to_mv(self.obj.gpuaddr, len(self._q) * 4).cast("I")
+    # the words of the kernargs, as runs at their byte offsets
+    runs:list[tuple[int, list]] = [(off, [UOp.const(val, dtypes.uint32 if sz == 4 else dtypes.uint16)]) for val,off,sz in data.consts_info]
+    runs.append((data.samp_off, data.samplers))
+    if data.NIR:
+      runs.append((data.buf_off, [b.getaddr(self.devs) for b in ubos]))
+      runs += [(data.buf_off + o, [v.ccast(dt)]) for v,(o,dt) in zip(vals, TinyELF.iter_sig(data.signature[len(bufs):], len(ubos)*8))]
+      if data.wgsz != 0xfc: runs.append((data.wgsz * 4, list(prg.arg.local_size)))
+    else:
+      runs += [(data.buf_offs[i], [b.getaddr(self.devs)]) for i, b in enumerate(ubos)]
+      runs += [(data.buf_offs[i+len(ubos)], [v.ccast(dt)]) for i,(v,(_,_,dt,_)) in enumerate(zip(vals, data.signature[len(bufs):]))]
 
-  def _submit(self, dev:QCOMDevice):
-    if self.binded_device == dev: submit_req = self.submit_req
-    else: submit_req, _ = self._build_gpu_command(dev)
-    dev.last_cmd = kgsl.IOCTL_KGSL_GPU_COMMAND(dev.fd, __payload=submit_req).timestamp
+    def _tex(b, ibo=False):
+      imgdt, shape, buf = b
+      pitch = shape[1] * 4 * imgdt.itemsize
+      fmt = mesa.FMT6_32_32_32_32_FLOAT if imgdt.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
+      return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt),
+              qreg.a6xx_tex_const_1(width=shape[1], height=shape[0]),
+              qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=pitch, pitchalign=ctz(pitch)-6), 0, buf.getaddr(self.devs),
+              qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0]
+    runs += [(data.tex_off, flatten(map(_tex, texs))), (data.ibo_off, flatten(map(functools.partial(_tex, ibo=True), ibos)))]
 
-  def exec(self, prg:QCOMProgram, args_state:QCOMArgsState, global_size, local_size):
-    self.bind_args_state(args_state)
+    # laid out as a linear in the cmdbuf tail, like amd's kernargs: the runs in order, zero bytes between them and after the last
+    out, end = [], 0
+    for off, run in sorted([r for r in runs if r[1]], key=lambda r: r[0]) + [(data.kernargs_alloc_size, [])]:
+      assert off >= end, f"kernargs run at {off} overlaps the one ending at {end}"
+      if off > end: out.append(UOp(Ops.BINARY, arg=bytes(off - end)))
+      out += (run:=[w if isinstance(w, UOp) else UOp.const(w, dtypes.uint32) for w in run])
+      end = off + sum(w.dtype.itemsize for w in run)
+    return UOp(Ops.LINEAR, src=tuple(out))
+
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = qcom_build_program(self.dev, prg, self.devs)
+    global_size, local_size = prg.arg.global_size, prg.arg.local_size
+    if data.max_threads < prod(local_size): raise RuntimeError("Too many resources requested for launch")
+    if any(g*l>mx for g,l,mx in zip(global_size, local_size, [65536, 65536, 65536])) and any(l>mx for l,mx in zip(local_size, [1024, 1024, 1024])):
+      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
 
     def cast_int(x, ceil=False): return (math.ceil(x) if ceil else int(x)) if isinstance(x, float) else x
     global_size_mp = [cast_int(g*l) for g,l in zip(global_size, local_size)]
 
+    args_addr, lib_addr = self.kernargs(call, prg, data).getaddr(self.devs), lib.getaddr(self.devs)
+    stack_addr = UOp.placeholder((data.hw_stack_offset * 4,), dtypes.uint8, 0, device=self.devs).rtag("stack").getaddr(self.devs)
+
     self.cmd(mesa.CP_SET_MARKER, qreg.a6xx_cp_set_marker_0(mode=mesa.RM6_COMPUTE))
     self.reg(mesa.REG_A6XX_SP_UPDATE_CNTL, qreg.a6xx_sp_update_cntl(cs_state=True, cs_uav=True))
     self.reg(mesa.REG_A6XX_SP_UPDATE_CNTL, 0x0)
     self.reg(mesa.REG_A6XX_SP_CS_TSIZE, qreg.a6xx_sp_cs_tsize(0x80)) # is this right? mesa uses 1
     self.reg(mesa.REG_A6XX_SP_CS_USIZE, qreg.a6xx_sp_cs_usize(0x40)) # mesa also uses 1
-    self.reg(mesa.REG_A6XX_SP_MODE_CNTL, qreg.a6xx_sp_mode_cntl(isammode=mesa.ISAMMODE_GL if prg.NIR else mesa.ISAMMODE_CL,
-                                                                constant_demotion_enable=prg.NIR))
+    self.reg(mesa.REG_A6XX_SP_MODE_CNTL, qreg.a6xx_sp_mode_cntl(isammode=mesa.ISAMMODE_GL if data.NIR else mesa.ISAMMODE_CL,
+                                                                constant_demotion_enable=data.NIR))
     self.reg(mesa.REG_A6XX_SP_PERFCTR_SHADER_MASK, qreg.a6xx_sp_perfctr_shader_mask(cs=True))
-    self.reg(mesa.REG_A6XX_TPL1_MODE_CNTL, qreg.a6xx_tpl1_mode_cntl(isammode=mesa.ISAMMODE_GL if prg.NIR else mesa.ISAMMODE_CL))
+    self.reg(mesa.REG_A6XX_TPL1_MODE_CNTL, qreg.a6xx_tpl1_mode_cntl(isammode=mesa.ISAMMODE_GL if data.NIR else mesa.ISAMMODE_CL))
     self.reg(mesa.REG_A6XX_TPL1_DBG_ECO_CNTL, 0)
     self.cmd(mesa.CP_WAIT_FOR_IDLE)
 
@@ -144,97 +161,68 @@ class QCOMComputeQueue(HWQueue):
              cast_int(global_size[0], ceil=True), cast_int(global_size[1], ceil=True), cast_int(global_size[2], ceil=True))
 
     self.reg(mesa.REG_A6XX_SP_CS_CNTL_0,
-             qreg.a6xx_sp_cs_cntl_0(threadsize=mesa.THREAD64, halfregfootprint=prg.hregs, fullregfootprint=prg.fregs, branchstack=prg.brnchstck),
-             qreg.a6xx_sp_cs_cntl_1(constantrammode=mesa.CONSTLEN_256, shared_size=prg.shared_size), # should this be CONSTLEN_512?
-             0, prg.prg_offset, *data64_le(prg.lib_gpu.va_addr),
-             qreg.a6xx_sp_cs_pvt_mem_param(memsizeperitem=prg.pvtmem_size_per_item), *data64_le(prg.dev._stack.va_addr),
-             qreg.a6xx_sp_cs_pvt_mem_size(totalpvtmemsize=prg.pvtmem_size_total))
+             qreg.a6xx_sp_cs_cntl_0(threadsize=mesa.THREAD64, halfregfootprint=data.hregs, fullregfootprint=data.fregs, branchstack=data.brnchstck),
+             qreg.a6xx_sp_cs_cntl_1(constantrammode=mesa.CONSTLEN_256, shared_size=data.shared_size), # should this be CONSTLEN_512?
+             0, data.prg_offset, lib_addr,
+             qreg.a6xx_sp_cs_pvt_mem_param(memsizeperitem=data.pvtmem_size_per_item), stack_addr,
+             qreg.a6xx_sp_cs_pvt_mem_size(totalpvtmemsize=data.pvtmem_size_total))
 
-    if prg.NIR and prg.wgsz != 0xfc: to_mv(int(args_state.buf.va_addr) + prg.wgsz * 4, 12)[:] = struct.pack("III", *local_size)
+    # the kernargs sit in the cmdbuf, so the const upload is sized to them (in vec4s) rather than to the whole constlen: it must not read past
     self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_CONSTANTS, state_src=mesa.SS6_INDIRECT,
-                                                             state_block=mesa.SB6_CS_SHADER, num_unit=1024 // 4),
-             *data64_le(args_state.buf.va_addr))
+                                                             state_block=mesa.SB6_CS_SHADER, num_unit=data.kernargs_alloc_size // 16), args_addr)
     self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_SHADER, state_src=mesa.SS6_INDIRECT,
-                                                             state_block=mesa.SB6_CS_SHADER, num_unit=ceildiv(prg.image_size, 128)),
-             *data64_le(prg.lib_gpu.va_addr))
+                                                             state_block=mesa.SB6_CS_SHADER, num_unit=ceildiv(data.image_size, 128)), lib_addr)
 
     self.reg(mesa.REG_A6XX_SP_REG_PROG_ID_0, 0xfcfcfcfc, 0xfcfcfcfc, 0xfcfcfcfc, 0xfc, qreg.a6xx_sp_cs_const_config(constlen=1024 // 4, enabled=True))
 
-    self.reg(mesa.REG_A6XX_SP_CS_PVT_MEM_STACK_OFFSET, qreg.a6xx_sp_cs_pvt_mem_stack_offset(prg.hw_stack_offset))
+    self.reg(mesa.REG_A6XX_SP_CS_PVT_MEM_STACK_OFFSET, qreg.a6xx_sp_cs_pvt_mem_stack_offset(data.hw_stack_offset))
     # image_size is in bytes, but INSTR_SIZE is measured in units of instruction groups (16 instructions, 8 bytes each)
     # https://elixir.bootlin.com/mesa/mesa-26.1.5/source/src/freedreno/ir3/ir3_shader.h#L719-L723
-    self.reg(mesa.REG_A6XX_SP_CS_INSTR_SIZE, qreg.a6xx_sp_cs_instr_size(ceildiv(prg.image_size, 128)))
+    self.reg(mesa.REG_A6XX_SP_CS_INSTR_SIZE, qreg.a6xx_sp_cs_instr_size(ceildiv(data.image_size, 128)))
 
-    if prg.samp_cnt > 0:
+    if data.samp_cnt > 0:
       self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_SHADER, state_src=mesa.SS6_INDIRECT,
-                                                               state_block=mesa.SB6_CS_TEX, num_unit=args_state.prg.samp_cnt),
-               *data64_le(args_state.buf.va_addr + args_state.prg.samp_off))
-      self.reg(mesa.REG_A6XX_SP_CS_SAMPLER_BASE, *data64_le(args_state.buf.va_addr + args_state.prg.samp_off))
-      self.reg(mesa.REG_A6XX_TPL1_CS_BORDER_COLOR_BASE, *data64_le(prg.dev.border_color_buf.va_addr))
+                                                               state_block=mesa.SB6_CS_TEX, num_unit=data.samp_cnt), args_addr + data.samp_off)
+      self.reg(mesa.REG_A6XX_SP_CS_SAMPLER_BASE, args_addr + data.samp_off)
+      self.reg(mesa.REG_A6XX_TPL1_CS_BORDER_COLOR_BASE,
+               UOp.placeholder((0x1000,), dtypes.uint8, 0, device=self.devs, tag="border_color").getaddr(self.devs))
 
-    if prg.tex_cnt > 0:
+    if data.tex_cnt > 0:
       self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_CONSTANTS, state_src=mesa.SS6_INDIRECT,
-                                                               state_block=mesa.SB6_CS_TEX, num_unit=min(16, args_state.prg.tex_cnt)),
-               *data64_le(args_state.buf.va_addr + args_state.prg.tex_off))
-      self.reg(mesa.REG_A6XX_SP_CS_TEXMEMOBJ_BASE, *data64_le(args_state.buf.va_addr + args_state.prg.tex_off))
+                                                               state_block=mesa.SB6_CS_TEX, num_unit=min(16, data.tex_cnt)), args_addr + data.tex_off)
+      self.reg(mesa.REG_A6XX_SP_CS_TEXMEMOBJ_BASE, args_addr + data.tex_off)
 
-    if prg.ibo_cnt > 0:
+    if data.ibo_cnt > 0:
       self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST6_UAV, state_src=mesa.SS6_INDIRECT,
-                                                               state_block=mesa.SB6_CS_SHADER, num_unit=args_state.prg.ibo_cnt),
-               *data64_le(args_state.buf.va_addr + args_state.prg.ibo_off))
-      self.reg(mesa.REG_A6XX_SP_CS_UAV_BASE, *data64_le(args_state.buf.va_addr + args_state.prg.ibo_off))
+                                                               state_block=mesa.SB6_CS_SHADER, num_unit=data.ibo_cnt), args_addr + data.ibo_off)
+      self.reg(mesa.REG_A6XX_SP_CS_UAV_BASE, args_addr + data.ibo_off)
 
-    self.reg(mesa.REG_A6XX_SP_CS_CONFIG,
-             qreg.a6xx_sp_cs_config(enabled=True, nsamp=args_state.prg.samp_cnt, ntex=args_state.prg.tex_cnt, nuav=args_state.prg.ibo_cnt))
+    self.reg(mesa.REG_A6XX_SP_CS_CONFIG, qreg.a6xx_sp_cs_config(enabled=True, nsamp=data.samp_cnt, ntex=data.tex_cnt, nuav=data.ibo_cnt))
 
-    if prg.NIR:
+    if data.NIR:
       self.reg(mesa.REG_A6XX_SP_CS_CONST_CONFIG_0,
-               qreg.a6xx_sp_cs_const_config_0(wgidconstid=prg.wgid, wgsizeconstid=prg.wgsz, wgoffsetconstid=0xfc, localidregid=prg.lid),
+               qreg.a6xx_sp_cs_const_config_0(wgidconstid=data.wgid, wgsizeconstid=data.wgsz, wgoffsetconstid=0xfc, localidregid=data.lid),
                qreg.a6xx_sp_cs_wge_cntl(linearlocalidregid=0xfc, threadsize=mesa.THREAD64))
       self.cmd(mesa.CP_EXEC_CS, 0,
                qreg.cp_exec_cs_1(ngroups_x=global_size[0]), qreg.cp_exec_cs_2(ngroups_y=global_size[1]), qreg.cp_exec_cs_3(_ngroups_z=global_size[2]))
     else: self.cmd(mesa.CP_RUN_OPENCL, 0)
 
     self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
-    return self
 
-class QCOMArgsState(HCQArgsState):
-  def __init__(self, buf:HCQBuffer, prg:QCOMProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    super().__init__(buf, prg, bufs, vals=vals)
-    ctypes.memset(int(self.buf.va_addr), 0, prg.kernargs_alloc_size)
+  def submit(self, cmdbuf:UOp) -> UOp:
+    ib, ib_off = unwrap_view(cmdbuf)
+    fd, ctxid = [UOp.variable(n, 0, 2**31 - 1, dtypes.int32, param=True) for n in ("kgsl_fd", "kgsl_ctx")]
+    obj = cstruct(kgsl.struct_kgsl_command_object, gpuaddr=ib.getaddr(self.devs) + ib_off, size=cmdbuf.max_numel(), flags=kgsl.KGSL_CMDLIST_IB)
+    req = cstruct(kgsl.struct_kgsl_gpu_command, cmdlist=obj.getaddr(HCQ_RUNTIME_DEV.value), cmdsize=ctypes.sizeof(kgsl.struct_kgsl_command_object),
+                  numcmds=1, context_id=ctxid)
+    ret = UOp.placeholder((1,), dtypes.int32, device=self.devs, volatile=True, tag="submit_ret")
 
-    ubos = [bufs[slot] for _,slot,_,shape in prg.signature if slot < len(bufs) and not is_image_shape(shape)]
-    uavs = [(dt,shape,bufs[slot]) for _,slot,dt,shape in prg.signature if slot < len(bufs) and is_image_shape(shape)]
-    # NIR can reorder images to different texture slots
-    ibos, texs = uavs[:prg.ibo_cnt], [uavs[prg.ibo_cnt + (prg.tex_to_image[i] if prg.NIR else i)] for i in range(prg.tex_cnt)]
-    for cnst_val,cnst_off,cnst_sz in prg.consts_info:
-      to_mv(cast(int, self.buf.va_addr) + cnst_off, cnst_sz)[:] = cnst_val.to_bytes(cnst_sz, byteorder='little')
+    idir, base, nr, struct_t = kgsl.IOCTL_KGSL_GPU_COMMAND.args
+    ioctl_cmd = (idir << 30) | (ctypes.sizeof(struct_t) << 16) | (base << 8) | nr
+    return ret.index(0).store(ccall(libc.dll.ioctl, fd, UOp.const(ioctl_cmd, dtypes.uint32), req.after(cmdbuf).index(0)))
 
-    if prg.samp_cnt > 0: to_mv(int(self.buf.va_addr) + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
-    if prg.NIR:
-      self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
-      for v,(o,dt) in zip(vals, TinyELF.iter_sig(prg.signature[len(bufs):], len(ubos)*8)):
-        self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_off + o)
-    else:
-      for i, b in enumerate(ubos): self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=prg.buf_offs[i])
-      for i,(v,(_,_,dt,_)) in enumerate(zip(vals, prg.signature[len(bufs):])):
-        self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_offs[i+len(ubos)])
-
-    def _tex(b, ibo=False):
-      imgdt, shape, buf = b
-      pitch = shape[1] * 4 * imgdt.itemsize
-      fmt = mesa.FMT6_32_32_32_32_FLOAT if imgdt.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
-      return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt),
-              qreg.a6xx_tex_const_1(width=shape[1], height=shape[0]),
-              qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=pitch, pitchalign=ctz(pitch)-6), 0, *data64_le(buf.va_addr),
-              qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0]
-
-    self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off)
-    self.bind_sints_to_buf(*flatten(map(functools.partial(_tex, ibo=True), ibos)), buf=self.buf, fmt='I', offset=prg.ibo_off)
-
-class QCOMProgram(HCQProgram['QCOMDevice']):
-  def __init__(self, dev: QCOMDevice, obj: TinyELF):
-    self.dev: QCOMDevice = dev
+class QCOMProgramData:
+  def __init__(self, dev:QCOMDevice, obj:TinyELF):
     self.signature, self.name, self.NIR = obj.signature, obj.name, isinstance(dev.renderer, IR3Renderer)
 
     if self.NIR:
@@ -259,26 +247,12 @@ class QCOMProgram(HCQProgram['QCOMDevice']):
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
     else: self._parse_lib(obj.lib)
 
-    self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
-    to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
-
     self.pvtmem_size_per_item: int = round_up(self.pvtmem, 512) >> 9
     self.pvtmem_size_total: int = self.pvtmem_size_per_item * 128 * 2
     self.hw_stack_offset: int = round_up(next_power2(round_up(self.pvtmem, 512)) * 128 * 16, 0x1000)
     self.shared_size: int = max(1, (self.shmem - 1) // 1024)
     self.max_threads = min(1024, ((384 * 32) // (max(1, (self.fregs + round_up(self.hregs, 2) // 2)) * 128)) * 128)
-    dev._ensure_stack_size(self.hw_stack_offset * 4)
-
-    kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100)
-    super().__init__(QCOMArgsState, self.dev, obj, kernargs_alloc_size=kernargs_alloc_size)
-    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
-
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
-               vals:tuple[int|None, ...]=(), wait=False, **kw):
-    if self.max_threads < prod(local_size): raise RuntimeError("Too many resources requested for launch")
-    if any(g*l>mx for g,l,mx in zip(global_size, local_size, [65536, 65536, 65536])) and any(l>mx for l,mx in zip(local_size, [1024, 1024, 1024])):
-      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
-    return super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
+    self.kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100)
 
   def _parse_lib(self, lib):
     # Extract image binary
@@ -326,38 +300,37 @@ class QCOMProgram(HCQProgram['QCOMDevice']):
     reg_desc_off = _read_lib(lib, 0x34)
     self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)
 
-class QCOMAllocator(HCQAllocatorBase):
+_qcom_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[QCOMProgramData, UOp]] = {}
+def qcom_build_program(dev:QCOMDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[QCOMProgramData, UOp]:
+  if (cached:=_qcom_program_cache.get(key:=(prg.src[3].arg, devs))) is None:
+    data = QCOMProgramData(dev, prg.to_elf())
+    image = bytes(data.image).ljust(round_up(len(data.image), 4), b"\x00")
+    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
+    cached = _qcom_program_cache[key] = (data, patch(buf, [], image))
+  return cached
+
+class QCOMAllocator(HCQAllocator['QCOMDevice']):
   def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
     return self.dev._gpu_map(opts.external_ptr, size) if opts.external_ptr else self.dev._gpu_alloc(size)
 
-  def _do_copy(self, src_addr, dest_addr, size, prof_text):
-    self.dev.synchronize()
-    with cpu_profile(prof_text, f"{self.dev.device}:COPY"): ctypes.memmove(dest_addr, src_addr, size)
-
-  def _copyin(self, dest:HCQBuffer, src:memoryview): self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, f"TINY -> {self.dev.device}")
-  def _copyout(self, dest:memoryview, src:HCQBuffer): self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, f"{self.dev.device} -> TINY")
-
-  def _as_buffer(self, src:HCQBuffer) -> memoryview: return to_mv(src.cpu_view().addr, src.size)
-
   def _do_free(self, opaque, options:BufferSpec): self.dev._gpu_free(opaque)
 
 def flag(nm, val): return (val << getattr(kgsl, f"{nm}_SHIFT")) & getattr(kgsl, f"{nm}_MASK")
 
-class QCOMDevice(HCQCompiled):
+class QCOMDevice(HCQ2Compiled):
+  timestamp_divider = 19.2
+  has_copy_queue = False
+  pm_encode = PatternMatcher([
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_qcom_compute", name="submit"), lambda ctx, submit: encode_submit(QCOMComputeQueue(ctx, submit))),
+  ])
+
   def __init__(self, device:str=""):
     self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
-    self.dummy_addr = int(self._gpu_alloc(0x1000).va_addr)
 
     flags = kgsl.KGSL_CONTEXT_PREAMBLE | kgsl.KGSL_CONTEXT_PWR_CONSTRAINT | kgsl.KGSL_CONTEXT_NO_FAULT_TOLERANCE | kgsl.KGSL_CONTEXT_NO_GMEM_ALLOC \
       | flag("KGSL_CONTEXT_PRIORITY", getenv("QCOM_PRIORITY", 8)) | flag("KGSL_CONTEXT_PREEMPT_STYLE", kgsl.KGSL_CONTEXT_PREEMPT_STYLE_FINEGRAIN)
     self.ctx = kgsl.IOCTL_KGSL_DRAWCTXT_CREATE(self.fd, flags=flags).drawctxt_id
-
-    self.cmd_buf = self._gpu_alloc(16 << 20)
-    self.cmd_buf_allocator = BumpAllocator(size=self.cmd_buf.size, base=int(self.cmd_buf.va_addr), wrap=True)
-
-    self.border_color_buf = self._gpu_alloc(0x1000, fill_zeroes=True)
-
-    self.last_cmd:int = 0
+    self._stack:Buffer|None = None # private-memory stack
 
     # Set max power
     struct.pack_into('IIQQ', pwr:=memoryview(bytearray(0x18)), 0, 1, self.ctx, mv_address(_:=memoryview(array.array('I', [1]))), 4)
@@ -374,9 +347,25 @@ class QCOMDevice(HCQCompiled):
     if PROFILE and self.gpu_id[:2] < (7, 3):
       System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
 
-    super().__init__(device, QCOMAllocator(self), [QCOMCLRenderer, IR3Renderer], QCOMProgram, QCOMSignal, functools.partial(QCOMComputeQueue, self),
+    super().__init__(device, QCOMAllocator(self), [QCOMCLRenderer, IR3Renderer], None,
                      arch=("a%d%d%d" + (",IMAGE_PITCH_ALIGNMENT=64" if IMAGE else "")) % self.gpu_id)
 
+    self.var_vals = {"kgsl_fd": self.fd.fd, "kgsl_ctx": self.ctx}
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.PARAM, tag="stack", name="b"), lambda ctx, b: ctx._ensure_stack_size(b.max_numel())),
+      (UPat(Ops.PARAM, tag="dummy"), lambda ctx: ctx.dummy),
+      (UPat(Ops.PARAM, tag="border_color"), lambda ctx: ctx.border_color),
+    ]) + self.pm_bufferize
+
+  @functools.cached_property
+  def dummy(self) -> Buffer: return Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True) # cache flush target
+
+  @functools.cached_property
+  def border_color(self) -> Buffer: # zeros: the samplers clamp to a black border
+    (b:=Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)) \
+      .as_memoryview(force_zero_copy=True)[:] = bytes(0x1000)
+    return b
+
   def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
     flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
     if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
@@ -404,12 +393,18 @@ class QCOMDevice(HCQCompiled):
       kgsl.IOCTL_KGSL_GPUOBJ_FREE(self.fd, id=mem.meta[0].id)
       FileIOInterface.munmap(mem.va_addr, mem.meta[0].mmapsize)
 
-  def _ensure_stack_size(self, sz):
-    if not hasattr(self, '_stack'): self._stack = self._gpu_alloc(sz)
-    elif self._stack.size < sz:
-      self.synchronize()
-      self._gpu_free(self._stack)
-      self._stack = self._gpu_alloc(sz)
+  def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
+    if sig[0] < value:
+      ts = kgsl.IOCTL_KGSL_CMDSTREAM_READTIMESTAMP_CTXTID(self.fd, context_id=self.ctx, type=kgsl.KGSL_TIMESTAMP_QUEUED).timestamp
+      with contextlib.suppress(OSError, RuntimeError):
+        kgsl.IOCTL_KGSL_DEVICE_WAITTIMESTAMP_CTXTID(self.fd, context_id=self.ctx, timestamp=ts, timeout=int(timeout or self.wait_timeout_ms))
+    super()._wait_signal(sig, value, timeout)
+
+  def _ensure_stack_size(self, sz:int) -> Buffer: # one stack for the device, grown to the deepest program's private memory
+    if self._stack is None or self._stack.nbytes < sz:
+      if self._stack is not None: self.synchronize()
+      self._stack = Buffer(self.device, sz, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+    return self._stack
 
   def _at_profile_finalize(self):
     super()._at_profile_finalize()

tinygrad/runtime/support/am/amdev.py

@@ -144,7 +144,7 @@ class AMMemoryManager(MemoryManager):
     self.dev.gmc.flush_tlb(ip='MM', vmid=0)
 
 class AMDev:
-  Version = 0xA0000008
+  Version = 0xA000000D
 
   def _disable_aspm(self):
     # L1 across retimers makes reads oscillate to 0xffffffff; power on defaults it enabled. Clearing the GPU endpoint
@@ -200,6 +200,7 @@ class AMDev:
         self.smu.mode1_reset()
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
       self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))
+    elif not self.is_vf: self.psp._tmr_init()
 
     # Booting done
     self.is_booting = False
@@ -213,6 +214,7 @@ class AMDev:
         self.smu.set_clocks(level=None)
       else: self.smu.set_clocks(level=-1) # last level, max perf.
       for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
+      self.reg("regSCRATCH_REG5").write(self.psp.tmr_size) # scratch registers are writable after GFX initialization
       self.reg("regSCRATCH_REG7").write(AMDev.Version)
       self.reg("regSCRATCH_REG6").write(1) # set initialized state.
 
@@ -222,7 +224,7 @@ class AMDev:
     self.smi_dev, self.is_err_state = smi_dev, False
 
     # Memory manager & firmware
-    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
+    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(3 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
       va_bits=48, first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base, reserve_ptable=not self.large_bar,
       palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
     self.fw = AMFirmware(self)

tinygrad/runtime/support/am/ip.py

@@ -603,10 +603,9 @@ class AM_PSP(AM_IP):
     self.ring_size = 0x10000
     self.ring_paddr = self.adev.mm.palloc(self.ring_size, zero=False, boot=True)
 
-    self.max_tmr_size, self.tmr_size = 0x1300000, 0
+    self.tmr_size, self.tmr_paddr = 0, 0
     self.boot_time_tmr = self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,14), (14,0,2), (14,0,3)}
     self.autoload_tmr = self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,14)}
-    self.tmr_paddr = self.adev.mm.palloc(self.max_tmr_size, align=am.PSP_TMR_ALIGNMENT, zero=False, boot=True) if not self.boot_time_tmr else 0
 
   def init_hw(self):
     spl_key = am.PSP_FW_TYPE_PSP_SPL if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0) else am.PSP_FW_TYPE_PSP_KDB
@@ -655,10 +654,13 @@ class AM_PSP(AM_IP):
     return self._wait_for_bootloader() if compid != am.PSP_BL__LOAD_SOSDRV else 0
 
   def _tmr_init(self):
-    # Load TOC and calculate TMR size
-    self._prep_msg1(fwm:=self.adev.fw.sos_fw[am.PSP_FW_TYPE_PSP_TOC])
-    self.tmr_size = self._load_toc_cmd(len(fwm)).resp.tmr_size
-    assert self.tmr_size <= self.max_tmr_size
+    if self.adev.partial_boot: self.tmr_size = self.adev.reg("regSCRATCH_REG5").read()
+    else:
+      # Load TOC and calculate TMR size
+      self._prep_msg1(fwm:=self.adev.fw.sos_fw[am.PSP_FW_TYPE_PSP_TOC])
+      self.tmr_size = self._load_toc_cmd(len(fwm)).resp.tmr_size
+    # First runtime allocation on both full and partial boots, so the resident TMR keeps the same address.
+    if not self.boot_time_tmr: self.tmr_paddr = self.adev.mm.pa_allocator.alloc(self.tmr_size, am.PSP_TMR_ALIGNMENT)
 
   def _ring_create(self):
     # If the ring is already created, destroy it

tinygrad/runtime/support/amd.py

@@ -30,7 +30,7 @@ class AMDIP:
 # https://github.com/torvalds/linux/blob/9207d47f966be9f4d52e7e0119ac2b7a7e366f3e/drivers/gpu/drm/amd/amdgpu/amdgpu_discovery.c#L3163
 def import_module(name:str, target:tuple[int, int, int], submod=""):
   # version overrides
-  target = {("smu", (13, 0, 7)): (13, 0, 0)}.get((name, target), target)
+  target = {("smu", (13, 0, 7)): (13, 0, 0), ("smu", (13, 0, 10)): (13, 0, 0)}.get((name, target), target)
   mod = getattr(tinygrad.runtime.autogen.am, submod) if submod else tinygrad.runtime.autogen.am
   if (children:=[c for c in mod.__all__ if c.startswith(name) and (v:=tuple(map(int, c.split('_')[1:])))[0] == target[0] and v <= target]):
     return getattr(mod, children[-1])

tinygrad/runtime/support/c.py

@@ -138,4 +138,5 @@ class DLL(ctypes.CDLL):
 
   def __getattr__(self, nm):
     if self.nm not in self._loaded_: raise AttributeError(f"failed to load library {self.nm}: {self.emsg}")
-    return super().__getattr__(nm)
+    (fn:=super().__getattr__(nm)).__module__ = f"tinygrad.runtime.autogen.{self.nm}"
+    return fn

tinygrad/runtime/support/compiler_qcom.py

@@ -14,7 +14,7 @@ class QCOMCompiler(Compiler):
     else:
       # extract once into the download cache, all processes share the rootfs (extract=True)
       self.arch, self.chip_id = arch, 0x6030001
-      fs, root = fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz', extract=True), pathlib.Path(__file__).parents[3]
+      fs, root = fetch('https://git.tinygrad.win/tinygrad/images/releases/download/v2/qcomcl.tar.gz', extract=True), pathlib.Path(__file__).parents[3]
       self.compiler_process = self.server(f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3" if (qemu:=shutil.which("qemu-aarch64-static"))
                                           else (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {root}:{root} "
                                                 f"-e PYTHONPATH={root} -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3"), arch)

tinygrad/runtime/support/hcq2.py

@@ -1,16 +1,16 @@
 from __future__ import annotations
 from typing import cast, TypeVar, Generic, Any, TYPE_CHECKING
-import functools, time, itertools, decimal, weakref, os, statistics
+import functools, time, itertools, decimal, weakref, statistics, ctypes, importlib
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ
-from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, perf_counter_us
-from tinygrad.device import Device, Buffer, MultiBuffer, BufferSpec, Compiled, LRUAllocator, DepsTracker
+from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2, cpu_profile, mv_address
+from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, perf_counter_us, DEV
+from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, DepsTracker
 from tinygrad.device import ProfileGraphEntry, ProfileGraphEvent, ProfileDeviceEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
-from tinygrad.dtype import dtypes, DType
+from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
 from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
+from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
 from tinygrad.engine.realize import lower_and_compile
 
 if TYPE_CHECKING: from tinygrad.runtime.support.hcq import HCQBuffer # TODO: remove that
@@ -20,7 +20,8 @@ if TYPE_CHECKING: from tinygrad.runtime.support.hcq import HCQBuffer # TODO: rem
 
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
-HCQ_DEVS = frozenset(("AMD", "CPU"))
+HCQ_CACHE_THRESH = ContextVar("HCQ_CACHE_THRESH", 64)
+HCQ_DEVS = frozenset(("NV", "QCOM", "CPU")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
 
 @dataclass(frozen=True)
 class HCQInfo:
@@ -35,6 +36,16 @@ class HCQInfo:
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
+def get_enqueue_devs(call:UOp) -> Any|None:
+  if call.src[0].op not in (Ops.PROGRAM, Ops.COPY): return None # only these bodies can be enqueued
+  if not (bufs:=get_call_arg_uops(call)) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
+  if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
+  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
+  # cpu has no queue (yet)
+  if not all_devices_in(devs, HCQ_DEVS) or to_tuple(devs)[0].startswith("CPU"): return None
+  # a device without a copy queue leaves copies to its allocator
+  return devs if call.src[0].op is not Ops.COPY or Device[to_tuple(devs)[0]].has_copy_queue else None
+
 def unwrap_view(v:UOp) -> tuple[UOp, int]: # look through views to (base, byte offset)
   if v.op in (Ops.BITCAST, Ops.AFTER): return unwrap_view(v.src[0])
   if v.op is not Ops.SHRINK: return v, 0
@@ -48,29 +59,55 @@ def to_name(*parts:str) -> str: return "_".join(parts).replace(":", "_").lower()
 def timeline(devs:tuple[str, ...]) -> UOp: return UOp.placeholder((2,), dtypes.uint64, 0, device=devs, volatile=True, tag="timeline")
 def timeline_value(devs:tuple[str, ...]) -> UOp: return timeline(devs).index(1).load()
 
+def rt_addr(b:UOp, dev) -> UOp:
+  base, off = unwrap_view(b)
+  return patch(UOp.placeholder((1,), dtypes.uint64, device=base.device, tag="addr"), [(0, base.getaddr(dev))]).index(0).load() + off
+
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   fn = to_name("submit", (devs:=to_tuple(devs))[0].split(":")[0], queue.split(":")[0])
   return UOp.custom_function(fn, UOp(Ops.LINEAR, src=tuple(cmds), arg=(devs, queue)))
 
-# *****************
-# 0.1. prep: replace buffers with params
+# C FFI
 
-def replace_call_buffers(ctx:tuple[list[UOp], dict[UOp, int]], call:UOp) -> UOp|None:
-  bufs, slots = ctx
-  for s in call.src[1:]:
-    if s.op is not Ops.PARAM and not s.is_bound_var and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s)
-  return call.replace(src=call.src[:1] + tuple(s if s.op is Ops.PARAM or s.is_bound_var else s.param_like(slots[s]) for s in call.src[1:]))
-pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])
+@functools.cache
+def cfunc_buf(lib:str, name:str) -> Buffer:
+  fn = getattr(importlib.import_module(f"tinygrad.runtime.autogen.{lib}").dll, name)
+  (b:=Buffer(HCQ_RUNTIME_DEV.value, 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = unwrap(ctypes.cast(fn, ctypes.c_void_p).value)
+  return b
+
+def ccall(fn:Any, *args:UOp|int) -> UOp:
+  ptr = UOp.placeholder((1,), dtypes.uint64, 0, device=HCQ_RUNTIME_DEV.value, tag=("cfunc", fn.__module__.split(".")[-1], fn.__name__))
+  ret = dtypes.void if fn.restype is None else dtypes.uint64 if fn.restype is ctypes.c_void_p else \
+    next(d for d in DTYPES_DICT.values() if d.fmt == fn.restype._type_)
+  cargs = [UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args]
+  return UOp.custom_function(fn.__name__, ptr.index(0).load()).call(*cargs, ret_dtype=ret)
+
+def cstruct(struct_t, **fields:UOp|int) -> UOp:
+  flds = {n: (o, {1: dtypes.uchar, 2: dtypes.ushort, 4: dtypes.uint, 8: dtypes.ulong}[ctypes.sizeof(t)]) for n, t, o, *_ in struct_t._real_fields_}
+  rows = [(flds[n][0], v.cast(flds[n][1]) if isinstance(v, UOp) else UOp.const(v, flds[n][1])) for n, v in fields.items()]
+  buf = UOp.placeholder((ctypes.sizeof(struct_t),), dtypes.uint8, device=HCQ_RUNTIME_DEV.value, volatile=True, tag=struct_t.__name__)
+  return patch(buf, rows, bytes(ctypes.sizeof(struct_t)))
+
+# *****************
+# 0.1. prep: eager buffers become tagged params
+
+def replace_buffer(ctx:tuple[bool, list[UOp], dict[UOp, int]], b:UOp) -> UOp:
+  use_rt, bufs, slots = ctx
+  if slots.setdefault(b, len(bufs)) == len(bufs): bufs.append(b)
+  param = UOp.param(slots[b], b.dtype, b.max_numel(), b.device)
+  return param if use_rt else param.replace(tag="lt_input")
+pm_replace_buffers = PatternMatcher([(UPat(Ops.BUFFER, name="b"), replace_buffer)])
 
 # *****************
 # 1.1. prep: staging copies
 
-STAGING_SIZE, STAGING_SLOTS = (4 if os.getenv("CI") else 128) << 20, 2 # the staging mlocks into the device: ci runners cap locked memory at 8mb
+STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
 
 @functools.cache
 def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
 
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS)
+def _need_staging(a, b):
+  return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS) and Device[to_tuple(a.device)[0]].has_copy_queue
 
 def stage_copy_ext(call:UOp) -> UOp|None:
   if (d:=next((d for b in call.src[1:] for d in to_tuple(b.device) if not d.startswith("CPU")), None)) is None: return None
@@ -87,47 +124,30 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
     copies += [src[off:off+n].copy_to_device("CPU").call(stage, src[off:off+n]), stage.copy_to_device(dst.device).call(dst[off:off+n], stage)]
   return UOp(Ops.LINEAR, src=tuple(copies))
 
+pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), stage_copy_ext),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
+])
+
 # *****************
 # 1.2. prep: one call per device: the args pick their lane, the DEVICE axis binds to it
 
 def unwrap_call(call:UOp) -> UOp|None:
-  if call.src[0].op not in (Ops.PROGRAM, Ops.COPY) or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None
+  if get_enqueue_devs(call) is None or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None
   dnum = UOp.variable("_device_num", 0, n - 1, dtypes.int)
   return UOp(Ops.LINEAR, src=tuple(call.replace(src=(call.src[0], *[a if a.is_bound_var else select_lane(a, i) for a in call.src[1:]], dnum.bind(i)))
                                    for i in range(n)))
 pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)])
 
-# *****************
-# 1.3. prep: kernel copies
-
-def _get_enqueue_devs(call:UOp) -> Any|None:
-  if call.src[0].op not in (Ops.PROGRAM, Ops.COPY): return None # only these bodies can be enqueued
-  if not (bufs:=get_call_arg_uops(call)) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
-  if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
-  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
-  # cpu has no queue (yet)
-  return devs if all_devices_in(devs, HCQ_DEVS) and not to_tuple(devs)[0].startswith("CPU") else None
-
-def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None:
-  if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
-  d, s = (UOp.param(i, dst.dtype, n:=dst.max_numel(), device=devs) for i in range(2))
-  ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
-  return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
-
-pm_insert_copy_staging = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), stage_copy_ext),
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), copy_with_kernel)
-])
-
 # *****************
 # 2. deps
 
 class HCQDepsTracker(DepsTracker):
-  # TODO: optimize
   @staticmethod
-  def _key(buf:Any) -> tuple[Any, int, int]:
-    return (buf, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
+  def _key(a:UOp) -> tuple[Any, int, int]: # (base, lane) and the byte range: overlapping views of one base depend
+    lane, view = (a.arg, a.src[0]) if a.op is Ops.MSELECT else (None, a)
+    base, off = unwrap_view(view)
+    return (base, lane), off, off + view.max_numel() * view.dtype.itemsize
 
 @dataclass
 class BatchCtx:
@@ -136,37 +156,38 @@ class BatchCtx:
   tracker:HCQDepsTracker = field(default_factory=HCQDepsTracker)
   queues:dict[str, list[str]] = field(init=False)
   first:dict[tuple[str, str], int] = field(init=False); last:dict[tuple[str, str], int] = field(init=False) # noqa: E702
+  prev:list[int|None] = field(init=False)
   signal_tags:set[int] = field(init=False)
   slots:dict[str, UOp] = field(init=False)
 
   def __post_init__(self):
-    self.queues, self.first, self.last = {}, {}, {}
+    self.queues, self.first, self.last, self.prev = {}, {}, {}, []
     for tag, (_, devs, q) in enumerate(self.batch):
       if q not in self.queues.setdefault(devs[0], []): self.queues[devs[0]].append(q)
       self.first.setdefault((devs[0], q), tag)
+      self.prev.append(self.last.get((devs[0], q)))
       self.last[(devs[0], q)] = tag
     self.signal_tags = {tag for (dev, q), tag in self.last.items() if q != self.epilogue_queue(dev)}
-    self.slots = {dev: UOp.placeholder((len(qs) + 1 + (2 * len(self.batch) if self.profile else 0),), dtypes.uint64, device=(dev,), volatile=True,
-                                       tag="slots") for dev, qs in self.queues.items()}
+    # a slot is [signal][timestamp], 16 bytes: the queue signals, the timeline, then two per call if profiling
+    self.slots = {dev: UOp.placeholder((2 * (len(qs) + 1 + (2 * len(self.batch) if self.profile else 0)),), dtypes.uint64, device=(dev,),
+                                       volatile=True, tag="slots") for dev, qs in self.queues.items()}
 
   def epilogue_queue(self, dev:str) -> str: return "COMPUTE:0" if len(self.queues[dev]) > 1 else self.queues[dev][0] # closes the device
 
-  def slot(self, devs:tuple[str, ...], i:int) -> UOp: return self.slots[devs[0]].shrink(((i, i + 1),)) # not [i:i+1]: the slice path is 10x the cost
+  def slot(self, devs:tuple[str, ...], i:int) -> UOp: return self.slots[devs[0]].shrink(((2 * i, 2 * i + 2),)) # not a slice: 10x the cost
   def queue_signal(self, devs:tuple[str, ...], queue:str) -> UOp: return self.slot(devs, self.queues[devs[0]].index(queue))
   def sched_timeline(self, devs:tuple[str, ...]) -> UOp: return self.slot(devs, len(self.queues[devs[0]]))
   def stamps(self, devs:tuple[str, ...], tag:int) -> tuple[int, ...]: return (st:=len(self.queues[devs[0]])+1+2*tag, st + 1) if self.profile else ()
 
-def _call_bufs(call:UOp) -> list[Any]:
-  def dep_buf(b:UOp) -> Any:
-    if (base:=(b.src[0] if b.op is Ops.MSELECT else b).storage_base).op is Ops.PARAM: return b if b.op is Ops.MSELECT else base
-    return cast(MultiBuffer, base.buffer).bufs[b.arg] if b.op is Ops.MSELECT else base.buffer
-  return [dep_buf(a) for a in get_call_arg_uops(call)]
-
 def _wait_ins(ctx:BatchCtx, call:UOp, device:str, queue:str, tag:int) -> list[UOp]:
-  bufs, write = _call_bufs(call), get_call_outs_ins(call)[0]
+  bufs, write = list(get_call_arg_uops(call)), get_call_outs_ins(call)[0]
   latest:dict[tuple[str, str], int] = {} # (producer device, queue) -> the latest submit tag to wait on, same-queue submits are fifo
   for d, q, t in ctx.tracker.access_resources(bufs, list(range(len(bufs)) if write is None else write), (device, queue, tag)):
     if t < tag and (d, q) != (device, queue): latest[(d, q)] = max(latest.get((d, q), 0), t)
+
+  # NV waits break QMD chaining, so also wait for the previous launch
+  if latest and device.split(":")[0] == "NV" and queue.startswith("COMPUTE") and (p:=ctx.prev[tag]) is not None: latest[(device, queue)] = p
+
   ctx.signal_tags |= set(latest.values())
   return [UOp(Ops.INS, arg=("wait", dtypes.void), src=(ctx.queue_signal((d,), q), UOp.const(t + 1, dtypes.uint64))) for (d, q), t in latest.items()]
 
@@ -186,7 +207,7 @@ def _emit_submits(ctx:BatchCtx, call_waits:list[list[UOp]]) -> tuple[list[UOp],
 
     # and make hcq call
     name, est = get_call_name(call, get_call_arg_uops(call)), estimate_uop(call)
-    kerns.append((devices, name, est, ctx.stamps(devices, tag), getattr(call.src[0].arg, "profile_key", None)))
+    kerns.append((devices, name, est, tuple(2 * s + 1 for s in ctx.stamps(devices, tag)), getattr(call.src[0].arg, "profile_key", None)))
 
     ts_ins = [UOp(Ops.INS, arg=("timestamp", dtypes.void), src=(ctx.slot(devices, i),)) for i in ctx.stamps(devices, tag)]
     q += ts_ins[:1] + [call] + ts_ins[1:]
@@ -210,13 +231,17 @@ def _finalize_batch(ctx:BatchCtx) -> UOp:
   submits += [_epilogue(ctx, dev) for dev in ctx.queues]
   fence = UOp.custom_function("hcq_fence", *[ctx.sched_timeline((dev,)) for dev in ctx.queues],
                               *[ctx.queue_signal((dev,), q) for dev, qs in ctx.queues.items() for q in qs])
-  merged = [m.replace(src=(*m.src, fence)) for m in _merge_queues(submits)]
+  merged:list[UOp] = [] # the submits in order, after the fence
+  for m in _merge_queues(submits): merged.append(m.after(fence, *merged[-1:]))
   estimates = sum((estimate_uop(call) for call, _, _ in ctx.batch), start=Estimates()).simplify()
-  return UOp.sink(*merged, arg=KernelInfo("hcq_submit", estimates=estimates), tag=1).call(aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns)))
+  sink = UOp.sink(*merged, arg=KernelInfo("hcq_submit", estimates=estimates), tag=1)
+  for pm in [Device[d].pm_batch for d in ctx.queues if Device[d].pm_batch is not None]: # a device adds its own work to the batch
+    if (r:=pm.rewrite(sink)) is not None: sink = r
+  return sink.call(aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns)))
 
 @rewrite_group(new_ctx=False)
 def sched_batches(l:UOp, profile:bool) -> UOp:
-  devs = [() if (d:=_get_enqueue_devs(c)) is None else tuple(Device.canonicalize(x) for x in to_tuple(d)) for c in l.src]
+  devs = [() if (d:=get_enqueue_devs(c)) is None else tuple(Device.canonicalize(x) for x in to_tuple(d)) for c in l.src]
   queues = ["COMPUTE:0" if c.src[0].op is Ops.PROGRAM else "COPY:0" for c in l.src]
   srcs:list[UOp] = []
   for hcq, grp in itertools.groupby(zip(l.src, devs, queues), key=lambda e: bool(e[1])):
@@ -231,13 +256,13 @@ class EncodeCtx:
   devs:tuple[str, ...]
   inputs:dict[tuple[UOp, str], int] = field(default_factory=dict)
   table:UOp = field(default_factory=lambda: UOp.placeholder((1,), dtypes.uint64, device="CPU", tag="inputs"))
-  lt_patches:list[UOp] = field(default_factory=list)
+  lt_patches:dict[UOp, list[UOp]] = field(default_factory=dict) # placeholder -> the stores into it that resolve when the linear links
 
 class HWQueue:
   q_rewrite:PatternMatcher
 
   def __init__(self, ctx:EncodeCtx, submit:UOp):
-    self.ctx, self.lin, self.deps = ctx, submit.src[0], list(submit.src[1:])
+    self.ctx, self.lin = ctx, submit.src[0]
     self.devs, self.queue = self.lin.arg
     self.dev = Device[self.devs[0]]
     self.blob, self.patches = bytearray(), list[tuple[int, UOp]]()
@@ -246,7 +271,8 @@ class HWQueue:
     for w in words:
       c = w
       while isinstance(c, UOp) and c.op is Ops.CAST: c = c.src[0]
-      if isinstance(c, UOp) and c.op is not Ops.CONST:
+      if isinstance(c, UOp) and c.op is Ops.BINARY: self.blob += c.arg
+      elif isinstance(c, UOp) and c.op is not Ops.CONST:
         self.patches.append((len(self.blob), w))
         self.blob += bytes(w.dtype.itemsize)
       else:
@@ -256,28 +282,79 @@ class HWQueue:
 
   def submit(self, cmdbuf:UOp) -> UOp: raise NotImplementedError("queues need a submit")
 
+# *****************
+# 3.1. hcq special functions
+
+def hcq_fence(ctx:EncodeCtx, f:UOp) -> UOp:
+  lasts, sigs = f.src[:len(ctx.devs)], f.src[len(ctx.devs):]
+  last:tuple[UOp, ...] = ()
+
+  # wait for prev schedule to not collide
+  # TODO: timeout?
+  for i, dev in enumerate(ctx.devs):
+    slots, off = unwrap_view(lasts[i])
+    slots = patch(slots, [], bytes(slots.max_numel() * slots.dtype.itemsize)) # zeroed at link
+    target = slots.after(*last, tv:=timeline_value((dev,))).index(off // slots.dtype.itemsize).load()
+    done = timeline((dev,)).after(target, loop:=UOp.loop(i)).index(0).load()
+    bumped = timeline((dev,)).after(done.end(loop, done < target)).index(1).store(nxt:=tv + UOp.const(1, dtypes.uint64))
+    last = (slots.after(bumped).index(off // slots.dtype.itemsize).store(nxt),)
+
+  # re-arm the signals
+  for sig in sigs:
+    base, off = unwrap_view(sig)
+    last = (base.after(*last).index(off // sig.dtype.itemsize).store(0),)
+  return last[0].barrier(*last[1:])
+
+pm_hcq_encode = PatternMatcher([
+  (UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name="f"), hcq_fence),
+
+  # after blocks are lowered, rechain stores saving original order
+  (UPat(Ops.AFTER, src=(UPat(dtype=dtypes.void, name="root"),), allow_any_len=True, name="a"),
+    lambda root, a: root.substitute({s.buf_uop: s.buf_uop.after(*a.src[1:]) for s in root.toposort() if s.op is Ops.STORE}, walk=True)),
+])
+
+# *****************
+# 3.2. split
+
+def _is_input_addr(g:UOp) -> bool:
+  base = unwrap_view(g.src[0])[0]
+  if base.op is Ops.MSELECT: base = unwrap_view(base.src[0])[0] # a lane of a view
+  return base.op is Ops.PARAM and base.tag is None
+
 def addrs_to_table(ctx:EncodeCtx, g:UOp) -> UOp|None:
+  if not _is_input_addr(g): return None
   base, off = unwrap_view(g.src[0])
-  param = base.src[0].base if base.op is Ops.MSELECT else base # unwrap mselects
-  if param.op is not Ops.PARAM or param.tag is not None: return None
   slot = ctx.inputs.setdefault((base, to_tuple(g.arg)[0]), len(ctx.inputs))
   return ctx.table.index(slot).load() + UOp.const(off, dtypes.uint64)
-pm_addrs_to_table = PatternMatcher([(UPat(Ops.GETADDR, name="g"), addrs_to_table)])
 
 def _is_link_patch(w:UOp) -> bool:
-  if w.op is Ops.GETADDR: return True
-  if w.op in {Ops.LOAD, Ops.INDEX, Ops.PARAM} or w.is_variable: return False
+  if w.op is Ops.GETADDR: return not _is_input_addr(w)
+  if w.op is Ops.PARAM: return w.tag is not None
+  if w.op is Ops.BUFFER: return w.addrspace is AddrSpace.GLOBAL # a register is written at runtime
+  if w.op in {Ops.LOAD, Ops.AFTER} or w.is_variable: return False
   return all(_is_link_patch(s) for s in w.src)
 
-def patch(buf:UOp, rows:list[tuple[int, UOp]], *deps:UOp) -> UOp: # the buffer after every row's word is stored at its byte offset
-  groups:dict[tuple[DType, int], list[tuple[int, UOp]]] = {}
-  for o, w in rows: groups.setdefault((w.dtype, o % w.dtype.itemsize), []).append((o, w))
+def hoist_links(ctx:EncodeCtx, a:UOp) -> UOp|None:
+  links, rest = partition(a.src[1:], lambda s: s.op is Ops.STORE and _is_link_patch(s))
+  if not links: return None
+  # nest the addr placeholders patches under their getaddr
+  ws = UOp.sink(*links)
+  sub = {g: g.replace(src=(g.src[0].after(*ctx.lt_patches[g.src[0]]),)) for g in ws.toposort() if g.op is Ops.GETADDR and g.src[0] in ctx.lt_patches}
+  ctx.lt_patches.setdefault(unwrap_view(a.src[0])[0], []).extend(ws.substitute(sub).src)
+  return a.src[0].after(*rest)
 
-  base, stores = buf.after(*deps), [] # the views hang off the buffer after its deps: the stores wait for them, the caller's after keeps the base flat
-  for (dt, phase), grp in groups.items():
+pm_patches = PatternMatcher([(UPat(Ops.GETADDR, name="g"), addrs_to_table), (UPat(Ops.AFTER, name="a"), hoist_links)])
+
+def patch(buf:UOp, rows:list[tuple[int, UOp]], blob:bytes|None=None) -> UOp:
+  groups:dict[tuple[DType, int, bool], list[tuple[int, UOp]]] = {} # split by: dtype, alignment, is_link (rt/lt can't share a store)
+  for offb, w in rows: groups.setdefault((w.dtype, offb % w.dtype.itemsize, _is_link_patch(w)), []).append((offb, w))
+
+  dep = [buf.store(UOp(Ops.BINARY, arg=blob).bitcast(buf.dtype))] if blob is not None else []
+  base, stores = buf.after(*dep), [] # keep buf.after to be sure that link applies patches after the blob
+  for (dt, phase, _), grp in groups.items():
     view = base[phase:phase + (buf.max_numel() - phase) // dt.itemsize * dt.itemsize].bitcast(dt)
     stores.append(view.index(UOp.stack(*[UOp.const((o - phase) // dt.itemsize) for o, _ in grp])).store(UOp.stack(*[w for _, w in grp])))
-  return buf.after(*deps, *stores)
+  return buf.after(*dep, *stores)
 
 def encode_submit(hq:HWQueue) -> UOp:
   # applying the rewrite
@@ -291,36 +368,8 @@ def encode_submit(hq:HWQueue) -> UOp:
 
   buf = UOp.placeholder((len(hq.blob),), dtypes.uint8, device=hq.devs, tag=to_name("cmdbuf", hq.queue))
 
-  words = UOp.sink(*[w for _, w in hq.patches]).substitute({l: buf[o:e] for l, (o, e) in views.items()})
-  words = graph_rewrite(words, pm_addrs_to_table, ctx=hq.ctx, name="addrs to table").src
-
-  links, runtime = partition(list(zip([o for o, _ in hq.patches], words)), lambda r: _is_link_patch(r[1]))
-  hq.ctx.lt_patches.append(patch(buf, links, buf.store(UOp(Ops.BINARY, arg=bytes(hq.blob)).bitcast(buf.dtype))))
-  return hq.submit(patch(buf, runtime, *hq.deps).shrink(((0, stream),)))
-
-# *****************
-# 3.1. hcq special functions
-
-def hcq_fence(ctx:EncodeCtx, f:UOp) -> UOp:
-  lasts, sigs = f.src[:len(ctx.devs)], f.src[len(ctx.devs):]
-  last:tuple[UOp, ...] = ()
-
-  # wait for prev schedule to not collide
-  # TODO: timeout?
-  for i, dev in enumerate(ctx.devs):
-    slots, off = unwrap_view(lasts[i])
-    ctx.lt_patches.append(slots.after(slots.store(UOp(Ops.BINARY, arg=bytes(slots.max_numel() * slots.dtype.itemsize)).bitcast(slots.dtype))))
-    done = timeline((dev,)).after(*last, loop:=UOp.loop(i)).index(0).load()
-    waited = done.end(loop, done < slots.index(off // slots.dtype.itemsize).load())
-    nxt = timeline_value((dev,)) + UOp.const(1, dtypes.uint64)
-    last = (timeline((dev,)).after(waited).index(1).store(nxt), slots.after(waited).index(off // slots.dtype.itemsize).store(nxt))
-
-  # re-arm the signals
-  for sig in sigs:
-    base, off = unwrap_view(sig)
-    last = (base.after(*last).index(off // sig.dtype.itemsize).store(0),)
-  return last[0].barrier(*last[1:])
-pm_hcq_encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name="f"), hcq_fence)])
+  words = UOp.sink(*[w for _, w in hq.patches]).substitute({l: buf[o:e] for l, (o, e) in views.items()}).src
+  return hq.submit(patch(buf, list(zip([o for o, _ in hq.patches], words)), bytes(hq.blob)).shrink(((0, stream),)))
 
 # *****************
 # 4. lower call
@@ -330,8 +379,9 @@ def lower_call(call:UOp) -> UOp|None:
 
   # encode bodies
   ctx = EncodeCtx(call.arg.aux.device)
-  pm = sum([Device[d].pm_encode for d in dedup([d.split(":")[0] for d in ctx.devs])], pm_hcq_encode)
-  body = graph_rewrite(call.src[0], pm, ctx=ctx, walk=True, name="encode body")
+  devs = [Device[d] for d in dedup([d.split(":")[0] for d in ctx.devs])]
+  body = graph_rewrite(call.src[0], sum([d.pm_encode for d in devs], PatternMatcher([])) + pm_hcq_encode, ctx=ctx, bpm=pm_patches, name="encode")
+  body = graph_rewrite(body, sum([d.pm_lower for d in devs if d.pm_lower is not None], PatternMatcher([])), ctx=ctx, bpm=pm_patches, name="lower")
 
   # resize table
   body = body.substitute({ctx.table: (table:=UOp.placeholder((len(ctx.inputs),), dtypes.uint64, device="CPU", tag="inputs"))})
@@ -347,10 +397,10 @@ def lower_call(call:UOp) -> UOp|None:
   # reenum ranges
   rngs = {r: r.replace(arg=(i,)+r.arg[1:]) for i, r in enumerate(sorted([u for u in tops if u.op is Ops.RANGE], key=lambda r: r.arg))}
   # and sub all of them
-  sink = body.substitute(params | vals | rngs)
+  sink = body.substitute(params | vals | rngs, enter_calls=True)
 
   # move all lt-patches to the args
-  patched = {p.src[0]: p for p in ctx.lt_patches}
+  patched = {b: b.after(*dedup(stores)) for b, stores in ctx.lt_patches.items()}
   args = [patched.get(b, b) for b in bufs]
 
   if VIZ: graph_rewrite(UOp.sink(*args), PatternMatcher([]), name="View Link-Time Patches")
@@ -361,46 +411,48 @@ def lower_call(call:UOp) -> UOp|None:
   return call.replace(src=(sink, *args), arg=replace(call.arg, aux=info))
 pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call", allow_any_len=True), lower_call)])
 
-hcq_compile_cache:dict[tuple[UOp, bool, bool], UOp] = {} # uops are hash-consed: the linear itself is the key, plus whether inputs bind
+hcq_compile_cache:dict[tuple[UOp, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
 
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
-  if input_uops is not None:
-    slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
-    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer")
+  if any(isinstance(getattr(c.without_after.arg, "aux", None), HCQInfo) for c in linear.src): return linear # compiled already
 
-  # TODO: this needs a cleanup
-  bufmap = {s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}
-  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear, profile, input_uops is None)))) is None:
-    linear = graph_rewrite(linear.substitute(bufmap, walk=True), pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, name="prep calls")
-    linear = sched_batches(linear, profile)
-    linear = graph_rewrite(linear, pm_encode, walk=True, name="encode")
-    with Context(EMULATED_DTYPES=""):
-      final_linear = hcq_compile_cache[cache_key] = lower_and_compile(linear).substitute({v: k for k, v in bufmap.items()}, walk=True)
-  return final_linear.substitute(bufmap, walk=True)
+  if input_uops is not None:
+    use_rt = len(linear.src) < HCQ_CACHE_THRESH # small schedules use runtime address patches so linked schedules can be cached without input buffers
+    slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
+    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(use_rt, input_uops, slots), walk=True, name="replace buffers")
+    if (cached:=hcq_compile_cache.get(key:=(linear, profile))) is not None: return cached
+  lin = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, name="prep calls")
+  lin = sched_batches(lin, profile)
+  lin = graph_rewrite(lin, pm_encode, walk=True, name="encode")
+  with Context(EMULATED_DTYPES=""): final_linear = lower_and_compile(lin)
+  if input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
+  return final_linear
 
 # *****************
-# 5. bufferize placeholders
+# 5. link
 
-def bufferize_buf(ctx:bool, b:UOp) -> UOp|None: # ctx: a kept link (the jit's) owns the linear's buffers, a one-shot borrows ring slots
+@dataclass
+class LinkCtx: inputs:dict[UOp, UOp]; use_rt:bool; refs:list[UOp] = field(default_factory=list) # noqa: E702
+
+def bufferize_buf(ctx:LinkCtx, b:UOp) -> UOp|None: # ctx: a kept link (the jit's) owns the linear's buffers, a one-shot borrows ring slots
   if b.tag is None: return None # a param, not a placeholder
 
   dev = cast(HCQ2Compiled, Device[to_tuple(b.device)[0]])
 
-  if b.arg.slot == 0 or b.tag == "program": r = cast(Buffer, unwrap(dev.pm_bufferize.rewrite(b, ctx=dev))) # device state and programs
-  elif ctx: r = Buffer(dev.device, b.max_numel(), b.dtype, options=BufferSpec(host=b.arg.volatile, uncached=True, cpu_access=True), preallocate=True)
+  # device owns the placeholders it names
+  if (r:=cast(Buffer|None, dev.pm_bufferize.rewrite(b, ctx=dev))) is not None: pass
+  elif not ctx.use_rt:
+    spec = BufferSpec(host=b.arg.volatile, uncached=b.arg.volatile, cpu_access=True)
+    r = Buffer(dev.device, b.max_numel(), b.dtype, options=spec, preallocate=True)
   else: r = dev.rt_view(b.max_numel() * b.dtype.itemsize, b.dtype, host=b.arg.volatile)
 
   return UOp.from_buffer(r, HCQ_RUNTIME_DEV.value)
-pm_bufferize_placeholders = PatternMatcher([(UPat(Ops.PARAM, name="b"), bufferize_buf)])
 
-# *****************
-# 6. link
-
-def resolve_getaddr(ctx:list[UOp], g:UOp) -> UOp|None:
+def resolve_getaddr(ctx:LinkCtx, g:UOp) -> UOp|None:
   buf, off = unwrap_view(g.src[0])
   if buf.op not in {Ops.BUFFER, Ops.MSELECT}: return None
-  ctx.append(buf) # add to refs
+  ctx.refs.append(buf) # add to refs
   return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]).va_addr + off, dtypes.uint64)
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
@@ -418,25 +470,31 @@ def fold_words(buf:UOp, offs:UOp, ws:UOp) -> UOp:
   return UOp(Ops.NOOP)
 
 pm_link = PatternMatcher([
+  (UPat(Ops.CAST, src=(UPat(Ops.CAST, src=(UPat.cvar(),), name="inner"),), name="c"), lambda c, inner: inner.src[0].cast(c.dtype)),
+  (UPat(Ops.PARAM, name="b"), lambda ctx, b: ctx.inputs[b] if b in ctx.inputs else bufferize_buf(ctx, b)),
   (UPat(Ops.GETADDR, name="g"), resolve_getaddr),
   (UPat(GroupOp.ALU, src=UPat.cvar().or_casted(), name="a"),
     lambda a: UOp.const(exec_alu(a.op, a.dtype, [s.val for s in a.src], False), a.dtype)),
   (UPat(name="buf").store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast())), fold_binary),
   (UPat(name="buf").index(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="offs")).store(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="ws")),
     fold_words),
-  (UPat(Ops.AFTER, name="a"), lambda a: None if a.is_bound_var else a.src[0] if all(s.op is Ops.NOOP for s in a.src[1:]) else
-   panic(RuntimeError, f"unresolved link words on {a.src[0].op}")),
+  (UPat(Ops.AFTER, name="a"), lambda a: None if a.is_bound_var or a.src[0].op is Ops.CALL else
+   a.src[0] if all(s.op is Ops.NOOP for s in a.src[1:]) else panic(RuntimeError, f"unresolved link words on {a.src[0].op}")),
 ])
 
 link_linear_cache:weakref.WeakKeyDictionary[UOp, UOp] = weakref.WeakKeyDictionary() # a baked link lives as long as its bound linear
 
-@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp, cache=True) -> UOp:
-  if (linked:=link_linear_cache.get(linear)) is not None: return linked
-  bufferized = graph_rewrite(linear, pm_bufferize_placeholders, ctx=cache, name="bufferize")
-  linked = graph_rewrite(bufferized, pm_link, ctx=(refs:=list[UOp]()), bottom_up=False, name="link")
-  if refs: linked = linked.replace(src=(linked.src[0].replace(src=linked.src[0].src + tuple(dedup(refs))), *linked.src[1:]))
-  if cache: link_linear_cache[linear] = linked
+@rewrite_group(lambda _,input_uops=None,allow_cache=True,ret=None: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+def hcq_link(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
+  if allow_cache and (linked:=link_linear_cache.get(linear)) is not None: return linked
+
+  # if we have any link time buffers, do not cache this linear
+  cache = allow_cache and not any(u.tag == "lt_input" for u in linear.toposort() if u.op is Ops.PARAM)
+
+  inputs = {UOp.param(i, b.dtype, b.max_numel(), b.device).replace(tag="lt_input"): b for i, b in enumerate(input_uops or ())}
+  linked = graph_rewrite(linear, pm_link, ctx=(ctx:=LinkCtx(inputs, use_rt=allow_cache and not cache)), walk=True, name="link")
+  if ctx.refs: linked = linked.replace(src=(linked.src[0].after(*dedup(ctx.refs)), *linked.src[1:])) # attach refs to linear
+  if cache and linked is not linear: link_linear_cache[linear] = linked
   return linked
 
 # *****************
@@ -445,8 +503,10 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
 class HCQ2Compiled(Compiled):
   timestamp_divider: float = 1000.0
   wait_timeout_ms: float = 30000.0
+  sleep_timeout_ms: int|None = None
   rt_nbytes: int = 64 << 20 # the pool every per-linear buffer is carved out of
   pm_encode: PatternMatcher = PatternMatcher([]) # the backend's own encode rules, matched by its submit names
+  var_vals: dict[str, int] = {}
 
   def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
     self.can_recover = can_recover
@@ -455,6 +515,7 @@ class HCQ2Compiled(Compiled):
       (UPat(Ops.PARAM, tag="timeline"), lambda ctx: ctx.timeline),
       (UPat(Ops.PARAM, tag="program", name="b"),
        lambda ctx, b: ctx.prog_bufs.setdefault(b, Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)))),
+      (UPat(Ops.PARAM, name="b"), lambda b: cfunc_buf(*b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == "cfunc" else None),
     ])
     super().__init__(device, allocator, compilers, runtime, None, arch=arch)
 
@@ -462,8 +523,10 @@ class HCQ2Compiled(Compiled):
     self.prof_ents:dict[tuple[Buffer, int], ProfileGraphEntry] = {} # (a batch's timestamps, start slot) -> entry, read at synchronize
 
   @functools.cached_property
-  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]: zeroed host memory
-    return Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]
+    buf = Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    buf._buf.cpu_view().view(fmt='B')[:16] = bytes(16)
+    return buf
 
   def collect_prof(self):
     if PROFILE:
@@ -494,7 +557,7 @@ class HCQ2Compiled(Compiled):
     return Buffer(self.device, self.rt_allocator(uncached, host).size, dtypes.uint8, options=spec, preallocate=True)
 
   def rt_view(self, nbytes:int, dtype:DType=dtypes.uint8, uncached:bool=True, host:bool=False) -> Buffer: # a slot of the ring, wraps silently
-    off = self.rt_allocator(uncached, host).alloc(max(nbytes, 1), alignment=128)
+    off = self.rt_allocator(uncached, host).alloc(max(nbytes, 1), alignment=256)
     return self.rt_buffer(uncached, host).view(nbytes // dtype.itemsize, dtype, off).ensure_allocated()
 
   def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
@@ -502,7 +565,8 @@ class HCQ2Compiled(Compiled):
     st, done = time.perf_counter(), sig[0]
     while done < value:
       if done != (done:=sig[0]): st = time.perf_counter()
-      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
+      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
+      elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
 
   def synchronize(self, timeout:int|None=None):
     self._wait_signal(tl:=self.timeline._buf.cpu_view().view(fmt='Q'), tl[1], timeout)
@@ -510,6 +574,9 @@ class HCQ2Compiled(Compiled):
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
 
+  def on_sleep(self):
+    if (iface:=getattr(self, "iface", None)) is not None and hasattr(iface, "sleep"): iface.sleep(self.sleep_timeout_ms)
+
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
@@ -529,8 +596,11 @@ class HCQ2Buffer:
     return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
 
 class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
-  def _as_buffer(self, buf:HCQBuffer) -> memoryview:
-    return unwrap(buf.view).mv
+  def _as_buffer(self, buf:HCQBuffer) -> memoryview|None: return buf.view.mv if buf.view is not None else None
+
+  def _copyout(self, dest:memoryview, src:HCQBuffer): # TODO: remove with memcpy on cpu worker?
+    self.dev.synchronize()
+    with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): ctypes.memmove(mv_address(dest), src.cpu_view().addr, dest.nbytes)
 
   def _map(self, buf:HCQBuffer) -> HCQBuffer: # a mapping lives on the opaque, like hcq1: the lru hands the same one to many Buffers
     if self.dev not in buf.mapped_devs:

tinygrad/runtime/support/system.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
-import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, itertools, struct, socket
-import subprocess, time, enum, atexit
-from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten, pluralize
+import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, struct, socket
+from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, DEBUG, pluralize
 from tinygrad.runtime.autogen import libc, pci, vfio
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
@@ -81,8 +80,7 @@ class _System:
 
   @functools.cache
   def list_devices(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None):
-    if getenv("REMOTE", ""): return [(functools.partial(RemotePCIDevice,sock=s), x) for s,x in RemotePCIDevice.remote_list(vendor,devices,base_class)]
-    return [(APLRemotePCIDevice if OSX else PCIDevice, x) for x in System.pci_scan_bus(vendor, devices, base_class)]
+    return [(PCIDevice, x) for x in System.pci_scan_bus(vendor, devices, base_class)]
 
   def pci_probe_device(self, device:str, dev_id:int, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None):
     try: cl, pcibus = (ds:=hcq_filter_visible_devices(self.list_devices(vendor, devices, base_class), device))[dev_id]
@@ -253,13 +251,12 @@ class PCIAllocationMeta: mapping:VirtMapping; has_cpu_mapping:bool; hMemory:int=
 class PCIIfaceBase:
   @property
   def peer_group(self) -> str: return getattr(self.pci_dev, 'peer_group', type(self.pci_dev).__name__)
-  def is_local(self) -> bool: return not isinstance(self.pci_dev, RemotePCIDevice)
   def is_bar_small(self) -> bool: return self.pci_dev.bar_info(self.vram_bar)[1] == (256 << 20)
 
   def __init__(self, dev, dev_id, vendor, devices:tuple[tuple[int, tuple[int, ...]], ...], vram_bar, va_start, va_size,
                dev_impl_t, base_class:int|None=None):
     self.pci_dev = System.pci_probe_device(dn:=dev.__class__.__name__[:-6], dev_id, vendor, devices, base_class=base_class)
-    if self.is_local(): System.reserve_va(va_start, va_size)
+    System.reserve_va(va_start, va_size)
     with contextlib.suppress(Exception): self.pci_dev.resize_bar(vram_bar)
     self.dev_impl = dev_impl_t(self.pci_dev)
     self.dev, self.vram_bar, self.count = dev, vram_bar, len(hcq_filter_visible_devices(System.list_devices(vendor, devices, base_class), dn))
@@ -283,15 +280,13 @@ class PCIIfaceBase:
   def free(self, b:HCQBuffer):
     if b.owner != self.dev: self.dev.iface.dev_impl.mm.unmap_range(b.va_addr, round_up(b.size, 0x1000))
     if b.owner == self.dev and b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
-    if b.owner == self.dev and self.is_local() and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
+    if b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
 
   def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
     return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
 
   def map(self, b:HCQBuffer):
     if b.owner is not None and b.owner._is_cpu():
-      if not self.is_local(): raise RuntimeError(f"P2P mapping not supported for remote devices: {b.owner} -> {self.dev}")
-
       System.lock_memory(int(b.va_addr), b.size)
       paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(int(b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
       snooped, uncached = True, True
@@ -305,143 +300,3 @@ class PCIIfaceBase:
 
     self.dev_impl.mm.map_range(int(b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
     return HCQBuffer(b.va_addr, b.size, meta=b.meta, owner=b.owner)
-
-# *** Remote PCI Devices
-
-class RemoteCmd(enum.IntEnum):
-  PROBE,MAP_BAR,MAP_SYSMEM_FD,CFG_READ,CFG_WRITE,RESET,MMIO_READ,MMIO_WRITE,MAP_SYSMEM,SYSMEM_READ,SYSMEM_WRITE,RESIZE_BAR,PING = range(13)
-
-class RemoteMMIOInterface(MMIOInterface):
-  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0, rd_cmd=RemoteCmd.MMIO_READ, wr_cmd=RemoteCmd.MMIO_WRITE):
-    self.dev, self.residx, self.nbytes, self.fmt, self.off, self.el_sz = dev, residx, nbytes, fmt, off, struct.calcsize(fmt)
-    self.rd_cmd, self.wr_cmd = rd_cmd, wr_cmd
-
-  def __getitem__(self, index):
-    sl = index if isinstance(index, slice) else slice(index, index + 1)
-    start, stop = (sl.start or 0) * self.el_sz, (sl.stop or len(self)) * self.el_sz
-    data = self.dev._bulk_read(self.rd_cmd, self.residx, self.off + start, stop - start)
-    result = data if self.fmt == 'B' else list(struct.unpack(f'<{(stop - start) // self.el_sz}{self.fmt}', data))
-    return result if isinstance(index, slice) else result[0]
-
-  def __setitem__(self, index, val):
-    start = (index.start or 0) * self.el_sz if isinstance(index, slice) else index * self.el_sz
-    data = (val if self.fmt == 'B' else struct.pack(f'<{len(val)}{self.fmt}', *val)) if isinstance(index, slice) else struct.pack(f'<{self.fmt}', val)
-    self.dev._bulk_write(self.wr_cmd, self.residx, self.off + start, data)
-
-  def view(self, offset:int=0, size:int|None=None, fmt=None):
-    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset, self.rd_cmd, self.wr_cmd)
-
-class RemotePCIDevice(PCIDevice):
-  _bulk_sent:int = 0
-  _bulk_recv:int = 0
-  _rpc_count:int = 0
-  _start_time:float = 0.0
-
-  @staticmethod
-  @functools.cache
-  def remote_sock(host:str, port:int) -> socket.socket:
-    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-    sock.settimeout(getenv("REMOTE_TIMEOUT", 3))
-    sock.connect((host, port))
-    sock.settimeout(None)
-    if DEBUG >= 1 and RemotePCIDevice._start_time == 0.0:
-      RemotePCIDevice._start_time = time.perf_counter()
-      def _print_stats():
-        dt = time.perf_counter() - RemotePCIDevice._start_time
-        sent_mb, recv_mb = RemotePCIDevice._bulk_sent / 1e6, RemotePCIDevice._bulk_recv / 1e6
-        print(f"remote: sent {sent_mb:,.2f} MB ({sent_mb/dt:,.2f} MB/s), recv {recv_mb:,.2f} MB ({recv_mb/dt:,.2f} MB/s), "
-              f"{RemotePCIDevice._rpc_count:,} roundtrips in {dt:.2f}s")
-      atexit.register(_print_stats)
-    return sock
-
-  @staticmethod
-  @functools.cache
-  def remote_list(vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None) -> list[tuple[socket.socket, str]]:
-    payload = array.array('I', itertools.chain.from_iterable((m, d) for m, ds in devices for d in ds)).tobytes()
-    def q(r:str) -> list[tuple[socket.socket, str]]:
-      sock = RemotePCIDevice.remote_sock((host:=r.strip().split(":")[0]), (port:=int(r.strip().split(":")[1]) if ":" in r else 6667))
-      data_len, _, _, _ = RemotePCIDevice._rpc(sock, 0, RemoteCmd.PROBE, base_class or 0, len(payload), vendor, payload=payload)
-      return [(sock, f"remote:{host}:{port}:{d}") for d in RemotePCIDevice._recvall(sock, data_len).decode().split('\n')]
-    return flatten([q(r) for r in getenv("REMOTE", "").split(",") if r.strip()])
-
-  @staticmethod
-  def _recvall(sock:socket.socket, n:int) -> bytes:
-    data = b''
-    while len(data) < n and (chunk:=sock.recv(n - len(data))): data += chunk
-    if len(data) < n: raise RuntimeError("Connection closed")
-    return data
-
-  @staticmethod
-  def _rpc(sock:socket.socket, dev_id:int, cmd:int, *args:int, bar:int=0, readout_size:int=0, payload:bytes=b'', has_fd=False):
-    sock.sendall(struct.pack('<BIIQQQ', cmd, dev_id, bar, *(*args, 0, 0, 0)[:3]) + payload)
-    if has_fd:
-      msg, anc, _, _ = sock.recvmsg(17, socket.CMSG_LEN(4))
-      fd = struct.unpack('<i', anc[0][2][:4])[0]
-    else: msg, fd = RemotePCIDevice._recvall(sock, 17), None
-    if (resp:=struct.unpack('<BQQ', msg))[0] != 0:
-      raise RuntimeError(f"RPC failed: {RemotePCIDevice._recvall(sock, resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
-    RemotePCIDevice._rpc_count += 1
-    return (resp[1], resp[2]) + ((RemotePCIDevice._recvall(sock, readout_size) if readout_size > 0 else None),) + (fd,)
-
-  def __init__(self, devpref:str, pcibus:str, sock:socket.socket):
-    self.sock, self.pcibus, self.dev_id = sock, pcibus, int(pcibus.split(':')[-1]) if ':' in pcibus else 0
-    self.peer_group = sock.getpeername()[0]
-    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
-
-    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
-
-  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes:
-    RemotePCIDevice._bulk_recv += size
-    return unwrap(self._rpc(self.sock, self.dev_id, cmd, offset, size, bar=idx, readout_size=size)[2])
-  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes):
-    RemotePCIDevice._bulk_sent += len(data)
-    self.sock.sendall(struct.pack('<BIIQQQ', cmd, self.dev_id, idx, offset, len(data), 0) + data)
-
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
-    paddrs_len, handle, _, _ = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM, size, int(contiguous))
-    paddrs = list(struct.unpack(f'<{paddrs_len // 8}Q', self._recvall(self.sock, paddrs_len)))
-    return RemoteMMIOInterface(self, handle, size, fmt='B', rd_cmd=RemoteCmd.SYSMEM_READ, wr_cmd=RemoteCmd.SYSMEM_WRITE), paddrs
-
-  def reset(self): self._rpc(self.sock, self.dev_id, RemoteCmd.RESET)
-  def read_config(self, offset:int, size:int): return self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_READ, offset, size)[0]
-  def write_config(self, offset:int, value:int, size:int): self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_WRITE, offset, size, value)
-
-  @functools.cache
-  def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_BAR, bar=bar_idx)[:2]
-  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
-    return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt)
-  def resize_bar(self, bar_idx:int): self._rpc(self.sock, self.dev_id, RemoteCmd.RESIZE_BAR, bar=bar_idx)
-
-class APLRemotePCIDevice(RemotePCIDevice):
-  APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"
-
-  @classmethod
-  def ensure_app(cls):
-    commit = "c0d024f9ff0e1dc8fdf217f255da7101d91e8323"
-    app_name = f"TinyGPU_{commit}.zip"
-    if (_ensure_downloads_dir() / app_name).is_file() and os.path.exists(cls.APP_PATH): return
-    print("Downloading TinyGPU.app...")
-    with contextlib.suppress(RuntimeError): system("pkill -f TinyGPU")
-    system(f"ditto -xk {fetch(f'https://github.com/tinygrad/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications")
-    print(system(f"{cls.APP_PATH} install"))
-
-  def __init__(self, devpref:str, pcibus:str):
-    self.ensure_app()
-    sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    for i in range(100):
-      with contextlib.suppress(ConnectionRefusedError, FileNotFoundError):
-        sock.connect(sock_path)
-        break
-      if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-      time.sleep(0.05)
-    else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
-    super().__init__(devpref, "usb4", sock=sock)
-
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
-    mapped_size, _, _, fd = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM_FD, size, int(contiguous), has_fd=True)
-    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
-
-    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
-    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
-    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

tinygrad/schedule/__init__.py

@@ -1,8 +1,8 @@
-import time, inspect
+import time, inspect, dataclasses
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo
 from tinygrad.uop.spec import type_verify, spec_tensor
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, dedup
 
 # **** schedule linearizer
 
@@ -11,67 +11,60 @@ def _unwrap_src(s: UOp) -> UOp:
   while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: s = s.src[0]
   return s
 
-# a buffer state is AFTER | BUFFER | PARAM. MSELECT/MSTACK join per-device states
+# unwrap per-device buffer arguments without dropping their ordering dependencies
 def _states(s: UOp) -> list[UOp]:
   s = _unwrap_src(s)
   if s.op in {Ops.MSELECT, Ops.MSTACK}: return [st for ss in s.src for st in _states(ss)]
   assert s.op in {Ops.AFTER, Ops.BUFFER, Ops.PARAM}, f"input to kernel must resolve to a buffer state, not {s.op}"
   return [s]
 
-def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
-  kernels, remaining = partition(after.src[1:], lambda s: s.op in {Ops.CALL, Ops.END})
-  deps, remaining = partition(remaining, lambda s: s.op is Ops.AFTER)
-  if invalid := [s for s in remaining if s.op is not Ops.STORE]:
-    raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
-  return tuple(kernels), tuple(deps)
-
 def create_schedule(sched_sink:UOp) -> UOp:
   with cpu_profile(TracingKey("toposort sched_sink")):
-    # build kernel dependency graph: edges from producer kernel to consumer kernels
+    afters = [u for u in sched_sink.toposort(gate_kernel_sink) if u.op is Ops.AFTER]
+    kernels = dict.fromkeys(k for u in afters for k in u.src[1:] if k.op in {Ops.CALL, Ops.END})
+    dependencies: dict[UOp, set[UOp]] = {}
+    writes: dict[UOp, set[UOp]] = {}
+    reads: list[tuple[UOp, UOp]] = []
+    ancestors: dict[UOp, set[UOp]] = {}
+    for k in kernels:
+      call = k.src[0] if k.op is Ops.END else k
+      states = [st for s in call.src[1:] for st in _states(s)]
+      for st in states:
+        if st not in ancestors: ancestors[st] = kernels.keys() & st.toposort(enter_calls=False).keys()
+      # AFTER supplies ordering dependencies, not evidence that its returned buffer was written.
+      dependencies[k] = set().union(*(ancestors[st] for st in states))
+      read_args, write_args = call.call_access()
+      reads += [(k, st) for s in read_args for st in _states(s)]
+      for s in write_args:
+        for st in _states(s): writes.setdefault(st.buf_uop, set()).add(k)
+    for u in afters:
+      for dep in (s for s in u.src[1:] if s.op is Ops.AFTER):
+        for k in (s for s in u.src[1:] if s in kernels):
+          dependencies[k].update(kernels.keys() & dep.toposort(enter_calls=False).keys() - {k})
+    # Tensor reads require the contents preceding writes absent from their argument ancestry (not an AFTER property).
+    for k, st in reads:
+      for writer in writes.get(st.buf_uop, set()):
+        if writer is not k and writer not in ancestors[st]: dependencies[writer].add(k)
     children: dict[UOp, list[UOp]] = {}
-    in_degree: dict[UOp, int] = {}
-    writes: dict[UOp, list[tuple[UOp, UOp, tuple[UOp, ...]]]] = {}  # buffer -> (AFTER, prior state, new kernels)
-    reads: list[tuple[UOp, UOp, UOp]] = []  # (reader AFTER, reader kernel, buffer state read)
-    for u in sched_sink.toposort(gate_kernel_sink):
-      if u.op is not Ops.AFTER: continue
-      kernels, after_deps = _split_after(u)
-      prev_state = _unwrap_src(u.src[0])
-      prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set()
-      writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels)))
-      for k in kernels:
-        in_degree.setdefault(k, 0)
-        if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
-        kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
-        read_states = [st for s in kernel_deps for st in _states(s)]
-        reads += [(u, k, st) for st in read_states]
-        # RAW deps: a kernel runs after the kernels that produced the states it reads or joins
-        for st in read_states + [st for s in after_deps for st in _states(s)]:
-          if st.op is Ops.AFTER:
-            for t in _split_after(st)[0]:
-              children.setdefault(t, []).append(k)
-              in_degree[k] += 1
-    # WAR deps: a kernel reading buffer state S must run before another write that supersedes S. an AFTER only
-    # supersedes its immediate prior state; join members already present in that prior state are ordering deps, not writes
-    for u, k, s in reads:
-      for a, prev_state, write_kernels in writes.get(s.buf_uop, []):
-        if a is u or prev_state is not s: continue
-        for t in write_kernels:
-          if t is not k and t not in k.backward_slice:
-            children.setdefault(k, []).append(t)
-            in_degree[t] += 1
+    in_degree = {k:len(deps) for k,deps in dependencies.items()}
+    for k, deps in dependencies.items():
+      for p in deps: children.setdefault(p, []).append(k)
 
   with cpu_profile(TracingKey("linearize schedule")):
     queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
     linearized: list[UOp] = []
     while len(queue):
       rk = queue.popleft()
-      if rk.op is Ops.LINEAR:
-        linearized.extend(rk.src)
-      else:
-        k = rk.src[0] if rk.op is Ops.END else rk
-        assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if not s.is_bound_var)
-        linearized.append(k.src[0].call(*buf_uops))
+      k = rk.src[0] if rk.op is Ops.END else rk
+      assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
+      buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if not s.is_bound_var)
+      body = k.src[0]
+      # Storage aliases may share a parameter now that their dependencies are in the schedule.
+      if body.op is Ops.SINK and len(set(buf_uops)) != len(buf_uops):
+        params = {p for p in body.toposort(enter_calls=False) if p.op is Ops.PARAM and p.arg.slot >= 0}
+        body = body.substitute({p:q for p in params
+                               if (q:=p.replace(arg=dataclasses.replace(p.arg, slot=buf_uops.index(buf_uops[p.arg.slot])))) in params})
+      linearized.append(body.call(*buf_uops))
       for x in children.get(rk, []):
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
@@ -120,8 +113,8 @@ schedule_cache: dict[bytes, UOp] = {}
 def lower_sink_to_linear(call:UOp) -> UOp|None:
   function = call.src[0]
   if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None
-  # value calls (with RETURNED outputs) are inlined positionally during prepare: their bodies are not programs to schedule
-  if any(x.unsharded_base.is_unbound for x in call.src[1:]): return None
+  # value calls (with unbound outputs) are inlined positionally during prepare: their bodies are not programs to schedule
+  if call.has_unbound_outputs: return None
   st = time.perf_counter()
   cache_key = function.key
   if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
@@ -188,6 +181,8 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
 
   # this recursively resolves the linear_call and allocates buffers
   linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
+  for call in linear.src:
+    if call.src[0].op is Ops.PROGRAM: call.call_access()
 
   # create copies
   linear = graph_rewrite(linear, pm_copy_from_store, name="create COPY kernels for SDMA")

tinygrad/schedule/multi.py

@@ -294,11 +294,11 @@ multi_pm = PatternMatcher([
     lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.arg, multi.src[1:])),
 
   # rewrite value-producing calls explicitly for UNSHARD
-  (UPat(Ops.CALL, name="call"), lambda call: rewrite_into_function(call) if call.num_returned else None),
+  (UPat(Ops.CALL, name="call"), lambda call: rewrite_into_function(call) if call.has_unbound_outputs else None),
   (UPat((Ops.CALL, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
   # just strip the UNSHARD from non-value-producing CALLs (custom kernels, etc.) — value-producing CALLs are handled by rewrite_into_function
   (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.UNSHARD])), lambda root:
-    UOp(root.op, src=tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), arg=root.arg) if root.num_returned == 0 else None),
+    UOp(root.op, src=tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), arg=root.arg) if not root.has_unbound_outputs else None),
   (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
         src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi),
   # STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile)