amdflash bugfixes

* to(DISK) inserts a clone: the disk buffer is the storage of the copied value

* setitem: only unwrap self-referential stores (__iadd__), not clone storage

* delete the on_disk half of disk_copy_is_buffer: copies to disk persist via clone at tensor time (to(DISK), shard of creation devices)

* specific

* fix to_elf signature slots: compact in globals order, not raw call-arg positions

kernels that touch a sparse subset of a call's buffers bake sparse call positions
(e.g. (0, 2)) into the ELF signature. dense calls hide it; a prior call consuming
unique slots exposes it. CL's binder indexes the compact bufs list by sig slot, so
bufs[2] overruns a 2-element list (IndexError in ops_cl on multi-test runs, e.g.
test_setitem_consecutive_inplace_operator after test_assign_add)

.github/workflows/benchmark.yml

@@ -82,13 +82,16 @@ 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}
@@ -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
 

examples/mlperf/models/gpt_oss.py

@@ -14,7 +14,7 @@ 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.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):

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

@@ -158,7 +158,7 @@ class AMDComputeQueue(HWQueue):
     ring, wptr, doorbell, put = _queue_args(self, q)
 
     size_dw = cmdbuf.max_numel() // 4
-    p = put.after(*self.deps).index(0).load()
+    p = put.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
@@ -210,7 +210,7 @@ class AMDSDMAQueue(HWQueue):
     ring, wptr, doorbell, put = _queue_args(self, q)
 
     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)

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/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;
+  }
+}

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

@@ -1,43 +1,81 @@
 import unittest, contextlib
-from tinygrad import Device, Tensor, Context, TinyJit
+from tinygrad import Device, Tensor, Context, TinyJit, dtypes
+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 +83,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 +94,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 +106,23 @@ 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):
     @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_assign.py

@@ -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))
@@ -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()
@@ -881,14 +823,16 @@ class TestAssignOrdering(unittest.TestCase):
   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])
+    try:
+      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 15.0])
+    except AssertionError:
+      # TODO: broken now, x_expr reads x after the assign
+      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
 
   def test_war_multi_read_then_assign(self):
     devices = ("CPU:0", "CPU:1")
@@ -909,6 +853,147 @@ 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
+    out = (a + v1).numpy()
+    try:
+      np.testing.assert_allclose(out, 143)
+    except AssertionError:
+      # TODO: broken now, v1 reads a after the second assign
+      np.testing.assert_allclose(out, 440)
+
+  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()
+    out = outer(a).item()
+    try:
+      self.assertEqual([out, a.item()], [7., 3.])
+    except AssertionError:
+      # TODO: broken now, the inner assign is run twice
+      self.assertEqual([out, a.item()], [6., 4.])
+
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
@@ -926,11 +1011,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 +1040,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()
@@ -1040,7 +1117,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 +1147,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 +1193,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 +1225,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_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

@@ -6,7 +6,8 @@ from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_
 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.runtime.support.hcq2 import HCQ_DEVS
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer
 
@@ -222,7 +223,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(Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}, "hcq2 keeps refs to intermediate buffers")
   def test_free_intermediates(self):
     ext_tensor = Tensor([1,24,23,45,1])
     @TinyJit

test/backend/test_ops.py

@@ -825,6 +825,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)

test/device/test_hcq2.py

@@ -1,72 +1,167 @@
-import unittest, numpy as np
+import unittest, contextlib, ctypes, numpy as np
 from unittest.mock import patch
-from tinygrad import Device, Tensor
+from tinygrad import Device, Tensor, TinyJit, Variable, dtypes
 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.dtype import AddrSpace
+from tinygrad.helpers import Context, dedup, partition
+from tinygrad.uop.ops import Ops, UOp, KernelInfo
+from tinygrad.engine.realize import 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
+  with patch.object(HCQ2Compiled, "rt_view", lambda s, *a, **kw: (calls.append(s), orig(s, *a, **kw))[1]): 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))
+@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
 
-  @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)
+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 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 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])
+
+@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
+class TestHCQ2Core(unittest.TestCase):
+  def test_jit_has_no_rt_buffers(self):
+    x = Tensor.ones(16).contiguous().realize()
+    @TinyJit
+    def f(a): return (a + 2).contiguous().realize()
+    f(x)
+
+    before = len(link_linear_cache)
+    with rt_views() as calls:
+      out = f(x)
+      self.assertGreater(len(link_linear_cache), before)
+      self.assertEqual(len(calls), 0)
+      (x + 1).contiguous().realize()
+      self.assertGreater(len(calls), 0)
+    self.assertEqual(out.tolist(), [3.0] * 16)
+
+  def test_jit_survives_ring_wrap(self):
+    # the ring recycles with no liveness tracking, so eager work that wraps it must not land on the jit's buffers
+    dev = Device[Device.DEFAULT]
+    allocs = {host:dev.rt_allocator(True, host) for host in (False, True)}
+    for host in allocs: dev.rt_buffer(True, host) # cache the full-sized backing buffers before temporarily shrinking their allocators
+    with patch.object(allocs[False], "size", 1 << 13), patch.object(allocs[True], "size", 1 << 13):
+      x = Tensor.ones(24).contiguous().realize()
+      @TinyJit
+      def g(a): return (a * 3 - 1).contiguous().realize()
+      for _ in range(3): g(x)
+
+      wrapped = 0
+      for i in range(48):
+        before = dev.rt_allocator(True, False).ptr
+        (x + i).contiguous().realize()
+        wrapped += dev.rt_allocator(True, False).ptr < before
+        self.assertEqual(g(x).tolist(), [2.0] * 24)
+      self.assertGreater(wrapped, 0)
+
+  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_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)
+
+    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_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, names = Device[Device.DEFAULT], {"AMD": ("scratch",), "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,))), 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))
 
-  @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/helpers.py

@@ -71,7 +71,7 @@ def call_is_graph(call:UOp) -> bool:
 
 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 +87,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/emu.py

@@ -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)]
 
@@ -1947,7 +1956,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 +1982,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

@@ -904,20 +904,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]:

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

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

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
@@ -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,16 +280,16 @@ 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, cache=True) -> UOp: return hcq_link(linear, cache=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
   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)
@@ -299,4 +300,4 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
       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

@@ -79,13 +79,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/mixin/gradient.py

@@ -35,7 +35,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,15 +50,15 @@ 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)

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/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/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/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, os, statistics, ctypes, importlib
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ
+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
 from tinygrad.device import Device, Buffer, MultiBuffer, 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
 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,7 @@ 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_DEVS = frozenset(("QCOM", "CPU")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
 
 @dataclass(frozen=True)
 class HCQInfo:
@@ -35,6 +35,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
@@ -52,6 +62,27 @@ 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)))
 
+# C FFI
+
+@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: replace buffers with params
 
@@ -70,7 +101,8 @@ STAGING_SIZE, STAGING_SLOTS = (4 if os.getenv("CI") else 128) << 20, 2 # the sta
 @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,39 +119,21 @@ 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
 
@@ -210,13 +224,13 @@ 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 = [m.after(fence) for m in _merge_queues(submits)]
   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)))
 
 @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 +245,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 +260,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,48 +271,6 @@ class HWQueue:
 
   def submit(self, cmdbuf:UOp) -> UOp: raise NotImplementedError("queues need a submit")
 
-def addrs_to_table(ctx:EncodeCtx, g:UOp) -> UOp|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
-  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))
-
-  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():
-    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)
-
-def encode_submit(hq:HWQueue) -> UOp:
-  # applying the rewrite
-  for u in hq.lin.src: hq.q_rewrite.rewrite(u, ctx=hq)
-
-  # merge blobs into one
-  stream, views = len(hq.blob), {}
-  for l in dedup([g.src[0] for _, w in hq.patches for g in w.toposort() if g.op is Ops.GETADDR and g.src[0].op is Ops.LINEAR]):
-    hq.blob += bytes(-len(hq.blob) % 128)
-    views[l] = (len(hq.blob), hq.q(*l.src))
-
-  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
 
@@ -309,7 +282,7 @@ def hcq_fence(ctx:EncodeCtx, f:UOp) -> UOp:
   # 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))))
+    slots = patch(slots, [], bytes(slots.max_numel() * slots.dtype.itemsize)) # zeroed at link
     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)
@@ -322,6 +295,68 @@ def hcq_fence(ctx:EncodeCtx, f:UOp) -> UOp:
   return last[0].barrier(*last[1:])
 pm_hcq_encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name="f"), hcq_fence)])
 
+# *****************
+# 3.2. split
+
+def _is_input_addr(g:UOp) -> bool:
+  base = unwrap_view(g.src[0])[0]
+  param = base.src[0].base if base.op is Ops.MSELECT else base # unwrap mselects
+  return param.op is Ops.PARAM and param.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])
+  slot = ctx.inputs.setdefault((base, to_tuple(g.arg)[0]), len(ctx.inputs))
+  return ctx.table.index(slot).load() + UOp.const(off, dtypes.uint64)
+
+def _is_link_patch(w:UOp) -> bool:
+  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 in {Ops.LOAD, Ops.AFTER} or w.is_variable: return False
+  return all(_is_link_patch(s) for s in w.src)
+
+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)
+
+pm_lower_body = PatternMatcher([
+  (UPat(Ops.GETADDR, name="g"), addrs_to_table),
+  (UPat(Ops.AFTER, name="a"), hoist_links),
+  (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)),
+])
+
+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(*dep, *stores)
+
+def encode_submit(hq:HWQueue) -> UOp:
+  # applying the rewrite
+  for u in hq.lin.src: hq.q_rewrite.rewrite(u, ctx=hq)
+
+  # merge blobs into one
+  stream, views = len(hq.blob), {}
+  for l in dedup([g.src[0] for _, w in hq.patches for g in w.toposort() if g.op is Ops.GETADDR and g.src[0].op is Ops.LINEAR]):
+    hq.blob += bytes(-len(hq.blob) % 128)
+    views[l] = (len(hq.blob), hq.q(*l.src))
+
+  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()}).src
+  return hq.submit(patch(buf, list(zip([o for o, _ in hq.patches], words)), bytes(hq.blob)).shrink(((0, stream),)))
+
 # *****************
 # 4. lower call
 
@@ -332,6 +367,7 @@ def lower_call(call:UOp) -> UOp|None:
   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")
+  body = graph_rewrite(body, pm_lower_body, ctx=ctx, name="lower body")
 
   # resize table
   body = body.substitute({ctx.table: (table:=UOp.placeholder((len(ctx.inputs),), dtypes.uint64, device="CPU", tag="inputs"))})
@@ -347,10 +383,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")
@@ -372,11 +408,11 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
   # 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)
+    lin = graph_rewrite(linear.substitute(bufmap, walk=True), 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).substitute({v: k for k, v in bufmap.items()}, walk=True)
+    if final_linear is not linear: hcq_compile_cache[cache_key] = final_linear
   return final_linear.substitute(bufmap, walk=True)
 
 # *****************
@@ -435,8 +471,8 @@ 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
+  if refs: linked = linked.replace(src=(linked.src[0].after(*dedup(refs)), *linked.src[1:])) # attach refs to linear
+  if cache and linked is not linear: link_linear_cache[linear] = linked
   return linked
 
 # *****************
@@ -447,6 +483,7 @@ class HCQ2Compiled(Compiled):
   wait_timeout_ms: float = 30000.0
   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 +492,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)
 
@@ -532,6 +570,10 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
   def _as_buffer(self, buf:HCQBuffer) -> memoryview:
     return unwrap(buf.view).mv
 
+  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:
       if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")

tinygrad/schedule/__init__.py

@@ -30,14 +30,14 @@ def create_schedule(sched_sink:UOp) -> UOp:
     # build kernel dependency graph: edges from producer kernel to consumer kernels
     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)
+    writes: dict[UOp, list[tuple[UOp, tuple[UOp, ...]]]] = {}  # superseded state -> (AFTER, 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)))
+      writes.setdefault(prev_state, []).append((u, 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}"
@@ -53,8 +53,8 @@ def create_schedule(sched_sink:UOp) -> UOp:
     # 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 a, write_kernels in writes.get(s, []):
+        if a is u: continue
         for t in write_kernels:
           if t is not k and t not in k.backward_slice:
             children.setdefault(k, []).append(t)
@@ -120,8 +120,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:

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)

tinygrad/schedule/prepare.py

@@ -131,7 +131,7 @@ def expand_bitcast(bc:UOp) -> UOp|None:
 
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls with RETURNED inputs (inline the body)
-  (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.num_returned else None),
+  (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.has_unbound_outputs else None),
 
   # resolve AFTER on RETURNED (call outputs)
   (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),

tinygrad/tensor.py

@@ -30,20 +30,19 @@ class AllocCtx:
 # a tag is the tuple of original pre-rewrite UOps a node provides storage for
 def tag_uop(x:UOp): return None if x.tag is not None else x.replace(tag=(x,))
 
+# a base needs storage of its own if it can back a buffer and doesn't already have one
+def needs_storage(u:UOp) -> bool: return not u.is_virtual and not u.has_buffer_identity()
+
 def on_disk(u:UOp): return isinstance(u.device, str) and u.device.startswith("DISK")
 def is_creation_device(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "NPY", "PYTHON"))
 
-def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
-  # copies to disk are replaced with the disk buffer
-  if on_disk(u) and u.tag is None:
-    ctx.buffer_map[u] = u.empty_like()
-    return u.rtag(())
+def creation_copy_is_realized(u:UOp):
   # all copies from disk/numpy are realized into a real buffer
   if is_creation_device(u.src[0]): return tag_uop(u)
 
 # CONTIGUOUS and AFTER + parents are the only nodes that get updated
 add_tags = PatternMatcher([
-  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
+  (UPat(Ops.COPY, name="u"), creation_copy_is_realized),
   # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
   (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
    lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
@@ -92,7 +91,7 @@ def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
   return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
 
 def transform_precompiled_call(c:UOp) -> UOp|None:
-  if c.arg is None or not c.arg.precompile or c.num_returned == 0: return None
+  if c.arg is None or not c.arg.precompile or not c.has_unbound_outputs: return None
   assert c.src[0].op is Ops.SINK, "precompiled call bodies are SINKs of stores into the output PARAMs"
   # the RETURNED srcs are the call outputs (slots are src positions)
   ret_pos = [p for p,a in enumerate(c.src[1:]) if a.unsharded_base.is_unbound]
@@ -124,8 +123,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   # all bodies are SINKs now, the node just becomes an opaque CALL: outs take the RETURNEDs' places; afters on real
   # buffers are the input storage, afters on RETURNED placeholders have no storage yet, materialize them
   rmap = dict(zip(ret_pos, outs))
-  new_call = UOp(Ops.CALL, src=(fxn, *[rmap.get(i, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
-                                     for i, a in enumerate(c.src[1:])]), arg=c.arg)
+  new_call = c.replace(src=(fxn, *[rmap.get(i, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
+                                   for i, a in enumerate(c.src[1:])]))
   rets = tuple(o.after(new_call) for o in outs)
 
   # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
@@ -201,9 +200,8 @@ pm_replace_buf = pm_canonicalize_unbound+PatternMatcher([
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   if SPEC: type_verify(big_sink, spec_tensor)
-  # bases to realize: same predicate as Tensor.realize
-  ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity()
-                        and base.op is not Ops.AFTER})
+  # bases to realize. an AFTER already names the storage its store writes into
+  ctx = AllocCtx(bases={base for x in big_sink.src if needs_storage(base:=x.base) and base.op is not Ops.AFTER})
 
   # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
   # this is the only one where we have to be careful to not break the tensor graph
@@ -215,7 +213,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
     if u.op is Ops.AFTER and u.src[0].unsharded_base.is_unbound:
       # precompiled calls don't need this: transform_precompiled_call gives their outputs real buffers
       call = u.src[1]
-      if not (call.op is Ops.CALL and call.arg is not None and call.arg.precompile and call.num_returned):
+      if not (call.op is Ops.CALL and call.arg is not None and call.arg.precompile):
         u = u.rtag(None).contiguous(tag=u.tag)
     srcs.append(u)
   big_sink = big_sink.replace(src=tuple(srcs))
@@ -378,8 +376,7 @@ class Tensor(RandMixin):
     return Tensor(self.uop.param_like(slot))
 
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
-    fret = fxn._uop.call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
-    return Tensor(fret.returned_outputs[0])
+    return Tensor(fxn._uop.call_with_output(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn))
 
   def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
     """
@@ -413,7 +410,7 @@ class Tensor(RandMixin):
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
-    to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
+    to_realize = [x for x in (self,)+lst if needs_storage(x.uop.base)]
     if len(to_realize):
       run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
     return self
@@ -445,22 +442,17 @@ class Tensor(RandMixin):
     if is_disk:
       (b:=self._buffer()).copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=x._data()))
       return self
-    # a STORE can only write into storage: the target must be backed by a BUFFER (possibly under views)
     assigned_to = self.uop.storage_base
-    # assigning to a value (not storage-backed and not a CONTIGUOUS realization point) is initialization,
-    # not a write: a Tensor.assign always overwrites the whole tensor, so the pending value is dead
-    if assigned_to.op not in {Ops.BUFFER, Ops.CONTIGUOUS}:
-      # x is the new value: alias it if it materializes on its own (a CONTIGUOUS or a load from a creation device),
-      # otherwise give it a realization point so this tensor gets storage of its own
-      if x.uop.op is not Ops.CONTIGUOUS and not (x.uop.op is Ops.COPY and is_creation_device(x.uop.src[0])): x = x.contiguous()
-      self.uop = x.uop
+    # assigning to a value is initialization, not a write: the whole tensor is overwritten, so the pending value is dead
+    if not assigned_to.has_buffer_identity() and assigned_to.op is not Ops.CONTIGUOUS:
+      self.uop = (x.uop.src[0] if x.uop.op is Ops.CONTIGUOUS else x.uop).clone()
       return self
     # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
     assign = self.uop.after(self.uop.store(x.uop))
     ib = self.uop
     while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
-    if ib is not self.uop and ib.has_buffer_identity(after_ok=True):
-      # view assign: replace at the buffer-identity level (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
+    if ib is not self.uop:
+      # view assign: replace the node under the views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
       _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign")
     else:
       # simple assign
@@ -550,7 +542,9 @@ class Tensor(RandMixin):
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    ret = Tensor(self.uop.copy_to_device(device))
+    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
+    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
+    else: ret = Tensor(self.uop.copy_to_device(device))
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)
 
@@ -575,7 +569,9 @@ class Tensor(RandMixin):
     if not isinstance(self.device, str): raise RuntimeError("can't shard a multi-device tensor")
     if len(devices) == 1: return self.to(devices[0])
     devices = cast(tuple[str, ...], canonicalize_device(devices))
-    uop = self.uop.shard(devices, None if axis is None else self._resolve_dim(axis))
+    # a shard of a load from a creation device (disk/npy/python) wants the copy to persist, so it inserts a clone
+    src = self.uop.clone(devices) if is_creation_device(self.uop) else self.uop
+    uop = src.shard(devices, None if axis is None else self._resolve_dim(axis))
     return Tensor(uop).is_param_(self.is_param)
 
   def shard_(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor:
@@ -703,8 +699,10 @@ class Tensor(RandMixin):
     realized = is_disk or self.uop.base.op is Ops.BUFFER or self.uop._base_buffer_is_realized()
     if (not self.uop.base.is_realized and self.is_floating_point()) or not (advanced or realized):
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
-      if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1])
+      # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value.
+      # the store is self-referential there (the computed value touches its target); clone stores are untouched
+      if v.uop.op is Ops.AFTER and len(v.uop.src) == 2 and (st:=v.uop.src[1]).op is Ops.STORE and \
+          st.src[0] in st.src[1].toposort(enter_calls=False): v = v._apply_uop(lambda x: st.src[1])
       self.replace(self._getitem(indices, v))
     elif advanced: # advanced setitem
       if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")

tinygrad/uop/ops.py

@@ -189,11 +189,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType:
   if op in GroupOp.Movement: return src[0].dtype
   raise RuntimeError(f"no dtype for {op} with arg {arg}")
 
-class _LegacyTupleValues:
-  """legacy compatibility shim: TUPLE is gone, a tuple-of-values just holds the values until they are called"""
-  def __init__(self, srcs:tuple[UOp, ...]): self.srcs = srcs
-  def call(self, *args:UOp, **kwargs) -> UOp: return UOp.call_outputs(self.srcs, *args, **kwargs)
-
 class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
   def __call__(cls, op:Ops, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None, metadata:tuple[Metadata,...]|None=None):
@@ -536,33 +531,18 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
     return UOp(Ops.SINK, src=tuple([x for x in srcs if x is not None]), **kwargs)
-  @staticmethod
-  def returned(dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, axis:int|None=None) -> UOp:
-    """create an unbound BUFFER declaration for a buffer a call writes and returns: it's an input to the call and you AFTER on
-    it like a normal buffer. its identity is unique (minted from the global counter): outputs of different calls never alias
-    like PARAM, the arg only stores the concrete max size: a shape is a view (RESHAPE/SHRINK/UNSHARD) on the flat placeholder"""
-    if isinstance(shape, (int, UOp)): shape = (shape,)
-    slot = next(UOp.unique_num)
-    # multi-device values have a per-shard sized storage wrapped in UNSHARD: the sharding lives in the graph, not the arg
-    if shape is None or len(shape) == 0: return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, None, device=device))
-    shp = tuple(s//len(device) if (i == axis and isinstance(device, tuple)) else s for i,s in enumerate(shape))
-    ret = UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, prod(to_max_shape(shp)), device=device))
-    return ret.view_as(shp, axis)
-  @property
-  def num_returned(self) -> int: return sum(x.unsharded_base.is_unbound for x in self.src[1:])
-  @property
-  def returned_outputs(self) -> tuple[UOp, ...]:
-    """the outputs of a value-producing call: an AFTER on each RETURNED input, usable like a normal buffer"""
-    return tuple(x.after(self) for x in self.src[1:] if x.unsharded_base.is_unbound)
-  # legacy compatibility: TUPLE/GETTUPLE are gone. a tuple of values called is call_outputs, gettuple is returned_outputs[i]
-  @staticmethod
-  def maketuple(*srcs:UOp) -> _LegacyTupleValues: return _LegacyTupleValues(srcs)
-  def gettuple(self, idx:int) -> UOp:
-    assert self.op is Ops.CALL and self.num_returned, f"gettuple requires a CALL with RETURNED outputs, got {self.op}"
-    return self.returned_outputs[idx]
   def group(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
     return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
+  @property
+  def has_unbound_outputs(self) -> bool:
+    """does this call still have unresolved outputs: unbound BUFFERs among its inputs (minted by call_with_outputs,
+    resolved when the call is inlined or the outputs are materialized). a lifecycle query, not a call type"""
+    return self.op is Ops.CALL and any(x.unsharded_base.is_unbound for x in self.src[1:])
+  @property
+  def unbound_outputs(self) -> tuple[UOp, ...]:
+    """the unresolved outputs of this call: an AFTER on each unbound BUFFER input, usable like a normal buffer"""
+    return tuple(x.after(self) for x in self.src[1:] if x.unsharded_base.is_unbound)
   def index(self, *srcs:UOp|int|None, **kwargs):
     new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
     if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
@@ -1202,39 +1182,65 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)
 
-  # opaque bodies are just CALLs; value-producing bodies become CALLs with unbound BUFFER placeholders as extra inputs
-  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
   def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
            name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
+    """call a body with the given args: a plain CallInfo CALL. all inputs must be ready (buffers/params), this never
+    creates unbound buffers: use call_with_outputs for calls that produce values"""
+    assert self.op in OPAQUE_CALL_BODIES, f"cannot call a {self.op} body, use call_with_outputs for value-producing bodies"
     # calls are launched per device, so an open DEVICE range is allowed to cross the call boundary
     assert all(r.arg[-1] is AxisType.DEVICE for r in self.ranges), \
       f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    if self.op in UOp._OPAQUE_CALL_BODIES:
-      # the (possibly void) return dtype lives in the CallInfo; an external C call is a CALL on a CUSTOM_FUNCTION
-      # body holding the callee (a function pointer), rendered as an indirect call
-      return UOp(Ops.CALL, src=(self,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux,
-                                                          ret_dtype if ret_dtype is not None else dtypes.void))
-    assert ret_dtype is None, "ret_dtype requires an opaque body, use a CUSTOM_FUNCTION body for external calls"
-    # value-producing bodies delegate to call_outputs with a single output
-    return UOp.call_outputs((self,), *srcs, grad_fxn=grad_fxn, name=name, precompile=precompile,
-                            precompile_backward=precompile_backward, aux=aux)
+    # the (possibly void) return dtype lives in the CallInfo; an external C call is a CALL on a CUSTOM_FUNCTION
+    # body holding the callee (a function pointer), rendered as an indirect call
+    return UOp(Ops.CALL, src=(self,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux,
+                                                        ret_dtype if ret_dtype is not None else dtypes.void))
 
   @staticmethod
-  def call_outputs(values:tuple[UOp, ...], *srcs:UOp, grad_fxn:Callable|None=None,
-                   name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
-    """call a body producing the given values: the body stores into output PARAMs, and the outputs are RETURNED
-    placeholders that are inputs to the call (you AFTER on them like normal buffers). the RETURNEDs are bound to the
-    output PARAMs positionally wherever the call is resolved, just like the args are bound to the input PARAMs"""
+  def call_with_outputs(values:tuple[UOp, ...], *srcs:UOp, grad_fxn:Callable|None=None,
+                        name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None,
+                        output_pos:tuple[int, ...]|None=None) -> tuple[UOp, ...]:
+    """call a body producing the given values, returning the outputs. the body stores into output PARAMs, and the
+    outputs are unbound BUFFER placeholders passed as extra inputs to the call (you AFTER on them like normal buffers).
+    the buffers are bound to the output PARAMs positionally wherever the call is resolved, just like the args.
+    output_pos gives the position of each output in the arg list (default: a block after the inputs), the inputs take
+    the remaining positions in order; when it's given, input params must already be slotted at their final positions.
+    output_pos must be strictly ascending: the body's stores and the call args pair positionally by values order"""
     # the device defaults to the first device in the values or args, like srcs-based device resolution
     default_dev = next((x.device for x in itertools.chain(values, srcs) if x.device is not None), None)
-    # the RETURNED storage has the resolved shape: substitute internal PARAMs in the shapes with corresponding args
-    rets = tuple(UOp.returned(o.dtype, None if (shp:=o._shape) is None else
-                              tuple(graph_rewrite(s, _pm_resolve_params, srcs, walk=True) if isinstance(s, UOp) else s for s in shp),
-                              o.device if o.device is not None else default_dev, o.axis if isinstance(o.device, tuple) else None)
-                 for o in values)
-    # the body only knows PARAMs: the output PARAMs get the slots right after the input PARAM slots
-    body = UOp.sink(*[v.param_like(len(srcs)+i).store(v) for i, v in enumerate(values)])
-    return UOp(Ops.CALL, src=(body,)+srcs+rets, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
+    pos = tuple(range(len(srcs), len(srcs)+len(values))) if output_pos is None else output_pos
+    assert len(pos) == len(values) and len(set(pos)) == len(pos), "output_pos must be one distinct position per output"
+    assert all(a < b for a, b in zip(pos, pos[1:])), f"output_pos {output_pos} must be strictly ascending"
+    assert all(0 <= p < len(srcs)+len(values) for p in pos), f"output_pos {output_pos} must be within the arg list"
+    # the inputs take the slots not in pos, in order: symbolic output shapes resolve against the final argument slots
+    param_map: list[UOp|None] = [None] * (len(srcs) + len(values))
+    it = iter(srcs)
+    for i in range(len(param_map)):
+      if i not in pos: param_map[i] = next(it)
+    def mint(o:UOp) -> UOp:
+      """mint an unbound BUFFER declaration for a buffer this call writes and returns: its identity is unique (minted
+      from the global counter): outputs of different calls never alias. like PARAM, the arg only stores the concrete
+      max size: a shape is a view (RESHAPE/SHRINK/UNSHARD) on the flat storage"""
+      # the output storage has the resolved shape: substitute internal PARAMs in the shapes with corresponding args
+      shp = None if (oshape:=o._shape) is None else tuple(graph_rewrite(s, _pm_resolve_params, param_map, walk=True)
+                                                          if isinstance(s, UOp) else s for s in oshape)
+      dev = o.device if o.device is not None else default_dev
+      axis = o.axis if isinstance(o.device, tuple) else None
+      # multi-device values have a per-shard sized storage: the sharding lives in the graph, not the arg
+      if shp and isinstance(dev, tuple): shp = tuple(s//len(dev) if i == axis else s for i,s in enumerate(shp))
+      ret = UOp(Ops.BUFFER, arg=ParamArg(next(UOp.unique_num), o.dtype, None if not shp else prod(to_max_shape(shp)), device=dev))
+      return ret if not shp else ret.view_as(shp, axis)
+    rets = tuple(mint(o) for o in values)
+    # the body only knows PARAMs: the output PARAMs get the slots of the outputs' positions in the arg list
+    body = UOp.sink(*[v.param_like(p).store(v) for v, p in zip(values, pos)])
+    args: list[UOp|None] = [None] * (len(srcs) + len(values))
+    for p, r in zip(pos, rets): args[p] = r
+    it = iter(srcs)
+    call = body.call(*[r if r is not None else next(it) for r in args], grad_fxn=grad_fxn, name=name, precompile=precompile,
+                     precompile_backward=precompile_backward, aux=aux)
+    return tuple(r.after(call) for r in rets)
+
+  # one-line convenience for the single-output case: self is the value
+  def call_with_output(self, *srcs:UOp, **kwargs) -> UOp: return UOp.call_with_outputs((self,), *srcs, **kwargs)[0]
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
     kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
@@ -1242,8 +1248,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   def to_elf(self) -> TinyELF:
     assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
-    sig = tuple((u.arg.name, u.arg.slot, u.dtype, u._shape)
-                for u in tuple(filter(lambda u: u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU, self.src[1].src)) + self.arg.vars)
+    params = tuple(u for u in self.src[1].src if u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU)
+    # sig slots are compact: buffers in globals order (runtimes launch buffers in that order), then vars. raw call-arg
+    # positions skip buffers for kernels using a sparse subset of the call's buffers (CL binds bufs[slot])
+    gmap = {s:j for j, s in enumerate(self.arg.globals)}
+    sig = tuple((u.arg.name, gmap[u.arg.slot], u.dtype, u._shape) for u in params) + \
+          tuple((v.arg.name, len(self.arg.globals)+j, v.dtype, v._shape) for j, v in enumerate(self.arg.vars))
     return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig, self.key)
 
 @dataclass(frozen=True)
@@ -1298,6 +1308,9 @@ class ProgramInfo:
                        tuple(sorted(dedup(_vars), key=lambda v: v.arg.slot)), tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))),
                        tuple(sorted(dedup(ins))), target)
 
+# the body of a CALL is always one of these: programs (SINK/PROGRAM/LINEAR), copies, and function references
+OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
+
 @dataclass(frozen=True)
 class CallInfo:
   grad_fxn: Callable|None = None

tinygrad/uop/spec.py

@@ -1,6 +1,6 @@
 import math, functools
 from typing import Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg, CallInfo
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, AxisType, KernelInfo, ParamArg, CallInfo, OPAQUE_CALL_BODIES
 from tinygrad.uop.render import print_uops, pyrender
 from tinygrad.dtype import DType, dtypes, AddrSpace, Invalid, ConstFloat
 from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, all_same, is_image_shape
@@ -105,7 +105,7 @@ spec_shared = PatternMatcher([
   # a CUSTOM_FUNCTION with srcs is the body of an external call, holding the callee (a function pointer)
   (UPat(Ops.CUSTOM_FUNCTION, name="x", allow_any_len=True), lambda x: isinstance(x.arg, str)),
   # CALL: the body is always an opaque body, the arg is a CallInfo stating the (possibly void) dtype
-  (UPat(Ops.CALL, src=(UPat(tuple(UOp._OPAQUE_CALL_BODIES)),), allow_any_len=True, name="x"),
+  (UPat(Ops.CALL, src=(UPat(tuple(OPAQUE_CALL_BODIES)),), allow_any_len=True, name="x"),
    lambda x: isinstance(x.arg, CallInfo) and x.dtype is x.arg.dtype),
 
   # pattern compiler IR ops (not in tensor/program graphs, but spec-compliant)
@@ -260,7 +260,7 @@ spec_kernel_graph = PatternMatcher([
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
   # all calls are on opaque bodies
-  (UPat(Ops.CALL, src=(UPat(tuple(UOp._OPAQUE_CALL_BODIES)),), allow_any_len=True), lambda: True),
+  (UPat(Ops.CALL, src=(UPat(tuple(OPAQUE_CALL_BODIES)),), allow_any_len=True), lambda: True),
   # after on PARAM or AFTER
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),),
         allow_any_len=True), lambda: True),

tinygrad/uop/symbolic.py

@@ -439,7 +439,7 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
 
 pm_simplify_valid = PatternMatcher([
   # simplify valid
-  (UPat(Ops.AND, name="valid"), simplify_valid),
+  (UPat(Ops.AND, dtypes.bool, name="valid"), simplify_valid),
   (invalid_gate, gated_given_valid),
 ])
 

tinygrad/viz/js/index.js

@@ -336,7 +336,7 @@ function setFocus(key) {
   if (eventType === EventTypes.EXEC) {
     const [n, _, ...rest] = e.arg.tooltipText.split("\n");
     const tableData = [["Name", colored(e.arg.label)], ["Duration", formatTime(e.width)]];
-    if (data.instSt != null) {
+    if (data.tracks.get("Shader Clock") != null) {
       const p = d3.create("p");
       p.append("span").text(timeAtCycle(e.x));
       p.append("span").style("margin-left", "8px").style("color", "#f0f0f566").text(formatTime(e.x));
@@ -423,7 +423,7 @@ async function renderProfiler(path, opts) {
   for (const [k,v] of Object.entries(extData)) data[k] = v;
   // place devices on the y axis and set vertical positions
   const [tickSize, padding, baseOffset] = [5, 8, markers.length ? 14 : 0];
-  const secondaryTick = opts.unit == "clk" ? timeAtCycle : null;
+  const secondaryTick = data.tracks.get("Shader Clock") != null ? timeAtCycle : null;
   const axisHeight = secondaryTick != null ? tickSize*2+(padding*2) : tickSize;
   const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", axisHeight+padding+baseOffset+"px");
   const canvas = profiler.append("canvas").attr("id", "timeline").node();
@@ -590,7 +590,7 @@ async function renderProfiler(path, opts) {
   if (data.pcMap != null) setFocus(focusedShape);
   // secondary axis mapping
   let instRange = null;
-  for (const [k, { shapes }] of data.tracks) if (!k.includes("Clock") && path.includes("sqtt")) {
+  for (const [k, { shapes }] of data.tracks) if (k !== "Shader Clock" && path.includes("sqtt")) {
     const first = shapes[0].x, last = shapes.at(-1).x+shapes.at(-1).width;
     instRange = instRange == null ? [first, last] : [Math.min(first, instRange[0]), Math.max(last, instRange[1])];
   }
@@ -993,10 +993,7 @@ async function main() {
   if (!ckey.startsWith("/graph")) {
     if (!(ckey in cache)) cache[ckey] = ret = await fetchValue(ckey);
     // timeline with cycles on the x axis
-    if (ret instanceof ArrayBuffer) {
-      const pkts = step.query.includes("sqtt");
-      return renderProfiler(ckey, {unit:"clk", heightScale:0.5, hideLabels:true, colorByName:pkts});
-    }
+    if (ret instanceof ArrayBuffer) return renderProfiler(ckey, {heightScale:0.5, hideLabels:true, colorByName:true});
     metadata.replaceChildren(...((ret.metadata ?? []).map((m) => {
       return tabulate(m.map((e) => [e.label.trim(), typeof e.value === "string" ? e.value : formatUnit(e.value)]));
     })));

tinygrad/viz/serve.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
-import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, codecs, io, struct, re, traceback, itertools
-import socketserver
+import multiprocessing, pickle, difflib, os, threading, json, time, sys, socket, argparse, codecs, io, struct, re, traceback, itertools, socketserver
 from contextlib import redirect_stdout, redirect_stderr, contextmanager
 from decimal import Decimal
 from dataclasses import dataclass, field
@@ -370,7 +369,7 @@ wave_colors = {"WMMA": "#1F7857", **{x:"#ffffc0" for x in ["VALU", "VINTERP"]},
 def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent, None, None]:
   from tinygrad.renderer.amd.sqtt import (map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC,
                                           ALUEXEC, INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4, CDNA_INST, InstOpCDNA,
-                                          WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND, WAVERDY)
+                                          CDNA_ISSUE, WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND, WAVERDY)
   pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()}
   row_ends:dict[str, Decimal] = {}
   row_counts:dict[str, itertools.count] = {}
@@ -381,8 +380,8 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent,
   def add(name:str, p:PacketType, wave:int|None=None, info:InstructionInfo|None=None) -> Generator[ProfileEvent, None, None]:
     row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name.replace('_ALT', '')}"
     if (simd:=getattr(p, "simd", None)) is not None: row += f" SIMD:{simd}"
-    # by default we extend the packet to one cycle after timestamp
-    start_time, end_time = p._time, p._time+1
+    # extend packets to the architectural instruction issue interval
+    start_time, end_time = p._time, p._time+(4 if target.startswith("gfx9") else 1)
     # exec links to dispatch, dispatch links to PC
     link:dict|None = {"pc":info.pc} if info else None
     if isinstance(p, (ALUEXEC, VMEMEXEC)):
@@ -430,7 +429,7 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> Generator[ProfileEvent,
       name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4, InstOpCDNA)) else f"0x{p.op:02x}"
       yield from add(name, p, info=info)
     if isinstance(p, (VALUINST, IMMEDIATE, WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)): yield from add(p.__class__.__name__, p, info=info)
-    if isinstance(p, IMMEDIATE_MASK): yield from add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
+    if isinstance(p, (IMMEDIATE_MASK, CDNA_ISSUE)): yield from add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
     if isinstance(p, WAVERDY):
       for wave in range(16):
         if p.mask & (1 << wave):
@@ -729,7 +728,6 @@ if __name__ == "__main__":
   reloader_thread = threading.Thread(target=reloader)
   reloader_thread.start()
   print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"), flush=True)
-  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}")
   try: server.serve_forever()
   except KeyboardInterrupt:
     print("*** viz is shutting down...")