Merge remote-tracking branch 'origin/master' into simpler_callify2

# Conflicts:
#	test/mockgpu/amd/emu.py

.github/workflows/benchmark.yml

@@ -97,7 +97,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -137,7 +137,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -185,7 +185,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -227,7 +227,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -272,7 +272,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code

conftest.py

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

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))).contiguous()
+    return (h + self.mlp(self.ln_2(h))).clone()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

extra/hcq2/ops_amd2.py

@@ -0,0 +1,716 @@
+from __future__ import annotations
+from typing import cast
+import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
+assert sys.platform != 'win32'
+from dataclasses import dataclass
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name
+from tinygrad.uop.ops import sint, UOp
+from tinygrad.device import BufferSpec, Buffer
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
+from tinygrad.helpers import ceildiv, unwrap, pluralize
+from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
+from tinygrad.renderer.llvmir import AMDLLVMRenderer
+from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
+from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
+from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
+from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
+from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
+from tinygrad.runtime.ops_amd import SQTT, PMC
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
+if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
+
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.uop.ops import Ops, UPat, PatternMatcher
+
+# *****************
+# PM4
+
+def _queue_args(hq:HWQueue, q) -> list[UOp]: # the ring and its pointers, tagged {name}_{queue} like the device's bufferize rules
+  shapes = [("ring", (q.ring.size,), q.ring.dtype)] + [(n, (1,), dtypes.uint64) for n in ("write_ptr", "doorbell", "put_value")]
+  return [UOp.placeholder(s, d, 0, device=hq.devs, volatile=True, tag=to_name(n, hq.queue)) for n, s, d in shapes]
+
+def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
+
+class AMDComputeQueue(HWQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))),
+     lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
+
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
+    self.pm4, self.gc, self.soc, self.nbio, self.target = self.dev.pm4, self.dev.gc, self.dev.soc, self.dev.nbio, self.dev.target
+
+  def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, _dw(vals) - 1), *vals)
+
+  def wreg(self, reg:AMDReg, *args:sint, **kwargs:int):
+    if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
+    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START
+    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
+      set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START
+    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
+    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
+
+  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
+    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
+                | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
+    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4)
+
+  def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
+    if self.target[0] != 9:
+      cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
+                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
+      return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
+    cp_coher_cntl = self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
+                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
+    return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
+
+  def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
+    if self.target[0] != 9:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
+                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
+      event_dw = self.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
+               | self.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | self.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
+                | self.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
+    else:
+      cache_flags_dw = 0 if not cache_flush else (self.pm4.EOP_TC_WB_ACTION_EN | self.pm4.EOP_TC_NC_ACTION_EN)
+      event_dw = self.pm4.EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | \
+                 self.pm4.EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
+      memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel)
+      ctxid = 0
+    addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64)
+    val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64)
+    self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid)
+
+  def memory_barrier(self):
+    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
+    self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
+                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
+    self.acquire_mem()
+
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = amd_build_program(self.dev, prg, self.devs)
+    info = prg.arg
+
+    # kernargs: a nested blob linear inside a getaddr, packed into the tail of the cmdbuf
+    ka_words = [get_call_arg_uops(call)[gi].getaddr(self.devs) for gi in info.globals] + \
+               [b.ccast(v.dtype) for v, b in zip(info.vars, get_call_var_uops(call, prg))] # a bound value is a bare const, the var has the width
+    pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
+    assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
+    ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4))
+
+    prog_addr = lib.getaddr(self.devs) + data.entry_point_offset
+    scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=self.devs).rtag("scratch").getaddr(self.devs)
+    args_addr = ka.getaddr(self.devs)
+
+    user_regs:list = []
+    if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000]
+    if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size]
+    user_regs += [args_addr]
+
+    dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
+      **({'cs_w32_en': int(data.wave32)} if self.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
+    self.acquire_mem(gli=0, gl2=0)
+    self.wreg(self.gc.regCOMPUTE_PGM_LO, prog_addr >> 8)
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
+    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
+    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
+    for xcc_id in range(self.dev.xccs):
+      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8)
+    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
+    self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
+    self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH")))
+    self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0)
+    self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *info.global_size, dispatch_init)
+    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
+
+  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.pm4.int_sel__mec_release_mem__none)
+
+  def signal(self, signal:UOp, value:UOp):
+    self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
+                     self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
+
+  def submit(self, cmdbuf:UOp) -> UOp:
+    q = self.dev.compute_queue
+
+    ring, wptr, doorbell, put = _queue_args(self, q)
+
+    size_dw = cmdbuf.max_numel() // 4
+    p = put.index(0).load()
+    i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf,))
+    copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
+    next_put = p + size_dw
+    flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
+    return doorbell.after(flush).index(0).store(next_put)
+
+# *****************
+# SDMA
+
+class AMDSDMAQueue(HWQueue):
+  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: ()),
+    (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 __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
+    self.sdma, self.target, self.max_copy_size = self.dev.sdma, self.dev.target, self.dev.max_copy_size
+
+  def copy(self, call:UOp):
+    sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
+    hdr = self.sdma.SDMA_OP_COPY | self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR)
+    for off in range(0, sz, self.max_copy_size):
+      self.q(hdr, self.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, self.max_copy_size)-1), 0,
+             *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(self.devs), call.src[1].getaddr(self.devs))))
+
+  def wait(self, signal:UOp, value:UOp):
+    op = self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
+       | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
+    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), 0xffffffff,
+           self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
+
+  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))
+
+  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)
+    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), self.sdma.SDMA_OP_TRAP, 0)
+
+  def submit(self, cmdbuf:UOp) -> UOp:
+    # sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail
+    q = unwrap(self.dev.sdma_queue(int(self.queue.split(":")[1])))
+
+    ring, wptr, doorbell, put = _queue_args(self, q)
+
+    rs, size_dw = q.ring.size, cmdbuf.max_numel() // 4
+    put_b = put.index(0).load()
+    tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
+    fits = (size_dw <= rs - tail).cast(dtypes.int)
+    start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
+    zi = UOp.range(zero_amt, 10, dtype=dtypes.int, src=(cmdbuf,))
+    zero_tail = ring.index(tail + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
+    i = UOp.range(size_dw, 11, dtype=dtypes.int, src=(cmdbuf,))
+    copy = ring.index(start_dw + i).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
+    next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype)
+    flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
+    return doorbell.after(flush).index(0).store(next_put)
+
+@dataclass(frozen=True)
+class AMDProgramData:
+  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
+  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
+  enable_dispatch_ptr:int; enable_private_segment_sgpr:int
+
+_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
+def amd_build_program(dev, prg:UOp, devs:tuple[str, ...]) -> tuple[AMDProgramData, UOp]:
+  # the image parses once per lib, each device set gets its own program buffer of it
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, devs))) is None:
+    data, image = _amd_program_image(dev, lib)
+    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
+    cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
+  return cached
+
+@functools.cache
+def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]:
+  image, sections, relocs = elf_loader(lib)
+  rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
+  for off, sym, typ, addent in relocs:
+    assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
+    image[off:off+8] = struct.pack('<q', sym - off + addent)
+  desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
+  if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
+    raise RuntimeError("Too many resources requested: group_segment_size")
+  edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
+
+  data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+    rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
+    rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
+    wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
+    kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
+    enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
+  return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
+
+class AMDAllocator(HCQAllocator['AMDDevice']):
+  def __init__(self, dev:AMDDevice):
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
+
+  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
+
+  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
+
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
+
+@dataclass
+class AMDQueueDesc:
+  ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
+  eop_buffer: Buffer|None = None; cwsr_buffer: Buffer|None = None; params: tuple|None = None  # noqa: E702
+
+class KFDIface:
+  kfd:FileIOInterface|None = None
+  event_page:HCQBuffer|None = None
+  gpus:list[FileIOInterface] = []
+  count:int = 0
+
+  def _is_usable_gpu(self, gpu_id):
+    with contextlib.suppress(OSError): return int(gpu_id.read()) != 0
+    return False
+
+  def __init__(self, dev, device_id):
+    self.dev = dev
+
+    kfd_topo_path = "/sys/devices/virtual/kfd/kfd/topology/nodes"
+
+    # Initialize KFD interface during first run
+    if KFDIface.kfd is None:
+      KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR)
+      gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))]
+      KFDIface.gpus = hcq_filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD")
+      KFDIface.count = len(KFDIface.gpus)
+
+    if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?")
+
+    self.gpu_id = int(FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/gpu_id").read())
+    self.props = {(p:=l.split())[0]: int(p[1]) for l in FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/properties").read().splitlines()}
+    self.dev_sysfs_path = f"/sys/class/drm/renderD{self.props['drm_render_minor']}/device"
+    ip_base = f"{self.dev_sysfs_path}/ip_discovery/die/0"
+    id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP}
+    ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip]
+    self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw}
+    self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)
+
+    self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version)
+    kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
+    if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
+
+    # Set these for our device.
+    if KFDIface.event_page is None:
+      KFDIface.event_page = self.alloc(0x8000, uncached=True)
+      kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
+    else: self.map(KFDIface.event_page)
+
+    # Event to wait for queues completion
+    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
+    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
+
+    # OS events to collect memory and hardware faults
+    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
+    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)
+
+    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
+      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
+    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
+    flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
+
+    if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
+    else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
+
+    # Make mapped cpu address to be uncachable
+    if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED
+
+    if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC
+
+    if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
+      buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
+    else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0)
+
+    try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, gpu_id=self.gpu_id, flags=flags, mmap_offset=buf)
+    except OSError as e:
+      if e.errno == errno.EINVAL and (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) and cpu_access:
+        raise MemoryError("Cannot allocate host-visible VRAM. Ensure the resizable BAR option is enabled on your system.") from e
+      if e.errno == errno.ENOMEM: raise MemoryError(f"Cannot allocate {size} bytes: no memory is available.") from e
+      raise
+
+    if not (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR):
+      buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset)
+      assert addr == buf == mem.va_addr
+
+    view = MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None
+    self.map(hcqbuf:=HCQBuffer(mem.va_addr, mem.size, meta=mem, view=view, owner=self.dev))
+    return hcqbuf
+
+  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
+    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
+    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+    assert stm.n_success == 1
+
+  def map(self, mem):
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped
+
+    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
+    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
+    assert stm.n_success == 1
+    return HCQBuffer(mem.va_addr, mem.size, meta=mem.meta, owner=mem.owner)
+
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
+    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
+      queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
+      eop_buffer_address=eop_buffer._buf.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer._buf.size if eop_buffer else 0,
+      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
+      write_pointer_address=gart._buf.va_addr+wptr, read_pointer_address=gart._buf.va_addr+rptr+8*xcc_id)
+
+    if not hasattr(self, 'doorbells'):
+      self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
+      self.doorbells = cast(FileIOInterface, KFDIface.kfd).mmap(0, 0x2000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, self.doorbells_base)
+
+    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
+    doorbell = Buffer("CPU", 1, dtypes.uint64,
+      options=BufferSpec(external_ptr=self.doorbells + queue.doorbell_offset - self.doorbells_base), preallocate=True)
+    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr+8*xcc_id).ensure_allocated(),
+      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer)
+
+  def sleep(self, tm:int):
+    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
+    if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang()
+
+  def on_device_hang(self):
+    def _str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
+
+    # try to collect fault info if not already set from sleep().
+    if not self.queue_event_arr[1].memory_exception_data.gpu_id and not self.queue_event_arr[2].hw_exception_data.gpu_id:
+      with contextlib.suppress(RuntimeError): self.sleep(tm=1)
+
+    report = []
+    if self.queue_event_arr[1].memory_exception_data.gpu_id:
+      report += [f"MMU fault: 0x{self.queue_event_arr[1].memory_exception_data.va:X} | {_str(self.queue_event_arr[1].memory_exception_data.failure)}"]
+    if self.queue_event_arr[2].hw_exception_data.gpu_id: report += [f"HW fault: {_str(self.queue_event_arr[2].hw_exception_data)}"]
+
+    raise RuntimeError("\n".join(report))
+
+  def require_profile_mode(self, can_set_mode=True):
+    if self.dev.target[0] == 9: return
+    fn = f'{self.dev_sysfs_path}/power_dpm_force_performance_level'
+    if (perflevel:=FileIOInterface(fn).read().strip()) != 'profile_standard':
+      if can_set_mode:
+        atexit.register(lambda: os.system(f"echo '{perflevel}' | sudo tee {fn} > /dev/null"))
+        os.system(f"echo 'profile_standard' | sudo tee {fn} > /dev/null")
+        self.require_profile_mode(can_set_mode=False)
+      else:
+        raise RuntimeError("PMC/SQTT requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
+
+  @functools.cached_property
+  def drm_dev_info(self) -> amdgpu_drm.struct_drm_amdgpu_info_device:
+    amdgpu_drm.DRM_IOCTL_AMDGPU_INFO(self.drm_fd, query=amdgpu_drm.AMDGPU_INFO_DEV_INFO,
+      return_pointer=ctypes.addressof(inf:=amdgpu_drm.struct_drm_amdgpu_info_device()), return_size=ctypes.sizeof(inf))
+    return inf
+  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return ((self.drm_dev_info.cu_bitmap[se % 4][sa + (se // 4) * 2] >> (2 * wgp)) & 0x3) == 0x3
+
+class PCIIface(PCIIfaceBase):
+  def __init__(self, dev, dev_id):
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
+      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
+    self._compute_props()
+
+  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
+    return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs)
+
+  def require_profile_mode(self): return True
+  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)
+
+  def _compute_props(self):
+    self.ip_versions = self.dev_impl.ip_ver
+
+    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
+    if self.dev_impl.gc_info.header.version_major == 2:
+      cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se
+    else:
+      cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)
+      max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se
+
+    array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs
+    self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count,
+      'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
+      'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
+      'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
+
+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
+    assert cwsr_buffer is None, "no cwsr buffer for am"
+
+    rcvr_params: tuple
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
+      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, idx)))
+    else:
+      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, eop_buffer._buf.va_addr, eop_buffer._buf.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
+
+    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
+    doorbell = Buffer("CPU", 1, dtypes.uint64, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True)
+    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr).ensure_allocated(),
+      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, params=rcvr_params)
+
+  def _collect_interrupts(self, reset=False, drain_only=False):
+    d = self.dev
+    if drain_only: d.iface.dev_impl.ih.drain()
+    else: d.iface.dev_impl.ih.interrupt_handler()
+
+    if reset and d.iface.dev_impl.recover(force=True):
+      cq = d.compute_queue
+      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
+      d.iface.dev_impl.gfx.setup_ring(*cq.params)
+      (tl:=d.timeline._buf.cpu_view().view(fmt='Q'))[0] = tl[1]
+
+  def sleep(self, timeout):
+    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
+      self.pci_dev.irq_fd.read(8 * events_cnt)
+    self._collect_interrupts()
+    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
+
+  def on_device_hang(self):
+    self._collect_interrupts(reset=True)
+    raise RuntimeError("Device hang detected")
+
+  def device_fini(self): self.dev_impl.fini()
+
+class USBIface(PCIIface):
+  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
+    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
+      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
+    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
+    self.dev_impl = AMDev(self.pci_dev)
+    self._compute_props()
+    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
+    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
+    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
+
+  def _dma_region(self, ctrl_addr, sys_addr, size):
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
+    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
+    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
+
+  def sleep(self, timeout): pass
+
+  # we don't own the sram region, so the buffer never frees it
+  @functools.cached_property
+  def usb_sram(self) -> Buffer:
+    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
+
+def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
+
+class AMDDevice(HCQ2Compiled):
+  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
+  max_scratch_psize = 0
+  pm_encode = PatternMatcher([
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_compute", name="submit"), lambda ctx, submit: encode_submit(AMDComputeQueue(ctx, submit))),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_copy", name="submit"), lambda ctx, submit: encode_submit(AMDSDMAQueue(ctx, submit))),
+  ])
+
+  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
+
+  def device_props(self): return self.iface.props
+
+  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
+
+  def __init__(self, device:str=""):
+    self.iface = self._select_iface(device)
+    self.is_usb = isinstance(self.iface, USBIface)
+    if self.is_usb: self.rt_nbytes = 4 << 20
+
+    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
+    self.arch = "gfx%d%x%x" % self.target
+    assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}"
+    if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
+
+    self.xccs = self.iface.props.get('num_xcc', 1)
+    self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs
+    self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs
+    self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu']
+    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
+
+    self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets")
+    self.soc = import_soc(self.target)
+    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}")
+    self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
+    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
+                    bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})
+
+    self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP],
+                      bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
+
+    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
+    if self.is_aql:
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
+
+    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
+    self.sdma_queues:dict = {}
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
+
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
+
+    # Scratch setup
+    self.max_private_segment_size = 0
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+
+    if self.is_usb:
+      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
+      raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
+
+    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
+    if self.pmc_enabled:
+      self.iface.require_profile_mode()
+
+      self.pmc_sched:list[PMCSample] = []
+      self.pmc_counters = import_pmc(self.target)
+
+      # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses
+      l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC")
+      pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS"
+      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
+        if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
+
+      raise NotImplementedError("PMC start not migrated to hcq2 yet")
+
+    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
+    self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
+    if self.sqtt_enabled:
+      self.iface.require_profile_mode()
+
+      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
+      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt * self.xccs)]
+      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
+      self.sqtt_next_cmd_id = itertools.count(0)
+
+  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
+    ring = Buffer(self.device, ring_size // 4, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
+    gart = Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
+
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL:
+      self.aql_gart = gart
+      self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
+        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+        max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1)
+      self.aql_gart._buf.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+    cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE)
+    cwsr_buffer = Buffer(self.device, cwsr_buffer_size, dtypes.uint8, preallocate=True) if ctx_save_restore_size else None
+    eop_buffer = Buffer(self.device, eop_buffer_size, dtypes.uint8, preallocate=True) if eop_buffer_size else None
+
+    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
+             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
+
+    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.PARAM, tag=to_name(name, qname)), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+    ]) + self.pm_bufferize
+
+    return queue
+
+  @functools.cached_property
+  def compute_queue(self) -> AMDQueueDesc:
+    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
+    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
+    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
+    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
+    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
+    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
+    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
+      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
+      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
+      debug_memory_size=round_up(self.wave_cnt * 32, 64))
+
+  def sdma_queue(self, idx:int):
+    if getenv("AMD_DISABLE_SDMA"): return None
+    if idx in self.sdma_queues: return self.sdma_queues[idx]
+    with contextlib.suppress(OSError):
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
+    return self.sdma_queues.get(idx, None)
+
+  def tmpring_size(self, private_segment_size):
+    private_segment_size = max(private_segment_size, 128)
+
+    lanes_per_wave = 64 # wave64
+    mem_alignment_size = 256 if self.target[0] != 9 else 1024
+    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
+    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+
+    # NOTE: xcc logic is correct only for GFX9.
+    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
+    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
+    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
+
+    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
+
+    if hasattr(self, 'aql_desc'):
+      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
+      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
+              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+
+      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
+      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
+      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
+        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
+        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
+      self.aql_desc.compute_tmpring_size = tmpring
+      self.aql_gart._buf.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
+
+    return tmpring
+
+  def scratch_buffer(self, private_segment_size):
+    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
+    if self.max_private_segment_size < private_segment_size:
+      lanes_per_wave = 64 # wave64
+      mem_alignment_size = 256 if self.target[0] != 9 else 1024
+      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
+      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+      self.max_private_segment_size = private_segment_size
+    return self.scratch
+
+  def on_device_hang(self): self.iface.on_device_hang()
+
+  def device_props(self): return self.iface.props

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))).contiguous().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).clone().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).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).clone().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):

test/amd/test_sqtt_profiler.py

@@ -1,6 +1,5 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit, dtypes
-from test.helpers import is_hcq2_device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Compiled, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
@@ -115,7 +114,8 @@ class TestSQTTProfiler(unittest.TestCase):
     kernel_name = sqtt[0]["name"]
     for i,e in enumerate(sqtt[1:], start=1): self.assertEqual(e["name"], f"{kernel_name} n{i+1}")
 
-  def test_jit_graph(self, kernel_count=3*(5 if is_hcq2_device() else 1)): # hcq2 traces the graphed kernels too
+  # TODO: can we trace SQTT for graphed kernels?
+  def test_jit_graph(self, kernel_count=3*1):
     @TinyJit
     def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
     t = Tensor.empty(32)

test/backend/test_assign.py

@@ -45,6 +45,16 @@ class TestAssign(unittest.TestCase):
     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])
+
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
     xs = X[2:4]
@@ -1014,10 +1024,10 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(c.tolist(), [[0,0],[0,0]])
 
-  def test_contiguous(self):
+  def test_clone(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
-    self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
+    c = t.permute(1,0).clone()
+    self.assertIs(c.uop.base.op, Ops.AFTER)
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(c.tolist(), [[1,1],[2,1]])
 
@@ -1032,6 +1042,16 @@ 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)
@@ -1043,10 +1063,10 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(d.tolist(), [[0,0],[0,0]])
 
-  def test_detach_contiguous(self):
+  def test_detach_clone(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
-    self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
+    d = t.permute(1,0).clone().detach()
+    self.assertIs(d.uop.base.op, Ops.AFTER)
     d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(d.tolist(), [[1,1],[2,1]])
 

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())
-    # contiguous folded const can still schedule
-    a = Tensor.empty(1, 0).sum().contiguous()
+    # An explicitly cloned folded constant still owns persistent storage.
+    a = Tensor.empty(1, 0).sum().clone()
     _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_ops.py

@@ -3107,13 +3107,6 @@ class TestOps(unittest.TestCase):
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])
 
-  def test_gather_bool_index(self):
-    helper_test_op(None, lambda x,y: x.gather(dim=0, index=y.bool().long()),
-                         lambda x,y: x.gather(dim=0, index=y.cast(dtypes.bool).cast(dtypes.int)),
-                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
-    helper_test_op(None, lambda x,y: x[y.bool().long()], lambda x,y: x[y.cast(dtypes.bool).cast(dtypes.int)],
-                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
-
   def test_scatter(self):
     b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
     a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)

test/backend/test_profiler.py

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

test/backend/test_schedule.py

@@ -115,7 +115,8 @@ 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))
-    sched = check_schedule(base, 4)
+    # The pending clone is already contiguous, so assign-back needs no separate contiguous buffer.
+    sched = check_schedule(base, 2)
     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

@@ -75,6 +75,11 @@ 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()

test/device/test_hcq2.py

@@ -1,10 +1,10 @@
 import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
-from tinygrad.device import Buffer, BufferSpec
+from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context, dedup, partition
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
+from tinygrad.uop.ops import Ops, UOp, KernelInfo
 from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.runtime.autogen import libc
@@ -42,25 +42,6 @@ def patch_words(batch:UOp) -> list[UOp]:
 def rt_params(batch:UOp) -> list[str]:
   return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
 
-class TestHCQ2Deps(unittest.TestCase):
-  def test_disjoint_write_preserves_dependencies(self):
-    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
-    for write in ([], [0]):
-      tracker = hcq2.HCQDepsTracker()
-      tracker.access_resources([b.shrink(((0, 4),))], write, 0)
-      self.assertEqual(tracker.access_resources([b.shrink(((4, 8),))], [0], 1), [])
-      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
-
-  def test_partial_write_preserves_dependencies(self):
-    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
-    for write in ([], [0]):
-      tracker = hcq2.HCQDepsTracker()
-      tracker.access_resources([b], write, 0)
-      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 1), [0])
-      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
-      self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
-      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
-
 @unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
 class TestHCQ2Core(unittest.TestCase):
   @staticmethod
@@ -209,9 +190,7 @@ class TestHCQ2Core(unittest.TestCase):
 
   def test_device_state_survives_as_link_refs(self):
     # a buffer the commands only address, never a param of the body, is kept by the linked call as a ref of what its getaddr resolved into
-    dev = Device[Device.DEFAULT]
-    names = {"AMD": () if getattr(dev, "is_aql", False) else ("scratch",), # the aql descriptor holds the scratch, nothing addresses it
-             "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
+    dev, names = Device[Device.DEFAULT], {"AMD": ("scratch",), "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
     @TinyJit
     def f(a): return (a * 2 + 1).contiguous().realize()
     x = Tensor.ones(16).contiguous().realize()
@@ -247,36 +226,6 @@ class TestHCQ2FFI(unittest.TestCase):
     got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
     self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
 
-  def test_device_lower_after_encode(self):
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = UOp.placeholder((1,), dtypes.int32, device="CPU", tag="result")
-      encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_encode"), lambda: UOp.custom_function("test_lower"))])
-      lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_lower"), lambda out=out: out.index(0).store(42))])
-      with patch.object(Device["CPU"], "pm_encode", encode), patch.object(Device["CPU"], "pm_lower", lower):
-        bufs = self._run(UOp.custom_function("test_encode"))
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 42)
-
-  def test_nested_cstruct_patches(self):
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      inner = hcq2.cstruct(init_c_struct_t(4, (("value", ctypes.c_uint32, 0),)), value=42)
-      outer = hcq2.cstruct(init_c_struct_t(8, (("ptr", ctypes.c_uint64, 0),)), ptr=inner.getaddr("CPU"))
-      out = UOp.placeholder((1,), dtypes.uint32, device="CPU", tag="result")
-      copied = hcq2.ccall(libc.memcpy, out.index(0), outer.bitcast(dtypes.uint64).index(0).load(), 4)
-      bufs = self._run(out.after(copied).index(0).load())
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32)._buf.cpu_view().view(fmt='I')[0], 42)
-
-
-class TestHCQ2Timeline(unittest.TestCase):
-  def test_reused_timeline_is_zeroed(self):
-    buf = Buffer("CPU", 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    addr = buf._buf.va_addr
-    buf._buf.cpu_view().view(fmt='B')[:] = b'\xff' * 16
-    buf.deallocate()
-    dev = HCQ2Compiled.__new__(HCQ2Compiled)
-    dev.device = "CPU"
-    self.assertEqual(dev.timeline._buf.va_addr, addr)
-    self.assertEqual(bytes(dev.timeline._buf.cpu_view()), bytes(16))
-
 
 if __name__ == "__main__":
   unittest.main()

test/mockgpu/amd/amdgpu.py

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

test/null/test_uop_graph.py

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

test/null/test_uop_vmin_vmax.py

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

test/null/test_uops.py

@@ -7,10 +7,28 @@ 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, type_verify
+from tinygrad.uop.spec import spec_program, spec_shared, spec_tensor, 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_validate_oob.py

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

test/null/test_viz.py

@@ -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).contiguous().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).clone().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)

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))[0].src[0].key
+def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop)).src[0].key
 
 class TestCall(unittest.TestCase):
   def test_call_plus(self):
@@ -370,7 +370,7 @@ class TestArgOrder(unittest.TestCase):
     x = Tensor.arange(3, dtype=dtypes.int).realize()
     call = self.make_intersperse_call(x, precompile=True)[0].src[1]
     # the transform must preserve the RETURNED's src position: its placeholder is at src 1, the input stays at src 2
-    from tinygrad.tensor import transform_precompiled_call
+    from tinygrad.schedule.prepare import transform_precompiled_call
     new = transform_precompiled_call(call)
     new_call = new.src[0].src[1].src[1]
     # the out buffer takes the RETURNED's position (src 1), the input value keeps its position (src 2)

test/unit/test_callify.py

@@ -1,5 +1,7 @@
 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):
@@ -107,6 +109,75 @@ 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()
@@ -114,6 +185,13 @@ class TestCallify(unittest.TestCase):
     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)

test/unit/test_llm_amd.py

@@ -28,12 +28,6 @@ class TestQ8Quantize(unittest.TestCase):
     # xsum holds the two per-16 sums per 32-wide group
     np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))
 
-  def test_quantize_rounding_ties(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    values = np.array([-127,127]+[i+0.5 for i in range(-15,15)],dtype=np.float32)
-    quant,_,_ = q8_quantize(Tensor(values),1,32)
-    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(32).numpy(),np.rint(values).astype(np.int8))
-
   def test_q6_linear_compiles_in_function(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
@@ -50,33 +44,22 @@ class TestQ8Quantize(unittest.TestCase):
     self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
     self.assertEqual(linear.weight.dtype, dtypes.uint32)
 
-  def test_q4_k_linear(self): self._test_quant_linear(12, 144)
-  def test_iq4_linear(self): self._test_quant_linear(23, 136)
-  def test_q5_linear(self): self._test_quant_linear(13, 176)
-
-  def _test_quant_linear(self, ggml_type, block_bytes):
+  def test_q4_k_linear(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
-    in_features, out_features = 2048, 64
-    packed = rng.integers(0, 256, (out_features*in_features//256, block_bytes), dtype=np.uint8)
-    packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
-    if ggml_type in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, out_features*in_features, ggml_type).reshape(out_features, in_features)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
+    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
     weight = decoded.numpy()
-    linear = Linear(in_features, out_features, bias=False)
-    linear.weight = decoded
-    for tokens in (1, 3, 32, 64, 128):
-      with self.subTest(tokens=tokens):
-        x = rng.normal(size=(tokens, in_features)).astype(np.float32 if tokens == 3 else np.float16)
-        reference_x = x.astype(np.float32)
-        if tokens < 16:
-          grouped = reference_x.reshape(tokens, -1, 32)
-          scale = np.maximum(np.abs(grouped).max(-1, keepdims=True) / 127, 1e-8)
-          reference_x = (np.clip(np.rint(grouped/scale), -127, 127)*scale).reshape(tokens, in_features)
-        reference_w = weight if tokens < 16 else weight.astype(np.float16).astype(np.float32)
-        np.testing.assert_allclose(linear(Tensor(x)).numpy(), reference_x @ reference_w.T, rtol=3e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, ggml_type)
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 12)
 
   def test_q6_linear_multiple_tokens(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -120,37 +103,6 @@ class TestQ8Quantize(unittest.TestCase):
     expected = q.scaled_dot_product_attention(cache[0, :, :, :3], cache[1, :, :, :3], enable_gqa=True)
     np.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=2e-3, atol=2e-3)
 
-  def test_flash_attention_decode_beyond_256_chunks(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    n = 257 * 64
-    q = Tensor.zeros(1, 1, 1, 32, dtype=dtypes.half).realize()
-    k = Tensor.zeros(1, 1, n, 32, dtype=dtypes.half)
-    v = Tensor.zeros(1, 1, n-64, 32, dtype=dtypes.half).cat(Tensor.ones(1, 1, 64, 32, dtype=dtypes.half), dim=2)
-    cache = Tensor.stack(k, v).contiguous().realize()
-    for valid, expected in ((1, 0), (n, 1/257)):
-      with self.subTest(valid=valid):
-        valid_kv_len = UOp.variable("valid_kv_len", 1, n).bind(valid)
-        assigned = Tensor(cache.uop.after(Tensor(valid_kv_len).uop))
-        np.testing.assert_allclose(flash_attention(q, assigned, valid_kv_len).numpy(), expected, rtol=2e-3, atol=2e-4)
-
-  def test_flash_attention_decode_long_context_random(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    Tensor.manual_seed(42)
-    n, valid = 257*64, 257*64 - 13  # past the old 256-chunk partial limit, with a ragged tail
-    q = Tensor.randn(1, 8, 1, 128, dtype=dtypes.half).realize()
-    cache = Tensor.randn(2, 1, 2, n, 128, dtype=dtypes.half).realize()
-    out = flash_attention(q, cache, valid).realize()
-    expected = q.scaled_dot_product_attention(cache[0, :, :, :valid], cache[1, :, :, :valid], enable_gqa=True)
-    np.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=2e-3, atol=2e-3)
-
-  def test_flash_attention_decode_chunk_round_accumulator_range(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    valid_kv_len, max_kv_len = 6749, 6784  # three chunk rounds, with a ragged tail
-    q = Tensor.zeros(1, 8, 1, 32, dtype=dtypes.half).realize()
-    cache = Tensor.stack(Tensor.zeros(1, 1, max_kv_len, 32, dtype=dtypes.half),
-                         Tensor.full((1, 1, max_kv_len, 32), 5500, dtype=dtypes.half)).contiguous().realize()
-    np.testing.assert_allclose(flash_attention(q, cache, valid_kv_len).numpy(), 5500, rtol=2e-3, atol=2e-3)
-
   def test_prefill_attention_unaligned_start(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)

tinygrad/device.py

@@ -280,13 +280,9 @@ class DepsTracker:
       if i in write:
         for dmap in [self.w_dependency_map, self.r_dependency_map]:
           kept = []
-          for entry in dmap[key]:
-            st, en, dep = entry
-            if st == en: continue
-            if en <= s or e <= st: kept.append(entry)
-            else:
-              if st < s: kept.append((st, s, dep))
-              if e < en: kept.append((e, en, dep))
+          for st,en,dep in dmap[key]:
+            if st < min(s, en): kept.append((st, min(s, en), dep))
+            if max(e, st) < en: kept.append((max(e, st), en, dep))
           dmap[key] = kept
         self.w_dependency_map[key].append((s, e, new_dependency))
       else: self.r_dependency_map[key].append((s, e, new_dependency))
@@ -341,9 +337,8 @@ class Compiled:
 
   has_copy_queue:bool = True
 
-  pm_batch:Any = None
-  pm_encode:Any = None
-  pm_lower:Any = None
+  pm_encode:Any = None # per queue kind: queue ops -> flat command words
+  pm_lower:Any = None # per queue kind: custom_function(submit, cmdbuf) -> the queue push
   pm_bufferize:Any = None
 
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):

tinygrad/function.py

@@ -26,9 +26,8 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
 def renumber_invalid_outputs(uret:UOp) -> UOp:
-  invalid = invalid_outputs(uret)
   return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
-                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid)})
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):

tinygrad/helpers.py

@@ -239,7 +239,7 @@ TRANSCENDENTAL = ContextVar("TRANSCENDENTAL", 1)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 1)
+VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)

tinygrad/llm/kernels/amd.py

@@ -75,7 +75,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.contiguous().reshape(nblocks * Q6_WORDS)
+      self.weight = padded.clone().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)))
@@ -101,18 +101,23 @@ class Linear(nn.Linear):
     return super().__call__(x)
 
 def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  return UOp(Ops.CUSTOMI, src=(a, b, c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32))
+  # int8 4-wide dot, widened to scalar multiply-adds (2% decode slower than the sudot4 builtin, but portable)
+  for i in range(4):
+    av = ((a >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
+    bv = ((b >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
+    c = c + av*bv
+  return c
 
 def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
   return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32))
 
-def _amd_load(ptr:UOp, lanes:int|None=None, stream:bool=False) -> UOp:
+def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
   assert ptr.op is Ops.INDEX
   # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
   if lanes is None: return ptr.load(arg="nontemporal")
   buf, coords = ptr.src[0], ptr.src[1:]
   idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(arg="nontemporal" if stream else None)
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
 
 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
@@ -148,19 +153,22 @@ def iq4_half_lut(device:str) -> Tensor:
 @functools.cache
 def _q8_quantize_kernel(q:UOp, scale:UOp, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
   groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, AxisType.GLOBAL), UOp.range(32, -1, AxisType.WARP)
+  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
   token, group = token_group//groups, token_group%groups
-  value = x.reshape(tokens, groups, 32)[token, group, lane].float()
-  # Quantize each input once, then pack four neighboring lanes into one word.
-  d = (warp_reduce(value.abs(), maximum=True, full_wave=True)/127).maximum(1e-8)
-  rounded = UOp(Ops.CUSTOM, src=(value/d,), arg=("__builtin_nearbyintf({0})", dtypes.float))
-  quant = rounded.clip(-127, 127).cast(dtypes.int8)
-  word = quant.cast(dtypes.uint8).cast(dtypes.uint32) << ((lane%4)*8).cast(dtypes.uint32)
-  for offset in (1, 2):
-    word |= UOp(Ops.CUSTOM, src=(word,), arg=(f"__builtin_amdgcn_ds_swizzle({{0}}, {0x1f | offset<<10})", dtypes.uint32))
-  stores = (q[token, group, (lane//4).valid((lane%4).eq(0))].store(word),
-            scale[token, group.valid(lane.eq(0))].store(d),
-            xsum[token, group, (lane//16).valid((lane%16).eq(0))].store(warp_reduce(quant.float())))
+  x = x.reshape(tokens, groups, 32)
+  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
+  word_lane = lane.minimum(7)
+  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
+  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
+  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
+  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
+  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
+  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
+  store_half = (lane & 4) >> 2
+  stores = (q[token, group, lane.valid(lane < 8)].store(word),
+            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
+                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
+                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
   return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
 
 def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
@@ -214,8 +222,8 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
     # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
     base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
     # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
-    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4, stream=True) for half in range(2))
-    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4, stream=True) for half in range(2))
+    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
+    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
     dots = [UOp.const(0, dtypes.int32)] * 2
     for word_idx in range(8):
       within = (subgroup*32 + word_idx*4)%128
@@ -395,25 +403,22 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
   _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
   assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DPL, WAVES, PARTIALS = H // H_KV, block_n, D // WARP_SIZE, waves, out.shape[2]
+  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
   assert CHUNK % WAVES == 0
   SEC = CHUNK // WAVES  # keys each wave scans independently
-  total_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  live_chunks = min(total_chunks, PARTIALS) if isinstance(total_chunks, int) else total_chunks.minimum(PARTIALS)
+  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
   block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
   lane, wave = UOp.range(WARP_SIZE, -1, axis_type=AxisType.WARP), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
   b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
   # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
   qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
   zerof = UOp.const(0, dtypes.float)
-  # Each block scans every PARTIALS-th chunk, keeping an online softmax across rounds.
-  chunk_round = UOp.range((total_chunks-1-block_chunk)//PARTIALS+1, 4, AxisType.REDUCE)
-  chunk_id = block_chunk + chunk_round*PARTIALS
   valids: list[UOp] = []
   scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
   vfrags: list[tuple[UOp, ...]] = [()]*SEC
   for j in range(SEC):
-    key = chunk_id*CHUNK + wave*SEC + j
+    key = block_chunk*CHUNK + wave*SEC + j
     valid = key < valid_kv_len
     valids.append(valid)
     kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
@@ -421,31 +426,23 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
     vfrags[j] = tuple(valid.where(v, zerof) for v in _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL))
     for h in range(G):
       s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
-      scores[j][h] = valid.where(s, UOp.const(-1e30, dtypes.float))
-  # A finite initial max keeps fully masked waves from computing exp(-inf - -inf).
-  acc_reg, max_reg, sum_reg = _reg((G, DPL), 2, 0), _reg((G,), 3, -1e30), _reg((G,), 4, 0)
-  prev_acc, prev_max, prev_sum = acc_reg.after(chunk_round), max_reg.after(chunk_round), sum_reg.after(chunk_round)
-  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), prev_max[h].load()) for h in range(G)]
-  # Rescale the previous rounds to the new max, then accumulate this round's keys.
-  alpha = [((prev_max[h].load()-row_max[h])*LOG2E).exp2() for h in range(G)]
-  accs = [[alpha[h]*prev_acc[h, i].load() for i in range(DPL)] for h in range(G)]
-  row_sums = [alpha[h]*prev_sum[h].load() for h in range(G)]
+      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
+  ninf = UOp.const(-math.inf, dtypes.float)
+  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
+  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
+  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
   for j in range(SEC):
     for h in range(G):
-      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), zerof)
+      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
       accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
       row_sums[h] = row_sums[h] + beta
-  update = UOp.group(acc_reg.store(UOp.stack(*(x for acc in accs for x in acc)).reshape(G, DPL)),
-                     max_reg.store(UOp.stack(*row_max)), sum_reg.store(UOp.stack(*row_sums))).end(chunk_round)
-  acc_reg, max_reg, sum_reg = acc_reg.after(update), max_reg.after(update), sum_reg.after(update)
   # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
   acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
   ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
-  # Normalize before fp16 to avoid overflow. Nonempty waves have sum >= 1; empty waves keep their zero accumulator.
-  stores = [lds_acc[wave, h, lane].store((acc_reg[h].load() / sum_reg[h].load().maximum(1)).cast(dtypes.half)) for h in range(G)]
+  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
   # NOTE: duplicate stores of the same value from every lane are harmless here
-  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((max_reg[h].load(), sum_reg[h].load()))]
+  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
   barrier = UOp.barrier(UOp.group(*stores))
   acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
   tid = wave*WARP_SIZE + lane
@@ -453,16 +450,14 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   for i in range(-(-G*D//(WAVES*WARP_SIZE))):
     flat = tid + i*WAVES*WARP_SIZE
     h, d = flat // D, flat % D
-    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)))
-    # LDS holds normalized values; restore each wave's sum before combining.
-    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * ml_lds[w, h, 1].load() * acc_lds[w, h, d].load().float()
-               for w in range(WAVES)), zerof)
+    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
+    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
     oidx = out[b, kv_head*G + h, block_chunk, d]
     if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
     final_stores.append(oidx.store(val))
   hstat = tid
-  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)))
-  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), zerof)
+  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
+  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
   q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
   final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
   return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
@@ -499,7 +494,7 @@ def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UO
 
 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
   B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(48, max_kv_len // 64)
+  chunks = min(256, max_kv_len // 64)
   partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
   stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
   fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)

tinygrad/mixin/elementwise.py

@@ -58,11 +58,14 @@ class ElementwiseMixin(CreationMixin):
 
   def contiguous(self, **kwargs) -> Self:
     """
-    Returns a contiguous tensor.
+    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.
     """
     if self.dtype in dtypes.weaks: return self
     uop = self._uop
-    if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(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)
     return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
 
   def contiguous_backward(self) -> Self:

tinygrad/mixin/rand.py

@@ -52,6 +52,7 @@ 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)
@@ -65,7 +66,8 @@ 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))
-    return cls._rand(key, counter, shape, dt, contiguous=contiguous)
+    out = cls._rand(key, counter, shape, dt, contiguous=False)
+    return cls._wrap_uop(out._uop.clone()) if contiguous else out
 
   def rand_like(self, **kwargs) -> Self:
     """
@@ -293,7 +295,8 @@ 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)
-    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)
+    mask = self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p
+    return self._wrap_uop(mask._uop.clone()).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/cstyle.py

@@ -67,8 +67,8 @@ base_rewrite = PatternMatcher([
 
   # call an external function: the CUSTOM_FUNCTION body holds the callee (a function pointer), the other srcs are the args
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, src=(UPat(name="fptr"),)),), allow_any_len=True, name="x"), lambda ctx,x,fptr:
-   f"((({ctx.abi}{ctx.render_dtype(x.dtype)}(*)({', '.join(ctx.render_type(y, ptr=True) for y in x.src[1:])}))({ctx[fptr]}))" +
-   f"({', '.join(f'({ctx.render_type(y, ptr=True)})({ctx[y]})' for y in x.src[1:])}))" + (";" if x.dtype is dtypes.void else "")),
+   f"((({ctx.abi}{ctx.render_dtype(x.dtype)}(*)({', '.join(ctx.render_type(y) for y in x.src[1:])}))({ctx[fptr]}))" +
+   f"({', '.join(f'({ctx.render_type(y)})({ctx[y]})' for y in x.src[1:])}))" + (";" if x.dtype is dtypes.void else "")),
 
   # custom passes through with format
   (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg[0].format(*[ctx[y] for y in x.src])),
@@ -187,8 +187,7 @@ class CStyleLanguage(Renderer):
       return prefix + self.type_map.get(dtype, dtype.name).replace(" ", "_") + str(sz) + suffix
     return prefix + self.type_map.get(dtype, dtype.name) + suffix
 
-  def render_type(self, u:UOp, ptr=False): # ptr: an address is a pointer whatever its addrspace (a register array passed to a function)
-    return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=ptr and u.addrspace not in (None, AddrSpace.ALU), shape=u._shape)
+  def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
   def render_ptr(self, u:UOp):
     # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type
     if u.max_numel() > 1 or u.dtype != u.src[0].dtype:

tinygrad/runtime/ops_cpu.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import platform, sys, ctypes, mmap, struct, time
-from typing import cast, Any
+from typing import cast
 from tinygrad.helpers import to_mv, from_mv, OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
-from tinygrad.device import BufferSpec, TinyELF, Program, Device, Buffer
+from tinygrad.device import BufferSpec, TinyELF, Program, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator
 from tinygrad.runtime.support.c import DLL
@@ -88,10 +88,6 @@ class CPUAllocator(HCQAllocator['CPUDevice']):
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
     dest[:] = self._as_buffer(src)[:len(dest)]
-  def map(self, buf:Buffer) -> HCQBuffer: # another device's buffer as cpu memory: an hcq buffer through its view, anything else through its bytes
-    if isinstance(buf._buf, HCQBuffer): return super().map(buf)
-    mv = cast(Any, Device[buf.device].allocator)._as_buffer(buf.ensure_allocated()._buf)
-    return HCQBuffer(addr:=mv_address(mv), mv.nbytes, meta=mv, view=MMIOInterface(addr, mv.nbytes, fmt='B'), owner=self.dev)
   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

@@ -95,6 +95,11 @@ class QMD:
 
 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 nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
@@ -122,6 +127,11 @@ class NVQueue(HWQueue):
     return doorbell.after(queued).index(0).store(UOp.const(fifo.token, dtypes.uint32))
 
 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 __init__(self, ctx, submit):
     super().__init__(ctx, submit)
 
@@ -180,6 +190,11 @@ class NVComputeQueue(NVQueue):
     self.prev_qmd = qmd
 
 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
+
   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)):

tinygrad/runtime/ops_qcom.py

@@ -53,6 +53,13 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 
 class QCOMComputeQueue(HWQueue):
   dev:QCOMDevice
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
   def cmd(self, opcode:int, *vals): self.q(pkt7_hdr(opcode, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 

tinygrad/runtime/support/hcq2.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, pe
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, DepsTracker
 from tinygrad.device import ProfileGraphEntry, ProfileGraphEvent, ProfileDeviceEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
-from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
+from tinygrad.dtype import dtypes, DType, DTYPES_DICT
 from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
@@ -35,15 +35,12 @@ class HCQInfo:
   slots:tuple[tuple[str, int], ...] = () # per device, the position of its batch slots in the args
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
-def cpu_bytes(d:Any) -> bool: # host memory: the cpu, numpy, a file. the cpu maps it as is, a queue device reaches it through that mapping
-  return not all_devices_in(d, HCQ_DEVS - {"CPU"}) and all(hasattr(Device[x].allocator, "_as_buffer") for x in to_tuple(d))
 
 def get_enqueue_devs(call:UOp) -> Any|None:
   if call.src[0].op not in (Ops.PROGRAM, Ops.COPY): return None # only these bodies can be enqueued
-  if not (bufs:=get_call_arg_uops(call)): return None
+  if not (bufs:=get_call_arg_uops(call)) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
   if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
-  if not all(all_devices_in(b.device, HCQ_DEVS) or cpu_bytes(b.device) for b in bufs): return None
-  devs = min(bufs, key=lambda b: not all_devices_in(b.device, HCQ_DEVS - {"CPU"})).device # prio to enqueue on a device with queues
+  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
   # cpu has no queue (yet)
   if not all_devices_in(devs, HCQ_DEVS) or to_tuple(devs)[0].startswith("CPU"): return None
   # a device without a copy queue leaves copies to its allocator
@@ -62,11 +59,6 @@ def to_name(*parts:str) -> str: return "_".join(parts).replace(":", "_").lower()
 def timeline(devs:tuple[str, ...]) -> UOp: return UOp.placeholder((2,), dtypes.uint64, 0, device=devs, volatile=True, tag="timeline")
 def timeline_value(devs:tuple[str, ...]) -> UOp: return timeline(devs).index(1).load()
 
-def rt_addr(b:UOp, dev=None) -> UOp: # the address of a view as a runtime value: a word on the runtime device the link patches
-  base, off = unwrap_view(b)
-  word = UOp.placeholder((1,), dtypes.uint64, device=HCQ_RUNTIME_DEV.value, tag="addr")
-  return patch(word, [(0, base.getaddr(dev or HCQ_RUNTIME_DEV.value))]).index(0).load() + off
-
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   fn = to_name("submit", (devs:=to_tuple(devs))[0].split(":")[0], queue.split(":")[0])
   return UOp.custom_function(fn, UOp(Ops.LINEAR, src=tuple(cmds), arg=(devs, queue)))
@@ -110,13 +102,17 @@ STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) <
 @functools.cache
 def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
 
-def _need_staging(a, b): # a queue device copying from memory nobody maps goes through the cpu
-  return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not (all_devices_in(b.device, HCQ_DEVS) or cpu_bytes(b.device)) \
-    and Device[to_tuple(a.device)[0]].has_copy_queue
+def _need_staging(a, b):
+  return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS) and Device[to_tuple(a.device)[0]].has_copy_queue
+
+def stage_copy_ext(call:UOp) -> UOp|None:
+  if (d:=next((d for b in call.src[1:] for d in to_tuple(b.device) if not d.startswith("CPU")), None)) is None: return None
+  return pm.rewrite(call) if (pm:=getattr(Device[d], "pm_stage_copy", None)) is not None else None
 
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
 
+  assert src.dtype.itemsize == dst.dtype.itemsize, "staged copies must be dtype-size matched"
   base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
   chunk = (STAGING_SIZE // STAGING_SLOTS) // it
   for i, off in enumerate(range(0, src.max_numel(), chunk)):
@@ -125,6 +121,7 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   return UOp(Ops.LINEAR, src=tuple(copies))
 
 pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), stage_copy_ext),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
 ])
 
@@ -230,13 +227,9 @@ def _finalize_batch(ctx:BatchCtx) -> UOp:
   submits += [_epilogue(ctx, dev) for dev in ctx.queues]
   fence = UOp.custom_function("hcq_fence", *[ctx.sched_timeline((dev,)) for dev in ctx.queues],
                               *[ctx.queue_signal((dev,), q) for dev, qs in ctx.queues.items() for q in qs])
-  merged:list[UOp] = [] # the submits in order, after the fence
-  for m in _merge_queues(submits): merged.append(m.after(fence, *merged[-1:]))
+  merged = [m.after(fence) for m in _merge_queues(submits)]
   estimates = sum((estimate_uop(call) for call, _, _ in ctx.batch), start=Estimates()).simplify()
-  sink = UOp.sink(*merged, arg=KernelInfo("hcq_submit", estimates=estimates), tag=1)
-  for pm in [Device[d].pm_batch for d in ctx.queues if Device[d].pm_batch is not None]: # a device adds its own work to the batch
-    if (r:=pm.rewrite(sink)) is not None: sink = r
-  return sink.call(aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns)))
+  return UOp.sink(*merged, arg=KernelInfo("hcq_submit", estimates=estimates), tag=1).call(aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns)))
 
 @rewrite_group(new_ctx=False)
 def sched_batches(l:UOp, profile:bool) -> UOp:
@@ -258,15 +251,7 @@ class EncodeCtx:
   lt_patches:dict[UOp, list[UOp]] = field(default_factory=dict) # placeholder -> the stores into it that resolve when the linear links
 
 class HWQueue:
-  q_rewrite = PatternMatcher([ # the ops of a queue: a queue defines the methods it supports
-    (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.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: ctx.memory_barrier()),
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("wait_eq", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val, eq=True)),
-    (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)),
-  ])
+  q_rewrite:PatternMatcher
 
   def __init__(self, ctx:EncodeCtx, submit:UOp):
     self.ctx, self.lin = ctx, submit.src[0]
@@ -287,7 +272,6 @@ class HWQueue:
         self.blob += (v & (1 << 8 * n) - 1).to_bytes(n, 'little')
     return len(self.blob)
 
-  def memory_barrier(self): pass # a copy queue has nothing to flush
   def submit(self, cmdbuf:UOp) -> UOp: raise NotImplementedError("queues need a submit")
 
 # *****************
@@ -302,24 +286,17 @@ def hcq_fence(ctx:EncodeCtx, f:UOp) -> UOp:
   for i, dev in enumerate(ctx.devs):
     slots, off = unwrap_view(lasts[i])
     slots = patch(slots, [], bytes(slots.max_numel() * slots.dtype.itemsize)) # zeroed at link
-    target = slots.after(*last, tv:=timeline_value((dev,))).index(off // slots.dtype.itemsize).load()
-    done = timeline((dev,)).after(target, loop:=UOp.loop(i)).index(0).load()
-    bumped = timeline((dev,)).after(done.end(loop, done < target)).index(1).store(nxt:=tv + UOp.const(1, dtypes.uint64))
-    last = (slots.after(bumped).index(off // slots.dtype.itemsize).store(nxt),)
+    done = timeline((dev,)).after(*last, loop:=UOp.loop(i)).index(0).load()
+    waited = done.end(loop, done < slots.index(off // slots.dtype.itemsize).load())
+    nxt = timeline_value((dev,)) + UOp.const(1, dtypes.uint64)
+    last = (timeline((dev,)).after(waited).index(1).store(nxt), slots.after(waited).index(off // slots.dtype.itemsize).store(nxt))
 
   # re-arm the signals
   for sig in sigs:
     base, off = unwrap_view(sig)
     last = (base.after(*last).index(off // sig.dtype.itemsize).store(0),)
   return last[0].barrier(*last[1:])
-
-pm_hcq_encode = PatternMatcher([
-  (UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name="f"), hcq_fence),
-
-  # after blocks are lowered, rechain stores saving original order
-  (UPat(Ops.AFTER, src=(UPat(dtype=dtypes.void, name="root"),), allow_any_len=True, name="a"),
-    lambda root, a: root.substitute({s.buf_uop: s.buf_uop.after(*a.src[1:]) for s in root.toposort() if s.op is Ops.STORE}, walk=True)),
-])
+pm_hcq_encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq_fence", name="f"), hcq_fence)])
 
 # *****************
 # 3.2. split
@@ -338,7 +315,6 @@ def addrs_to_table(ctx:EncodeCtx, g:UOp) -> UOp|None:
 def _is_link_patch(w:UOp) -> bool:
   if w.op is Ops.GETADDR: return not _is_input_addr(w)
   if w.op is Ops.PARAM: return w.tag is not None
-  if w.op is Ops.BUFFER: return w.addrspace is AddrSpace.GLOBAL # a register is written at runtime
   if w.op in {Ops.LOAD, Ops.AFTER} or w.is_variable: return False
   return all(_is_link_patch(s) for s in w.src)
 
@@ -351,7 +327,12 @@ def hoist_links(ctx:EncodeCtx, a:UOp) -> UOp|None:
   ctx.lt_patches.setdefault(unwrap_view(a.src[0])[0], []).extend(ws.substitute(sub).src)
   return a.src[0].after(*rest)
 
-pm_patches = PatternMatcher([(UPat(Ops.GETADDR, name="g"), addrs_to_table), (UPat(Ops.AFTER, name="a"), hoist_links)])
+pm_lower_body = PatternMatcher([
+  (UPat(Ops.GETADDR, name="g"), addrs_to_table),
+  (UPat(Ops.AFTER, name="a"), hoist_links),
+  (UPat(Ops.AFTER, src=(UPat(dtype=dtypes.void, name="root"),), allow_any_len=True, name="a"),
+    lambda root, a: root.substitute({s.buf_uop: s.buf_uop.after(*a.src[1:]) for s in root.toposort() if s.op is Ops.STORE}, walk=True)),
+])
 
 def patch(buf:UOp, rows:list[tuple[int, UOp]], blob:bytes|None=None) -> UOp:
   groups:dict[tuple[DType, int, bool], list[tuple[int, UOp]]] = {} # split by: dtype, alignment, is_link (rt/lt can't share a store)
@@ -387,9 +368,9 @@ def lower_call(call:UOp) -> UOp|None:
 
   # encode bodies
   ctx = EncodeCtx(call.arg.aux.device)
-  devs = [Device[d] for d in dedup([d.split(":")[0] for d in ctx.devs])]
-  body = graph_rewrite(call.src[0], sum([d.pm_encode for d in devs], PatternMatcher([])) + pm_hcq_encode, ctx=ctx, bpm=pm_patches, name="encode")
-  body = graph_rewrite(body, sum([d.pm_lower for d in devs if d.pm_lower is not None], PatternMatcher([])), ctx=ctx, bpm=pm_patches, name="lower")
+  pm = sum([Device[d].pm_encode for d in dedup([d.split(":")[0] for d in ctx.devs])], pm_hcq_encode)
+  body = graph_rewrite(call.src[0], pm, ctx=ctx, walk=True, name="encode body")
+  body = graph_rewrite(body, pm_lower_body, ctx=ctx, name="lower body")
 
   # resize table
   body = body.substitute({ctx.table: (table:=UOp.placeholder((len(ctx.inputs),), dtypes.uint64, device="CPU", tag="inputs"))})
@@ -397,7 +378,6 @@ def lower_call(call:UOp) -> UOp|None:
   # the placeholders become the body's params in visit order, variables bind by name after them, the ranges renumber
   tops = body.toposort()
   bufs, alus = partition([u for u in tops if u.op is Ops.PARAM], lambda u: u.tag is not None)
-  bufs += [b for b in ctx.lt_patches if b not in bufs] # a patched placeholder is an arg, the link applies its patches through the args
   names = dedup([a.arg.name for a in alus])
   # bufs to params
   params = {b: UOp.param(i, b.dtype, b.shape, HCQ_RUNTIME_DEV.value, volatile=b.arg.volatile, name=f"{b.arg.name}_{i}") for i, b in enumerate(bufs)}
@@ -452,8 +432,7 @@ def bufferize_buf(ctx:LinkCtx, b:UOp) -> UOp|None: # ctx: a kept link (the jit's
   # device owns the placeholders it names
   if (r:=cast(Buffer|None, dev.pm_bufferize.rewrite(b, ctx=dev))) is not None: pass
   elif not ctx.use_rt:
-    spec = BufferSpec(host=b.arg.volatile, uncached=b.arg.volatile, cpu_access=True)
-    r = Buffer(dev.device, b.max_numel(), b.dtype, options=spec, preallocate=True)
+    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)
@@ -532,10 +511,8 @@ class HCQ2Compiled(Compiled):
     self.prof_ents:dict[tuple[Buffer, int], ProfileGraphEntry] = {} # (a batch's timestamps, start slot) -> entry, read at synchronize
 
   @functools.cached_property
-  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]
-    buf = Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    buf._buf.cpu_view().view(fmt='B')[:16] = bytes(16)
-    return buf
+  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]: zeroed host memory
+    return Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
 
   def collect_prof(self):
     if PROFILE:
@@ -611,8 +588,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
     self.dev.synchronize()
     with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): ctypes.memmove(mv_address(dest), src.cpu_view().addr, dest.nbytes)
 
-  def map(self, buf:Buffer) -> HCQBuffer: # another device's buffer: an hcq buffer as is, anything else through its cpu mapping
-    return self._map(buf._buf if hasattr(buf._buf, "va_addr") else buf.get_buf("CPU"))
   def _map(self, buf:HCQBuffer) -> HCQBuffer: # a mapping lives on the opaque, like hcq1: the lru hands the same one to many Buffers
     if self.dev not in buf.mapped_devs:
       if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")

tinygrad/runtime/support/usb.py

@@ -1,10 +1,11 @@
 import ctypes, struct, time, functools, itertools
-from tinygrad.runtime.autogen import libusb, libc
-from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv, to_tuple
-from tinygrad.dtype import dtypes, DType, AddrSpace
+from typing import Any, cast
+from tinygrad.runtime.autogen import libusb
+from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv, unwrap, to_tuple
+from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
-from tinygrad.device import Buffer, BufferSpec
-from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV, ccall, patch, rt_addr, unwrap_view, cpu_bytes
+from tinygrad.device import Buffer, BufferSpec, Device
+from tinygrad.runtime.support.hcq2 import HCQInfo, make_submit, HCQ_RUNTIME_DEV
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
 
@@ -226,6 +227,7 @@ class USBMMIOInterface(MMIOInterface):
     return (index * self.el_sz, self.el_sz)
 
   def __getitem__(self, index):
+    Device[HCQ_RUNTIME_DEV.value].synchronize() # one driver on the link: drain the compiled submits before python touches it
     off, sz = self._off_from_index(index)
     if self.pcimem:
       assert sz % 4 == 0 and off % 4 == 0, f"pcie_mem_read requires 4-byte aligned access, got off={off}, sz={sz}"
@@ -234,6 +236,7 @@ class USBMMIOInterface(MMIOInterface):
     return data if isinstance(index, slice) else int.from_bytes(data, "little")
 
   def __setitem__(self, index, data):
+    Device[HCQ_RUNTIME_DEV.value].synchronize()
     off, _ = self._off_from_index(index)
     data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
     if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
@@ -246,233 +249,75 @@ class USBMMIOInterface(MMIOInterface):
     return USBMMIOInterface(self.usb, self.addr+offset, self.nbytes-offset if size is None else size, fmt=fmt or self.fmt, pcimem=self.pcimem)
 
 # *****************
-# hcq2: the host program drives the board over libusb. an access to device memory becomes a transfer on the link, the buffer holding the
-# libusb handle. every transfer is h.after(call), so the link threads through the program in order: h is always the link after all before it.
-# a copy between the host and vram goes through the controller's sram in chunks: the queue moves each chunk, the host streams it
 
-# *****************
-# 0. helpers
+# TODO: unported to the hcq2 rewrite, keeps the old signal placeholder helper alive
+def make_buf(devs, slot:int=0, tag:str="signal") -> UOp: return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True, tag=tag)
 
-HALF, CHUNK = 0x40000, 0x40000 - 512 # the sram has two halves, a chunk fills one: its payload, then the block holding its sentinel
-PENDING = 0xff # a transfer status libusb never sets: the program marks a transfer before submitting it, reaps it once the status changed
-USB_HOST = ("usb_link", "usb_stage", "usb_xfer", "usb_zeros", "put_value", "cmdbuf_copy") # the device's placeholders that live in host memory
+def _libusb(devs, dep:tuple[UOp, ...], fn:str, *args) -> UOp:
+  # the CUSTOM_FUNCTION body holds the callee (the loaded function pointer), the call args are plain dataflow
+  fptr = make_buf(devs, tag=f"func:{fn}").after(*dep).index(0).load()
+  return UOp.custom_function(fn, fptr).call(make_buf(devs, tag="usb_handle").index(0).load(),
+    *[UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args], ret_dtype=dtypes.void)
 
-def usb_buf(dev, tag:str, n:int=1, dt:DType=dtypes.uint8) -> UOp: return UOp.placeholder((n,), dt, 0, device=to_tuple(dev)[0], tag=f"usb_{tag}")
-def usb_link(dev) -> UOp: return usb_buf(dev, "link", 3, dtypes.uint64) # [the libusb handle, its context, the chunks of the last batch]
-def usb_xfer(x:UOp, field:str) -> UOp: # a field of a libusb transfer struct
-  f = getattr(libusb.struct_libusb_transfer, field)
-  return x[f.offset:f.offset + f.size].bitcast({4: dtypes.int32, 8: dtypes.uint64}[f.size]).index(0)
+def usb_bulk(devs, dep, endpoint:int, data:UOp, length, timeout:int=1000) -> UOp: # NULL actual_length out param
+  return _libusb(devs, dep, "libusb_bulk_transfer", endpoint, data, length, UOp.const(0, dtypes.uint64), timeout)
 
-def usb_reg(dt:DType, *vals:UOp|int) -> UOp: # an array on the program's stack holding vals: what a transfer reads or writes
-  r = UOp.placeholder((max(1, len(vals)),), dt, addrspace=AddrSpace.REG)
-  return r.after(*[r.index(i).store(v.cast(dt) if isinstance(v, UOp) else UOp.const(v, dt)) for i, v in enumerate(vals)])
+def usb_stream(devs, dep:tuple[UOp, ...], addr:UOp, data:UOp, nbytes:int, write:bool) -> UOp:
+  hdr = UOp.placeholder((2,), dtypes.uint64, device=devs, tag="usb_scratch").after(*dep)
+  arm = _libusb(devs, (hdr.index(0).store(addr), hdr.index(1).store(UOp.const(nbytes // 4, dtypes.uint64))), "libusb_control_transfer",
+                0x40, 0xF0, (0x60 if write else 0x20) | (0x0F << 8), 1 if write else 2, hdr.index(0), 12, 5000)
+  return usb_bulk(devs, (arm,), 0x02 if write else 0x81, data, nbytes)
 
-def _addr(b:UOp, idx:UOp, dt:DType) -> UOp: return rt_addr(b) + (idx * dt.itemsize).cast(dtypes.uint64) # of an element of a view
-def _host(b:UOp) -> bool: return b.device is None or cpu_bytes(b.device) # the stack, or memory the host program reads in place
+def usb_load(b:UOp, idx:UOp, dt) -> UOp:
+  got = UOp.placeholder((1,), dt, device=(devs:=to_tuple(b.device)), tag="usb_scratch")
+  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*dt.itemsize).cast(dtypes.uint64)
+  return got.after(usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, got.index(0), dt.itemsize, False)).index(0).load()
 
-# *****************
-# 1. transfers
+def usb_write(b:UOp, idx:UOp, v:UOp) -> UOp:
+  val = (s:=UOp.placeholder((1,), v.dtype, device=(devs:=to_tuple(b.device)), tag="usb_scratch")).after(s.index(0).store(v))
+  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*v.dtype.itemsize).cast(dtypes.uint64)
+  return usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, val.index(0), v.dtype.itemsize, True)
 
-def usb_ctrl(h:UOp, rtype:int, req:int, val:UOp|int, idx:UOp|int, data:UOp, n:UOp|int, timeout:int=1000) -> UOp:
-  return h.after(ccall(libusb.libusb_control_transfer, h.index(0).load(), rtype, req, val, idx, data, n, timeout))
+def usb_idle(devs) -> UOp:
+  v = usb_load(make_buf(devs, tag="timeline_signal").after(loop:=UOp.loop(0)), UOp.const(0, dtypes.int), dtypes.uint64)
+  return v.end(loop, v + 1 < make_buf(devs, tag="timeline_value").index(0).load())
 
-def usb_bulk(h:UOp, ep:int, data:UOp, n:UOp|int, timeout:int=10000) -> UOp: # NULL actual_length
-  return h.after(ccall(libusb.libusb_bulk_transfer, h.index(0).load(), ep, data, n, UOp.const(0, dtypes.uint64), timeout))
+def usb_scsi(devs, read:bool, nbytes:int) -> UOp:
+  return _libusb(devs, (usb_idle(devs),), "libusb_control_transfer", 0x40, 0xF2, ceildiv(nbytes, 512) | (0x8000 if read else 0),
+                 (ceildiv(nbytes, 0x4000) & 0xFF) << 8, UOp.const(0, dtypes.uint64), 0, 1000)
 
-def usb_poke(h:UOp, addr:UOp, val:UOp) -> UOp: # 0xF0 mode 0: a dword in one control transfer. the header: the address, then a dword
-  return usb_ctrl(h, 0x40, 0xF0, 0x60 | 0x0F00, 0, usb_reg(dtypes.uint64, addr, val.bitcast(dtypes.uint32).cast(dtypes.uint64)).index(0), 12, 5000)
+def usb_stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if (cin:=to_tuple(src.device)[0].startswith("CPU")) == to_tuple(dst.device)[0].startswith("CPU"): return None
 
-def usb_stream(h:UOp, addr:UOp, data:UOp, n:UOp|int, write:bool) -> UOp: # 0xF0 mode 1/2: the header, then the payload on the bulk endpoint
-  hdr = usb_reg(dtypes.uint64, addr, n // 4)
-  h = usb_ctrl(h, 0x40, 0xF0, (0x60 if write else 0x20) | 0x0F00, 1 if write else 2, hdr.index(0), 12, 5000)
-  return usb_bulk(h, 0x02 if write else 0x81, data, n)
+  total, ops, win = dst.nbytes(), [], cast(Any, Device[(devs:=to_tuple((dst if cin else src).device))[0]]).iface.usb_sram
+  for off in range(0, total, win.size): # off and nb are bytes, the two ends of the copy can have different dtypes
+    sram = UOp.from_buffer(win)[0:(nb:=min(win.size, total - off))]
+    s, d = src[off // src.dtype.itemsize:(off + nb) // src.dtype.itemsize], dst[off // dst.dtype.itemsize:(off + nb) // dst.dtype.itemsize]
+    if cin:
+      push = usb_bulk(devs, (usb_scsi(devs, False, nb),), 0x02, s.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
+      ops += [UOp.custom_function("hcq", push.sink()).call(sram, s, name="hcq_copyin", aux=HCQInfo(devs)),
+              sram.copy_to_device(d.device).call(d, sram)]
+    else:
+      pad = UOp.new_buffer("CPU", round_up(nb, 512), dtypes.uint8)[0:nb]
+      submit = make_submit(s.copy_to_device(sram.device).call(sram, s), devs=devs, queue="COPY:0")
+      pull = usb_bulk(devs, (submit,), 0x81, pad.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
+      ops += [UOp.custom_function("hcq", pull.sink()).call(pad, sram, s, name="hcq_copyout", aux=HCQInfo(devs)),
+              pad.copy_to_device("CPU").call(d, pad)]
+  return UOp(Ops.LINEAR, src=tuple(ops))
+pm_usb_stage = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), usb_stage_copy)])
 
-# *****************
-# 2. batch: a copy between the host and vram goes through the sram in chunks, a batch numbers them from 0 in queue order. the queue side moves
-# each chunk between the sram and vram and clears what it waited for, the host side streams it: a run of copies one way is a loop over a table
-# of its chunks. the fence counts the chunks the queue is done with: the host resets it once the batch before is done
-
-def usb_wire(size:UOp|int) -> UOp|int: return (size + 4 + 511) // 512 * 512 # a chunk on the wire: the payload, then its sentinel in the last dword
-def usb_sentinel(g:UOp) -> UOp: return ((g & 0xFFFFFF) | 0x51000000).cast(dtypes.uint32)
-def _staged(call:UOp) -> bool: return call.op is Ops.CALL and call.src[0].op is Ops.COPY and _host(call.src[1]) != _host(call.src[2])
-def _chunks(call:UOp) -> list[tuple[UOp, int, int]]: # (the host view, byte offset, bytes) per chunk of a copy
-  host, win = (call.src[2], CHUNK) if _host(call.src[2]) else (call.src[1], 2 * HALF)
-  return [(host, off, min(win, host.nbytes() - off)) for off in range(0, host.nbytes(), win)]
-
-def ins(name:str, dst:UOp, val:UOp) -> UOp: return UOp(Ops.INS, arg=(name, dtypes.void), src=(dst, val))
-def usb_chunks(lin:UOp, first:dict[UOp, tuple[int, int]]) -> UOp: # the queue side: per chunk wait for the host, move it, release the sram
-  dev, zero = lin.arg[0][0], UOp.const(0, dtypes.uint32) # first: the number of a copy's first chunk and of its run's
-  def chunks(call:UOp) -> list[UOp]:
-    cin, sram, ops = _host(call.src[2]), usb_buf(dev, "sram", 2 * HALF), [] # into vram from the host, or out of it
-    vram, (n0, k0) = call.src[1] if cin else call.src[2], first[call]
-    for n, (_, off, nb) in enumerate(_chunks(call), start=n0): # off and nb are bytes: the two ends of a copy can differ in dtype
-      wo = ((n - k0) & 1) * HALF if cin else 0 # a run alternates the halves from half 0
-      w, v = sram[wo:wo + nb], vram[off // vram.dtype.itemsize:(off + nb) // vram.dtype.itemsize]
-      if cin: # the host streamed it, the sentinel is cleared for the next batch
-        ops.append(ins("wait_eq", sentinel:=sram[(o:=wo + usb_wire(nb) - 4):o + 4].bitcast(dtypes.uint32), usb_sentinel(UOp.const(n, dtypes.uint32))))
-        ops += [call.replace(src=(call.src[0], v, w)), ins("store", sentinel, zero)]
-      else: # the host armed a read of the window
-        ops += [ins("wait", go:=usb_buf(dev, "go", 1, dtypes.uint32), UOp.const(n + 1, dtypes.uint32)), ins("store", go, zero)]
-        ops += [call.replace(src=(call.src[0], w, v)), ins("store", usb_buf(dev, "cq", 0x1000)[12:16].bitcast(dtypes.uint32), zero)] # send it
-      ops.append(ins("store", usb_buf(dev, "fence", 1, dtypes.uint32), UOp.const(n + 1, dtypes.uint32))) # the chunk is done with the sram
-    return ops
-  return lin.replace(src=tuple(u for call in lin.src for u in (chunks(call) if call in first else [call])))
-
-def usb_batch(s:UOp) -> UOp|None: # the copies of the batch in queue order: the queue side into each submit, the host side after the last one
-  lins = [submit.without_after.src[0] for submit in s.src]
-  runs, first, k = [], {}, 0 # the chunks numbered in queue order, a run of copies one way at a time
-  for cin, grp in itertools.groupby([call for lin in lins for call in lin.src if _staged(call)], key=lambda call: _host(call.src[2])):
-    chunks:list = []
-    for call in grp: first[call], chunks = (k + len(chunks), k), chunks + _chunks(call)
-    runs.append((cin, k, chunks))
-    k += len(chunks)
-  if not runs: return None
-  s = s.substitute({lin: usb_chunks(lin, first) for lin in lins})
-  h = usb_link(lins[0].arg[0][0]).after(s.src[-1])
-  h = usb_drained(h, h.index(2).load() + 1) # the batch before is done with the sram: its chunks, then the fence restarts
-  h = usb_ctrl(h, 0x40, 0xE5, rt_addr(usb_buf(h.device, "fence", 1, dtypes.uint32)), 0, UOp.const(0, dtypes.uint64), 0) # 0xE5 writes a byte
-  for cin, k0, chunks in runs: h = (usb_copyin if cin else usb_copyout)(h, chunks, k0)
-  return s.replace(src=(*s.src, h.index(2).store(UOp.const(k, dtypes.uint64))))
-pm_usb_batch = PatternMatcher([(UPat(Ops.SINK, name="s"), usb_batch)])
-
-def usb_table(chunks:list[tuple[UOp, int, int]]) -> UOp: # [host address, bytes] per chunk: the link patches it, a jit input at runtime
-  table = UOp.placeholder((2 * len(chunks),), dtypes.uint64, device=HCQ_RUNTIME_DEV.value, tag="usb_table")
-  rows = [(16 * i, (v:=unwrap_view(host))[0].getaddr(HCQ_RUNTIME_DEV.value) + (v[1] + off)) for i, (host, off, _) in enumerate(chunks)]
-  return patch(table, rows + [(16 * i + 8, UOp.const(nb, dtypes.uint64)) for i, (_, _, nb) in enumerate(chunks)])
-
-def usb_reap(h:UOp, xfer:UOp) -> UOp: # poll the event loop until the async transfer is done
-  loop = UOp.range(UOp(Ops.NOOP), next(UOp.unique_num), dtype=dtypes.void, src=(h,)) # a loop on the link, unique by it
-  events = ccall(libusb.libusb_handle_events_timeout, h.after(loop).index(1).load(), usb_reg(dtypes.uint64, 0, 0).index(0))
-  status = usb_xfer(xfer.after(events), "status").load()
-  return h.after(status.end(loop, status.eq(PENDING)))
-
-def usb_drained(h:UOp, need:UOp) -> UOp: # wait until the queue is done with the chunks before need: it may lag one chunk, the other half
-  loop, slot = UOp.range(UOp(Ops.NOOP), next(UOp.unique_num), dtype=dtypes.void, src=(h,)), usb_reg(dtypes.uint32)
-  h = usb_ctrl(h.after(loop), 0xC0, 0xE4, rt_addr(usb_buf(h.device, "fence", 1, dtypes.uint32)), 0, slot.index(0), 1) # 0xE4 reads controller memory
-  fence = slot.after(h).index(0).load().cast(dtypes.uint64) # one byte can't tear, the count is compared mod 256
-  return h.after(fence.end(loop, ((need - fence) & 0xff) > 1))
-
-def usb_chunk(h:UOp, table:UOp, i:UOp, half:int, k0:int) -> UOp: # stream chunk i of the run into a half of the sram: an async bulk after its arm
-  xfer, stage = usb_buf(h.device, f"xfer{half}", 64), usb_buf(h.device, "stage", 2 * HALF)
-  g = (i + k0).cast(dtypes.uint64) # chunk k0 + i of the batch
-  addr, size = table.index(2 * i).load(), table.index(2 * i + 1).load().cast(dtypes.int)
-  wire = usb_wire(size)
-  h = usb_reap(h, xfer) # the transfer that used this half before
-  h = h.after(ccall(libc.memcpy, stage.after(h).index(half * HALF), addr, size.cast(dtypes.uint64)))
-  h = h.after(stage.after(h).bitcast(dtypes.uint32).index((half * HALF + wire - 4) // 4).store(usb_sentinel(g)))
-  h = usb_drained(h, g) # the queue is done with the chunk that used this half before
-  h = usb_ctrl(h, 0x40, 0xF2, wire // 512, half * 16 | ((wire + 0x3fff) // 0x4000 << 8), UOp.const(0, dtypes.uint64), 0) # arm the write
-  xfer = xfer.after(h)
-  xfer = xfer.after(usb_xfer(xfer, "status").store(PENDING), usb_xfer(xfer, "length").store(wire),
-                    usb_xfer(xfer, "buffer").store(rt_addr(stage) + half * HALF))
-  return h.after(ccall(libusb.libusb_submit_transfer, xfer.index(0)))
-
-def usb_copyin(h:UOp, chunks:list, k0:int) -> UOp: # pairs of chunks on the two halves, two transfers in flight. k0: the chunks before the run
-  table, n = usb_table(chunks), len(chunks)
-  pairs = n // 2 if n // 2 > 1 else 0 # a single pair is unrolled: the linearizer places a one trip loop with the code around it
-  h = usb_drained(h, UOp.const(k0 + 1, dtypes.uint64)) # the sram is free: the run starts on half 0
-  if pairs:
-    j = UOp.range(pairs, next(UOp.unique_num), dtype=dtypes.int)
-    h = h.after(usb_chunk(usb_chunk(h.after(j), table, j * 2, 0, k0), table, j * 2 + 1, 1, k0).end(j))
-  for i in range(pairs * 2, n): h = usb_chunk(h, table, UOp.const(i, dtypes.int), i & 1, k0)
-  for half in range(2): h = usb_reap(h, usb_buf(h.device, f"xfer{half}", 64)) # nothing in flight after a run: a read may come next
-  return h
-
-def usb_copyout(h:UOp, chunks:list, k0:int) -> UOp: # per chunk: arm a read of the sram, release the queue to fill it, pull it
-  table, stage = usb_table(chunks), usb_buf(h.device, "stage", 2 * HALF)
-  h = usb_drained(h, UOp.const(k0 + 1, dtypes.uint64)) # the sram is free and nothing else writes the controller's memory
-  i = UOp.range(len(chunks), next(UOp.unique_num), dtype=dtypes.int)
-  addr, size = table.index(2 * i).load(), table.index(2 * i + 1).load().cast(dtypes.int)
-  wire = (size + 511) // 512 * 512
-  hi = usb_ctrl(h.after(i), 0x40, 0xF2, (wire // 512) | 0x8000, (wire + 0x3fff) // 0x4000 << 8, UOp.const(0, dtypes.uint64), 0) # arm the read
-  hi = usb_poke(hi, rt_addr(usb_buf(h.device, "go", 1, dtypes.uint32)), (i + k0 + 1).cast(dtypes.uint32))
-  hi = usb_bulk(hi, 0x81, stage.index(0), wire)
-  hi = hi.after(ccall(libc.memcpy, addr, stage.after(hi).index(0), size.cast(dtypes.uint64)))
-  return h.after(hi.end(i))
-
-# *****************
-# 3. lower: the host's accesses to device memory. a load streams the value into a register, a store pokes it, a loop of stores streams the source
-
-def _remote(b:UOp) -> bool: return (p:=unwrap_view(b)[0]).op is Ops.PARAM and not _host(p) and not str(p.tag).startswith(USB_HOST)
-def _deps(b:UOp) -> tuple[UOp, ...]: # what a buffer view is after
-  return (b.src[1:] if b.op is Ops.AFTER else ()) + (_deps(b.src[0]) if b.op in (Ops.BITCAST, Ops.SHRINK, Ops.AFTER) else ())
-
-def _affine(idx:UOp, r:UOp) -> UOp|None: # idx = base + r: the base, None if idx doesn't walk r with unit stride
-  if idx is r: return UOp.const(0, r.dtype)
-  if idx.op is not Ops.ADD or r not in idx.src: return None
-  base = idx.src[1] if idx.src[0] is r else idx.src[0]
-  return base if r not in base.ranges else None
-
-def usb_load(b:UOp, idx:UOp, ld:UOp) -> UOp:
-  slot = usb_reg(ld.dtype)
-  h = usb_stream(usb_link(b.device).after(*_deps(b)), _addr(b, idx, ld.dtype), slot.index(0), ld.dtype.itemsize, False)
-  return slot.after(h).index(0).load()
-
-def usb_store(b:UOp, idx:UOp, v:UOp) -> UOp:
-  if idx.op is Ops.STACK: # a patch: word by word, each after the one before
-    h = usb_store(b, idx.src[0], v.src[0])
-    for i, w in zip(idx.src[1:], v.src[1:]): h = usb_store(b.after(h), i, w)
-    return h
-  h, addr = usb_link(b.device).after(*_deps(b)), _addr(b, idx, v.dtype)
-  if v.dtype.itemsize == 4: return usb_poke(h, addr, v)
-  return usb_poke(usb_poke(h, addr, v.cast(dtypes.uint32)), addr + 4, (v >> 32).cast(dtypes.uint32))
-
-def usb_copy(dst:UOp, di:UOp, v:UOp, r:UOp) -> UOp|None: # a loop of unit stride stores from host memory: one stream
-  if not _remote(dst): return None
-  if v.op is Ops.LOAD and not _remote(sb:=v.src[0].src[0]): s0, deps = _affine(v.src[0].src[1], r), _deps(sb) # from host memory
-  elif v.vmin == v.vmax == 0: sb, s0, deps = usb_buf(dst.device, "zeros", r.src[0].vmax * v.dtype.itemsize), UOp.const(0, dtypes.int), ()
-  else: return usb_store(dst, di, v).end(r)
-  if s0 is None or (d0:=_affine(di, r)) is None: return usb_store(dst, di, v).end(r)
-  h, cnt = usb_link(dst.device).after(*_deps(dst), *deps, *r.src[1:]), r.src[0]
-  # the firmware can't stream zero bytes: an empty loop streams one element into a scratch word instead
-  addr = (cnt > 0).where(_addr(dst, d0, v.dtype), rt_addr(usb_buf(dst.device, "scratch", 1, dtypes.uint32)))
-  return usb_stream(h, addr, sb.index(s0.minimum(sb.max_numel() - 1)), (cnt * v.dtype.itemsize).maximum(v.dtype.itemsize), True)
-
-pm_usb_lower = PatternMatcher([ # a store inside a loop is left to the loop's END
-  (UPat.var("dst").index(UPat.var("di")).store(UPat.var("v")).end(UPat(Ops.RANGE, name="r")), usb_copy),
-  (UPat.var("b").index(UPat.var("idx")).store(UPat.var("v")), lambda b, idx, v: None if idx.ranges or not _remote(b) else usb_store(b, idx, v)),
-  (UPat.var("b").index(UPat.var("idx")).load(name="ld"), lambda b, idx, ld: usb_load(b, idx, ld) if _remote(b) else None),
-])
-
-# *****************
-# 4. bufferize: the device's usb state, bound once for the life of the device
-
-def _init(b:Buffer, data:bytes) -> Buffer: # the buffer, allocated and holding data
-  b.ensure_allocated()._buf.cpu_view().view(fmt='B')[:len(data)] = data
-  return b
-def _region(dev, b) -> Buffer: # a window of the controller's memory: not ours to free
-  return Buffer(dev.device, b.size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
-
-@functools.cache
-def _link(dev) -> Buffer: # the libusb handle, its context, the chunks of the last batch
-  handles = [ctypes.addressof(x.contents) for x in (dev.iface.pci_dev.usb.usb.handle, USB3.ctx())]
-  return _init(Buffer("CPU", 3, dtypes.uint64), struct.pack('QQQ', *handles, 0))
-@functools.cache
-def _xfer(dev, tag:str) -> Buffer: # an async bulk out: the program sets its buffer and length, then reaps its status
-  t = libusb.libusb_alloc_transfer(0).contents
-  t.dev_handle, t.endpoint, t.type, t.timeout = dev.iface.pci_dev.usb.usb.handle, 0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK, 10000
-  return Buffer("CPU", ctypes.sizeof(t), dtypes.uint8, options=BufferSpec(external_ptr=ctypes.addressof(t), nolru=True), preallocate=True)
-@functools.cache
-def _cpu(dev, tag:str, n:int) -> Buffer: return Buffer("CPU", n, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
-@functools.cache
-def _word(dev, tag:str) -> Buffer: # a dword in vram
-  return _init(Buffer(dev.device, 1, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True, nolru=True)), bytes(4))
-@functools.cache
-def _fence(dev) -> Buffer: return _init(_region(dev, dev.iface.sys_buf).view(1, dtypes.uint32, 0x800), bytes(4))
-@functools.cache
-def _sram(dev) -> Buffer: return _init(_region(dev, dev.iface.sram), bytes(2 * HALF)) # no stale sentinel from an earlier process
-@functools.cache
-def _cq(dev) -> Buffer: return _region(dev, dev.iface.cq_buf)
+USB_HOST_TAGS = {"signal", "timeline_signal"}
+pm_usb_hostio = PatternMatcher([
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))),),
+        name="ld"), lambda b, idx, ld: usb_load(b, idx, ld.dtype)),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))), UPat(name="v"))), usb_write)])
 
 pm_usb_bufferize = PatternMatcher([
-  (UPat(Ops.PARAM, tag="usb_link"), lambda ctx: _link(ctx)),
-  (UPat(Ops.PARAM, tag={"usb_xfer0", "usb_xfer1"}, name="b"), lambda ctx, b: _xfer(ctx, b.tag)),
-  (UPat(Ops.PARAM, tag={"usb_stage", "usb_zeros"}, name="b"), lambda ctx, b: _cpu(ctx, b.tag, b.max_numel())),
-  (UPat(Ops.PARAM, tag={"usb_go", "usb_scratch"}, name="b"), lambda ctx, b: _word(ctx, b.tag)), # words in vram the host pokes
-  (UPat(Ops.PARAM, tag="usb_fence"), lambda ctx: _fence(ctx)),
-  (UPat(Ops.PARAM, tag="usb_sram"), lambda ctx: _sram(ctx)),
-  (UPat(Ops.PARAM, tag="usb_cq"), lambda ctx: _cq(ctx)),
-  (UPat(Ops.PARAM, name="b"), lambda b: Buffer("CPU", b.max_numel(), b.dtype, preallocate=True) if str(b.tag).startswith("cmdbuf_copy") else None),
-]) # the sdma cmdbuf streams into the ring from the host
+  (UPat(Ops.PARAM, tag={"systems", "runtime", "inputs", "usb_scratch"}, name="b"),
+   lambda ctx, b: Buffer("CPU", b.max_numel(), b.dtype, options=BufferSpec(nolru=True), preallocate=True)),
+  (UPat(Ops.PARAM, tag="usb_handle", name="b"), lambda ctx, b: ctx.signal(b.tag, ctx.iface.usb_handle, device="CPU")),
+  (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if not isinstance(b.tag, str) or not b.tag.startswith("func:") else
+   ctx.signal(b.tag, unwrap(ctypes.cast(getattr(libusb.dll, b.tag[5:]), ctypes.c_void_p).value), device="CPU")),
+])
 
 if DEV.interface.startswith("MOCK"): from test.mockgpu.usb import MockUSB3 as USB3  # type: ignore  # noqa: F811

tinygrad/schedule/__init__.py

@@ -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
+from tinygrad.schedule.prepare import prepare_rangeify, prepare_call_views
 from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
@@ -122,6 +122,8 @@ 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:

tinygrad/schedule/prepare.py

@@ -8,6 +8,109 @@ 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),
+])
+
+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}
+  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()))
+
+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),
+])
+
+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}
+
+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])
+
 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])
   if u.op is Ops.AFTER and (b:=walk_mop(u.src[0])) is not u.src[0]: return b.after(*u.src[1:])
@@ -24,6 +127,8 @@ 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
@@ -129,13 +234,10 @@ 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+PatternMatcher([
-  # resolve calls with RETURNED inputs (inline the body)
+earliest_rewrites = mop_cleanup+pm_resolve_call_outputs+PatternMatcher([
+  # Inline calls with unbound outputs.
   (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),
 
@@ -215,7 +317,9 @@ pm_copy_to_store = PatternMatcher([
 def prepare_rangeify(sink:UOp) -> UOp:
   # prepare for rangeify
   tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
-  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
+  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, 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,246 +1,55 @@
 # 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 dataclass, field, replace
+from dataclasses import 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 resolve_returned_after, 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.multi import multi_pm
+from tinygrad.schedule.prepare import buffer_view_subs, prepare_to_call, on_disk
 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 ***
 
-@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)
+@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
 
-# a tag is the tuple of original pre-rewrite UOps a node provides storage for
-def tag_uop(x:UOp): return None if x.tag is not None else x.replace(tag=(x,))
-
-# a base needs storage of its own if it can back a buffer and doesn't already have one
-def needs_storage(u:UOp) -> bool: return not u.is_virtual and not u.has_buffer_identity()
-
-def on_disk(u:UOp): return isinstance(u.device, str) and u.device.startswith("DISK")
 def is_creation_device(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "NPY", "PYTHON"))
 
-def creation_copy_is_realized(u:UOp):
-  # all copies from disk/numpy are realized into a real buffer
-  if is_creation_device(u.src[0]): return tag_uop(u)
-
-# CONTIGUOUS and AFTER + parents are the only nodes that get updated
-add_tags = PatternMatcher([
-  (UPat(Ops.COPY, name="u"), creation_copy_is_realized),
-  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
-   lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
-  (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
-
-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"
-  # 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: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))
-
-  # 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, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
-                                   for i, a in enumerate(c.src[1:])]))
-  rets = tuple(o.after(new_call) for o in outs)
-
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # 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)))
-
-  # 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)])
-
-# NOTE: adding rules to here is bad. these all need to run before the schedule cache
-pm_early_transform_tensor_graph = PatternMatcher([
-  # transform precompiled value-producing calls into opaque CALLs (outputs become real buffers)
-  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
-
-  # resolve AFTER on RETURNED placeholders (for precompiled calls)
-  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
-
-  # 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. an AFTER already names the storage its store writes into
-  ctx = AllocCtx(bases={base for x in big_sink.src if needs_storage(base:=x.base) and base.op is not Ops.AFTER})
-
-  # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
-  # this is the only one where we have to be careful to not break the tensor graph
-  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
-  srcs:list[UOp] = []
-  for u in big_sink.src:
-    if u.op is Ops.AFTER and u.src[0].unsharded_base.is_unbound:
-      # precompiled calls don't need this: transform_precompiled_call gives their outputs real buffers
-      call = u.src[1]
-      if not (call.op is Ops.CALL and call.arg is not None and call.arg.precompile):
-        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) -> None:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, *, tensors:list[Tensor]|None=None) -> 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)
-    scope_tensors: list[Tensor] = [t for tref in list(all_tensors) if (t:=tref()) is not None and t.uop.topovisit(visitor, in_scope)]
+    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)]
 
     # 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])
@@ -251,10 +60,15 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
       if s is ns: continue
       t.uop = ns
 
-def _tensor_holds(u:UOp) -> bool: return any((t:=tref()) is not None and t.uop is u for tref in list(all_tensors))
-
 # **** Tensor helper functions ****
 
+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]
+
 def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
 
 def _fromnp(x: 'numpy.ndarray') -> UOp:
@@ -313,7 +127,9 @@ 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 None or data.device == _device else data.copy_to_device(_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)
     # 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)
 
@@ -386,10 +202,33 @@ 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:
-    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")
+    """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")
     return self
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
@@ -397,9 +236,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")
-    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
+    call, becomes_map = self._prepare_call(*lst)
     _apply_map_to_tensors(becomes_map, name="buffers")
-    return create_linear_with_vars(big_sink)
+    return create_linear_with_vars(call)
 
   def schedule_linear(self, *lst:Tensor) -> UOp:
     """Creates the schedule needed to realize these Tensor(s)."""
@@ -410,7 +249,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 needs_storage(x.uop.base)]
+    to_realize = [x for x in (self,)+lst if not (b:=x.uop.base).is_virtual and not b.has_buffer_identity()]
     if len(to_realize):
       run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
     return self
@@ -425,6 +264,12 @@ 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)
@@ -442,21 +287,26 @@ class Tensor(RandMixin):
     if is_disk:
       (b:=self._buffer()).copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=x._data()))
       return self
-    assigned_to = self.uop.storage_base
-    # assigning to a value is initialization, not a write: the whole tensor is overwritten, so the pending value is dead
-    if not assigned_to.has_buffer_identity() and assigned_to.op is not Ops.CONTIGUOUS:
-      self.uop = (x.uop.src[0] if x.uop.op is Ops.CONTIGUOUS else x.uop).clone()
+    # 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()
       return self
-    # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
-    assign = self.uop.after(self.uop.store(x.uop))
-    ib = self.uop
-    while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
-    if ib is not self.uop:
-      # view assign: replace the node under the views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      _apply_map_to_tensors({ib: ib.after(assign)}, name="Embed View Assign")
-    else:
-      # simple assign
-      self.uop = assign
+    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
     return self
 
   def _buffer(self) -> Buffer:
@@ -529,7 +379,8 @@ class Tensor(RandMixin):
 
   def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor:
     """
-    Creates a clone of this tensor allocating a separate buffer for the data.
+    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.
     If `device` is specified, the clone is placed on that device.
     """
     ret = Tensor(self.uop.clone(device=device))
@@ -538,12 +389,13 @@ class Tensor(RandMixin):
 
   def to(self, device:str|tuple[str, ...]|None) -> Tensor:
     """
-    Moves the tensor to the given device.
+    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.
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
-    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
+    # 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))
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)
@@ -686,12 +538,20 @@ 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)
@@ -699,10 +559,7 @@ class Tensor(RandMixin):
     realized = is_disk or self.uop.base.op is Ops.BUFFER or self.uop._base_buffer_is_realized()
     if (not self.uop.base.is_realized and self.is_floating_point()) or not (advanced or realized):
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value.
-      # the store is self-referential there (the computed value touches its target); clone stores are untouched
-      if v.uop.op is Ops.AFTER and len(v.uop.src) == 2 and (st:=v.uop.src[1]).op is Ops.STORE and \
-          st.src[0] in st.src[1].toposort(enter_calls=False): v = v._apply_uop(lambda x: st.src[1])
+      if (rhs:=_inplace_rhs(v.uop)) is not None: v = v._apply_uop(lambda _, rhs=rhs: rhs)
       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

@@ -458,8 +458,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def ended_ranges(self) -> tuple[UOp, ...]:
-    if self.op is Ops.CALL and self.src[0].op is Ops.CUSTOM_FUNCTION and self.src[0].src: return ()
-    if self.op is Ops.END: return tuple(r for r in self.src[1:] if r.op is Ops.RANGE)
     if self.op in range_start: return self.src[range_start[self.op]:]
     if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
     # UNSHARD ends the DEVICE range: its src is per-device index math, the device axis is carried by the axis metadata
@@ -800,7 +798,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # *** uop Buffer stuff ***
 
-  unique_num = itertools.count(0)
+  # Fresh storage IDs decrease from -1; canonical slots are numbered from 0 within their scope.
+  unique_num = itertools.count(-1, -1)
 
   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
@@ -818,8 +817,11 @@ 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=self.commit_dtype() if dtype is None else dtype, device=device)
+    ret = UOp.empty(self.shard_shape if axis is not None else self.shape, dtype=dt, 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:
@@ -833,11 +835,13 @@ 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.copy_to_device(device)
+    return ret if ret.device == device else ret.clone(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:
@@ -1095,9 +1099,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       trunc = truncate.get(self.dtype) if dtypes.is_float(self.dtype) else math.trunc if dtypes.is_int(self.dtype) else None
       if trunc is not None and all(math.isfinite(v) for v in (smin, smax)): smin, smax = trunc(smin), trunc(smax)
       if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax
-      # a signed or float destination holds the part of the source that overlaps it: overflow is undefined, a nan bound overlaps nothing
-      if self.dtype in dtypes.floats+dtypes.sints+dtypes.weaks and smin <= self.dtype.max and self.dtype.min <= smax:
-        return max(self.dtype.min, smin), min(smax, self.dtype.max)
+      if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, smin), min(smax, self.dtype.max)
     return self.dtype.min, self.dtype.max
 
   @functools.cached_property

tinygrad/uop/spec.py

@@ -21,10 +21,12 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
   # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
   if 0<=idx.vmin and idx.vmax<sz: return True
 
-  # TODO: validate STACK, z3 can't model vectors
+  # TODO: validate these
+  # WEBGPU has a BITCAST in the index, PTX casts pointer to long
+  # VECTORIZE can't be properly modeled in z3 since it doesn't support vectors
   # don't descend into PARAM shape metadata; only the PARAM value participates in index arithmetic
   for x in idx.toposort(gate=lambda x: x.op is not Ops.PARAM) | gate.toposort(gate=lambda x: x.op is not Ops.PARAM):
-    if x.op is Ops.STACK: return True
+    if x.op in {Ops.BITCAST, Ops.STACK}: return True
 
   # if all is good and CHECK_OOB=1, validate with z3
   from tinygrad.uop.validate import validate_index_with_z3
@@ -91,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 on Movement Op, PARAM, BUFFER, CONTIGUOUS, RETURNED, or another AFTER
+  # AFTER preserves its target view.
   (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),
@@ -122,10 +124,9 @@ 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: 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
+  # STORE targets storage (or an AFTER/BITCAST/view of it). 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, Ops.CONTIGUOUS} else None if b.op is Ops.INDEX else False),
+   True if (b:=x.storage_base).op in {Ops.BUFFER, Ops.PARAM} 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),
@@ -174,7 +175,10 @@ 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)),
 
-  # CONTIGUOUS ensures the source UOp realizes
+  # 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.
   (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

tinygrad/uop/validate.py

@@ -29,34 +29,36 @@ z3_alu: dict[Ops, Callable[..., z3.ExprRef]] = python_alu | {Ops.CMOD: lambda a,
   Ops.FLOORMOD: lambda a,b: a-z3_floordiv(a,b)*b,
   Ops.SHR: lambda a,b: a/(2**b.as_long()), Ops.SHL: lambda a,b: a*(2**b.as_long()),
   Ops.AND: z3_and, Ops.WHERE: z3.If, Ops.XOR: z3_xor, Ops.MAX: lambda a,b: z3.If(a<b, b, a),}
-
-def create_bounded(name:str, vmin:int|z3.ArithRef, vmax:int|z3.ArithRef, solver:z3.Solver) -> z3.ArithRef:
-  solver.add((vmin <= (s:=z3.Int(name, ctx=solver.ctx)))&(s <= vmax))
-  return s
-def create_var(x:UOp, ctx:tuple[z3.Solver, dict[UOp, z3.ExprRef]]) -> z3.ExprRef:
-  name = x.arg.name if x.op in {Ops.PARAM, Ops.BUFFER} else f"{x.op.name.lower()}{len(ctx[1])}"
-  return z3.Bool(name, ctx=ctx[0].ctx) if x.dtype == dtypes.bool else create_bounded(name, x.vmin, x.vmax, ctx[0])
-# z3 does not model widths: a cast only converts between bool and int
-def z3_cast(c:UOp, x:z3.ExprRef) -> z3.ExprRef:
-  if (c.src[0].dtype == dtypes.bool) == (c.dtype == dtypes.bool): return x
-  return x != 0 if c.dtype == dtypes.bool else z3.If(x, 1, 0)
+def create_bounded(name:str, vmin:int, vmax:int, z3ctx:z3.Context) -> tuple[z3.ArithRef, z3.BoolRef]:
+  return (s:=z3.Int(name, ctx=z3ctx)), (vmin <= s)&(s <= vmax)
 
 z3_renderer = PatternMatcher([
-  # the valid condition is a constraint
-  (UPat.var("cond").where(UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), lambda x,cond,ctx: ctx[0].add(ctx[1][cond]) or ctx[1][x]),
+  (UPat.var("cond").where(UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), lambda x,cond,ctx: (ctx[1][x], ctx[1][cond])),
   # variables
-  (UPat((Ops.SPECIAL, Ops.RANGE), name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
-  # unknown values are variables bounded by their vmin/vmax: params, loads (non-pointer INDEX is a LOAD) and anything from floats
-  (UPat((Ops.PARAM, Ops.BUFFER, Ops.LOAD, Ops.INDEX), name="x"), create_var),
-  (UPat((Ops.CAST, Ops.BITCAST)+tuple(GroupOp.Comparison), src=UPat(dtype=dtypes.floats), name="x"), create_var),
-  # a bitcast between ints wraps into the target range, z3 ints are unbounded
-  (UPat(Ops.BITCAST, dtypes.ints, src=(UPat.var("x", dtypes.ints),), name="c"),
-   lambda c,x,ctx: (ctx[1][x]-c.dtype.min) % 2**(8*c.dtype.itemsize) + c.dtype.min),
+  (UPat(Ops.SPECIAL, name="x"), lambda x,ctx: create_bounded(x.arg, 0, ctx[1][x.src[0]]-1, ctx[0])),
+  (UPat(Ops.PARAM, name="x"), lambda x,ctx: create_bounded(x.arg.name, x.vmin, x.vmax, ctx[0])),
+  (UPat(Ops.BUFFER, name="x"), lambda x,ctx: create_bounded(x.arg.name, x.vmin, x.vmax, ctx[0]) if x.is_variable else None),
+  (UPat(Ops.RANGE, name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
+  # loads are variables bounded by the min/max of the dtype. non-pointer INDEX is also a LOAD
+  (UPat((Ops.LOAD, Ops.INDEX), dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx:
+    create_bounded(f"load{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
+  (UPat((Ops.LOAD, Ops.INDEX), dtypes.bool), lambda ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
   # constants
-  (UPat(Ops.CONST, arg=Invalid), lambda ctx: z3.Int("Invalid", ctx=ctx[0].ctx)),
-  (UPat(Ops.CONST, name="x"), lambda x,ctx: z3.BoolVal(x.val, ctx=ctx[0].ctx) if x.dtype == dtypes.bool else z3.IntVal(x.val, ctx=ctx[0].ctx)),
-  (UPat(Ops.CAST, src=(UPat.var("x"),), name="c"), lambda c,x,ctx: z3_cast(c, ctx[1][x])),
-  (UPat(GroupOp.ALU, name="x"), lambda x,ctx: z3_alu[x.op](*(ctx[1][s] for s in x.src))),
+  (UPat(Ops.CONST, arg=Invalid), lambda ctx: (z3.Int("Invalid", ctx=ctx[0]), None)),
+  (UPat(Ops.CONST, dtypes.weakint, name="x"), lambda x,ctx: (z3.IntVal(x.val, ctx=ctx[0]), None)),
+  (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.val, ctx=ctx[0]), None)),
+  # casts from floats create new variables
+  (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
+    create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
+  # A comparison between floats introduces a new bool variable
+  (UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats)), lambda ctx: (z3.Bool(f"float_cmp{len(ctx[1])}", ctx=ctx[0]), None)),
+  # a same-dtype cast states a width, which z3 does not model: identity. must precede the rules below (bool->bool)
+  (UPat(Ops.CAST, name="x"), lambda x,ctx: (ctx[1][x.src[0]], None) if x.dtype == x.src[0].dtype else None),
+  # casts from bool/int to int/bool
+  (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,),src=(UPat.var("x", dtypes.bool),)), lambda x,ctx: (z3.If(ctx[1][x], 1, 0), None)),
+  (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat.var("x", dtypes.ints+(dtypes.weakint,)),)), lambda x,ctx: (ctx[1][x], None)),
+  (UPat(Ops.CAST, dtypes.bool, name="x"), lambda x,ctx: (ctx[1][x.src[0]]!=0, None)),
+  (UPat(GroupOp.ALU, name="x"), lambda x,ctx: (z3_alu[x.op](*(ctx[1][s] for s in x.src)), None)),
 ])
 
 def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
@@ -67,8 +69,11 @@ def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
   for u in lst:
     # NOTE: we skip STACK here, it can't actually be accessed
     if u.op is Ops.STACK: continue
-    if (z3_rewritten:=z3_renderer.rewrite(u, ctx=(solver, z3map))) is None: raise NotImplementedError(f"{u.op} is not supported by z3")
-    z3map[u] = z3_rewritten
+    z3_rewritten: tuple[z3.ExprRef, z3.BoolRef|None]|None = z3_renderer.rewrite(u, ctx=(solver.ctx, z3map))
+    if z3_rewritten is None: raise NotImplementedError(f"{u.op} is not supported by z3")
+    new_u, constraint = z3_rewritten
+    if constraint is not None: solver.add(constraint)
+    z3map[u] = new_u
   assert all(u in z3map for u in uops), "UOp failed to rewrite to z3!"
   return [z3map[u] for u in uops]
 

tinygrad/uop/weak.py

@@ -43,16 +43,9 @@ pm_commit_weak = PatternMatcher([
 # consumers absorb the weak CAST off their srcs and default underivable consts; dtype-producing ops settle here.
 # a weakfloat Unary (sin/exp2/...) must resolve before the transcendental decomposition.
 _lower_weak_ops = GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}
-
-# only within the kind is a weak CAST a width statement: across kinds it converts the value, so it commits unless u recasts its srcs anyway
-def absorb_weak_src(u:UOp, s:UOp) -> UOp:
-  if s.op is not Ops.CAST or s.dtype not in dtypes.weaks: return s
-  if u.op in _lower_weak_ops or u.op is Ops.CAST or weak_dtype(s.src[0].dtype) is s.dtype: return s.src[0]
-  return s.src[0].cast(s.commit_dtype(dtypes.int))
-
 def lower_weak_node(u:UOp) -> UOp|None:
   if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None  # a committed const, not a consumer
-  src = tuple(absorb_weak_src(u, s) for s in u.src)
+  src = tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
   if derived_dtypes(u, src) is None:
     src = tuple(s.ccast(s.commit_dtype(dtypes.int)) if s.op is Ops.CONST and s.dtype in dtypes.weaks else s for s in src)
   if src == u.src: return None