Merge branch 'master' into migrate_calL_expand

.github/workflows/benchmark.yml

@@ -319,7 +319,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 11
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -436,7 +436,13 @@ jobs:
     - name: UsbGPU tiny tests
       run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    - name: UsbGPU (USB4/TB) install script
+      run: sh extra/setup_tinygpu_osx.sh
+    - name: UsbGPU (USB4/TB) boot time
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU (USB4/TB) tiny tests
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
   testcomma:
     strategy:
@@ -556,7 +562,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -623,6 +629,16 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
+        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
+          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
+        fi
+        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 

.github/workflows/platform.yml

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

.github/workflows/test.yml

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

AGENTS.md

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

examples/gpt2.py

@@ -57,7 +57,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
     h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).clone()
+    return (h + self.mlp(self.ln_2(h))).contiguous()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

examples/mlperf/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
+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, router_mfma
 
@@ -305,25 +305,10 @@ class GPTOSS:
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     h_normed = self.norm(h)
-
-    if getenv("FP8_LMHEAD", 0) and ASM_GEMM:
-      pad = (-self.dim) % 256
-      h2 = h_normed.reshape(-1, self.dim).pad(((0, 0), (0, pad)))
-      w2 = self.output.pad(((0, 0), (0, pad)))
-      hq, he8, hsi = quantize_mxfp8(h2)
-      oq, oe8, _ = quantize_mxfp8(w2)
-      if hsi is not None and can_use_asm_gemm(hq, oq.T):
-        logits = asm_gemm(hq, oq.T, mx=True, mx_scales=(hsi, he8, mx_pack(oe8), oe8), mx_w_stored=False)
-        logits = logits.reshape(bsz, seqlen, self.vocab_size).cast(dtypes.bfloat16)
-      else:
-        logits = h_normed @ self.output.T
-    elif ASM_GEMM:
-      pad = (-self.dim) % 256
-      h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
-      logits = asm_gemm(h_padded, w_padded.T) if can_use_asm_gemm(h_padded, w_padded.T) and getenv("VOCAB_ASM", 1) else h_normed @ self.output.T
-    else:
-      logits = h_normed @ self.output.T
-
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

extra/amdflash/README.md

@@ -1,116 +0,0 @@
-# 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

@@ -1,53 +0,0 @@
-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

@@ -1,292 +0,0 @@
-#!/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

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

extra/amdpci/am_smi.py

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

extra/gemm/simple_matmul.py

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

extra/hcq2/ops_amd2.py

@@ -145,7 +145,7 @@ class AMDComputeQueue(HWQueue):
   def wait(self, signal:UOp, value:UOp): self.wait_reg_mem(value.cast(dtypes.uint32), mem=signal.getaddr(self.devs))
 
   def timestamp(self, signal:UOp):
-    self.release_mem(signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+    self.release_mem(signal.getaddr(self.devs), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
                      self.pm4.int_sel__mec_release_mem__none)
 
   def signal(self, signal:UOp, value:UOp):
@@ -197,7 +197,7 @@ class AMDSDMAQueue(HWQueue):
 
   def timestamp(self, signal:UOp):
     self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
-           signal.getaddr(self.devs) + UOp.const(8, dtypes.uint64))
+           signal.getaddr(self.devs))
 
   def signal(self, signal:UOp, value:UOp): # a fence packet then a trap
     op = self.sdma.SDMA_OP_FENCE | (self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.target[0] != 9 else 0)

extra/models/llama.py

@@ -139,7 +139,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
     h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).clone().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -201,7 +201,7 @@ class Transformer:
     self.tok_embeddings = embedding(vocab_size, dim)
     self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
     self.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).clone().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().is_param_(False)
     self.forward_jit = TinyJit(self.forward) if jit else None
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):

extra/remote/bench.py

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

extra/remote/hcq1_remote.py

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

extra/remote/serve.py

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

test/amd/hw/test_ds.py

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

test/amd/test_emu2_pcode.py

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

test/amd/test_pdf.py

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

test/backend/test_arange.py

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

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu, is_hcq2_device
+from test.helpers import assert_kernel_count, needs_second_gpu
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -43,17 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
-
-  def test_assign_copy_retained_uses(self):
-    for use in (lambda x: x.reshape(1, 3), lambda x: x + 1):
-      with self.subTest(use=use):
-        x = Tensor([1., 2, 3], device="PYTHON").to(None)
-        retained = use(x)
-        dest = Tensor.empty(3).assign(x)
-        del x
-        dest.realize().assign(0).realize()
-        self.assertEqual(retained.tolist(), [[1., 2, 3]] if retained.ndim == 2 else [2., 3, 4])
+    assert_kernel_count(1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -509,10 +499,14 @@ class TestAssign(unittest.TestCase):
     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
+    # assign to a shape-changing bitcast view (only works on DISK currently)
     a = Tensor([0]*8, dtype=dtypes.uint8).realize()
     a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
+    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)
 
   def test_assign_dtype_mismatch(self):
     # assign should not implicitly cast dtypes - this can lose precision
@@ -629,7 +623,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(4 if is_hcq2_device() else 2)  # TODO: first copy is dead, could be 1
+    assert_kernel_count(2)  # TODO: first copy is dead, could be 1
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -639,17 +633,9 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
+    assert_kernel_count(1)
     self.assertEqual(base.tolist(), [1,4,3])
 
-  def test_assign_temporary_copy_reshape(self):
-    a = Tensor([[1., 2], [3, 4]], device="PYTHON")
-    c = Tensor.empty(2, 2).assign(a.to(None))
-    GlobalCounters.reset()
-    c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
-    self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
-
 class TestAssignOrdering(unittest.TestCase):
   """Tests for complex assign orderings that could differ between lazy and eager execution.
 
@@ -983,15 +969,20 @@ class TestAssignOrdering(unittest.TestCase):
     b.assign(a)
     b.assign(Tensor.zeros(4))
     b.realize()
-    self.assertListEqual(a.tolist(), [7., 7., 7., 7.])
+    try:
+      self.assertListEqual(a.tolist(), [7., 7., 7., 7.])
+    except AssertionError:
+      # TODO: broken now, b shares a's buffer, so the second assign to b overwrites a
+      self.assertListEqual(a.tolist(), [0., 0., 0., 0.])
 
   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.])
+    with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):  # TODO: broken now, raises
+      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
@@ -1024,12 +1015,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(c.tolist(), [[0,0],[0,0]])
 
-  def test_clone(self):
+  def test_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    c = t.permute(1,0).clone()
-    self.assertIs(c.uop.base.op, Ops.AFTER)
+    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())
-    self.assertEqual(c.tolist(), [[1,1],[2,1]])
+    try:
+      self.assertEqual(c.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      self.assertEqual(c.tolist(), [[1,3],[2,4]])
 
   def test_contiguous_backward(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
@@ -1042,16 +1037,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(cb.tolist(), [[1,2],[3,4]])
 
-  def test_detach_buffer_assignment(self):
-    for realized in (False, True):
-      with self.subTest(realized=realized):
-        base = Tensor([1., 2., 3.])
-        if realized: base.realize()
-        detached = base.detach()
-        detached.assign(detached + 1).realize()
-        self.assertEqual(detached.tolist(), [2., 3., 4.])
-        self.assertEqual(base.tolist(), [2., 3., 4.])
-
   def test_detach_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
     d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
@@ -1063,12 +1048,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(d.tolist(), [[0,0],[0,0]])
 
-  def test_detach_clone(self):
+  def test_detach_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    d = t.permute(1,0).clone().detach()
-    self.assertIs(d.uop.base.op, Ops.AFTER)
+    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())
-    self.assertEqual(d.tolist(), [[1,1],[2,1]])
+    try:
+      self.assertEqual(d.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      self.assertEqual(d.tolist(), [[1,3],[2,4]])
 
   def test_alu(self):
     a = Tensor([1,2,3,4]).contiguous().realize()
@@ -1114,16 +1103,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
-  def test_detach_assignment_preserves_earlier_update(self):
-    x = Tensor([1., 2.]).detach()
-    state = Tensor([0., 0.]).detach()
-    state.assign(state + x * 2)
-    result = state + 1
-    x.assign(x + 1).realize(state, result)
-    self.assertEqual(x.tolist(), [2., 3.])
-    self.assertEqual(state.tolist(), [2., 4.])
-    self.assertEqual(result.tolist(), [3., 5.])
-
 class TestPartialAssignToSharedBuffer(unittest.TestCase):
   def test_five_slices(self):
     big = Tensor.zeros(50).contiguous().realize()

test/backend/test_const_folding.py

@@ -86,8 +86,8 @@ class TestReduceOpsConstFolding(unittest.TestCase):
   def test_zero_size_realize_folded(self):
     # non contiguous folded output doesn't realize
     _check_ast_count(0, Tensor.empty(1, 0).sum())
-    # An explicitly cloned folded constant still owns persistent storage.
-    a = Tensor.empty(1, 0).sum().clone()
+    # contiguous folded const can still schedule
+    a = Tensor.empty(1, 0).sum().contiguous()
     _check_ast_count(2, a+2)
     self.assertIs(a.uop.base.op, Ops.BUFFER)
     np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)

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, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
 
@@ -19,8 +19,7 @@ settings.load_profile("my_profile")
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
 
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  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)]
+  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)]
   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 []
@@ -72,14 +71,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:
@@ -90,10 +89,10 @@ class TestDType(unittest.TestCase):
       if dtype.itemsize > self.DTYPE.itemsize:
         _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
 
-  def test_downcast_ops(self):
+  def test_upcast_to_ops(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
       if dtype.itemsize < self.DTYPE.itemsize:
-        _test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
+        _test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
 
   def test_bitcast(self):
     if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
@@ -113,7 +112,12 @@ 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))
@@ -284,10 +288,14 @@ 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):
-  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())
+  @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_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -393,9 +401,6 @@ class TestEmulatedFp8e5m2(TestFp8e5m2):
   @classmethod
   def tearDownClass(cls): cls.stack.close()
 
-class TestFp8e4m3fnuz(TestDType): DTYPE = dtypes.fp8e4m3fnuz
-class TestFp8e5m2fnuz(TestDType): DTYPE = dtypes.fp8e5m2fnuz
-
 class TestImplicitFunctionTypeChange(unittest.TestCase):
   def test_functions(self):
     result = []

test/backend/test_jit.py

@@ -2,11 +2,12 @@
 import unittest
 import numpy as np
 
-from test.helpers import is_hcq2_device, assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
+from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
 from tinygrad.helpers import JIT, DEV, GlobalCounters
+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(is_hcq2_device(), "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_multitensor.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variab
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
@@ -76,7 +76,7 @@ class TestMultiTensor(unittest.TestCase):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
     out = (X + X)
-    linear = lower_and_compile(out.schedule_linear())
+    linear = compile_linear(out.schedule_linear())
     uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
     run_linear(linear)
     self.assertEqual(len(set(uops)), 1, "function was relinearized")

test/backend/test_ops.py

@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -815,6 +816,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
+  # TODO: x86 PARAM dtype fails SPEC=2
+  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)

test/backend/test_schedule.py

@@ -115,8 +115,7 @@ class TestSchedule(unittest.TestCase):
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
     base.assign(flat_base.reshape(4, 4))
-    # The pending clone is already contiguous, so assign-back needs no separate contiguous buffer.
-    sched = check_schedule(base, 2)
+    sched = check_schedule(base, 4)
     run_linear(*sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v

test/backend/test_setitem.py

@@ -1,4 +1,4 @@
-import unittest, operator
+import unittest
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
 from tinygrad.helpers import Context
 import numpy as np
@@ -75,11 +75,6 @@ class TestSetitem(unittest.TestCase):
     t.detach()[1, 2] = 5
     self.assertEqual(t[1, 2].item(), 5.0)
 
-  def test_setitem_detach_whole(self):
-    t = Tensor.zeros((3, 3)).realize()
-    t.detach()[:] = 5
-    np.testing.assert_equal(t.numpy(), np.full((3, 3), 5.))
-
   def test_setitem_permute(self):
     # setitem on permuted tensor should modify original
     t = Tensor.zeros((2, 3)).contiguous().realize()
@@ -383,32 +378,6 @@ class TestWithGrad(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       y[0] = 99.0
 
-  def test_unrealized_inplace_keeps_storage(self):
-    x = Tensor([1., 2.]).clone()
-    view = x[:1]
-    x += 3
-    x.realize()
-    self.assertEqual(x.tolist(), [4., 5.])
-    self.assertEqual(view.tolist(), [4.])
-
-  def test_unrealized_view_inplace_keeps_storage(self):
-    x = Tensor([1., 2.]).clone()
-    view = x[:1]
-    view += 3
-    view.realize()
-    self.assertEqual(x.tolist(), [4., 2.])
-    self.assertEqual(view.tolist(), [4.])
-
-  def test_set_augmented_backward(self):
-    for op, expected in ((operator.isub, [-1., -1.]), (operator.imul, [1., 2.]), (operator.itruediv, [-0.01, -0.005])):
-      with self.subTest(op=op.__name__):
-        z = Tensor([1.0, 2.0, 3.0, 4.0])
-        x = Tensor([10.0, 20.0])
-        z[:2] = op(z[:2], x)
-        z.sum().backward()
-        np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
-        np.testing.assert_allclose(x.grad.numpy(), expected)
-
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):
     N = 10

test/backend/test_uops.py

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

test/backend/test_wait_loop.py

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

test/device/test_hcq2.py

@@ -1,11 +1,11 @@
-import unittest, contextlib, ctypes, gc, numpy as np
+import unittest, contextlib, ctypes, numpy as np
 from unittest.mock import patch
-from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
+from tinygrad import Device, Tensor, TinyJit, Variable, dtypes
 from tinygrad.device import Buffer
 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 compile_linear, link_linear, lower_and_compile, run_linear
+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
@@ -16,14 +16,7 @@ from test.helpers import call_is_hcq
 @contextlib.contextmanager
 def rt_views():
   calls, orig = [], HCQ2Compiled.rt_view
-  def track(dev, *args, **kwargs):
-    calls.append(dev)
-    return orig(dev, *args, **kwargs)
-  with patch.object(HCQ2Compiled, "rt_view", track): yield calls
-
-def chain(x:Tensor, n:int) -> Tensor:
-  for _ in range(n): x = (x + 1).contiguous()
-  return x
+  with patch.object(HCQ2Compiled, "rt_view", lambda s, *a, **kw: (calls.append(s), orig(s, *a, **kw))[1]): yield calls
 
 @contextlib.contextmanager
 def encoded_batches():
@@ -31,10 +24,6 @@ def encoded_batches():
   with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
     yield batches
 
-def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
-  for _ in range(n): x = (x + 1).contiguous()
-  return x.realize()
-
 def patch_words(batch:UOp) -> list[UOp]:
   return [w for s in batch.src[0].toposort() if s.op is Ops.STORE and s.src[0].op is Ops.INDEX and s.src[0].src[1].op is Ops.STACK
           and s.src[1].op is Ops.STACK for w in s.src[1].src]
@@ -44,69 +33,39 @@ def rt_params(batch:UOp) -> list[str]:
 
 @unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
 class TestHCQ2Core(unittest.TestCase):
-  @staticmethod
-  def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
-
-  def compiled(self, n:int, jit=False):
-    x, inputs = self.input(), []
-    if jit:
-      f = TinyJit(lambda a: chain(a, n).realize())
-      f(x)
-      return f(x), f.captured._linear, [x.uop.base]
-    out = chain(x, n)
-    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
-
   def test_jit_has_no_rt_buffers(self):
+    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]
-    rings = [dev.rt_buffer(True, host) for host in (False, True)]
-    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
-    for n in (1, 65):
-      with self.subTest(kernels=n):
-        x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
-        for _ in range(2): f(x)
-        for u in f.captured.linear.toposort():
-          if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
-            addr = buf._buf.va_addr
-            self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
+    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)
 
-  def test_small_eager_cached(self):
-    _, compiled, inputs = self.compiled(1)
-    linked = link_linear(compiled, input_uops=inputs)
-    self.assertIs(link_linear(compiled, input_uops=inputs), linked)
-
-  def test_large_eager_not_cached(self):
-    _, compiled, inputs = self.compiled(65)
-    linked = link_linear(compiled, input_uops=inputs)
-    self.assertIsNot(link_linear(compiled, input_uops=inputs), linked)
-    self.assertNotIn(compiled, link_linear_cache)
-
-  def test_double_compile(self):
-    for n in (1, 65):
-      for jit in (False, True):
-        with self.subTest(kernels=n, jit=jit):
-          out, compiled, inputs = self.compiled(n, jit=jit)
-          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
-          before = tuple(inputs)
-          with rt_views() as borrowed:
-            for linear in (compiled, linked):
-              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
-          self.assertEqual(tuple(inputs), before)
-          self.assertFalse(borrowed)
-          run_linear(linked, input_uops=inputs, jit=True, wait=True)
-          self.assertEqual(out.tolist(), [2 + n] * 4)
-
-  def test_double_link(self):
-    for n in (1, 65):
-      for jit in (False, True):
-        with self.subTest(kernels=n, jit=jit):
-          out, compiled, inputs = self.compiled(n, jit=jit)
-          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
-          with rt_views() as borrowed:
-            again = link_linear(linked, input_uops=inputs, allow_cache=not jit)
-          self.assertIs(again, linked)
-          self.assertFalse(borrowed)
-          run_linear(again, input_uops=inputs, jit=True, wait=True)
-          self.assertEqual(out.tolist(), [2 + n] * 4)
+      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
@@ -126,13 +85,6 @@ class TestHCQ2Core(unittest.TestCase):
       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_map_cpu_buffer_preserves_contents(self):
-    src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
-    data = bytes(range(16))
-    src.as_memoryview(force_zero_copy=True)[:] = data
-    src.get_buf(Device.DEFAULT)
-    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
-
   def test_staged_copy_roundtrip(self):
     # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
     stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
@@ -150,7 +102,6 @@ class TestHCQ2Core(unittest.TestCase):
       @TinyJit
       def f(a): return (a.sin() * 3).contiguous().realize()
       for _ in range(3): f(x)
-      eager_chain(x)
 
     jit, eager = partition(batches, lambda c: c.arg.aux.table >= 0)
     self.assertTrue(jit and eager, f"want both kinds of batch, got {len(jit)} jit and {len(eager)} eager")
@@ -173,24 +124,9 @@ class TestHCQ2Core(unittest.TestCase):
       return max(c.arg.aux.nargs for c in batches)
     self.assertEqual(nargs(2), nargs(12))
 
-  def test_caches_hold_no_buffers(self):
-    # an eager template caches without its buffers and the jit's linear compiles once uncached: freeing the tensors frees the device memory
-    def step(i):
-      x = Tensor(np.full(1024, i, np.float32)).to(Device.DEFAULT).realize()
-      @TinyJit
-      def f(a): return (a * 2 + 1).contiguous().realize()
-      for _ in range(3): out = f(x)
-      self.assertEqual(out.tolist(), [2.0 * i + 1] * 1024)
-    step(1) # warms the programs, templates and rings
-    gc.collect()
-    used = GlobalCounters.mem_used
-    for i in range(2, 5): step(i)
-    gc.collect()
-    self.assertEqual(GlobalCounters.mem_used, used)
-
   def test_device_state_survives_as_link_refs(self):
     # a buffer the commands only address, never a param of the body, is kept by the linked call as a ref of what its getaddr resolved into
-    dev, names = Device[Device.DEFAULT], {"AMD": ("scratch",), "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
+    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()
@@ -206,7 +142,7 @@ class TestHCQ2FFI(unittest.TestCase):
   def _run(body:UOp) -> list[Buffer]:
     call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
     assert call is not None
-    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), allow_cache=False)
+    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]
 

test/external/external_test_usb_asm24.py

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

test/external/speed_v_theoretical.py

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

test/helpers.py

@@ -65,10 +65,6 @@ def assert_kernel_count(expected:int):
   got = GlobalCounters.kernel_count
   if got != expected: raise KernelCountException(expected, got)
 
-def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
-  from tinygrad.runtime.support.hcq2 import HCQ_DEVS
-  return Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}
-
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"

test/mockgpu/amd/emu.py

@@ -69,7 +69,7 @@ from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp, Inst
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
 from test.amd.helpers import decode_dpp16
-from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32, _bitreverse, _countbits
+from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32
 
 MASK32 = 0xFFFFFFFF
 
@@ -1566,19 +1566,9 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)  # type: ignore[union-attr]
 
-  if is_lds and op_name == 'DS_SWIZZLE_B32':
-    # The manual's reverse_bits operates on five-bit lane indices; thread indices wrap within the wave.
-    funcs = {'reverse_bits': lambda x: _bitreverse(x, 32) >> _c(27), 'count_ones': _countbits,
-             'thread_in': lambda x: ctx.rvgpr_dyn(addr_reg, x & _c(ctx.wave_size - 1)),
-             'thread_valid': lambda x: _lane_active(exec_mask, x & _c(ctx.wave_size - 1))}
-    result, _ = parse_pcode(pcode, {'offset0': offset0.cast(dtypes.uint8), 'offset1': offset1.cast(dtypes.uint8)}, funcs)
-    values = [result[f'thread_out@{i}'] for i in range(ctx.wave_size)]
-    # Snapshot every source before writing: destination and source registers may be identical.
-    reads = UOp(Ops.STACK, src=tuple(values))
-    return UOp.sink(*(ctx.wvgpr_dyn(vdst_reg, _c(i), val, exec_mask, after=reads) for i, val in enumerate(values)), *ctx.inc_pc())
-
   # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
   if is_lds and 'PERMUTE' in op_name:
+    pcode = get_pcode(inst.op)
     srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
             'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
     _, assigns = parse_pcode(pcode, srcs)

test/mockgpu/amd/pcode.py

@@ -630,10 +630,6 @@ class Parser:
         self.eat('DOT')
         dt_name = self.eat('IDENT').val
         return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
-      if name in self.funcs and self.try_eat('LBRACKET'):
-        index = self.parse()
-        self.eat('RBRACKET')
-        return self.funcs[name](index)
       if name == 'VGPR' and self.at('LBRACKET'):
         self.eat('LBRACKET')
         lane = self.parse()
@@ -1010,24 +1006,20 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
     # for loop
     if first == 'for':
-      # C-style loops use an exclusive bound; for/in loops use an inclusive bound.
-      if m := re.fullmatch(r'for\s*\(\s*(\w+)\s*=\s*(\d+);\s*\1\s*<\s*(\d+);\s*\1\s*(\+\+|\+=\s*\d+)\s*\)', line):
-        loop_var, start_val, end_val = m[1], int(m[2]), int(m[3]) - 1
-        step = 1 if m[4] == '++' else int(m[4][2:])
-      else:
-        p = Parser(toks, env, funcs)
-        p.eat_val('for', 'IDENT')
-        loop_var = p.eat('IDENT').val
-        p.eat_val('in', 'IDENT')
-        def parse_bound():
-          if p.at('NUM') and p.peek(1).type == 'QUOTE':
-            p.eat('NUM')
-            p.eat('QUOTE')
-          if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-          return int(p.parse())
-        start_val = parse_bound()
-        p.eat('COLON')
-        end_val, step = parse_bound(), 1
+      # Parse: for VAR in [SIZE']START : [SIZE']END do
+      p = Parser(toks, env, funcs)
+      p.eat_val('for', 'IDENT')
+      loop_var = p.eat('IDENT').val
+      p.eat_val('in', 'IDENT')
+      def parse_bound():
+        if p.at('NUM') and p.peek(1).type == 'QUOTE':
+          p.eat('NUM')
+          p.eat('QUOTE')
+        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
+        return int(p.parse())
+      start_val = parse_bound()
+      p.eat('COLON')
+      end_val = parse_bound()
       # Collect body
       i += 1
       body_lines: list[str] = []
@@ -1043,7 +1035,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       has_break = any('break' in bl.lower() for bl in body_lines)
       found_var = f'_found_{next(_break_var_ids)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
-      for loop_i in range(start_val, end_val + 1, step):
+      for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
         _, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
         if has_break:
@@ -1232,9 +1224,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       var = toks[0].val
       j, idx_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'EQUALS':
-        idx_expr = parse_tokens(idx_toks, env, funcs)
         # Static index: var[NUM] = value
-        if isinstance(idx := _single_value(idx_expr), int):
+        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
+          idx = int(idx_toks[0].val.rstrip('UuLl'))
           val = parse_tokens(toks[j+1:], env, funcs)
           existing = block_assigns.get(var, env.get(var))
           if existing is not None and isinstance(existing, UOp):
@@ -1246,6 +1238,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         # Dynamic index: var[expr] = value where var has @-elements
         elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
         if elems:
+          idx_expr = parse_tokens(idx_toks, env, funcs)
           val = parse_tokens(toks[j+1:], env, funcs)
           for elem_idx_str, old_elem in elems:
             elem_idx = int(elem_idx_str)
@@ -1414,29 +1407,16 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
 
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dict | None = None) -> tuple[dict, list]:
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
   env: dict = srcs.copy() if srcs else {}
   assigns: list[tuple[str, UOp]] = []
+  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
   lines: list[str] = []
-  blocks: list[str] = []
-  for raw in pcode.splitlines():
-    line = raw.split('//')[0].strip().rstrip(';')
-    if not line: continue
-    # Both block syntaxes share the same parser; braces supply the implicit end markers.
-    if line.startswith('}') and blocks:
-      end = blocks.pop()
-      line = line[1:].strip()
-      if not line.startswith(('elsif', 'else')): lines.append(end)
-    if m := re.match(r'(if|elsif|else|for)\b.*\{$', line):
-      blocks.append('endfor' if m[1] == 'for' else 'endif')
-      line = line[:-1].rstrip()
-      if m[1] in ('if', 'elsif'): line += ' then'
-    if not line: continue
-    line = re.sub(r'=\s*(\w+):(\w+)$', r'= {\1, \2}', line)
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] += ' ' + line
-    else: lines.append(line)
-  assert not blocks, "unclosed pcode block"
-  _, final, _ = parse_block(lines, 0, env, {**_FUNCS, **funcs} if funcs else None, assigns=assigns)
+  for l in raw_lines:
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
+    else: lines.append(l)
+  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
   sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
   for var, val in final.items():
     if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):

test/mockgpu/nv/nvdriver.py

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

test/mockgpu/nv/nvgpu.py

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

test/null/test_uops.py

@@ -7,28 +7,10 @@ from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.weak import pm_lower_weak
-from tinygrad.uop.spec import spec_program, spec_shared, spec_tensor, type_verify
+from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
 
-class TestStorageSpec(unittest.TestCase):
-  def test_contiguous_is_not_store_target(self):
-    value = (Tensor.empty(4).uop + 1).contiguous()
-    for target in (value, value.reshape(2, 2), value.detach()):
-      with self.subTest(op=target.op), self.assertRaises(RuntimeError):
-        type_verify(target.store(target), spec_tensor)
-
-  def test_contiguous_can_depend_on_other_storage_writes(self):
-    buf = Tensor.empty(4).uop
-    type_verify((buf + 1).contiguous().after(buf.store(buf + 1)), spec_tensor)
-
-  def test_detached_storage_can_carry_writes(self):
-    buf = Tensor.empty(4).uop
-    detached = buf.detach()
-    type_verify(detached.after(detached.store(buf + 1)), spec_tensor)
-    with self.assertRaises(RuntimeError):
-      type_verify((buf + 1).detach().after(buf.store(buf + 1)), spec_tensor)
-
 class TestDTypeFromUOp(unittest.TestCase):
   def test_broadcastable_promotion(self):
     self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)), None), dtypes.float32)

test/null/test_viz.py

@@ -454,7 +454,7 @@ class TestVizIntegration(unittest.TestCase):
   def test_jit(self):
     with save_viz():
       @TinyJit
-      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).clone().assign(a.to(c.device)), b.assign(c.to(b.device))
+      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
       a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
       for _ in range(3): Tensor.realize(*f(a, b, c))
     out = load_profile(cpu_events)
@@ -1073,11 +1073,10 @@ class TestCLI(unittest.TestCase):
       out = run_cli(*files, "-s", "NULL")
       aggregate = run_cli(*files, "-s", "NULL", "-t")
     self.assertEqual(len(out), 3*2)
-    # Operation count increases with N; FLOPS is a rate and also depends on the measured duration.
+    # flops increases as N gets larger
     gflops = [row["fmt"]["FLOPS"] for row in out]
-    flops = [rate * row["dur_ms"] * 1e-3 for rate, row in zip(gflops, out)]
-    self.assertGreater(flops[4], flops[2])
-    self.assertGreater(flops[5], flops[3])
+    self.assertGreater(gflops[4], gflops[2])
+    self.assertGreater(gflops[5], gflops[3])
     # aggregate flops
     self.assertEqual(len(aggregate), 2)
     agg_gflops = [row["fmt"]["FLOPS"] for row in aggregate]

test/unit/test_call.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.tensor import transform_to_call
 
-def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop)).src[0].key
+def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop))[0].src[0].key
 
 class TestCall(unittest.TestCase):
   def test_call_plus(self):
@@ -164,11 +164,16 @@ class TestCallSchedule(unittest.TestCase):
     np.testing.assert_equal(big.numpy(), ref)
 
   def test_precompile_full_buffer_assign(self):
+    from test.helpers import assert_kernel_count
+    from tinygrad.helpers import GlobalCounters
     @function(precompile=True)
     def f(x:Tensor) -> Tensor: return x * 2 + 1
     a = Tensor.arange(8).float().realize()
     cache = Tensor.zeros(8).realize()
+    GlobalCounters.reset()
+    # the call writes the whole assign target directly: it binds into the output position, no copy kernel
     cache.assign(f(a)).realize()
+    assert_kernel_count(1)
     np.testing.assert_equal(cache.numpy(), np.arange(8)*2+1)
 
   def test_reshape_precompile(self):

test/unit/test_callify.py

@@ -1,7 +1,5 @@
 import unittest
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.tensor import transform_to_call
 
 class TestCallify(unittest.TestCase):
   def test_basic(self):
@@ -109,96 +107,5 @@ class TestCallify(unittest.TestCase):
     self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
     self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
 
-  def test_only_replace_inputs(self):
-    x = Tensor.empty(4)
-    body = UOp.sink((x.uop + 1).contiguous().copy_to_device("CPU:1"))
-    call = transform_to_call(body)
-    self.assertEqual(call.src[1:], (x.uop,))
-    self.assertIs(call.src[0], body.substitute({x.uop: x.uop.param_like(0)}))
-
-  def test_existing_params_do_not_alias_buffers(self):
-    x = Tensor.empty(4)
-    param = UOp.param(0, x.dtype, x.shape, device=x.device)
-    body = UOp.sink(x.uop + param)
-    call = transform_to_call(body)
-    self.assertEqual(set(call.src[1:]), {x.uop, param})
-    params = [u for u in call.src[0].toposort() if u.op is Ops.PARAM]
-    self.assertEqual({u.arg.slot for u in params}, {0, 1})
-    self.assertIs(call.src[0].substitute({u: call.src[1+u.arg.slot] for u in params}, walk=True), body)
-
-  def test_scalar_param_binding_survives_renumbering(self):
-    from tinygrad.schedule import create_linear_with_vars
-    from tinygrad.engine.realize import run_linear
-    x = Tensor([1, 2, 3]).realize()
-    out = Tensor.empty_like(x)
-    binding = UOp.variable("amount", 1, 10, dtypes.int).bind(4)
-    param = binding.param_like(7)
-    call = transform_to_call(UOp.sink(out.uop.after(out.uop.store(x.uop + param))))
-    call = call.replace(src=(call.src[0], *(binding if arg is param else arg for arg in call.src[1:])))
-    run_linear(*create_linear_with_vars(call))
-    self.assertEqual(out.tolist(), [5, 6, 7])
-
-  def test_nested_params_keep_their_scope(self):
-    x = Tensor.empty(4)
-    param = UOp.param(7, x.dtype, x.shape, device=x.device)
-    nested_body = UOp.sink(param + 1)
-    nested = nested_body.call(*([x.uop] * 8))
-    call = transform_to_call(UOp.sink(x.uop + param, nested))
-    self.assertIs(call.src[0].src[1].src[0], nested_body)
-    self.assertEqual(set(call.src[1:]), {x.uop, param})
-
-  def test_fresh_slots_are_negative_and_canonical_slots_are_dense(self):
-    x = Tensor.empty(4)
-    param = UOp.placeholder((4,), x.dtype, device=x.device)
-    inner = x.uop.param_like(0)
-    outputs = UOp.call_with_outputs((inner + 1, inner + 2), x.uop)
-    fresh = [x.uop.arg.slot, param.arg.slot, *(out.src[0].arg.slot for out in outputs)]
-    self.assertLess(fresh[0], 0)
-    self.assertTrue(all(a > b for a, b in zip(fresh, fresh[1:])))
-    call = transform_to_call(UOp.sink(*outputs, param))
-    unbound = [u.arg.slot for u in call.src[0].toposort() if u.is_unbound]
-    self.assertEqual(unbound, list(range(len(outputs))))
-    params = [u.arg.slot for u in call.src[0].toposort(enter_calls=False) if u.op is Ops.PARAM]
-    self.assertEqual(params, list(range(len(call.src)-1)))
-    self.assertIn(param, call.src[1:])
-
-  def test_unbound_renumbering_preserves_distinct_outputs(self):
-    def output(): return UOp.call_with_outputs((Tensor(1., dtype=dtypes.float, device="CPU").uop,))[0]
-    canonical = transform_to_call(UOp.sink(output())).src[0].src[0]
-    body = UOp.sink(canonical, output())
-    call = transform_to_call(body)
-    self.assertEqual(len([u for u in call.src[0].toposort() if u.is_unbound]), 2)
-    self.assertIs(transform_to_call(call.src[0]).src[0], call.src[0])
-
-  def test_intermediate_contiguous_stays_a_value(self):
-    x = (Tensor([1, 2, 3]).realize() + 1).contiguous()
-    original = x.uop
-    y = (x * 2).realize()
-    self.assertIs(x.uop, original)
-    self.assertIs(x.uop.op, Ops.CONTIGUOUS)
-    self.assertEqual(y.tolist(), [4, 6, 8])
-
-  def test_intermediate_clone_persists(self):
-    x = (Tensor([1, 2, 3]).realize() + 1).clone()
-    y = (x * 2).realize()
-    self.assertTrue(x.uop.has_buffer_identity())
-    self.assertEqual(x.tolist(), [2, 3, 4])
-    self.assertEqual(y.tolist(), [4, 6, 8])
-
-  def test_creation_copy_has_storage(self):
-    x = Tensor([1, 2, 3], device="PYTHON").to("CPU")
-    self.assertTrue(x.uop.has_buffer_identity(after_ok=True))
-    y = Tensor.empty(3, dtype=dtypes.int, device=x.device).assign(x).realize()
-    y.assign(0).realize()
-    self.assertEqual(x.tolist(), [1, 2, 3])
-
-  def test_zero_size_cat_with_rng(self):
-    # Empty outputs must not replay a pending RNG counter update.
-    a = Tensor.rand(2, 2)
-    b = Tensor.rand(2, 0)
-    t = a.cat(b, dim=1).realize()
-    self.assertEqual(t.shape, (2, 2))
-    self.assertListEqual(t.tolist(), a.tolist())
-
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

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

tinygrad/codegen/late/regalloc.py

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

tinygrad/device.py

@@ -187,7 +187,7 @@ class Buffer:
     # zero copy with as_memoryview (disabled by default due to use after free)
     if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'):
       if not no_sync: self.allocator.dev.synchronize()
-      if (mv:=self.allocator._as_buffer(self._buf)) is not None: return mv
+      return self.allocator._as_buffer(self._buf)
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv

tinygrad/engine/jit.py

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

tinygrad/engine/realize.py

@@ -130,7 +130,7 @@ class ExecContext:
   cache: bool = True
 
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.MSELECT, Ops.SHRINK, Ops.BITCAST): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
   if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
   return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -283,18 +283,17 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
   linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
-def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
-  return hcq_link(linear, input_uops=input_uops, allow_cache=allow_cache)
+def 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), input_uops=inputs)
+  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.without_after, perf_counter_us(), pm_exec.rewrite(call.without_after, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), allow_cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
   while True:
     if clear_l2:
       if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()

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 = 23
+VERSION = 22
 _db_connection = threading.local()
 def db_connection():
   if (conn:=getattr(_db_connection, "conn", None)) is None:

tinygrad/llm/kernels/amd.py

@@ -6,7 +6,6 @@ from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
-from tinygrad.renderer.cstyle import HIPRenderer
 
 BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
@@ -32,7 +31,7 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
   if device is None or device.split(":")[0] != "AMD": return False
   # @function contexts set ALLOW_DEVICE_USAGE=0 (scheduling must not open devices); the device is always open here
   with Context(ALLOW_DEVICE_USAGE=1):
-    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11 and isinstance(Device[device].renderer, HIPRenderer)
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
 
 def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
   for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
@@ -75,7 +74,7 @@ class Linear(nn.Linear):
       nbytes, nblocks = raw.max_numel(), raw.max_numel() // Q6_BYTES
       byte_view = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(nbytes, dtypes.uint8, raw_offset)))
       padded = byte_view.reshape((nblocks, Q6_BYTES)).pad_to((nblocks, Q6_PADDED)).bitcast(dtypes.uint32)
-      self.weight = padded.clone().reshape(nblocks * Q6_WORDS)
+      self.weight = padded.contiguous().reshape(nblocks * Q6_WORDS)
     else:
       self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
         .view(raw.max_numel() * raw.dtype.itemsize // dtypes.uint32.itemsize, dtypes.uint32, raw_offset)))
@@ -541,8 +540,7 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
   q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
   k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
-  # All waves must finish reading Q/K before their shared memory is reused for P/V.
-  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk).barrier()
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
   S_reg = S_reg.after(qk_done, S_reg.store(S_reg * SCALE))
   rm, rn = UOp.range(TM, 250), UOp.range(TN, 251)
   q_idx = q_base + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m

tinygrad/mixin/elementwise.py

@@ -58,14 +58,11 @@ class ElementwiseMixin(CreationMixin):
 
   def contiguous(self, **kwargs) -> Self:
     """
-    Requests a contiguous layout for this value when it is computed.
-    This does not reserve independent storage or retain an intermediate result across realizations; use `clone()` for that.
+    Returns a contiguous tensor.
     """
     if self.dtype in dtypes.weaks: return self
     uop = self._uop
-    src = uop
-    while src.op in {Ops.DETACH, Ops.CONTIGUOUS_BACKWARD}: src = src.src[0]
-    if uop.op is Ops.CONTIGUOUS or self.device is None or src.has_buffer_identity(after_ok=True): return self._wrap_uop(uop)
+    if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
     return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
 
   def contiguous_backward(self) -> Self:

tinygrad/mixin/rand.py

@@ -52,7 +52,6 @@ class RandMixin(OpMixin):
     Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[0, 1)`.
 
     You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
-    By default, the random values get persistent storage when computed. `contiguous=False` leaves them as an expression.
 
     ```python exec="true" source="above" session="tensor" result="python"
     Tensor.manual_seed(42)
@@ -66,8 +65,7 @@ class RandMixin(OpMixin):
     if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
     device = cast(str, canonicalize_device(device))
     key, counter = cls._next_counter(device, ceildiv(prod(shape) * dt.itemsize, 4))
-    out = cls._rand(key, counter, shape, dt, contiguous=False)
-    return cls._wrap_uop(out._uop.clone()) if contiguous else out
+    return cls._rand(key, counter, shape, dt, contiguous=contiguous)
 
   def rand_like(self, **kwargs) -> Self:
     """
@@ -295,8 +293,7 @@ class RandMixin(OpMixin):
     if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]")
     if not TRAINING or p == 0: return self
     if p == 1: return self.const_like(0)
-    mask = self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p
-    return self._wrap_uop(mask._uop.clone()).where(self, 0) / (1.0 - p)
+    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)
 
   def scaled_dot_product_attention(self, key:Self, value:Self, attn_mask:Self|None=None, dropout_p:float=0.0,
                                    is_causal:bool=False, enable_gqa:bool=False) -> Self:

tinygrad/renderer/amd/generate.py

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

tinygrad/renderer/isa/__init__.py

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

tinygrad/renderer/isa/x86.py

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

tinygrad/runtime/autogen/am/__init__.py

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

tinygrad/runtime/autogen/am/regs.py

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

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

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

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

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

tinygrad/runtime/ops_amd.py

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

tinygrad/runtime/ops_cpu.py

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

tinygrad/runtime/ops_nv.py

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

tinygrad/runtime/ops_python.py

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

tinygrad/runtime/support/am/amdev.py

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

tinygrad/runtime/support/am/ip.py

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

tinygrad/runtime/support/amd.py

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

tinygrad/runtime/support/compiler_qcom.py

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

tinygrad/runtime/support/hcq2.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 from typing import cast, TypeVar, Generic, Any, TYPE_CHECKING
-import functools, time, itertools, decimal, weakref, statistics, ctypes, importlib
+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, HCQ2, cpu_profile, mv_address
-from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, perf_counter_us, DEV
-from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, DepsTracker
+from tinygrad.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, DTYPES_DICT
@@ -20,8 +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_CACHE_THRESH = ContextVar("HCQ_CACHE_THRESH", 64)
-HCQ_DEVS = frozenset(("NV", "QCOM", "CPU")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
+HCQ_DEVS = frozenset(("QCOM", "CPU")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
 
 @dataclass(frozen=True)
 class HCQInfo:
@@ -85,19 +84,19 @@ def cstruct(struct_t, **fields:UOp|int) -> UOp:
   return patch(buf, rows, bytes(ctypes.sizeof(struct_t)))
 
 # *****************
-# 0.1. prep: eager buffers become tagged params
+# 0.1. prep: replace buffers with params
 
-def replace_buffer(ctx:tuple[bool, list[UOp], dict[UOp, int]], b:UOp) -> UOp:
-  use_rt, bufs, slots = ctx
-  if slots.setdefault(b, len(bufs)) == len(bufs): bufs.append(b)
-  param = UOp.param(slots[b], b.dtype, b.max_numel(), b.device)
-  return param if use_rt else param.replace(tag="lt_input")
-pm_replace_buffers = PatternMatcher([(UPat(Ops.BUFFER, name="b"), replace_buffer)])
+def replace_call_buffers(ctx:tuple[list[UOp], dict[UOp, int]], call:UOp) -> UOp|None:
+  bufs, slots = ctx
+  for s in call.src[1:]:
+    if s.op is not Ops.PARAM and not s.is_bound_var and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s)
+  return call.replace(src=call.src[:1] + tuple(s if s.op is Ops.PARAM or s.is_bound_var else s.param_like(slots[s]) for s in call.src[1:]))
+pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])
 
 # *****************
 # 1.1. prep: staging copies
 
-STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
+STAGING_SIZE, STAGING_SLOTS = (4 if os.getenv("CI") else 128) << 20, 2 # the staging mlocks into the device: ci runners cap locked memory at 8mb
 
 @functools.cache
 def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
@@ -139,11 +138,10 @@ pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)])
 # 2. deps
 
 class HCQDepsTracker(DepsTracker):
+  # TODO: optimize
   @staticmethod
-  def _key(a:UOp) -> tuple[Any, int, int]: # (base, lane) and the byte range: overlapping views of one base depend
-    lane, view = (a.arg, a.src[0]) if a.op is Ops.MSELECT else (None, a)
-    base, off = unwrap_view(view)
-    return (base, lane), off, off + view.max_numel() * view.dtype.itemsize
+  def _key(buf:Any) -> tuple[Any, int, int]:
+    return (buf, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
 @dataclass
 class BatchCtx:
@@ -152,38 +150,37 @@ class BatchCtx:
   tracker:HCQDepsTracker = field(default_factory=HCQDepsTracker)
   queues:dict[str, list[str]] = field(init=False)
   first:dict[tuple[str, str], int] = field(init=False); last:dict[tuple[str, str], int] = field(init=False) # noqa: E702
-  prev:list[int|None] = field(init=False)
   signal_tags:set[int] = field(init=False)
   slots:dict[str, UOp] = field(init=False)
 
   def __post_init__(self):
-    self.queues, self.first, self.last, self.prev = {}, {}, {}, []
+    self.queues, self.first, self.last = {}, {}, {}
     for tag, (_, devs, q) in enumerate(self.batch):
       if q not in self.queues.setdefault(devs[0], []): self.queues[devs[0]].append(q)
       self.first.setdefault((devs[0], q), tag)
-      self.prev.append(self.last.get((devs[0], q)))
       self.last[(devs[0], q)] = tag
     self.signal_tags = {tag for (dev, q), tag in self.last.items() if q != self.epilogue_queue(dev)}
-    # a slot is [signal][timestamp], 16 bytes: the queue signals, the timeline, then two per call if profiling
-    self.slots = {dev: UOp.placeholder((2 * (len(qs) + 1 + (2 * len(self.batch) if self.profile else 0)),), dtypes.uint64, device=(dev,),
-                                       volatile=True, tag="slots") for dev, qs in self.queues.items()}
+    self.slots = {dev: UOp.placeholder((len(qs) + 1 + (2 * len(self.batch) if self.profile else 0),), dtypes.uint64, device=(dev,), volatile=True,
+                                       tag="slots") for dev, qs in self.queues.items()}
 
   def epilogue_queue(self, dev:str) -> str: return "COMPUTE:0" if len(self.queues[dev]) > 1 else self.queues[dev][0] # closes the device
 
-  def slot(self, devs:tuple[str, ...], i:int) -> UOp: return self.slots[devs[0]].shrink(((2 * i, 2 * i + 2),)) # not a slice: 10x the cost
+  def slot(self, devs:tuple[str, ...], i:int) -> UOp: return self.slots[devs[0]].shrink(((i, i + 1),)) # not [i:i+1]: the slice path is 10x the cost
   def queue_signal(self, devs:tuple[str, ...], queue:str) -> UOp: return self.slot(devs, self.queues[devs[0]].index(queue))
   def sched_timeline(self, devs:tuple[str, ...]) -> UOp: return self.slot(devs, len(self.queues[devs[0]]))
   def stamps(self, devs:tuple[str, ...], tag:int) -> tuple[int, ...]: return (st:=len(self.queues[devs[0]])+1+2*tag, st + 1) if self.profile else ()
 
+def _call_bufs(call:UOp) -> list[Any]:
+  def dep_buf(b:UOp) -> Any:
+    if (base:=(b.src[0] if b.op is Ops.MSELECT else b).storage_base).op is Ops.PARAM: return b if b.op is Ops.MSELECT else base
+    return cast(MultiBuffer, base.buffer).bufs[b.arg] if b.op is Ops.MSELECT else base.buffer
+  return [dep_buf(a) for a in get_call_arg_uops(call)]
+
 def _wait_ins(ctx:BatchCtx, call:UOp, device:str, queue:str, tag:int) -> list[UOp]:
-  bufs, write = list(get_call_arg_uops(call)), get_call_outs_ins(call)[0]
+  bufs, write = _call_bufs(call), get_call_outs_ins(call)[0]
   latest:dict[tuple[str, str], int] = {} # (producer device, queue) -> the latest submit tag to wait on, same-queue submits are fifo
   for d, q, t in ctx.tracker.access_resources(bufs, list(range(len(bufs)) if write is None else write), (device, queue, tag)):
     if t < tag and (d, q) != (device, queue): latest[(d, q)] = max(latest.get((d, q), 0), t)
-
-  # NV waits break QMD chaining, so also wait for the previous launch
-  if latest and device.split(":")[0] == "NV" and queue.startswith("COMPUTE") and (p:=ctx.prev[tag]) is not None: latest[(device, queue)] = p
-
   ctx.signal_tags |= set(latest.values())
   return [UOp(Ops.INS, arg=("wait", dtypes.void), src=(ctx.queue_signal((d,), q), UOp.const(t + 1, dtypes.uint64))) for (d, q), t in latest.items()]
 
@@ -203,7 +200,7 @@ def _emit_submits(ctx:BatchCtx, call_waits:list[list[UOp]]) -> tuple[list[UOp],
 
     # and make hcq call
     name, est = get_call_name(call, get_call_arg_uops(call)), estimate_uop(call)
-    kerns.append((devices, name, est, tuple(2 * s + 1 for s in ctx.stamps(devices, tag)), getattr(call.src[0].arg, "profile_key", None)))
+    kerns.append((devices, name, est, ctx.stamps(devices, tag), getattr(call.src[0].arg, "profile_key", None)))
 
     ts_ins = [UOp(Ops.INS, arg=("timestamp", dtypes.void), src=(ctx.slot(devices, i),)) for i in ctx.stamps(devices, tag)]
     q += ts_ins[:1] + [call] + ts_ins[1:]
@@ -303,8 +300,8 @@ pm_hcq_encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name
 
 def _is_input_addr(g:UOp) -> bool:
   base = unwrap_view(g.src[0])[0]
-  if base.op is Ops.MSELECT: base = unwrap_view(base.src[0])[0] # a lane of a view
-  return base.op is Ops.PARAM and base.tag is None
+  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
@@ -400,47 +397,46 @@ def lower_call(call:UOp) -> UOp|None:
   return call.replace(src=(sink, *args), arg=replace(call.arg, aux=info))
 pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call", allow_any_len=True), lower_call)])
 
-hcq_compile_cache:dict[tuple[UOp, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
+hcq_compile_cache:dict[tuple[UOp, bool, bool], UOp] = {} # uops are hash-consed: the linear itself is the key, plus whether inputs bind
 
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
-  if any(isinstance(getattr(c.without_after.arg, "aux", None), HCQInfo) for c in linear.src): return linear # compiled already
-
   if input_uops is not None:
-    use_rt = len(linear.src) < HCQ_CACHE_THRESH # small schedules use runtime address patches so linked schedules can be cached without input buffers
     slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
-    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(use_rt, input_uops, slots), walk=True, name="replace buffers")
-    if (cached:=hcq_compile_cache.get(key:=(linear, profile))) is not None: return cached
-  lin = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, name="prep calls")
-  lin = sched_batches(lin, profile)
-  lin = graph_rewrite(lin, pm_encode, walk=True, name="encode")
-  with Context(EMULATED_DTYPES=""): final_linear = lower_and_compile(lin)
-  if input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
-  return final_linear
+    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer")
+
+  # 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:
+    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)
 
 # *****************
-# 5. link
+# 5. bufferize placeholders
 
-@dataclass
-class LinkCtx: inputs:dict[UOp, UOp]; use_rt:bool; refs:list[UOp] = field(default_factory=list) # noqa: E702
-
-def bufferize_buf(ctx:LinkCtx, b:UOp) -> UOp|None: # ctx: a kept link (the jit's) owns the linear's buffers, a one-shot borrows ring slots
+def bufferize_buf(ctx:bool, b:UOp) -> UOp|None: # ctx: a kept link (the jit's) owns the linear's buffers, a one-shot borrows ring slots
   if b.tag is None: return None # a param, not a placeholder
 
   dev = cast(HCQ2Compiled, Device[to_tuple(b.device)[0]])
 
-  # device owns the placeholders it names
-  if (r:=cast(Buffer|None, dev.pm_bufferize.rewrite(b, ctx=dev))) is not None: pass
-  elif not ctx.use_rt:
-    r = Buffer(dev.device, b.max_numel(), b.dtype, options=BufferSpec(host=b.arg.volatile, uncached=True, cpu_access=True), preallocate=True)
+  if b.arg.slot == 0 or b.tag == "program": r = cast(Buffer, unwrap(dev.pm_bufferize.rewrite(b, ctx=dev))) # device state and programs
+  elif ctx: r = Buffer(dev.device, b.max_numel(), b.dtype, options=BufferSpec(host=b.arg.volatile, uncached=True, cpu_access=True), preallocate=True)
   else: r = dev.rt_view(b.max_numel() * b.dtype.itemsize, b.dtype, host=b.arg.volatile)
 
   return UOp.from_buffer(r, HCQ_RUNTIME_DEV.value)
+pm_bufferize_placeholders = PatternMatcher([(UPat(Ops.PARAM, name="b"), bufferize_buf)])
 
-def resolve_getaddr(ctx:LinkCtx, g:UOp) -> UOp|None:
+# *****************
+# 6. link
+
+def resolve_getaddr(ctx:list[UOp], g:UOp) -> UOp|None:
   buf, off = unwrap_view(g.src[0])
   if buf.op not in {Ops.BUFFER, Ops.MSELECT}: return None
-  ctx.refs.append(buf) # add to refs
+  ctx.append(buf) # add to refs
   return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]).va_addr + off, dtypes.uint64)
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
@@ -458,30 +454,24 @@ def fold_words(buf:UOp, offs:UOp, ws:UOp) -> UOp:
   return UOp(Ops.NOOP)
 
 pm_link = PatternMatcher([
-  (UPat(Ops.CAST, src=(UPat(Ops.CAST, src=(UPat.cvar(),), name="inner"),), name="c"), lambda c, inner: inner.src[0].cast(c.dtype)),
-  (UPat(Ops.PARAM, name="b"), lambda ctx, b: ctx.inputs[b] if b in ctx.inputs else bufferize_buf(ctx, b)),
   (UPat(Ops.GETADDR, name="g"), resolve_getaddr),
   (UPat(GroupOp.ALU, src=UPat.cvar().or_casted(), name="a"),
     lambda a: UOp.const(exec_alu(a.op, a.dtype, [s.val for s in a.src], False), a.dtype)),
   (UPat(name="buf").store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast())), fold_binary),
   (UPat(name="buf").index(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="offs")).store(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="ws")),
     fold_words),
-  (UPat(Ops.AFTER, name="a"), lambda a: None if a.is_bound_var or a.src[0].op is Ops.CALL else
-   a.src[0] if all(s.op is Ops.NOOP for s in a.src[1:]) else panic(RuntimeError, f"unresolved link words on {a.src[0].op}")),
+  (UPat(Ops.AFTER, name="a"), lambda a: None if a.is_bound_var else a.src[0] if all(s.op is Ops.NOOP for s in a.src[1:]) else
+   panic(RuntimeError, f"unresolved link words on {a.src[0].op}")),
 ])
 
 link_linear_cache:weakref.WeakKeyDictionary[UOp, UOp] = weakref.WeakKeyDictionary() # a baked link lives as long as its bound linear
 
-@rewrite_group(lambda _,input_uops=None,allow_cache=True,ret=None: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
-  if allow_cache and (linked:=link_linear_cache.get(linear)) is not None: return linked
-
-  # if we have any link time buffers, do not cache this linear
-  cache = allow_cache and not any(u.tag == "lt_input" for u in linear.toposort() if u.op is Ops.PARAM)
-
-  inputs = {UOp.param(i, b.dtype, b.max_numel(), b.device).replace(tag="lt_input"): b for i, b in enumerate(input_uops or ())}
-  linked = graph_rewrite(linear, pm_link, ctx=(ctx:=LinkCtx(inputs, use_rt=allow_cache and not cache)), walk=True, name="link")
-  if ctx.refs: linked = linked.replace(src=(linked.src[0].after(*dedup(ctx.refs)), *linked.src[1:])) # attach refs to linear
+@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+def hcq_link(linear:UOp, cache=True) -> UOp:
+  if (linked:=link_linear_cache.get(linear)) is not None: return linked
+  bufferized = graph_rewrite(linear, pm_bufferize_placeholders, ctx=cache, name="bufferize")
+  linked = graph_rewrite(bufferized, pm_link, ctx=(refs:=list[UOp]()), bottom_up=False, name="link")
+  if refs: linked = linked.replace(src=(linked.src[0].after(*dedup(refs)), *linked.src[1:])) # attach refs to linear
   if cache and linked is not linear: link_linear_cache[linear] = linked
   return linked
 
@@ -491,7 +481,6 @@ def hcq_link(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UO
 class HCQ2Compiled(Compiled):
   timestamp_divider: float = 1000.0
   wait_timeout_ms: float = 30000.0
-  sleep_timeout_ms: int|None = None
   rt_nbytes: int = 64 << 20 # the pool every per-linear buffer is carved out of
   pm_encode: PatternMatcher = PatternMatcher([]) # the backend's own encode rules, matched by its submit names
   var_vals: dict[str, int] = {}
@@ -543,7 +532,7 @@ class HCQ2Compiled(Compiled):
     return Buffer(self.device, self.rt_allocator(uncached, host).size, dtypes.uint8, options=spec, preallocate=True)
 
   def rt_view(self, nbytes:int, dtype:DType=dtypes.uint8, uncached:bool=True, host:bool=False) -> Buffer: # a slot of the ring, wraps silently
-    off = self.rt_allocator(uncached, host).alloc(max(nbytes, 1), alignment=256)
+    off = self.rt_allocator(uncached, host).alloc(max(nbytes, 1), alignment=128)
     return self.rt_buffer(uncached, host).view(nbytes // dtype.itemsize, dtype, off).ensure_allocated()
 
   def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
@@ -551,8 +540,7 @@ class HCQ2Compiled(Compiled):
     st, done = time.perf_counter(), sig[0]
     while done < value:
       if done != (done:=sig[0]): st = time.perf_counter()
-      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
-      elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
+      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
 
   def synchronize(self, timeout:int|None=None):
     self._wait_signal(tl:=self.timeline._buf.cpu_view().view(fmt='Q'), tl[1], timeout)
@@ -560,9 +548,6 @@ class HCQ2Compiled(Compiled):
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
 
-  def on_sleep(self):
-    if (iface:=getattr(self, "iface", None)) is not None and hasattr(iface, "sleep"): iface.sleep(self.sleep_timeout_ms)
-
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
@@ -582,7 +567,8 @@ class HCQ2Buffer:
     return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
 
 class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
-  def _as_buffer(self, buf:HCQBuffer) -> memoryview|None: return buf.view.mv if buf.view is not None else None
+  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()

tinygrad/runtime/support/system.py

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

tinygrad/schedule/__init__.py

@@ -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, tuple[UOp, ...]]]] = {}  # superseded state -> (AFTER, new kernels)
+    writes: dict[UOp, list[tuple[UOp, UOp, tuple[UOp, ...]]]] = {}  # buffer -> (AFTER, prior state, new kernels)
     reads: list[tuple[UOp, UOp, UOp]] = []  # (reader AFTER, reader kernel, buffer state read)
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       kernels, after_deps = _split_after(u)
       prev_state = _unwrap_src(u.src[0])
       prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set()
-      writes.setdefault(prev_state, []).append((u, tuple(k for k in kernels if k not in prev_kernels)))
+      writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels)))
       for k in kernels:
         in_degree.setdefault(k, 0)
         if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
@@ -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, write_kernels in writes.get(s, []):
-        if a is u: continue
+      for a, prev_state, write_kernels in writes.get(s.buf_uop, []):
+        if a is u or prev_state is not s: continue
         for t in write_kernels:
           if t is not k and t not in k.backward_slice:
             children.setdefault(k, []).append(t)
@@ -80,7 +80,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
-from tinygrad.schedule.prepare import prepare_rangeify, prepare_call_views
+from tinygrad.schedule.prepare import prepare_rangeify
 from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
@@ -122,14 +122,15 @@ def lower_sink_to_linear(call:UOp) -> UOp|None:
   if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): 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
-  call = prepare_call_views(call)
-  function = call.src[0]
   st = time.perf_counter()
   cache_key = function.key
   if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
     if SPEC: type_verify(function, spec_tensor)
+    # prepare can introduce new opaque CALLs (transform_precompiled_call turns precompiled calls opaque here): lower
+    # nested calls first so create_schedule doesn't look through their unscheduled bodies
     # support recursive CALLs
-    linear = create_schedule(get_kernel_graph(prepare_rangeify(function)))
+    prepared = graph_rewrite(prepare_rangeify(function), pm_schedule, name="schedule nested calls", enter_calls=True)
+    linear = create_schedule(get_kernel_graph(prepared))
     if SCACHE: schedule_cache[cache_key] = linear
   else:
     # schedule cache hit

tinygrad/schedule/prepare.py

@@ -8,108 +8,56 @@ from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.schedule.allreduce import create_allreduce_function
 from tinygrad.schedule.multi import multi_pm
 
-def on_disk(u:UOp): return isinstance(u.device, str) and u.device.startswith("DISK")
-
-def contiguous_mops_to_view(ctx:list[UOp]|None, c:UOp, src:UOp):
-  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
-  # A list holds CALL arguments; None rewrites views in the live Tensor graph.
-  # Ordinary copies keep their source graph so JIT can substitute its input buffer.
-  if ctx is None and c.op is Ops.COPY and not on_disk(src): return None
-  buf = src.base
-  while buf.op is Ops.BITCAST: buf = buf.src[0].base
-  # no symbolic shape
-  if buf.op not in {Ops.BUFFER, Ops.PARAM, Ops.UNSHARD} or not all_int(c.shape): return None
-
-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then view the resolved shard
-  unshard = None
-  if buf.op is Ops.UNSHARD:
-    if isinstance(c.device, str): return None
-    if (unshard := graph_rewrite(src, multi_pm, name="multi_buffer_view")).op is not Ops.UNSHARD: return None
-    src = unshard.src[0]
-
-  # offset the base buffer by the collapsed movement ops and view it
-  if (cv := src.contiguous_view()) is None or (buf := cv[0]).op not in {Ops.BUFFER, Ops.PARAM}: return None
-  view = buf[cv[1]:cv[1] + src.max_numel() * src.element_size() // buf.element_size()].bitcast(src.dtype)
-  if ctx is not None and view.op in {Ops.SHRINK, Ops.BITCAST}:
-    arg = view.substitute({u: ctx[u.arg.slot] for u in view.toposort() if u.op is Ops.PARAM and u.arg.slot >= 0})
-    if arg not in ctx: ctx.append(arg)
-    view = view.param_like(ctx.index(arg))
-  elif on_disk(buf) and buf.op is Ops.BUFFER and not buf.is_unbound: view = UOp.from_buffer(view.buffer, device=buf.device)
-  view = view.reshape(src.shape).unshard(unshard.arg, unshard.src[1:]) if unshard is not None else view.reshape(c.shape)
-  return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
-
-# Fold contiguous movement operations into buffer views.
-pm_mops_to_view = PatternMatcher([
-  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
-  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
-
-  # remove contiguous on movement ops before a copy on disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
-   copy.replace(src=(x,), tag=None) if on_disk(x) else None),
-  # push copy past movement ops on disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
-   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) if on_disk(x) else None),
-])
+# args without storage: fold movement ops that collapse to a contiguous range over a buffer into a static view of
+# the base buffer (like contiguous_mops_to_view); everything else materializes with CONTIGUOUS
+def arg_with_storage(a:UOp) -> UOp:
+  if a.has_buffer_identity(after_ok=True): return a
+  if all_int(a.shape) and len(a.shape) > 0 and (cv := a.contiguous_view()) is not None and (b := cv[0]).op in {Ops.BUFFER, Ops.PARAM}:
+    return b[cv[1]:cv[1] + a.max_numel() * a.element_size() // b.element_size()].bitcast(a.dtype).reshape(a.shape)
+  return a.contiguous()
 
 def transform_precompiled_call(c:UOp) -> UOp|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"
-  # Bind output storage at the existing argument positions.
-  outs = {p: a.empty_like() for p,a in enumerate(c.src[1:]) if a.unsharded_base.is_unbound}
+  # 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]
+  srcs = tuple(st.src[1] for st in c.src[0].src if st.op is Ops.STORE)
+
+  # add the outputs to the call
+  outs = tuple(c.src[1+p].empty_like() for p in ret_pos)
+  targets = [o.param_like(p).shrink_to(s.shape) for p,o,s in zip(ret_pos, outs, srcs)]
+
+  # how each stored value lands in its output PARAM target: a CONTIGUOUS materializes straight into the target and
+  # a real buffer/UNSHARD rebinds its storage to the target (once per unique value); everything else is copied into it
   placed:dict[UOp, UOp] = {}
-  items = []
-  for st in c.src[0].src:
-    value = st.src[1]
-    while value.op is Ops.AFTER: value = value.src[0]
-    # A custom kernel's output buffer can be the call output directly. Rebind each buffer only once.
-    if value.op in {Ops.BUFFER, Ops.UNSHARD} and value.has_buffer_identity() and value not in placed:
-      placed[value] = st.src[0]
-      items.append(st.src[1])
-    else: items.append(st.src[0].after(st))
-  body = UOp.sink(*items).substitute(placed)
-  call = c.replace(src=(body, *(outs.get(i, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
-                                   for i, a in enumerate(c.src[1:]))))
-  return UOp.sink(*(c.src[1+p].store(o.after(call).shrink_to(c.src[1+p].shape)) for p,o in outs.items()))
+  items:list[UOp] = []
+  for s, t in zip(srcs, targets):
+    deps:list[UOp] = []
+    while s.op is Ops.AFTER:
+      deps.extend(s.src[1:])
+      s = s.src[0]
+    if s not in placed:
+      if s.op is Ops.CONTIGUOUS: placed[s] = t.after(t.store(s.src[0]))
+      elif s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): placed[s] = t
+      if s in placed:
+        items.append(s.after(*deps))
+        continue
+    items.append(t.after(t.store(s.after(*deps))))
+  # swap every placed value for its target storage, also inside other stores' AFTER deps
+  fxn = UOp.sink(*(x.substitute(placed) for x in items))
 
-pm_resolve_call_outputs = PatternMatcher([
-  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
-  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
-])
+  # 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 = c.replace(src=(fxn, *[rmap.get(i, arg_with_storage(a)) for i, a in enumerate(c.src[1:])]))
+  rets = tuple(o.after(new_call) for o in outs)
 
-def buffer_view_subs(sink:UOp) -> dict[UOp, UOp]:
-  # Include intermediate nodes so every Tensor sharing a pending write receives the same view rewrite.
-  nodes = list(sink.toposort(enter_calls=False))
-  rewritten = graph_rewrite(UOp.sink(*nodes), pm_mops_to_view, bottom_up=True, name="fold buffer views")
-  return {u: v for u, v in zip(nodes, rewritten.src) if u is not v}
+  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
+  # NOTE: must use the resolved shapes of the RETURNED placeholders (which substitute PARAMs with external args), not raw body shapes
+  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, (c.src[1+p] for p in ret_pos)))
 
-def prepare_call_views(call:UOp) -> UOp:
-  # Lift contiguous views into call arguments, preserving their buffer/offset graph for JIT input substitution.
-  args = list(call.src[1:])
-  body = graph_rewrite(call.src[0], pm_mops_to_view, ctx=args, bottom_up=True, name="prepare call views")
-  return call.replace(src=(body, *args))
-
-def prepare_to_call(sink:UOp, tensor_roots:tuple[UOp, ...]) -> UOp:
-  # A copy used only to initialize another buffer can write directly into that destination.
-  # Include live Tensor graphs so retained copies and aliases keep their independent storage.
-  users:dict[UOp, set[UOp]] = {}
-  for u in UOp.sink(sink, *tensor_roots).toposort(enter_calls=False):
-    for src in u.src: users.setdefault(src, set()).add(u)
-  subs = {}
-  for store in sink.toposort(enter_calls=False):
-    if store.op is not Ops.STORE: continue
-    value = store.src[1]
-    if value.op is not Ops.AFTER or len(value.src) != 2: continue
-    buf, init = value.src
-    if init.op is not Ops.STORE or len(init.src) != 2 or init.src[0] is not buf or init.src[1].op is not Ops.COPY: continue
-    # Only this assignment may consume the copy, and only the initialization may use its storage.
-    if users.get(value) != {store} or users.get(buf) != {value, init}: continue
-    while buf.op is Ops.RESHAPE and users.get(buf.src[0]) == {buf}: buf = buf.src[0]
-    if buf.op is not Ops.BUFFER or buf.is_unbound or buf.buffer.is_allocated(): continue
-    subs[value] = init.src[1]
-  sink = sink.substitute(subs, walk=True)
-  sink = graph_rewrite(sink, pm_resolve_call_outputs, bottom_up=True, name="resolve call outputs")
-  return UOp.sink(*[u for u in sink.toposort(enter_calls=False)
-                   if u.op is Ops.AFTER and not u.is_bound_var and not u.src[0].unsharded_base.is_unbound])
+  # the AFTER outputs resolve against this: stores of each real output into its RETURNED placeholder
+  return UOp.sink(*[c.src[1+p].store(v) for p, v in zip(ret_pos, rets)])
 
 def walk_mop(u:UOp):
   if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD, Ops.BITCAST}: return walk_mop(u.src[0])
@@ -127,8 +75,6 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
   ctx[x] = after
 
 # *** fold moved AFTERs (hack for openpilot) ***
-# These temporary stores exist only in the schedule; they do not persist Tensor intermediates.
-pm_contiguous_to_store = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="c"), lambda c: c.clone())])
 pm_fold_moved_after = PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
   # replace ALU sources with AFTER versions found above
@@ -234,10 +180,16 @@ def expand_bitcast(bc:UOp) -> UOp|None:
   parts = [tmp>>8*i*ns for i in range(os//ns)]
   return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
 
-earliest_rewrites = mop_cleanup+pm_resolve_call_outputs+PatternMatcher([
-  # Inline calls with unbound outputs.
+earliest_rewrites = mop_cleanup+PatternMatcher([
+  # transform precompiled value-producing calls into opaque CALLs (outputs become real buffers)
+  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
+
+  # resolve calls with RETURNED inputs (inline the body)
   (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),
+
   # resolve allreduce (must be bottom up)
   (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
 
@@ -313,13 +265,52 @@ pm_copy_to_store = PatternMatcher([
   (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
 ])
 
+def bind_call_outputs(sink:UOp) -> UOp:
+  """fresh outputs of opaque precompiled calls directly stored into final storage: bind the storage into the call's
+  output positions and drop the stores; the call then writes the storage directly, no copy kernel (like
+  transform_precompiled_call giving outputs real buffers). done for all consumers of a call at once (all or
+  nothing) so the call stays a single kernel instance.
+
+  NOTE: output positions aren't recorded on the opaque call, so this treats any arg with an AFTER(call) consumer as
+  an output. that's sound here: the only AFTERs on call args come from the transform's resolved outputs — inputs
+  are read through their own states"""
+  parents: dict[UOp, list[UOp]] = {}
+  uses: dict[UOp, int] = {}
+  for u in sink.toposort(enter_calls=False):
+    for s in u.src:
+      parents.setdefault(s, []).append(u)
+      uses[s] = uses.get(s, 0)+1
+  subs: dict[UOp, UOp] = {}
+  for call in [u for u in uses if u.op is Ops.CALL and u.arg is not None and u.arg.precompile and not u.has_unbound_outputs]:
+    # every AFTER consumer of the call's outputs must be a single-use direct store into bindable final storage
+    items: list[tuple[UOp, UOp, UOp, UOp]] = []
+    for after in [a for a in parents.get(call, []) if a.op is Ops.AFTER and a.src[0] in call.src[1:]]:
+      out = after.src[0]
+      # dst must have buffer identity: kernel args bind via buf_uop (base buffer), which would drop a view's offset
+      ok = out.has_buffer_identity() and uses.get(out, 0) == 2 and len(pars:=parents.get(after, [])) == 1 and pars[0].op is Ops.STORE \
+        and pars[0].src[1] is after and (dst:=pars[0].src[0]).has_buffer_identity(after_ok=True) and not dst.unsharded_base.is_unbound \
+        and dst.dtype is out.dtype and dst.device == out.device and dst._shape is not None and dst._shape == out._shape
+      if not ok: break
+      items.append((pars[0], dst, after, out))
+    else:
+      if not items: continue
+      new_call = call.substitute({out: dst for _, dst, _, out in items}, enter_calls=False)
+      for st, dst, after, out in items: subs[st] = after.substitute({out: dst, call: new_call}, enter_calls=False)
+  return sink.substitute(subs, name="bind call outputs") if subs else sink
+
+# resolve precompiled call outputs to real buffers before anything else sees the calls (like multi_pm)
+pm_resolve_call_outputs = PatternMatcher([
+  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
+  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
+])
+
 @rewrite_group(new_ctx=False)
 def prepare_rangeify(sink:UOp) -> UOp:
+  # transform precompiled value-producing calls into opaque CALLs (outputs become real buffers)
+  tsink = bind_call_outputs(graph_rewrite(sink, pm_resolve_call_outputs, bottom_up=True, name="resolve call outputs"))
   # prepare for rangeify
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
-  if OPENPILOT_HACKS:
-    tsink = graph_rewrite(tsink, pm_contiguous_to_store, bottom_up=True, name="materialize contiguous")
-    tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
+  tsink = graph_rewrite(tsink, multi_pm, name="multi_pm")
+  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
   tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
   return tsink

tinygrad/tensor.py

@@ -1,55 +1,201 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
-from dataclasses import replace
+from dataclasses import dataclass, field, replace
 from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, _from_np_dtype, _to_np_dtype, PyConst, AddrSpace
 from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize, SPEC
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, graph_rewrite, rewrite_group
+from tinygrad.uop.ops import remove_all_tags
 from tinygrad.uop.spec import type_verify, spec_tensor
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
-from tinygrad.schedule.prepare import buffer_view_subs, prepare_to_call, on_disk
+from tinygrad.schedule.multi import multi_pm
 from tinygrad.device import Buffer, canonicalize_device
 from tinygrad.engine.realize import run_linear
 
 # *** callify: transform a tensor graph into a CALL UOp such that all state is properly scoped ***
 
-@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret.src)-1)}")
-def transform_to_call(big_sink:UOp) -> UOp:
-  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
-  if SPEC: type_verify(big_sink, spec_tensor)
-  # Storage declarations have unique global IDs; canonicalize them, including declarations inside nested calls.
-  unbound = [u for u in big_sink.toposort() if u.is_unbound]
-  body = big_sink.substitute({u: u.replace(arg=replace(u.arg, slot=i)) for i,u in enumerate(unbound)},
-                             enter_calls=True, walk=True, name="renumber buffers")
-  # PARAMs belong to the enclosing scope. Nested call bodies keep their own positional PARAMs.
-  inputs = [u for u in body.toposort(enter_calls=False)
-            if (u.op is Ops.PARAM and (u.addrspace is not AddrSpace.ALU or u.arg.slot >= 0)) or u.is_bound_var or
-               (u.op is Ops.BUFFER and u.addrspace is AddrSpace.GLOBAL and not u.is_unbound)]
-  params = {u: u.replace(arg=replace(u.arg, slot=i, name=f"p{i}" if u.addrspace is AddrSpace.ALU else u.arg.name))
-            if u.op is Ops.PARAM else u.param_like(i) for i,u in enumerate(inputs)}
-  ret = body.substitute(params, walk=True, name="replace inputs").call(*inputs)
-  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
-  return ret
+@dataclass
+class AllocCtx:
+  buffer_map: dict[UOp, UOp] = field(default_factory=dict)
+  bases: set[UOp] = field(default_factory=set)
+  stores: list[UOp] = field(default_factory=list)
+  replacements: list[UOp] = field(default_factory=list)
+  unbound: dict[UOp, UOp] = field(default_factory=dict)
+  views: set[UOp] = field(default_factory=set)
 
+# 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,))
+
+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(())
+  # 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),
+  # 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),
+  (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop),
+  (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(x) if x in ctx.bases else None),
+])
+
+def replace_contig_with_store_after(u:UOp):
+  # can't allocate a buffer for a virtual value
+  if u.is_virtual: return None
+  # if size is 0, remove the contig
+  if 0 in u.shape: return u.src[0]
+  # no real contig for DISK tensors, they are left alone
+  if on_disk(u): return u.rtag(None)
+  buf = u.empty_like()
+  return buf.after(buf.store(u.src[0])).rtag(u.tag)
+
+def wrap_tagged_in_contig(x:UOp):
+  if x.tag is None: return None          # untouched
+  # empty tag from rtag(()): a COPY already handled via buffer_map or merged into a parent AFTER.
+  # () is falsy but not None, so it isn't re-tagged like a bare (tag=None) node would be; just strip it here
+  if not x.tag: return x.rtag(None)
+  return x.rtag(None).contiguous(tag=x.tag)  # the tag moves onto the wrapping CONTIGUOUS
+
+def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
+  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
+  buf = src.base
+  while buf.op is Ops.BITCAST: buf = buf.src[0].base
+  # no symbolic shape
+  if buf.op not in {Ops.BUFFER, Ops.UNSHARD} or not all_int(c.shape): return None
+
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then view the resolved shard
+  unshard = None
+  if buf.op is Ops.UNSHARD:
+    if isinstance(c.device, str): return None
+    if (unshard := graph_rewrite(src, multi_pm, name="multi_buffer_view")).op is not Ops.UNSHARD: return None
+    src = unshard.src[0]
+
+  # offset the base buffer by the collapsed movement ops and view it
+  if (cv := src.contiguous_view()) is None or (buf := cv[0]).op is not Ops.BUFFER: return None
+  # NB: make offset a UOp.variable here to do the offset computation in the kernels
+  view = buf[cv[1]:cv[1] + src.max_numel() * src.element_size() // buf.element_size()].bitcast(src.dtype)
+  ctx.views.add(view)
+  if unshard is not None: return view.reshape(src.shape).unshard(unshard.arg, unshard.src[1:])
+  view = view.reshape(c.shape)
+  return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
+
+# NOTE: adding rules to here is bad. these all need to run before the schedule cache
+pm_early_transform_tensor_graph = PatternMatcher([
+  # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
+  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
+  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
+
+  # remove contiguous on movement ops before a copy on disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
+   copy.replace(src=(x,), tag=None) if on_disk(x) else None),
+  # push copy past movement ops to disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
+   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) if on_disk(x) else None),
+
+  # add CONTIGUOUS to tagged UOps
+  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"), wrap_tagged_in_contig),
+  # remove extra CONTIGUOUS on AFTER (only when target is contiguous)
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"),
+   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
+  # replace CONTIGUOUS with STORE+AFTER
+  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after),
+  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+])
+
 # a store's storage keeps the views and drops AFTERs (they only sequence stores)
 pm_drop_after = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: a.src[0])])
 
+def replace_input_buffer(ctx:AllocCtx, b:UOp):
+  ctx.replacements.append(b)
+  return b.param_like(len(ctx.replacements)-1)
+
+# unbound BUFFERs get canonical scope-local id slots here so structurally identical calls hash identically for the
+# schedule cache (fresh slots are all positive from the global counter; negative slots are already canonical)
+def canonicalize_unbound_buffer(ctx:AllocCtx, b:UOp):
+  if b.arg.slot >= 0 and b not in ctx.unbound: ctx.unbound[b] = b.replace(arg=replace(b.arg, slot=-1-len(ctx.unbound)))
+  return ctx.unbound.get(b)
+
+def canonicalize_call_body(ctx:AllocCtx, c:UOp):
+  body = graph_rewrite(c.src[0], pm_canonicalize_unbound, ctx=ctx, bottom_up=True)
+  return c.replace(src=(body,)+c.src[1:]) if body is not c.src[0] else None
+
+pm_canonicalize_unbound = PatternMatcher([
+  (UPat(Ops.CALL, name="c"), canonicalize_call_body),
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b: canonicalize_unbound_buffer(ctx, b) if b.is_unbound else None),
+])
+
+pm_replace_buf = pm_canonicalize_unbound+PatternMatcher([
+  # replace BUFFER with PARAM for cache key normalization (ALU addrspace buffers are Variables, they stay, and unbound BUFFERs too)
+  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
+   replace_input_buffer(ctx, b) if b.addrspace is AddrSpace.GLOBAL and not b.is_unbound else None),
+  # replace buffer views (SHRINK/BITCAST) with PARAM (only the views created by contiguous_mops_to_view)
+  (UPat((Ops.SHRINK, Ops.BITCAST), name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b in ctx.views else None),
+  # strip the stored value from bound Variables for cache key normalization, so different values hit same cache
+  (UPat(Ops.AFTER, name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.is_bound_var else None),
+])
+
+@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
+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})
+
+  # 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
+  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="add tags")
+
+  # final outputs of value calls materialize with fresh storage
+  # precompiled calls resolve to opaque CALLs in prepare (after callify), not by inlining: their outputs need the
+  # storage minted now so the call survives in the call body (a CONTIGUOUS can't be minted on the device-less AFTER).
+  # no tag: the tensor keeps the AFTER(RETURNED, CALL) node; backward and later realizes resolve through it again
+  srcs:list[UOp] = []
+  for u in big_sink.src:
+    if u.op is Ops.AFTER and u.src[0].unsharded_base.is_unbound:
+      if (call:=u.src[1]).op is Ops.CALL and call.arg is not None and call.arg.precompile:
+        buf = (val:=u.rtag(None)).empty_like()
+        u = buf.after(buf.store(val))
+      else: u = u.rtag(None).contiguous(tag=u.tag)
+    srcs.append(u)
+  big_sink = big_sink.replace(src=tuple(srcs))
+
+  # here we can break the tensor graph. tags propagate through replaces so we can still find the original UOps
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
+
+  # collect the stores (never entering call bodies) and map tagged AFTERs to their storage; tags are stripped at the end
+  # copies to disk are stores to the disk buffer; bound Variables are call inputs and RETURNEDs are call outputs
+  for u in big_sink.toposort(enter_calls=False):
+    if (u.op is Ops.COPY and on_disk(u)) or (u.op is Ops.AFTER and not u.is_bound_var and not u.src[0].unsharded_base.is_unbound):
+      ctx.stores.append(u)
+      if u.tag: ctx.buffer_map.update({t:graph_rewrite(u.src[0], pm_drop_after).shrink_to(t.shape) for t in u.tag})
+  ret = graph_rewrite(UOp.sink(*ctx.stores), pm_replace_buf+remove_all_tags, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
+  assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
+  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
+  return ret, ctx.buffer_map
+
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, *, tensors:list[Tensor]|None=None) -> None:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
     def visitor(node: UOp) -> bool: return True if node in applied_map else any(in_scope.get(s, False) for s in node.src)
-    if tensors is None: tensors = [t for tref in list(all_tensors) if (t:=tref()) is not None]
-    scope_tensors = [t for t in tensors if t.uop.topovisit(visitor, in_scope)]
+    scope_tensors: list[Tensor] = [t for tref in list(all_tensors) if (t:=tref()) is not None and t.uop.topovisit(visitor, in_scope)]
 
     # get all Tensors and apply the map. always walk: replace exactly the nodes the map names, values are final
     sink = UOp.sink(*[t.uop for t in scope_tensors])
@@ -60,14 +206,9 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, *, tensors:list[
       if s is ns: continue
       t.uop = ns
 
-# **** Tensor helper functions ****
+def _tensor_holds(u:UOp) -> bool: return any((t:=tref()) is not None and t.uop is u for tref in list(all_tensors))
 
-def _inplace_rhs(update:UOp) -> UOp|None:
-  # Recover the computed value of a read-modify-write; ordinary clone stores are not self-referential.
-  if update.op is not Ops.AFTER or len(update.src) != 2: return None
-  store = update.src[1]
-  if store.op is not Ops.STORE or store.src[0] not in store.src[1].toposort(enter_calls=False): return None
-  return store.src[1]
+# **** Tensor helper functions ****
 
 def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
 
@@ -127,9 +268,7 @@ class Tensor(RandMixin):
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
 
     # data might be on a different device
-    self.uop:UOp = data
-    if data.device is not None and data.device != _device:
-      self.uop = data.clone(_device) if is_creation_device(data) else data.copy_to_device(_device)
+    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device)
     # cast on the target device, the source may not hold the dtype (numpy has no fp8/bfloat16) or be able to compute it (DISK)
     if _dtype is not None: self.uop = self.uop.cast(_dtype)
 
@@ -202,33 +341,10 @@ class Tensor(RandMixin):
     """
     return [Tensor(u) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
 
-  def _prepare_call(self, *lst:Tensor) -> tuple[UOp, dict[UOp, UOp]]:
-    outs = (self,)+lst
-    _apply_map_to_tensors(buffer_view_subs(UOp.sink(*[x.uop for x in outs])), name="fold buffer views")
-    # Only requested outputs acquire storage. Intermediate values persist only when explicitly cloned.
-    bases = set()
-    for x in outs:
-      base = x.uop.base
-      while base.op is Ops.CONTIGUOUS_BACKWARD: base = base.src[0].base
-      bases.add(base)
-    subs:dict[UOp, UOp] = {}
-    for u in UOp.sink(*bases).toposort(enter_calls=False):
-      if u not in bases or u.is_virtual or on_disk(u): continue
-      if u.has_buffer_identity(after_ok=True) or u.storage_base.has_buffer_identity(): continue
-      if u.op is Ops.AFTER and u.src[1].op is Ops.CALL and u.src[1].arg.precompile: continue
-      subs[u] = u.substitute(subs, walk=True).clone()
-    _apply_map_to_tensors(subs, name="materialize")
-    sink = UOp.sink(*[x.uop for x in outs])
-    becomes_map = {u: graph_rewrite(u.src[0], pm_drop_after).shrink_to(u.shape)
-                   for u in sink.toposort(enter_calls=False)
-                   if u.op is Ops.AFTER and not u.is_bound_var and not u.src[0].unsharded_base.is_unbound}
-    tensor_roots = tuple(t.uop for ref in list(all_tensors) if (t:=ref()) is not None)
-    return transform_to_call(prepare_to_call(sink, tensor_roots)), becomes_map
-
   def callify(self, *lst:Tensor) -> Tensor:
-    """Groups the computation for these tensors into a deferred call. Returns `self` without executing the call."""
-    call, becomes_map = self._prepare_call(*lst)
-    _apply_map_to_tensors({x:y.after(call) for x,y in becomes_map.items()}, name="callify")
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    big_sink, buffer_map = transform_to_call(big_sink)
+    _apply_map_to_tensors({x:y.after(big_sink) for x,y in buffer_map.items()}, name="callify")
     return self
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
@@ -236,9 +352,9 @@ class Tensor(RandMixin):
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    call, becomes_map = self._prepare_call(*lst)
+    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
     _apply_map_to_tensors(becomes_map, name="buffers")
-    return create_linear_with_vars(call)
+    return create_linear_with_vars(big_sink)
 
   def schedule_linear(self, *lst:Tensor) -> UOp:
     """Creates the schedule needed to realize these Tensor(s)."""
@@ -249,7 +365,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 (b:=x.uop.base).is_virtual and not b.has_buffer_identity()]
+    to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
     if len(to_realize):
       run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
     return self
@@ -264,12 +380,6 @@ class Tensor(RandMixin):
     return self
 
   def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
-    """
-    Assigns `x` to this tensor and returns `self`. `x` must broadcast to this tensor's shape.
-    Tensor inputs must match its dtype and device, except that disk tensors accept inputs from any device.
-    Updates existing storage, or creates storage if this tensor is a computed value.
-    The write is deferred until realization, except for disk tensors.
-    """
     if self.dtype in dtypes.weaks: self.uop = self.uop.clone()
     is_disk = on_disk(self.uop)
     if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype)
@@ -287,26 +397,26 @@ class Tensor(RandMixin):
     if is_disk:
       (b:=self._buffer()).copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=x._data()))
       return self
-    # Assigning to a value initializes new storage; assigning to a buffer updates its storage.
-    if not self.uop.storage_base.has_buffer_identity():
-      self.uop = x.uop.clone()
+    # 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
       return self
-    update = self.uop.after(self.uop.store(x.uop))
-    base = self.uop
-    # Direct assignments need no alias search. A held reshape of a buffer also owns its update.
-    if not base.has_buffer_identity() and base.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH}:
-      tensors = [t for ref in list(all_tensors) if (t:=ref()) is not None]
-      held = {t.uop for t in tensors}
-      # Find the owning Tensor's buffer or pending write, preserving its shape for function argument substitution.
-      while base.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH}:
-        if base.has_buffer_identity() and base in held: break
-        base = base.src[0]
-      if base.has_buffer_identity(after_ok=True):
-        # Detach shares storage, but an assignment through it must not rewrite earlier computations using that storage.
-        if self.uop.op is Ops.DETACH: tensors = [t for t in tensors if t.uop.storage_base is base.storage_base]
-        _apply_map_to_tensors({base: base.after(update)}, name="Embed View Assign", tensors=tensors)
-        return self
-    self.uop = update
+    # 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
+      _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign")
+    else:
+      # simple assign
+      self.uop = assign
     return self
 
   def _buffer(self) -> Buffer:
@@ -379,8 +489,7 @@ class Tensor(RandMixin):
 
   def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor:
     """
-    Creates a tensor with independent storage, populated lazily when its value is needed.
-    Use this to retain an intermediate result across realizations or to modify it independently.
+    Creates a clone of this tensor allocating a separate buffer for the data.
     If `device` is specified, the clone is placed on that device.
     """
     ret = Tensor(self.uop.clone(device=device))
@@ -389,14 +498,11 @@ class Tensor(RandMixin):
 
   def to(self, device:str|tuple[str, ...]|None) -> Tensor:
     """
-    Returns this tensor on the given device, transferring its data lazily. Returns `self` if the device already matches.
-    Use `clone(device)` when the result needs independent, persistent storage.
+    Moves the tensor to the given device.
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # Copies from creation devices and copies to disk own persistent storage.
-    if is_creation_device(self.uop) or (isinstance(device, str) and device.startswith("DISK")): ret = Tensor(self.uop.clone(device))
-    else: ret = Tensor(self.uop.copy_to_device(device))
+    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)
 
@@ -421,9 +527,7 @@ 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))
-    # 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))
+    uop = self.uop.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:
@@ -538,20 +642,12 @@ class Tensor(RandMixin):
     if isinstance(v, Tensor):
       if v.dtype in dtypes.weaks: v = v.cast(least_upper_dtype(self.dtype, v.dtype))
       if v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
-    # Augmented view assignment may already have embedded its STORE in the parent. Undo that dependency
-    # before the functional setitem below, while retaining the computed RHS for autograd.
-    if isinstance(v, Tensor) and self.is_floating_point() and not self.uop._base_buffer_is_realized():
-      a = self.uop
-      if a.op is Ops.AFTER and len(a.src) == 2 and a.src[1] in v.uop.backward_slice and (view_rhs:=_inplace_rhs(a.src[1])) is not None:
-        _apply_map_to_tensors({a: a.src[0]}, name="functional setitem")
-        v = v._apply_uop(lambda _: view_rhs)
     # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
     v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
     if self.uop.op_in_backward_slice_with_self(Ops.BUFFER):
       shared = self.uop.base if self.uop.base.is_realized else None
       if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors
              if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
-        self._getitem(indices)  # invalid indices take precedence over the mutation restriction
         raise RuntimeError("can't setitem on a tensor with other uses")
     idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
     is_disk = on_disk(self.uop)
@@ -559,7 +655,8 @@ 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)
-      if (rhs:=_inplace_rhs(v.uop)) is not None: v = v._apply_uop(lambda _, rhs=rhs: rhs)
+      # __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])
       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

@@ -798,8 +798,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # *** uop Buffer stuff ***
 
-  # Fresh storage IDs decrease from -1; canonical slots are numbered from 0 within their scope.
-  unique_num = itertools.count(-1, -1)
+  unique_num = itertools.count(0)
 
   def getaddr(self, device=None) -> UOp:
     if self.without_after.op not in {Ops.BUFFER, Ops.SHRINK, Ops.BITCAST, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM, Ops.LINEAR}: return self
@@ -817,11 +816,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return UOp(Ops.BUFFER, arg=ParamArg(-id(opaque), opaque.dtype, size=opaque.size, device=device or opaque.device, buffer=opaque))
   def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None) -> UOp:
     device = canonicalize_device(self.device if device is None else device)
-    dt = self.commit_dtype() if dtype is None else dtype
-    if self.op is Ops.UNSHARD and isinstance(device, tuple):  # mirror the sharding on the fresh storage
-      return UOp.empty(self.src[0].shape, dtype=dt, device=device).unshard(self.arg, self.src[1:])
     axis = self.axis if isinstance(device, tuple) else None
-    ret = UOp.empty(self.shard_shape if axis is not None else self.shape, dtype=dt, device=device)
+    ret = UOp.empty(self.shard_shape if axis is not None else self.shape, dtype=self.commit_dtype() if dtype is None else dtype, device=device)
     return ret.unshard(axis) if axis is not None else ret
   @staticmethod
   def _frompy(x:list|tuple|bytes, dtype:DType, device:str|tuple[str, ...]|None=None) -> UOp:
@@ -835,13 +831,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       data = struct.pack(f"{prod(shape)}{bdtype.fmt}", *[truncate[bdtype](bdtype.const(xi)) for xi in fully_flatten(x)])
     ret.buffer.allocate(memoryview(bytearray(data))) # fake realize. buffer storage must be writable, and bytes isn't
     if ret.dtype != dtype: ret = ret.cast(dtype)
-    return ret if ret.device == device else ret.clone(device)
+    return ret if ret.device == device else ret.copy_to_device(device)
   def clone(self, device=None) -> UOp:
     device = device or self.device
     ret = self.empty_like(device=device)
     src = self if self.device is None or self.device == device else self.copy_to_device(device)
-    # The clone's STORE already materializes the value; a separate CONTIGUOUS is redundant.
-    if src.op is Ops.CONTIGUOUS: src = src.src[0]
     return ret.after(ret.store(src.cast(ret.dtype)))
   @recursive_property
   def device(self) -> str|tuple[str, ...]|None:
@@ -1254,12 +1248,8 @@ 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"
-    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))
+    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)
     return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig, self.key)
 
 @dataclass(frozen=True)
@@ -1417,7 +1407,6 @@ class UPat(RandMixin):
   @staticmethod
   def any(*src): return UPat(src=src, is_any=True)
   def or_casted(self, name:str|None=None): return UPat.any(self if name is None else self.named(name), UPat(Ops.CAST, name=name, src=(self,)))
-  def or_bitcasted(self, name:str|None=None): return UPat.any(self if name is None else self.named(name), UPat(Ops.BITCAST, name=name, src=(self,)))
   def or_after(self, name:str|None=None):
     return UPat.any(self if name is None else self.named(name), UPat(Ops.AFTER, name=name, src=(self,), allow_any_len=True))
   @staticmethod

tinygrad/uop/spec.py

@@ -93,7 +93,7 @@ spec_shared = PatternMatcher([
   # GROUP of stores (or groups, or NOOPs)
   (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END))), lambda: True),
 
-  # AFTER preserves its target view.
+  # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, RETURNED, or another AFTER
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.INDEX,
                                                      Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS})),),
         allow_any_len=True), lambda: True),
@@ -124,9 +124,10 @@ spec_shared = PatternMatcher([
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat()), validate_index),
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
 
-  # STORE targets storage (or an AFTER/BITCAST/view of it). INDEX stores are checked above.
+  # STORE: the target must be storage or a CONTIGUOUS realization point (or an AFTER/BITCAST/view of one);
+  # CONTIGUOUS targets are written into the buffer the CONTIGUOUS creates. INDEX stores are checked above
   (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x:
-   True if (b:=x.storage_base).op in {Ops.BUFFER, Ops.PARAM} else None if b.op is Ops.INDEX else False),
+   True if (b:=x.storage_base).op in {Ops.BUFFER, Ops.PARAM, Ops.CONTIGUOUS} else None if b.op is Ops.INDEX else False),
 
   # WMMA has a <a, b, acc>
   (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 5),
@@ -175,10 +176,7 @@ spec_tensor = PatternMatcher([
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
   (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)),
 
-  # Detached storage may carry pending writes in the Tensor graph.
-  (UPat(Ops.AFTER, src=(UPat(Ops.DETACH, name="x"),), allow_any_len=True), lambda x: x.storage_base.op in {Ops.BUFFER, Ops.PARAM}),
-
-  # Layout and autograd markers preserve the source value.
+  # CONTIGUOUS ensures the source UOp realizes
   (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), src=(UPat(),), arg=None), lambda: True),
 
   # TODO: this should not be here. STAGE is transformed to BUFFER later
@@ -200,8 +198,8 @@ spec_program = PatternMatcher([
   (UPat(GroupOp.All, name="x"), lambda x: False if x.op is not Ops.CAST and any(s.op is Ops.CONST for s in x.src) else None),
   (UPat(GroupOp.All-{Ops.CONST}, dtypes.weaks), lambda: False),
 
-  # allow special SHRINK of a buffer or its bitcast
-  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)).or_bitcasted(), UPat(), UPat.cvar().or_casted())), lambda: True),
+  # allow special SHRINK
+  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST).or_casted())), lambda: True),
 
   # movement ops are not allowed in programs
   (UPat(GroupOp.Movement), lambda: False),

tinygrad/uop/weak.py

@@ -4,8 +4,8 @@ from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype,
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, dtype_from_uop, promo_dtype
 
 # the decomps and float emulation commit bare consts at a dtype another src already states
-def commit_weak_consts(u:UOp, dt:DType|None) -> UOp:
-  return u if dt is None else u.replace(src=tuple(s.ccast(dt) if s.op is Ops.CONST and s.dtype in dtypes.weaks else s for s in u.src))
+def commit_weak_consts(u:UOp, dt:DType|None) -> UOp|None:
+  return None if dt is None else u.replace(src=tuple(s.ccast(dt) if s.op is Ops.CONST and s.dtype in dtypes.weaks else s for s in u.src))
 
 # the concrete dtypes u commits its srcs at: the operands' meet and u's own derived dtype, None if either is weak
 def derived_dtypes(u:UOp, src:tuple[UOp, ...]) -> tuple[DType, DType]|None:
@@ -81,11 +81,16 @@ def uncast_const(u:UOp) -> UOp|None:
 
 pm_uncast_const = PatternMatcher([(UPat(GroupOp.Broadcastable, name="u"), uncast_const)])
 
+def cast_const(u:UOp, s:UOp) -> UOp:
+  if s.op is not Ops.CONST or s.is_invalid: return s  # Invalid never commits
+  # bool is the one strong bare dtype: cconst, since .cast(bool) would fold at construction
+  if s.dtype is dtypes.bool: return UOp.cconst(s.val, s.dtype)
+  # commit at the dtype its consumer derives
+  return s.ccast(dts[0]) if (dts:=derived_dtypes(u, u.src)) is not None else s
+
 # commit every remaining bare const, keyed on the consumer: "bare" is a property of the edge
 def cast_consts(u:UOp) -> UOp|None:
   if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None  # a committed const's CONST is its value, not an edge
-  if (dts:=derived_dtypes(u, u.src)) is not None: u = commit_weak_consts(u, dts[0])
-  # bool is the one strong bare dtype: .cast(bool) would fold at construction. Invalid never commits.
-  return u.replace(src=tuple(UOp.cconst(s.val, s.dtype) if s.op is Ops.CONST and s.dtype is dtypes.bool and not s.is_invalid else s for s in u.src))
+  return None if (src:=tuple(cast_const(u, s) for s in u.src)) == u.src else u.replace(src=src)
 
-pm_cast_const = PatternMatcher([(UPat(GroupOp.All, name="u", custom_early_reject={Ops.CONST}), cast_consts)])
+pm_cast_const = PatternMatcher([(UPat(GroupOp.All, name="u"), cast_consts)])