Merge branch 'master' into invalid_try_3

staying weak is okay

AGENTS.md

@@ -3,3 +3,4 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
+- Read `./tinygrad/viz/README` for profiling

examples/mlperf/model_train.py

@@ -1285,7 +1285,7 @@ def train_llama3():
   from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
 
   INITMLPERF = getenv("INITMLPERF")
   RUNMLPERF = getenv("RUNMLPERF")
@@ -1482,7 +1482,8 @@ def train_llama3():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)
@@ -1667,7 +1668,7 @@ def train_llama3():
 def train_gptoss():
   from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1775,7 +1776,8 @@ def train_gptoss():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)

examples/mlperf/optim.py

@@ -21,6 +21,12 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
   return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
 
+def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
+  for g in grads: g.assign(g / grad_acc)
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
+  return total_norm
+
 class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)
@@ -44,38 +50,21 @@ class GradAccClipAdamW(Optimizer):
     n, sz = len(t.device), t.shape[0] // len(t.device)
     return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
 
-  def fstep(self, grads:list[Tensor]):
-    if self.fused:
-      out, extra = self._step([], grads)
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
-    else:
-      updates, extra = self._step([], grads)
+  def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
+    updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
     fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
     fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
 
-    Tensor.realize(*to_realize)
-    return extra[-1]
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
     grads = list(grads)
 
     for i in range(len(grads)):
       if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
-
-    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
-      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
-    else:
-      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
-
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
@@ -88,7 +77,7 @@ class GradAccClipAdamW(Optimizer):
       v_hat = v_new / (1.0 - self.b2_t)
       up = m_hat / (v_hat.sqrt() + self.eps)
       ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v
 
   def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
     w = master if master is not None else t
@@ -109,7 +98,7 @@ class GradAccClipAdamW(Optimizer):
         if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(new_w.shape)
+        ret = w_q.reshape(t.shape)
         return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
       if IMMEDIATE_SCALE:

extra/gemm/moe_gemm.py

@@ -0,0 +1,111 @@
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
+
+@functools.cache
+def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
+  M, K = A.shape
+  E, N, K2 = B.shape
+  assert K == K2, f"{A.shape} {B.shape}"
+  assert E == n_experts, f"{E} != {n_experts}"
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special((M // 256) * (N // 256), "gidx0")
+  sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, extra[0].base, extra[1].base, extra[2].base, threads, workgroups)
+  sink = UOp.sink(*sink_inputs,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_gemm_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*K+E*N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_gemm.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
+                                 f"-DGEMM_E={E}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def custom_hk_grouped_mxfp8_wgrad(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, expert_off:UOp, *, dname:str, n_experts:int) -> UOp:
+  N, M = A.shape
+  K, M2 = B.shape
+  assert M == M2, f"{A.shape} {B.shape}"
+  E = n_experts
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special(E * (N // 256) * (K // 256), "gidx0")
+  sink = UOp.sink(C.base, A.base, B.base, scale_A.base, scale_B.base, expert_off.base, threads, workgroups,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_wgrad_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(N*M+K*M)*A.dtype.itemsize+E*N*K*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_wgrad.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DWGRAD_M={M}", f"-DWGRAD_N={N}", f"-DWGRAD_K={K}",
+                                 f"-DWGRAD_E={E}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=lib)))
+
+def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> Tensor:
+  from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
+  M, N = g.shape
+  M2, K = xg.shape
+  assert M == M2, f"{g.shape} {xg.shape}"
+  assert M % 128 == 0 and N % 256 == 0 and K % 256 == 0, f"wgrad needs M%128,N%256,K%256, got {g.shape} {xg.shape}"
+  gT, _, g_si = transpose_quantize_mxfp8(g.contiguous())
+  xT, _, x_si = transpose_quantize_mxfp8(xg.contiguous())
+  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
+  is_multi = isinstance(g.device, tuple)
+  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
+  out = Tensor(inv.uop.multi(0), device=g.device) if is_multi else inv
+  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
+                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
+  out = out.sum(0) if is_multi else out.squeeze(0)
+  return out.reshape(n_experts, N, K)
+
+def mx_pack_3d(e8:Tensor) -> Tensor:
+  E, rows, scale_K = e8.shape
+  return e8.reshape(E, rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(E, rows, scale_K // 4).permute(0, 2, 1).contiguous()
+
+@functools.cache
+def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) -> tuple:
+  inputs = kernel.src[1:]
+  aq = Tensor(inputs[1], device=inputs[1].device)
+  bq = Tensor(inputs[2], device=inputs[2].device)
+  ae8 = Tensor(inputs[5], device=inputs[5].device)
+  be8 = Tensor(inputs[6], device=inputs[6].device)
+  E, N = bq.shape[0], bq.shape[1]
+  M, K = aq.shape
+  g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
+  x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
+  w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
+  expert_off = Tensor(inputs[7], device=inputs[7].device)
+  grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
+  grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
+  grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
+  grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
+  return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
+
+_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
+
+def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
+  if (pre_quantized := isinstance(w, tuple)):
+    w_q, w_e8 = w
+    E, N, K2 = w_q.shape
+  else:
+    E, N, K2 = w.shape
+  M, K = x.shape
+  assert K == K2, f"shape mismatch {x.shape} {w.shape}"
+  assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
+  dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
+  x_q, x_e8, x_si = quantize_mxfp8(x)
+  if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
+  w_si = mx_pack_3d(w_e8)
+  xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
+  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
+    ndev = len(x.device)
+    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
+                                 dtype=dtypes.bfloat16, device=x.device).uop.multi(row_axis), device=x.device)
+  else:
+    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
+  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
+                              fxn=functools.partial(custom_hk_grouped_mxfp8_gemm, dname=dname, n_experts=E),
+                              grad_fxn=(_grouped_bw_stored if pre_quantized else custom_grouped_mx_gemm_bw))[0]

extra/hcq2/hcq2.py

@@ -14,11 +14,11 @@ from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
 from tinygrad.engine.jit import DepsTracker
 
-HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
-
 # *****************
 # 0. helpers
 
+HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
+
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
 
 HCQ_DEVS = frozenset(("AMD",))
@@ -37,25 +37,11 @@ class HCQInfo:
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
-def unwrap_after(uop):
-  while uop.op is Ops.AFTER: uop = uop.src[0]
-  return uop
-
 def unwrap_mstack(u):
   return tuple(x for s in u.src for x in unwrap_mstack(s)) if u.op is Ops.MSTACK else (unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,))
 
-def make_getaddr(u, device=None):
-  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM): return u
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(u,), arg=device or to_tuple(u.device)[0])
-
-def make_ins(op, *srcs):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs))
-
-def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
-  return UOp.param(next(UOp.unique_num) if unique else 0, dtype, shape=(size,), device=devs).rtag(name or "temp")
-
-def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
-  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(dtype or buf.dtype))
+def make_patch(buf:UOp, off:sint, val:UOp) -> UOp:
+  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(buf.dtype))
 
 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
   data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
@@ -67,20 +53,24 @@ def make_cmdbuf(lin, devs):
   for s in (s for ins in lin.src for s in ins.src):
     if (ssimp:=s.simplify()).op is not Ops.CONST: patches.append((len(blob), ssimp))
     blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0)
-  cmdbuf = make_placeholder(devs, len(blob) // 4, dtypes.uint32, name="cmdbuf")
+  cmdbuf = UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
   return cmdbuf.after(make_binary_patch(cmdbuf, blob), *[make_patch(cmdbuf, off, s) for off, s in patches])
 
-def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
-
-def make_signal(devs, queue=None, sentinel=False):
-  return make_placeholder(devs, 1, dtypes.uint64, "sentinel_signal" if sentinel else (queue, "timeline_signal") if queue else "timeline_signal", unique=False)
-def make_signal_value(devs, queue=None):
-  return make_placeholder(devs, 1, dtypes.uint64, (queue, "timeline_value") if queue else "timeline_value", unique=False)
+def make_signal(devs, queue="COMPUTE:0", sentinel=False):
+  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal")
+def make_signal_value(devs, queue="COMPUTE:0"):
+  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value")
 
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
 def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf")
 
+def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
+  data, info = prg.arg
+  buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
+  words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars)
+  return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)])
+
 # *****************
 # 0.1. prep: replace buffers with params
 
@@ -131,8 +121,8 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str,
 
   waits = []
   for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = make_mstack([make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
-    val = make_mstack([make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
+    sig = UOp.mstack(*[make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
+    val = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
     waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag)))
   return waits, {dtag for _, _, dtag in deps}
 
@@ -154,9 +144,10 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
     waited |= cur_waited
 
     # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
-    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero)))
+    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero) + n))
     submit = make_submit(*waits, store, devs=devs, queue="COMPUTE:0")
-    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), n) for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])]
+    upd = [(tl, n + 1)] + [(make_signal_value(devs, queue=qn), n)
+                           for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) if qn != "COMPUTE:0"]
     bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
     finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
   return finalizers, waited
@@ -179,19 +170,18 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
   waited |= finalizer_waited
 
   src = []
-  for tag, ((call, _), (devices, queue), cmds) in enumerate(zip(batch, batch_info, call_waits)):
+  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
     # first queue use, sync prior device work with main signal
     if batch_info.index((devices, queue)) == tag:
-      epoch = UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))
-      cmds = [UOp(Ops.INS, arg="barrier", src=()), epoch] + cmds
+      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))] + q
 
     # and make hcq call
     info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
-    cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))]
+    q += [call.replace(arg=replace(call.arg, aux=info))]
 
     # signal queue timeline if someone waits for us
-    if tag in waited: cmds += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))]
-    src.append(UOp.custom_function("hcq", make_submit(*cmds, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
+    if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))]
+    src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
   return src + finalizers
 
 def sched_hcq_batches(l:UOp) -> UOp:
@@ -235,15 +225,6 @@ def merge_queues(linear:UOp) -> UOp:
   return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
 pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
 
-# *****************
-# 4.1. hcq lowering: programs
-
-def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
-  data, info = prg.arg
-  buf = make_placeholder(devs, data.kernargs_alloc_size // 4, dtypes.uint32, name="kernargs")
-  words = [w for gi in info.globals for w in data64_le(make_getaddr(get_call_arg_uops(call)[gi], devs))] + list(info.vars)
-  return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)])
-
 # *****************
 # 4.2. hcq lowering: ops to ir
 
@@ -286,16 +267,17 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION
 # *****************
 
 def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
-  bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}
+  bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
 
   order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
-  slots, table = {g:i for i,g in enumerate(order)}, make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)
+  slots = {g:i for i,g in enumerate(order)}
+  table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
 
   reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
   return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()
 
 def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
-  bufs = {b: make_placeholder(call.arg.aux.device, b.max_numel(), b.dtype, "template") for b in blobs}
+  bufs = {b: UOp.placeholder((b.max_numel(),), b.dtype, next(UOp.unique_num), device=call.arg.aux.device).rtag("template") for b in blobs}
   return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items())
 
 def rm_rt_uops(call:UOp) -> UOp|None:
@@ -322,7 +304,7 @@ def replace_params(call:UOp) -> UOp|None:
   by_root = {p.src[0]: p for p in patched}
   c_args = [by_root.get(a, a) for a in args]
 
-  sub = {unwrap_after(u): UOp.param(i, u.dtype, shape=unwrap_after(u).shape, device=u.device) for i,u in enumerate(c_args)} | \
+  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=u.device) for i,u in enumerate(c_args)} | \
         {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
   info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
   return call.replace(src=(body.substitute(sub), *c_args, *refhold), arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
@@ -332,7 +314,7 @@ pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTIO
 
 def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
   base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  itemsize = bv.src[0].dtype.itemsize if unwrap_after(bv.src[0]).op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
+  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
   return UOp(Ops.GETADDR, dtypes.uint64, src=(base,), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
 
 pm_early_simplify = PatternMatcher([
@@ -353,7 +335,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
       offs[b] = round_up(sizes.get(b.tag, 0), 128 // b.dtype.itemsize)
       sizes[b.tag] = offs[b] + b.max_numel()
   counts = collections.Counter(b.tag for b in bufs)
-  bases = {b.tag:make_placeholder(b.device, sizes[b.tag], b.dtype, b.tag) for b in bufs if counts[b.tag] > 1}
+  bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
   subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
   return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
@@ -400,7 +382,8 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
 
 def bufferize_buf(ctx:bool, buf:UOp) -> UOp|None:
   if buf.tag is None: return None
-  return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), HCQ_RUNTIME_DEV.value) for dev in to_tuple(buf.device)))
+  return UOp.mstack(*(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), HCQ_RUNTIME_DEV.value)
+                      for dev in to_tuple(buf.device)))
 pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 
 # *****************
@@ -474,8 +457,6 @@ class HCQ2Compiled(Compiled):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].timeline_signal()),
-      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].timeline_value()),
       (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)),
       (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, jit=ctx[1]))
     ])
@@ -491,13 +472,13 @@ class HCQ2Compiled(Compiled):
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache
-  def timeline_signal(self, queue:str|None=None, init_value:int=0) -> Buffer:
+  def timeline_signal(self, queue:str="COMPUTE:0", init_value:int=0) -> Buffer:
     buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
     buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf
 
   @functools.cache
-  def timeline_value(self, queue:str|None=None, init_value:int=1) -> Buffer:
+  def timeline_value(self, queue:str="COMPUTE:0", init_value:int=1) -> Buffer:
     buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
     buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf

extra/hcq2/ops_amd2.py

@@ -3,7 +3,7 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_getaddr, make_ins, make_cmdbuf, make_placeholder
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
 from extra.hcq2.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
@@ -36,7 +36,9 @@ class PM4Ops(FastEnum):
   SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto()  # noqa: E702
   RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702
 
-def pkt3(ctx, op:PM4Ops, *vals): return make_ins(op, ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)
+def pkt3(ctx, op:PM4Ops, *vals):
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x)
+    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
 
 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
   if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
@@ -90,26 +92,26 @@ def memory_barrier(ctx):
                  reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
     acquire_mem(ctx)))
 
-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=make_getaddr(dst, ctx.devs))
+def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))
 
 def pm4_barrier(ctx): return memory_barrier(ctx)
 
 def pm4_store(ctx, dst, val):
   if val.op is Ops.BINARY: return None
-  return release_mem(ctx, make_getaddr(dst, ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
+  return release_mem(ctx, dst.getaddr(ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
                      ctx.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
 
 def pm4_timestamp(ctx, dst):
-  return release_mem(ctx, make_getaddr(dst, ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+  return release_mem(ctx, dst.getaddr(ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
                      ctx.pm4.int_sel__mec_release_mem__none)
 
 def pm4_program(ctx, call, prg):
   data, info = prg.arg
   lib_gpu = prg.src[0]
   args = encode_kernargs_clike(call, prg, ctx.devs)
-  prog_addr = make_getaddr(lib_gpu, ctx.devs) + data.entry_point_offset
-  scratch_addr = make_getaddr(make_placeholder(ctx.devs, data.private_segment_size, dtypes.uint8, "scratch", unique=False), ctx.devs)
-  args_addr = make_getaddr(args, ctx.devs)
+  prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
+  scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
+  args_addr = args.getaddr(ctx.devs)
 
   user_regs = []
   if data.enable_private_segment_sgpr:
@@ -149,7 +151,7 @@ def pm4_submit(cmdbuf, devs):
 
   # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
   for d in devs: q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COMPUTE:0", name), unique=False)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # place the cmdbuf at the ring's write offset, wrapping the ring
@@ -176,36 +178,36 @@ pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
 class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
 
 def sdma_copy(ctx, call):
-  dst, src = call.src[1], call.src[2]
-  sz = src.max_numel() * src.dtype.itemsize
-  src_addr, dst_addr = make_getaddr(src, ctx.devs), make_getaddr(dst, ctx.devs)
-  return UOp(Ops.LINEAR, dtypes.void, tuple([make_ins(SDMAOps.COPY,
-     ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-     ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz - off, ctx.max_copy_size) - 1), 0,
-     *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)]))
+  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
+  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
+    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))
 
-def sdma_wait(ctx, dst, val):
+def sdma_wait(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
      | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(dst, ctx.devs)), val, 0xffffffff,
-    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
+    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
+    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
 
-def sdma_store(ctx, dst, val):
+def sdma_store(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
-  return UOp(Ops.LINEAR, dtypes.void, (
-    make_ins(SDMAOps.FENCE, op, *data64_le(make_getaddr(dst, ctx.devs)), val), make_ins(SDMAOps.TRAP, ctx.sdma.SDMA_OP_TRAP, 0)))
+  return UOp(Ops.LINEAR, src=(
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
 
-def sdma_timestamp(ctx, dst):
+def sdma_timestamp(ctx, ins, dst):
   op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return make_ins(SDMAOps.TIMESTAMP, op, *data64_le(make_getaddr(dst, ctx.devs)))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
 
 pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
   (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
-  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
-  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
 ])
 
 def sdma_submit(cmdbuf, devs):
@@ -214,7 +216,7 @@ def sdma_submit(cmdbuf, devs):
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COPY:0", name), unique=False)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
@@ -279,7 +281,7 @@ def amd_build_program(prg:UOp) -> UOp:
       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)
-    buf = make_placeholder(prg.device, len(image), dtypes.uint8, "program")
+    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
     cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
   return cached
 
@@ -628,10 +630,10 @@ class AMDDevice(HCQ2Compiled):
 
     qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
     ] + [
-      (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
-      (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx[0].timeline_value(q)),
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)),
     ]) + self.pm_bufferize
 
     return queue

extra/thunder/amd/grouped_mxfp8_gemm.cpp

@@ -0,0 +1,158 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef GEMM_M
+constexpr int GEMM_M = 8192;
+#endif
+#ifndef GEMM_N
+constexpr int GEMM_N = 8192;
+#endif
+#ifndef GEMM_K
+constexpr int GEMM_K = 8192;
+#endif
+#ifndef GEMM_E
+constexpr int GEMM_E = 8;
+#endif
+
+// Kernel
+constexpr int NUM_WARPS  = 8;
+constexpr int WARPS_ROW  = 2;
+constexpr int WARPS_COL  = 4;
+constexpr int BLOCK_ROW  = 256;
+constexpr int BLOCK_COL  = 256;
+constexpr int BLOCK_K    = 128;
+constexpr int HALF_ROW   = BLOCK_ROW / 2;
+constexpr int HALF_COL   = BLOCK_COL / 2;
+constexpr int REG_M      = BLOCK_ROW / WARPS_ROW / 2;
+constexpr int REG_N      = BLOCK_COL / WARPS_COL / 2;
+
+using G = kittens::group<NUM_WARPS>;
+
+__global__ __launch_bounds__(512, 2) void grouped_mxfp8_gemm_kernel(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, fp8e8m0 *scale_A_ptr, fp8e8m0 *scale_B_ptr,
+    const uint8_t *__restrict__ a_e8_unused,
+    const uint8_t *__restrict__ b_e8_unused,
+    const int *__restrict__ expert_off) {
+    constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K, E = GEMM_E;
+
+    kittens::gl<fp8e4m3, 1, 1, M, K>     A{A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e4m3, 1, 1, E * N, K> B{B_ptr, nullptr, nullptr, nullptr, nullptr};  // all experts stacked on rows
+    kittens::gl<bf16, 1, 1, M, N>        C{C_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    constexpr int k_iters      = K / BLOCK_K;
+    constexpr int NUM_THREADS  = NUM_WARPS * WARP_THREADS;
+
+    kittens::gl<fp8e8m0, k_iters * (M / BLOCK_ROW), 1, 16, 64>     scale_A_gl{scale_A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e8m0, E * k_iters * (N / BLOCK_COL), 1, 16, 64> scale_B_gl{scale_B_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    using ST_A     = st_fp8e4m3<HALF_ROW, BLOCK_K, st_16x128_s>;
+    using ST_B     = st_fp8e4m3<HALF_COL, BLOCK_K, st_16x128_s>;
+    using ST_Scale = st<fp8e8m0, 16, 64, st_16x64_s>;
+    using RT_A     = rt_fp8e4m3<REG_M, BLOCK_K>;
+    using RT_B     = rt_fp8e4m3<REG_N, BLOCK_K>;
+    using RT_C     = rt_fl<REG_M, REG_N, col_l, rt_16x16_s>;
+
+    __shared__ ST_A As[2][2];
+    __shared__ ST_B Bs[2][2];
+    __shared__ ST_Scale scale_A_smem[2], scale_B_smem[2];
+
+    RT_A a;
+    RT_B b0, b1;
+    RT_C cA, cB, cC, cD;
+    zero(cA); zero(cB); zero(cC); zero(cD);
+
+    constexpr int tiles_M  = M / BLOCK_ROW;
+    constexpr int tiles_N  = N / BLOCK_COL;
+    const int NUM_XCDS     = 8;
+    const int WGM          = 8;
+    int wgid = chiplet_transform_chunked(blockIdx.x, gridDim.x, NUM_XCDS, WGM * WGM);
+    int num_wgid_in_group = WGM * tiles_N;
+    int group_id     = wgid / num_wgid_in_group;
+    int first_pid_m  = group_id * WGM;
+    int group_size_m = min(tiles_M - first_pid_m, WGM);
+    int block_row    = first_pid_m + ((wgid % num_wgid_in_group) % group_size_m);
+    int block_col    = (wgid % num_wgid_in_group) / group_size_m;
+    int block_m      = block_row * BLOCK_ROW;
+    int block_n      = block_col * BLOCK_COL;
+
+    int e = 0;
+    #pragma unroll
+    for (int i = 1; i < E; i++) e += (expert_off[i] <= block_row * BLOCK_ROW);
+    e = __builtin_amdgcn_readfirstlane(e);
+    const int bcol_base = e * (N / HALF_COL);          // expert base in B row-tile (128-row) units
+    const int sb_base   = e * (k_iters * tiles_N);     // expert base into scale_B batches
+
+    int warp_m = warpid() / WARPS_COL;
+    int warp_n = warpid() % WARPS_COL;
+
+    using T = fp8e4m3;
+    constexpr int bpt      = ST_A::underlying_subtile_bytes_per_thread;
+    constexpr int bpm      = bpt * NUM_THREADS;
+    constexpr int copies_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm;
+    constexpr int copies_B = HALF_COL * BLOCK_K * sizeof(T) / bpm;
+    uint32_t sw_A[copies_A], sw_B[copies_B];
+    G::prefill_swizzled_offsets(As[0][0], A, sw_A);
+    G::prefill_swizzled_offsets(Bs[0][0], B, sw_B);
+
+    const T *a_base = (const T *)&A[{0, 0, 0, 0}];
+    const T *b_base = (const T *)&B[{0, 0, 0, 0}];
+    const int a_row_stride = A.template stride<2>() * sizeof(T);
+    const int b_row_stride = B.template stride<2>() * sizeof(T);
+    i32x4 a_srd = make_srsrc(a_base, (uint32_t)((uint64_t)M * a_row_stride), a_row_stride);
+    i32x4 b_srd = make_srsrc(b_base, (uint32_t)((uint64_t)E * N * b_row_stride), b_row_stride);
+
+    const int wid = warpid() % NUM_WARPS;
+    constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS;
+    uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+
+    int a_row_h0 = warp_m * REG_M;
+    int a_row_h1 = HALF_ROW + warp_m * REG_M;
+    int b_row_h0 = warp_n * REG_N;
+    int b_row_h1 = HALF_COL + warp_n * REG_N;
+
+
+    uint32_t a_lds[2][2] = {{a_lds_00, a_lds_01}, {a_lds_10, a_lds_11}};
+    uint32_t b_lds[2][2] = {{b_lds_00, b_lds_01}, {b_lds_10, b_lds_11}};
+
+    #pragma unroll 1
+    for (int kk = 0; kk < k_iters; kk++) {
+        G::load(As[0][0], A, {0, 0, block_row * 2,     kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds[0][0]));
+        G::load(As[0][1], A, {0, 0, block_row * 2 + 1, kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds[0][1]));
+        G::load(Bs[0][0], B, {0, 0, bcol_base + block_col * 2,     kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds[0][0]));
+        G::load(Bs[0][1], B, {0, 0, bcol_base + block_col * 2 + 1, kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds[0][1]));
+        G::load(scale_A_smem[0], scale_A_gl, {kk * tiles_M + block_row, 0, 0, 0});
+        G::load(scale_B_smem[0], scale_B_gl, {sb_base + kk * tiles_N + block_col, 0, 0, 0});
+        asm volatile("s_waitcnt vmcnt(0)");
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        __builtin_amdgcn_s_barrier();
+
+        fp8e8m0_4 sa_h0 = pack_scales(scale_A_smem[0].data, a_row_h0);
+        fp8e8m0_4 sa_h1 = pack_scales(scale_A_smem[0].data, a_row_h1);
+        fp8e8m0_4 sb_h0 = pack_scales(scale_B_smem[0].data, b_row_h0);
+        fp8e8m0_4 sb_h1 = pack_scales(scale_B_smem[0].data, b_row_h1);
+
+        auto bs0 = subtile_inplace<REG_N, BLOCK_K>(Bs[0][0], {warp_n, 0}); load(b0, bs0);
+        auto bs1 = subtile_inplace<REG_N, BLOCK_K>(Bs[0][1], {warp_n, 0}); load(b1, bs1);
+        auto as0 = subtile_inplace<REG_M, BLOCK_K>(As[0][0], {warp_m, 0}); load(a, as0);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cA, a, b0, cA, &sa_h0, &sb_h0);
+        mma_ABt_scaled(cB, a, b1, cB, &sa_h0, &sb_h1);
+        auto as1 = subtile_inplace<REG_M, BLOCK_K>(As[0][1], {warp_m, 0}); load(a, as1);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cC, a, b0, cC, &sa_h1, &sb_h0);
+        mma_ABt_scaled(cD, a, b1, cD, &sa_h1, &sb_h1);
+        __builtin_amdgcn_s_barrier();
+    }
+
+    store(C, cA, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + warp_n});
+    store(C, cB, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+    store(C, cC, {0, 0, block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m, block_col * WARPS_COL * 2 + warp_n});
+    store(C, cD, {0, 0, block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+}

extra/thunder/amd/grouped_mxfp8_wgrad.cpp

@@ -0,0 +1,141 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef WGRAD_M
+constexpr int WGRAD_M = 8192;
+#endif
+#ifndef WGRAD_N
+constexpr int WGRAD_N = 8192;
+#endif
+#ifndef WGRAD_K
+constexpr int WGRAD_K = 8192;
+#endif
+#ifndef WGRAD_E
+constexpr int WGRAD_E = 8;
+#endif
+
+constexpr int NUM_WARPS  = 8;
+constexpr int WARPS_ROW  = 2;
+constexpr int WARPS_COL  = 4;
+constexpr int BLOCK_ROW  = 256;
+constexpr int BLOCK_COL  = 256;
+constexpr int BLOCK_K    = 128;
+constexpr int HALF_ROW   = BLOCK_ROW / 2;
+constexpr int HALF_COL   = BLOCK_COL / 2;
+constexpr int REG_M      = BLOCK_ROW / WARPS_ROW / 2;
+constexpr int REG_N      = BLOCK_COL / WARPS_COL / 2;
+
+using G = kittens::group<NUM_WARPS>;
+
+__global__ __launch_bounds__(512, 2) void grouped_mxfp8_wgrad_kernel(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr,
+    fp8e8m0 *scale_A_ptr, fp8e8m0 *scale_B_ptr,
+    const int *__restrict__ expert_off) {
+    constexpr int M = WGRAD_M, N = WGRAD_N, K = WGRAD_K, E = WGRAD_E;
+
+    kittens::gl<fp8e4m3, 1, 1, N, M>     A{A_ptr, nullptr, nullptr, nullptr, nullptr};  // g^T
+    kittens::gl<fp8e4m3, 1, 1, K, M>     B{B_ptr, nullptr, nullptr, nullptr, nullptr};  // x^T
+    kittens::gl<bf16, 1, 1, E * N, K>    C{C_ptr, nullptr, nullptr, nullptr, nullptr};  // grad_w, experts stacked
+
+    constexpr int m_blocks    = M / BLOCK_K;      // 128-wide blocks along the contraction (token) axis
+    constexpr int tiles_N     = N / BLOCK_ROW;
+    constexpr int tiles_K     = K / BLOCK_COL;
+    constexpr int NUM_THREADS = NUM_WARPS * WARP_THREADS;
+
+    kittens::gl<fp8e8m0, m_blocks * tiles_N, 1, 16, 64> scale_A_gl{scale_A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e8m0, m_blocks * tiles_K, 1, 16, 64> scale_B_gl{scale_B_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    using ST_A     = st_fp8e4m3<HALF_ROW, BLOCK_K, st_16x128_s>;
+    using ST_B     = st_fp8e4m3<HALF_COL, BLOCK_K, st_16x128_s>;
+    using ST_Scale = st<fp8e8m0, 16, 64, st_16x64_s>;
+    using RT_A     = rt_fp8e4m3<REG_M, BLOCK_K>;
+    using RT_B     = rt_fp8e4m3<REG_N, BLOCK_K>;
+    using RT_C     = rt_fl<REG_M, REG_N, col_l, rt_16x16_s>;
+
+    __shared__ ST_A As[2];
+    __shared__ ST_B Bs[2];
+    __shared__ ST_Scale scale_A_smem, scale_B_smem;
+
+    RT_A a;
+    RT_B b0, b1;
+    RT_C cA, cB, cC, cD;
+    zero(cA); zero(cB); zero(cC); zero(cD);
+
+    const int wg        = blockIdx.x;
+    const int e         = wg / (tiles_N * tiles_K);
+    const int rem       = wg % (tiles_N * tiles_K);
+    const int block_row = rem / tiles_K;   // over grad_w rows (N)
+    const int block_col = rem % tiles_K;   // over grad_w cols (K)
+
+    const int o0  = __builtin_amdgcn_readfirstlane(expert_off[e]);
+    const int kk0 = o0 / BLOCK_K;
+    const int nk  = (__builtin_amdgcn_readfirstlane(expert_off[e + 1]) - o0) / BLOCK_K;
+
+    const int warp_m = warpid() / WARPS_COL;
+    const int warp_n = warpid() % WARPS_COL;
+
+    using T = fp8e4m3;
+    constexpr int bpt      = ST_A::underlying_subtile_bytes_per_thread;
+    constexpr int bpm      = bpt * NUM_THREADS;
+    constexpr int copies_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm;
+    constexpr int copies_B = HALF_COL * BLOCK_K * sizeof(T) / bpm;
+    uint32_t sw_A[copies_A], sw_B[copies_B];
+    G::prefill_swizzled_offsets(As[0], A, sw_A);
+    G::prefill_swizzled_offsets(Bs[0], B, sw_B);
+
+    const T *a_base = (const T *)&A[{0, 0, 0, 0}];
+    const T *b_base = (const T *)&B[{0, 0, 0, 0}];
+    const int a_row_stride = A.template stride<2>() * sizeof(T);
+    const int b_row_stride = B.template stride<2>() * sizeof(T);
+    i32x4 a_srd = make_srsrc(a_base, (uint32_t)((uint64_t)N * a_row_stride), a_row_stride);
+    i32x4 b_srd = make_srsrc(b_base, (uint32_t)((uint64_t)K * b_row_stride), b_row_stride);
+
+    const int wid = warpid() % NUM_WARPS;
+    constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS;
+    uint32_t a_lds_0 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_1 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_0 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_1 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1].data[0]) + wid * elem_per_warp * sizeof(T)));
+
+    const int a_row_h0 = warp_m * REG_M;
+    const int a_row_h1 = HALF_ROW + warp_m * REG_M;
+    const int b_row_h0 = warp_n * REG_N;
+    const int b_row_h1 = HALF_COL + warp_n * REG_N;
+
+    #pragma unroll 1
+    for (int t = 0; t < nk; t++) {
+        const int kk = kk0 + t;
+        G::load(As[0], A, {0, 0, block_row * 2,     kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds_0));
+        G::load(As[1], A, {0, 0, block_row * 2 + 1, kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds_1));
+        G::load(Bs[0], B, {0, 0, block_col * 2,     kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds_0));
+        G::load(Bs[1], B, {0, 0, block_col * 2 + 1, kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds_1));
+        G::load(scale_A_smem, scale_A_gl, {kk * tiles_N + block_row, 0, 0, 0});
+        G::load(scale_B_smem, scale_B_gl, {kk * tiles_K + block_col, 0, 0, 0});
+        asm volatile("s_waitcnt vmcnt(0)");
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        __builtin_amdgcn_s_barrier();
+
+        fp8e8m0_4 sa_h0 = pack_scales(scale_A_smem.data, a_row_h0);
+        fp8e8m0_4 sa_h1 = pack_scales(scale_A_smem.data, a_row_h1);
+        fp8e8m0_4 sb_h0 = pack_scales(scale_B_smem.data, b_row_h0);
+        fp8e8m0_4 sb_h1 = pack_scales(scale_B_smem.data, b_row_h1);
+
+        auto bs0 = subtile_inplace<REG_N, BLOCK_K>(Bs[0], {warp_n, 0}); load(b0, bs0);
+        auto bs1 = subtile_inplace<REG_N, BLOCK_K>(Bs[1], {warp_n, 0}); load(b1, bs1);
+        auto as0 = subtile_inplace<REG_M, BLOCK_K>(As[0], {warp_m, 0}); load(a, as0);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cA, a, b0, cA, &sa_h0, &sb_h0);
+        mma_ABt_scaled(cB, a, b1, cB, &sa_h0, &sb_h1);
+        auto as1 = subtile_inplace<REG_M, BLOCK_K>(As[1], {warp_m, 0}); load(a, as1);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cC, a, b0, cC, &sa_h1, &sb_h0);
+        mma_ABt_scaled(cD, a, b1, cD, &sa_h1, &sb_h1);
+        __builtin_amdgcn_s_barrier();
+    }
+
+    const int crow_base = e * (N / REG_M);   // grad_w rows are experts stacked; store coord is in REG_M units
+    store(C, cA, {0, 0, crow_base + block_row * WARPS_ROW * 2 + warp_m,              block_col * WARPS_COL * 2 + warp_n});
+    store(C, cB, {0, 0, crow_base + block_row * WARPS_ROW * 2 + warp_m,              block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+    store(C, cC, {0, 0, crow_base + block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m,  block_col * WARPS_COL * 2 + warp_n});
+    store(C, cD, {0, 0, crow_base + block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m,  block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+}

extra/torch_backend/backend.py

@@ -3,7 +3,7 @@
 # A002 Function argument `input` is shadowing a Python builtin
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, GroupOp
 from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
@@ -18,12 +18,16 @@ def _to_torch_device(device: str): return torch.device("tiny", int(device.partit
 
 import torch.utils.cpp_extension
 mod = torch.utils.cpp_extension.load(name="custom_device_extension", sources=[str(pathlib.Path(__file__).parent / "wrapped_tensor.cpp")])
+# TODO: this assumes a contiguous source, so PERMUTE/EXPAND/PAD/FLIP are wrong. UOp.contiguous_view_offset does it
+# properly, but it needs a device (these are deviceless)
+alias_ops = GroupOp.Movement | {Ops.BITCAST, Ops.DETACH, Ops.AFTER}
 def calculate_storage_offset(x: Tensor) -> int:
-  offset = 0
-  for u in x.uop.toposort():
-    if u.op == Ops.SHRINK:
+  offset, u = 0, x.uop
+  while u.op in alias_ops:
+    if u.op is Ops.SHRINK:
       u_strides = strides_for_shape(u.src[0].shape)
       for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+    u = u.src[0]
   return offset
 def wrap(x: Tensor, dev: torch.device|None=None) -> torch.Tensor:
   x._strides = strides_for_shape(x.shape) # always recalculate
@@ -220,7 +224,7 @@ def _as_strided(tensor:Tensor, size, stride, storage_offset=0):
 
 @torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
-  storage_offset = storage_offset or tensor.storage_offset()
+  if storage_offset is None: storage_offset = tensor.storage_offset()
   return _as_strided(tensor, size, stride, storage_offset)
 
 @torch.library.impl("aten::_reshape_alias", "privateuseone")
@@ -228,16 +232,16 @@ def _reshape_alias(tensor:torch.Tensor, size, stride):
   return _as_strided(tensor, size, stride)
 
 @torch.library.impl("aten::empty_strided", "privateuseone")
-def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False):
+def empty_strided(size, stride, dtype=None, layout=None, device=None, pin_memory=False):
   if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype), device=_from_torch_device(device)).contiguous()
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
   # TODO: should return with requested strides
   return wrap(ret)
 
 @torch.library.impl("aten::empty.memory_format", "privateuseone")
 def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
   if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device)).contiguous()
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
   return wrap(ret)
 
 @torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
@@ -652,35 +656,6 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.unfold": Tensor.unfold,
 }}
 
-# operations that need inplace treatment (use _inplace_op instead of wrap_fxn) AKA return original tensor
-inplace_ops = {
-  "aten.zero_",
-  "aten.fill_.Scalar",
-  "aten.fill_.Tensor",
-  "aten.add_.Tensor",
-  "aten.add_.Scalar",
-  "aten.mul_.Tensor",
-  "aten.mul_.Scalar",
-  "aten.floor_divide_.Tensor",
-  "aten.__ilshift__.Scalar",
-  "aten.__irshift__.Scalar",
-  "aten.relu_",
-  "aten.random_",
-  "aten.random_.from",
-  "aten.uniform_",
-  "aten.normal_",
-  "aten.logical_or_",
-  "aten.masked_fill_.Scalar",
-  "aten.masked_fill_.Tensor",
-}
-
-inplace_view_ops = {
-  "aten.squeeze_.dim",
-  "aten.unsqueeze_",
-  "aten.transpose_",
-  "aten.t_",
-}
-
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
@@ -694,7 +669,7 @@ def wrap_fxn(k,f):
     else: raise RuntimeError(f"unknown output type {type(out)}")
   return nf
 
-def wrap_inplace(k,f):
+def wrap_inplace(f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -702,7 +677,7 @@ def wrap_inplace(k,f):
     return orig
   return nf
 
-def wrap_inplace_view_op(k,f):
+def wrap_inplace_view_op(f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -735,11 +710,17 @@ def wrap_inplace_view_op(k,f):
     return orig
   return nf
 
+# the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
+# and a writable out arg must have come from tiny_backend_out so that wrap_out was applied
 for k,v in tiny_backend.items():
-  if k in inplace_view_ops: wrapper = wrap_inplace_view_op
-  elif k in inplace_ops: wrapper = wrap_inplace
-  else: wrapper = wrap_fxn
-  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
+  name, _, overload = k.removeprefix("aten.").partition(".")
+  op = getattr(getattr(aten, name), overload or "default")
+  writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
+  if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
+  elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
+  elif not writes or (writes == ["out"] and k in tiny_backend_out): fxn = wrap_fxn(k, v)
+  else: raise RuntimeError(f"{k} writes {writes}: expected an inplace first arg, or an out arg with {k} in tiny_backend_out")
+  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)
 
 @torch.library.impl("aten::equal", "privateuseone")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -775,21 +756,17 @@ def native_batch_norm(input, weight, bias, running_mean, running_var, training,
 @torch.library.impl("aten::native_batch_norm_backward", "privateuseone")
 def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask):
   grad_out_t, input_t = unwrap(grad_out), unwrap(input)
-  weight_t = unwrap(weight) if weight is not None else None
-  save_mean_t = unwrap(save_mean)
-  save_invstd_t = unwrap(save_invstd)
-  out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t)
-  targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m]
-  if targets:
-    grads = out.gradient(*targets, gradient=grad_out_t)
-    grad_input = grads.pop(0) if output_mask[0] else None
-    grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None
-  else:
-    grad_input, grad_weight = None, None
-  grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None
-  return (wrap(grad_input) if grad_input is not None else None,
-          wrap(grad_weight) if grad_weight is not None else None,
-          wrap(grad_bias) if grad_bias is not None else None)
+  dims, shape = tuple(x for x in range(input_t.ndim) if x != 1), (1, -1) + (1,)*(input_t.ndim-2)
+  # training differentiates the batch stats it was given, eval treats the running stats as constants
+  if train: mean, invstd = unwrap(save_mean), unwrap(save_invstd)
+  else: mean, invstd = unwrap(running_mean), unwrap(running_var).add(eps).rsqrt()
+  xhat = (input_t - mean.reshape(shape)) * invstd.reshape(shape)
+  grad_bias, grad_weight = grad_out_t.sum(axis=dims), (grad_out_t * xhat).sum(axis=dims)
+  grad_input = grad_out_t if not train else \
+    grad_out_t - (grad_bias.reshape(shape) + xhat * grad_weight.reshape(shape)) / (input_t.numel() // input_t.shape[1])
+  grad_input = grad_input * invstd.reshape(shape) * (unwrap(weight).reshape(shape) if weight is not None else 1)
+  return (wrap(grad_input) if output_mask[0] else None, wrap(grad_weight) if output_mask[1] else None,
+          wrap(grad_bias) if output_mask[2] else None)
 
 # _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
 # we need torch.autograd.Function with explicit AutogradPrivateUse1 registration

extra/torch_backend/test.py

@@ -71,6 +71,33 @@ class TestTorchBackend(unittest.TestCase):
     a = a.as_strided((1,1,5,5), (50,50,7,1), storage_offset=21)
     np.testing.assert_equal(a.cpu().numpy().sum(-1), [[[115,150,185,220,255]]])
 
+  def test_storage_offset_of_computed_tensor(self):
+    # a computed result owns its storage, so a slice anywhere in its history must not shift the offset
+    a = torch.arange(8., device=device)
+    self.assertEqual((a[3:]+1).storage_offset(), 0)
+
+  def test_storage_offset_through_aliases(self):
+    a = torch.arange(8., device=device)[3:]
+    self.assertEqual(a.detach().storage_offset(), 3)
+    self.assertEqual(a.view(torch.int32).storage_offset(), 3)
+    torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
+    self.assertEqual(a.detach().storage_offset(), 3)
+
+  @unittest.expectedFailure  # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
+  def test_storage_offset_non_contiguous_source(self):
+    a = torch.arange(12., device=device).reshape(3,4)
+    self.assertEqual(a.permute(1,0)[1:].storage_offset(), 1)
+    self.assertEqual(a.flatten()[3:].flip(0).storage_offset(), 0)
+
+  def test_as_strided_explicit_zero_offset(self):
+    # storage_offset=0 is a real offset, not "unspecified": it must not fall back to the input's own offset
+    a = torch.arange(6., device=device)
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,), 0).cpu().numpy(), [0,1])
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,)).cpu().numpy(), [3,4])
+
+  def test_empty_strided_default_dtype(self):
+    self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())
+
   def test_plus_inplace(self):
     a = torch.ones(4, device=device)
     b = torch.ones(4, device=device)
@@ -316,6 +343,21 @@ class TestTorchBackend(unittest.TestCase):
     assert b.shape == (4, 2, 3)
     np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))
 
+  def test_batchnorm_backward_realized_stats(self):
+    # the saved stats are a function of input in training, so grad_input must flow through them even when handed in realized.
+    # the backward eps is unused in training: torch differentiates the save_invstd it was given
+    x0, g0 = torch.randn(8, 4, 3, 3), torch.randn(8, 4, 3, 3)
+    def run(dev, bwd_eps):
+      x, go = x0.to(dev), g0.to(dev)
+      w, b = torch.linspace(0.5, 2.0, 4).to(dev), torch.zeros(4, device=dev)
+      rm, rv = torch.zeros(4, device=dev), torch.ones(4, device=dev)
+      out, sm, si = torch.ops.aten.native_batch_norm(x, w, b, rm, rv, True, 0.1, 1e-5)
+      grads = torch.ops.aten.native_batch_norm_backward(go, x, w, rm, rv, sm.clone().detach(), si.clone().detach(),
+                                                        True, bwd_eps, [True,True,True])
+      return [t.cpu().numpy() for t in grads]
+    for bwd_eps in [1e-5, 0.3]:
+      for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
+
   def test_batchnorm_unsqueeze(self):
     bn = torch.nn.BatchNorm2d(4).to(device)
     x = torch.randn(8, 4, 3, 3, device=device)
@@ -742,6 +784,20 @@ class TestTorchBackend(unittest.TestCase):
 
 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):
+  def test_unwrap_rejects_foreign_tensor(self):
+    # unwrap casts to the tiny impl, so a tensor from another backend must be refused rather than reinterpreted
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.unwrap(torch.ones(4))
+
+  def test_update_metadata_rejects_foreign_tensor(self):
+    # resizing a tensor we don't own would expose memory past its allocation
+    t = torch.ones(4)
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.mod.update_metadata(t, [8], [1], 0)
+    self.assertEqual(t.shape, (4,))
+
+  def test_unwrap_parameter_and_detached(self):
+    # nn.Parameter and detach rebuild the base OpaqueTensorImpl, which unwrap still has to accept
+    extra.torch_backend.backend.unwrap(torch.nn.Parameter(torch.ones(4, device="tiny")))
+    extra.torch_backend.backend.unwrap(torch.ones(4, device="tiny").detach())
 
   def test_calculate_storage_offset_no_shrink(self):
     t = Tensor.ones(3, 4)

extra/torch_backend/wrapped_tensor.cpp

@@ -124,16 +124,21 @@ at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceInd
     sizes, strides, storage_offset);
 }
 
+// shallow_copy_and_detach (nn.Parameter, aten.detach) rebuilds the base OpaqueTensorImpl, so that is the type every tiny tensor has
+at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>> *tiny_impl(const at::Tensor &tensor) {
+  auto* impl = dynamic_cast<at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(tensor.unsafeGetTensorImpl());
+  TORCH_CHECK(impl != nullptr, "expected a tiny tensor, got a ", tensor.device().str(), " one. move it with .to(\"tiny\") first");
+  return impl;
+}
+
 py::object unwrap_tensor(const at::Tensor &tensor) {
-  auto* impl = tensor.unsafeGetTensorImpl();
-  auto* opaque_impl = static_cast<at::TinyOpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(impl);
-  std::shared_ptr<c10::SafePyObject> tiny = opaque_impl->opaque_handle();
+  std::shared_ptr<c10::SafePyObject> tiny = tiny_impl(tensor)->opaque_handle();
   return py::reinterpret_borrow<py::object>(tiny->ptr(getPyInterpreter()));
 }
 
 void update_metadata(const at::Tensor &tensor, const std::vector<int64_t> &sizes,
                      const std::vector<int64_t> &strides, int64_t storage_offset) {
-  auto* impl = tensor.unsafeGetTensorImpl();
+  auto* impl = tiny_impl(tensor);
   impl->set_allow_tensor_metadata_change(true);
   impl->set_sizes_and_strides(sizes, strides, storage_offset);
 }

multi_migration.md

@@ -0,0 +1,86 @@
+# Multi-device op migration: MULTI/MSELECT/MSTACK → PAD / WHERE / STACK+INDEX
+
+## Status (updated)
+
+- **Stage 0 — DONE.** Internal `Ops.PAD` fills **Invalid** (`schedule/indexing.py:104`, bool keeps 0-fill); external `Tensor.pad`/`pad_to` always emit an explicit fill mask (`mixin/op.py:289`, `mixin/movement.py:267`) — required because a bare Invalid-pad leaks through elementwise ALU (`pad(x)+1` would read 0 instead of 1 in pad regions). REDUCE inputs with Invalid contribute the reduce identity (`pm_invalid_reduce_identity` in `uop/symbolic.py`, run in `get_kernel_graph` after gate lifting in `schedule/rangeify.py`) — only WHERE-alt gates whose condition involves a reduce range are rewritten, so gather-with-Invalid-index still poisons whole lanes. Same-condition nested where collapse rule added (`c?(c?t:f):f2 -> c?t:f2`) so the mask form folds to a single gate. All suites green (`test/unit`, `test/null`, `test/backend`, `test/external/external_test_schedule_scaling.py`, mypy, ruff).
+- **Stage 1 — representation in place behind `SYMBOLIC_MULTI`.** `symbolic_multi_pm` (`schedule/multi.py`) converts `MULTI→_unshard` (raw Invalid pad), `MSELECT→dnum.eq(i).where(x, Invalid)`, `MSTACK→STACK.index(dnum)` + INDEX(STACK,var)→nested-where lowering. `_unshard` uses the raw Invalid pad; `_unshard_fill` (0-fill) is used for the ALU allreduce in `copy_multi` because gated stores leave stale pad regions (the ALU-sum path can't use the identity rule). Basic shard ops work; full parity is Stage 2.
+- **Stage 2 — remaining.** Buffer level, reduce/allreduce split for shard-axis reduces, API surface.
+- **Stage 3 — remaining.**
+
+Notes: `test_schedule.py:test_pad_reduce_unsafe_multiview_st` went 4→5 kernels (pad now materializes an explicit mask; the mask form is also what makes the previously-wrong masked-pad+hazard case correct). `test_jit_footguns.py:test_symbolic_pad_view_frozen` went 2→4: the explicit mask recomputes from the symbolic shape, fixing the frozen-pad footgun. Also fixed a latent infinite loop: `(x+y) !=/< c → x !=/< c-y` collapse rules in `codegen/simplify.py` now only fire when the remaining side still contains the range (they previously shuffled constants forever when both sides were range-free).
+
+## Goal
+
+Replace the three multi-device UOps with a symbolic `_device_num` representation:
+
+| Old op | New form |
+|---|---|
+| `MULTI(x, axis)` | `x._unshard(axis)` — PAD with `_device_num`-dependent bounds back to full shape (helper already exists at `tinygrad/uop/ops.py:704-707`) |
+| `MSELECT(x, i)` | `dnum.eq(i).where(x, x.const_like(Invalid))` |
+| `MSTACK(s0..sn)` | `UOp(Ops.STACK, src=srcs).index(dnum)` — leading device axis, indexed per-device |
+
+where `dnum = UOp.variable("_device_num", 0, ndev-1)`. The per-device specialization mechanism already exists: `unwrap_multi` (`tinygrad/engine/realize.py:148-153`) binds `_device_num` per device at exec time.
+
+**Key semantic decision (approved):** internal `Ops.PAD` produces **Invalid** in padded regions; external `Tensor.pad` API still pads with 0. Staged migration: introduce the new representation first, keep old ops working, migrate call sites incrementally, delete old ops last.
+
+## Background: current design
+
+- `Ops.MULTI(src, axis)` (`tinygrad/uop/__init__.py:100`) — per-shard graph marker. Eliminated by `multi_pm` (`tinygrad/schedule/multi.py:162-195`) as the first step of `get_kernel_graph` (`tinygrad/schedule/rangeify.py:548`). Shape/axis tracking: `UOp.axis`/`UOp.bounds` (`tinygrad/uop/ops.py:667-702`).
+- `Ops.MSELECT(x, i)` / `Ops.MSTACK(srcs)` (`__init__.py:96`) — buffer-level ops. Spec at `tinygrad/uop/spec.py:181-184`; device prop `ops.py:816-819`; per-kernel PARAMs via debuf (`rangeify.py:474`); per-device dependency states (`tinygrad/schedule/__init__.py:11-17`); `MultiBuffer` (`ops.py:904-930`, `tinygrad/device.py:88-99`); only MSTACK can be `realized` (`ops.py:920-930`).
+- `_shard`/`_unshard` (`ops.py:704-714`) already emit symbolic SHRINK/PAD bounds with `_device_num`.
+- Naive allreduce already uses the target pattern: `dnum.eq(i).where(buf, state)` (`tinygrad/schedule/allreduce.py:27-33`).
+
+## Existing Invalid machinery (rely on this)
+
+- `pm_data_invalid` (`tinygrad/uop/symbolic.py:71-92`): Invalid poisons ALU (ops move inside the gate); gated LOAD folds to alt/0, gated STORE folds to NOOP.
+- `pm_remove_invalid` (`symbolic.py:94-96`): leftover Invalid → 0 in final codegen (`codegen/__init__.py:345`). Spec forbids Invalid in final programs (`spec.py:217`), so materialized Invalid regions read as 0.
+- STORE of CONST(Invalid) → NOOP (`rangeify.py:423-424`).
+- `identity_element(op, dtype)` exists (`ops.py:51`): ADD→0, MUL→1, MAX→dtype.min.
+- `found_after` (`rangeify.py:26`) already matches `WHERE(cond, PAD(x), Invalid)`.
+
+## Stage 0 — internal PAD = Invalid; external pad = explicit 0
+
+1. `tinygrad/schedule/indexing.py:100-104` (`convert_pad_to_where_to_keep_behavior_local`): fill value `0` → `UOp.const(x.dtype, Invalid)`, **except `dtypes.bool` keeps 0-fill** (False is the bool-reduce identity, and the external-pad mask below needs it).
+2. `tinygrad/mixin/op.py:282-290` (`_pad_constant`): **remove the `if value == 0: return base` shortcut** — always emit `pad(bool_ones).where(base, value)`. Required because bare Invalid-pad leaks through elementwise ALU: `pad(x)+1` gate-lifts to `where(valid, x+1, Invalid)` and reads 0 instead of 1 in pad regions. The mask lowers to a pure index expression (`valid.where(1,0)`), no extra kernel. External behavior unchanged for all `value`.
+3. **New rule**: `REDUCE(where(c, x, Invalid), op)` → `REDUCE(where(c, x, identity_element(op, dtype)), op)`. Must fire in rangeify/symbolic *before* codegen builds the accumulator loop — otherwise `pm_data_invalid` gate-lifts `acc + where(c,x,Invalid)` into `where(c, acc+x, Invalid)` and one invalid lane poisons the whole reduction. Placement (symbolic.py vs the reduce path in indexing.py) TBD at implementation; verify with `Tensor.pad(...).sum()/max()` tests.
+4. Audit: schedule tests with kernel counts involving pads; circular/reflect/replicate pads don't use PAD fill (verified, `op.py:292-312`) — unaffected; `allreduce.py:59,76` usum-of-padded-chunks gets *more* correct (disjoint regions).
+
+## Stage 1 — new representation behind env flag
+
+New `symbolic_multi_pm` PatternMatcher (in `schedule/multi.py` or new file), gated by env (e.g. `SYMBOLIC_MULTI`), run in `get_kernel_graph` right after `multi_pm`:
+
+- `MULTI(x, axis)` → `x._unshard(axis)`
+- `MSELECT(x, i)` → `dnum.eq(i).where(x, x.const_like(Invalid))` (Invalid from `tinygrad.dtype`)
+- `MSTACK(srcs)` → `STACK(*srcs).index(dnum)`, plus new lowering `INDEX(STACK(vals), var)` → nested `var.eq(k).where(src_k, Invalid)` (analogous to `convert_stack_to_where`, `indexing.py:113-121`; must fire before `validate_index` spec, `spec.py:118-122`)
+
+Flag off = zero behavior change; flag on = new forms flow through rangeify and specialize per device at exec.
+
+## Stage 2 — migrate producers/consumers (one commit each, independently testable)
+
+1. `UOp.shard` (`ops.py:715-717`): emit symbolic `_shard`+`_unshard` full-shape form directly instead of `.multi(axis)`; delete movement-op `multi_pm` rules that PAD subsumes (`pad_multi`, `permute_multi`, `expand_multi`, `reshape_multi`, `flip_multi`, `shrink_multi` — `multi.py:93-125`).
+2. ALU/STACK: `alu_multi`/`shard_srcs`/`stack_multi` (`multi.py:55-78,127-131`) become plain elementwise on full-shape padded tensors. `reduce_multi` (`multi.py:80-91`) keeps the shard-axis → local-reduce + ALLREDUCE split; Invalid-pad + identity rule replaces neutral-pad-value reasoning.
+3. allreduce (`schedule/allreduce.py`): naive path already matches; migrate ring/all2all MSELECT/MSTACK scratch-buffer assembly (lines 35-76) to WHERE/STACK+INDEX forms.
+4. Buffer level: debuf (`rangeify.py:474`), `_states`/`_unwrap_src` (`schedule/__init__.py:11-17`), `_collect_bufs` (`schedule/memory.py:9`), `unwrap_multi` (`realize.py:148-153`), JIT (`jit.py:127-130, 237`), callify (`callify.py:52-95`), `buffer`/`realized`/`buf_uop`/`has_buffer_identity` (`ops.py:841-930`).
+5. API surface: `UOp.multi/mselect/mstack` (`ops.py:662-725`), `Tensor.shard` (`tensor.py:333-347`), gradient (`mixin/gradient.py:72`), `_multi_like` (`mixin/creation.py:16-20`), embedding backward (`nn/__init__.py:309-354`), `copy_to_device(arg=)` MSELECT path (`ops.py:719-723`).
+
+## Stage 3 — removal
+
+Delete `Ops.MULTI/MSELECT/MSTACK` from the enum (`uop/__init__.py:96,100`), spec rules, viz colors (`viz/serve.py:51,56`), `UOp.axis`/`bounds` machinery (`ops.py:667-702`), remaining `multi_pm` rules, and `MultiBuffer` if fully subsumed. Flip flag default-on, then delete the flag.
+
+## Open implementation details
+
+- REDUCE-identity rule placement (must precede codegen accumulator construction).
+- INDEX(STACK, var) spec timing — the value-STACK INDEX violates the pointer-INDEX spec until lowered.
+- Whether `MultiBuffer`/tuple-`device` survives as the runtime container, or buffers become single-device with the device axis explicit in shape — decides how much of Stage 2.4 is rewrite vs delete.
+- Bool carve-out in Stage 0.1: verify no internal consumer needs Invalid-filled bool pads.
+
+## Verification (run at each stage)
+
+```bash
+python -m pytest test/unit/test_multitensor.py test/unit/test_allreduce.py test/null/test_multitensor.py test/unit/test_call.py -x -q -n12
+python -m pytest test/external/external_test_schedule_scaling.py -x -q   # test_concat_scaling
+python -m mypy tinygrad/
+python -m ruff check .
+```
+
+Also pad/reduce numeric tests after Stage 0 (`test_ops` pad tests, `Tensor.pad(...).sum()/max()`).

test/amd/hw/helpers.py

@@ -169,11 +169,10 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
   return parse_output(bytes(out_buf), n_lanes)
 
 def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
-  """Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  """Run instructions on real AMD hardware via HIPCompiler and the AMD runtime."""
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import flat_mv
+  from tinygrad.helpers import Target, flat_mv
 
   dev = Device["AMD"]
   compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -223,7 +222,7 @@ amdhsa.kernels:
 """
 
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
 
   buf_sz = _out_bytes(n_lanes)
   out_gpu = dev.allocator.alloc(buf_sz)

test/amd/hw/test_cdna_vop3.py

@@ -5,7 +5,7 @@ gfx950 hardware when USE_HW=1.
 """
 import ctypes, struct, unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -42,8 +42,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> int:
   return out_buf[0]
 
 def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -86,7 +85,7 @@ amdhsa.kernels:
 ...
 .end_amdgpu_metadata
 """
-  prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
+  prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
   prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
   out = bytearray(LANES * 4)
   dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)

test/amd/hw/test_dpp.py

@@ -6,7 +6,7 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
@@ -37,8 +37,7 @@ def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
   return list(out_buf)
 
 def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -84,7 +83,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
   out_gpu = dev.allocator.alloc(WAVE64 * 4)
   prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
   out = bytearray(WAVE64 * 4)

test/amd/hw/test_rdna4_permlane_var.py

@@ -5,7 +5,7 @@ real RDNA4 hardware when USE_HW=1.
 """
 import ctypes, unittest
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -36,8 +36,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
   return list(out_buf)
 
 def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device['AMD']
@@ -85,7 +84,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, 'test', lib)
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
   out_gpu = dev.allocator.alloc(LANES * 4)
   prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
   out = bytearray(LANES * 4)

test/amd/test_asm_kernel.py

@@ -36,7 +36,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
   threads = UOp.special(A.numel(), "lidx0")
-  var = UOp.param(2, dtypes.weakint, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
+  var = UOp.param(2, dtypes.int, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
   insts = [
     s_load_b128(s[4:7], s[0:1]),
     s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable

test/amd/test_mockgpu_invalid.py

@@ -7,15 +7,16 @@ class TestMockGPUInvalidInstruction(unittest.TestCase):
     """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
     test_code = '''
 import struct
+from dataclasses import replace
 from tinygrad import Device, Tensor
 from tinygrad.engine.realize import compile_linear
-from tinygrad.runtime.ops_amd import AMDProgram
 
 dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
 linear = compile_linear(b.schedule_linear())
-lib = bytearray(linear.src[-1].src[0].src[3].arg)
+compiled_prg = linear.src[-1].src[0]
+lib = bytearray(compiled_prg.src[3].arg)
 
 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
@@ -27,7 +28,7 @@ for i in range(0, len(lib) - 4, 4):
     break
 assert found, "s_endpgm not found"
 
-patched_prg = AMDProgram(dev, "patched", bytes(lib))
+patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
 b.uop.buffer.allocate()
 patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
 dev.synchronize()

test/backend/test_custom_kernel.py

@@ -323,7 +323,7 @@ class TestCustomKernel(unittest.TestCase):
   def test_multi_invalids_custom_kernel_no_copy(self):
     devs = ("CPU:0", "CPU:1")
     a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
-    c = Tensor(UOp.const(dtypes.float, Invalid, shape=(2, 4)).clone(device=devs).multi(0), device=devs)
+    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.multi(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     c.realize()

test/backend/test_linearizer.py

@@ -59,7 +59,6 @@ class TestLinearizer(unittest.TestCase):
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_ranges, "test inspects ranges, which are rewritten to loops on this renderer")
   def test_late_bias_load(self):
     img = Tensor.empty(1, 3, 16, 16)
     w = Tensor.empty(16, 3, 3, 3)
@@ -239,7 +238,6 @@ class TestLinearizer(unittest.TestCase):
         helper_arg_acc_dtype(d.conv2d(w, dtype=acc_dtype), expected_dtype)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_ranges, "test inspects ranges, which are rewritten to loops on this renderer")
   def test_simple_unroll_no_between_phi_dependencies(self):
     x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
     r = (x@y).relu()
@@ -301,7 +299,6 @@ class TestLinearizer(unittest.TestCase):
     program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
     assert not any(u.op == Ops.WHERE for u in tuple(program.src[1].src)), "found where where where should be folded"
 
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_ranges, "test inspects ranges, which are rewritten to loops on this renderer")
   def test_phi_simplification(self):
     def helper(t, max_ops=0):
       ast = helper_linearizer_opt(t)

test/backend/test_ops.py

@@ -339,6 +339,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.ones(256,256).max(1)[0], lambda: Tensor.ones(256,256).max(1), forward_only=True)
 
   def test_where(self):
+    helper_test_op([], lambda: torch.where(torch.tensor([True, False]), 1, 3).type(torch.int32),
+                   lambda: Tensor([True, False]).where(1, 3), forward_only=True)
     helper_test_op(
       [(100,)],
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32),
@@ -525,6 +527,7 @@ class TestOps(unittest.TestCase):
 
   def test_add(self):
     helper_test_op([(45,68), (45,68)], lambda x,y: x+y, Tensor.add)
+    helper_test_op([], lambda: torch.tensor(1)+0.5, lambda: Tensor(1)+0.5, forward_only=True)
     helper_test_op([(45,68), (45,68)], lambda x,y: x+y)
     helper_test_op([(), ()], lambda x,y: x+y)
   def test_add3(self):
@@ -903,6 +906,7 @@ class TestOps(unittest.TestCase):
   @unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
   def test_cos(self):
     helper_test_op([(45,65)], lambda x: x.cos())
+    helper_test_op([], lambda: torch.tensor(2.0).cos(), lambda: Tensor(2).cos(), forward_only=True)
     helper_test_op([()], lambda x: x.cos())
     if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
@@ -981,6 +985,8 @@ class TestOps(unittest.TestCase):
 
   def test_exp(self):
     helper_test_op([(45,65)], torch.exp, Tensor.exp)
+    # int scalar input: the python coefficients in the decomposition must stay float
+    helper_test_op([], lambda: torch.tensor(2.0).exp(), lambda: Tensor(2).exp(), forward_only=True)
     helper_test_op(None, torch.exp, Tensor.exp, vals=[[math.inf, -math.inf, math.nan]])
     helper_test_op([()], torch.exp, Tensor.exp)
   def test_exp2(self):

test/backend/test_rangeify.py

@@ -1,7 +1,7 @@
 import unittest
-from tinygrad import Tensor, nn, Device
+from tinygrad import Tensor, nn, Device, dtypes, Variable
 from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
-from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -66,6 +66,18 @@ class TestRangeifyAssign(unittest.TestCase):
     self.assertListEqual(lst2, B.permute(1, 0).tolist())
 
 class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_variable_stack_data(self):
+    # a bound-Variable STACK used as data gets ranges from its graph position
+    v = Variable("v", 0, 10).bind(3)
+    t = Tensor(UOp.stack(v, v+1).cast(dtypes.int)) + Tensor.arange(2)
+    self.assertListEqual(t.tolist(), [3, 5])
+
+  def test_variable_data_and_shape(self):
+    # the same Variable has a data edge through CAST and a structural edge through SHRINK
+    v = Variable("shared_v", 1, 10).bind(3)
+    t = Tensor.ones(10)[:v] * Tensor(v.cast(dtypes.float))
+    self.assertEqual(t.sum().item(), 9)
+
   def test_matmul_relu_cat(self):
     a = Tensor.ones(100, 512).contiguous().realize()
     c = Tensor.ones(1, 512).contiguous().realize()

test/backend/test_schedule.py

@@ -366,6 +366,7 @@ class TestCopyFolding(unittest.TestCase):
     x = y.one_hot(10)
     check_schedule(x, 3, filter_sink=False)
 
+  @unittest.skip("no longer supported")
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).clone().realize()
     zeros = Tensor.zeros(3, buffer=False).realize()
@@ -406,6 +407,7 @@ class TestCopyFolding(unittest.TestCase):
     assert t.uop.is_realized, f"didn't realize Tensor {t}"
     self.assertListEqual(t.tolist(), [1.,1.,1.,1.])
 
+  @unittest.skip("same-device copies are no-ops")
   def test_self_assign_same_device_copy(self):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph

test/backend/test_tensor.py

@@ -464,6 +464,7 @@ class TestTinygrad(unittest.TestCase):
     data = data + [-x for x in data]
     np.testing.assert_allclose(Tensor(data, dtype=dtypes.int32).numpy(), np.array(data).astype(np.int32))
 
+  @unittest.skip("list elements are python scalars: numpy values in a list have no inferred dtype, use Tensor(np.array(data))")
   def test_tensor_list_ndarray(self):
     data = [np.array([1, 2, 3]), np.array([1, 2, 3]), np.array([1, 2, 3])]
     np.testing.assert_equal(Tensor(data).numpy(), np.array(data))

test/backend/test_tensor_variable.py

@@ -18,13 +18,16 @@ class TestTensorVariable(unittest.TestCase):
       self.assertListEqual((vv * t).tolist(), [2, 2, 2])
     except RuntimeError: pass
 
+  # TODO: a Variable PARAM lowers to int32, so a bound value that doesn't fit int32 truncates or fails to bind
+  @unittest.expectedFailure
   def test_large_range_variable(self):
-    vv = Variable("b", 0, 2**40).bind(2**35)
-    # TODO: pm_lower_index_dtype lowers ALU PARAM to int32 unconditionally
-    try:
-      self.assertEqual(Tensor(vv).item(), 2**35)
-    except AssertionError:
-      pass
+    self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35)).item(), 2**35)
+
+  def test_variable_defers_like_a_literal(self):
+    vv = Variable("a", 1, 10).bind(2)
+    self.assertEqual(Tensor(vv).dtype, dtypes.weakint)
+    self.assertEqual((Tensor(vv) + Tensor([1], dtype=dtypes.int8)).dtype, dtypes.int8)  # takes the concrete side, no widening
+    self.assertEqual(Tensor(vv).item(), 2)                                              # a read commits at default_int
 
   def test_variable_tensor_dtype_arg(self):
     vv = Variable("a", 1, 10).bind(2)

test/backend/test_wait_loop.py

@@ -3,7 +3,6 @@ from tinygrad import Tensor, UOp
 from tinygrad.device import Device, Buffer, BufferSpec
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.engine.realize import run_linear
-from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.uop.ops import Ops, KernelInfo
 
 def wait_loop_kernel(C:UOp) -> UOp:
@@ -82,7 +81,6 @@ def loop_in_loop_kernel(C:UOp) -> UOp:
 
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=()))
 
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "loops are not supported in X86")
 class TestWaitLoop(unittest.TestCase):
   def test_wait_loop(self):
     c = Tensor.empty(1, dtype=dtypes.int)
@@ -109,7 +107,6 @@ class TestWaitLoop(unittest.TestCase):
     self.assertEqual(c.item(), 12)
 
 @unittest.skipUnless(Device.DEFAULT in ("CPU", "AMD", "NV"), "need proper uncached=True handling")
-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "loops are not supported in X86")
 class TestVolatileLoops(unittest.TestCase):
   def test_async_wait_ext(self):
     sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)

test/device/test_metal.py

@@ -1,7 +1,8 @@
 import unittest
-from tinygrad.device import CompileError, Device, BufferSpec
+from tinygrad.device import CompileError, Device, BufferSpec, TinyELF
+from tinygrad.helpers import Target
 if Device.DEFAULT=="METAL":
-  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
+  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
 class TestMetal(unittest.TestCase):
   def test_alloc_oom(self):
@@ -48,7 +49,7 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
     with self.assertRaises(RuntimeError):
       compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
+      device.runtime(TinyELF(compiled, "r_5", Target("METAL"), ()))
 
   def test_free(self):
     size = 2**16

test/device/test_ocl.py

@@ -1,18 +1,20 @@
 import unittest
 from unittest.mock import patch
 from tinygrad import Device
-from tinygrad.device import Buffer
+from tinygrad.device import Buffer, TinyELF
 from tinygrad.dtype import dtypes
-from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
+from tinygrad.helpers import Target
+from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLCompileCache(unittest.TestCase):
   def test_compile_cached(self):
     device = Device[Device.DEFAULT]
     src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
-    CLProgram(device, name="cached_test", lib=src.encode())
+    obj = TinyELF(src.encode(), "cached_test", Target("CL"), ())
+    device.runtime(obj)
     with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
-      CLProgram(device, name="cached_test", lib=src.encode())
+      device.runtime(obj)
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
@@ -27,7 +29,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
+      device.runtime(TinyELF(b"__kernel void test(__global int* a) { a[0] = 1; }", "", Target("CL"), ()))
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

test/external/external_test_am_fault_recovery.py

@@ -16,9 +16,9 @@ def _run(code:str, timeout:float=15.0) -> subprocess.CompletedProcess:
   return subprocess.run(["python", "-c", code], env={**os.environ, "AMD": "1"}, capture_output=True, text=True, timeout=timeout)
 
 def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
-  return _run('from tinygrad.device import Device; from tinygrad.runtime.ops_amd import AMDProgram; '
+  return _run('from tinygrad.device import Device, TinyELF; from tinygrad.helpers import Target; '
               'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
-              f'AMDProgram(dev, "test", HIPCompiler(dev.arch).compile("""{asm_src}"""))('
+              f'dev.runtime(TinyELF(HIPCompiler(dev.arch).compile("""{asm_src}"""), "test", Target("AMD", arch=dev.arch), ()))('
               'dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)')
 
 def _verify_recovery() -> subprocess.CompletedProcess:

test/external/external_test_gpu_crash.py

@@ -5,7 +5,8 @@ These tests intentionally cause GPU faults to verify error handling.
 Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
 """
 import unittest, re, importlib
-from tinygrad.device import Device
+from tinygrad.device import Device, TinyELF
+from tinygrad.helpers import Target
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
 RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
@@ -42,8 +43,8 @@ class TestGPUCrash(unittest.TestCase):
       self.fail("Device not working before test")
 
   def _run(self, code: str):
-    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
+    prg = self.dev.runtime(TinyELF(self.compiler.compile(assemble(code, is_cdna=self.is_cdna)), "test",
+                                   Target("AMD", arch=self.dev.arch), ()))
     prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
 
   def _run_insts(self, insts: list[Inst]):

test/helpers.py

@@ -15,7 +15,7 @@ from tinygrad.codegen.late.linearizer import linearize
 
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
-from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
+from tinygrad.runtime.ops_python import PythonRenderer
 
 def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp:
   if ren is None: ren = Renderer(Target())
@@ -83,14 +83,15 @@ def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
   return ret, (time.perf_counter_ns()-st)*1e-6
 
 def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple[int, ...]=()):
-  allocator = Device['PYTHON'].allocator
+  dev = Device['PYTHON']
+  allocator = dev.allocator
   bufs = []
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, (1,))
   prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src[2].arg))
+  prog = dev.runtime(prg.to_elf())
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 

test/mockgpu/amd/pcode.py

@@ -837,16 +837,16 @@ class Parser:
       idx = addr
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         val = _u32(0).cast(dtypes.uint64)
-        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i)).load().cast(dtypes.uint64) << _u64(i * 8))
+        for i in range(8): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint64) << _u64(i * 8))
       elif dt in (dtypes.uint8, dtypes.int8):
         val = mindex(idx).load().cast(dt)
       elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
         lo = mindex(idx).load().cast(dtypes.uint32)
-        hi = mindex(idx + _const(dtypes.int, 1)).load().cast(dtypes.uint32)
+        hi = mindex(idx + _const(adt, 1)).load().cast(dtypes.uint32)
         val = (lo | (hi << _u32(8))).cast(dt)
       else:
         val = _u32(0)
-        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i)).load().cast(dtypes.uint32) << _u32(i * 8))
+        for i in range(4): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint32) << _u32(i * 8))
     else:
       idx = addr >> _const(addr.dtype, 2)
       val = mindex(idx)

test/null/test_dtype_spec.py

@@ -205,7 +205,8 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
     assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
     assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
-    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.uint64
+    # uint64 has no common integer supertype with any signed int (JAX JEP), they all defer up to weakfloat
+    for st in dtypes.sints: assert least_upper_dtype(st, dtypes.uint64) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
     assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
 
@@ -224,8 +225,9 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
 
   def test_weakint_promo(self):
-    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.weakint)
-    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.int8)
+    assert least_upper_dtype(dtypes.weakint, dtypes.weakint) == dtypes.weakint
+    assert least_upper_dtype(dtypes.bool, dtypes.weakint) == dtypes.weakint
+    assert least_upper_dtype(dtypes.weakint, dtypes.int8) == dtypes.int8
 
   def test_weakfloat_promo(self):
     # weakfloat is a float, but is not one of dtypes.floats

test/null/test_llm_server.py

@@ -1,5 +1,5 @@
 import unittest, threading, time, json
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 class TestLLMServer(unittest.TestCase):
   """Integration tests using the real OpenAI client."""
@@ -66,20 +66,18 @@ class TestLLMServer(unittest.TestCase):
       self.assertEqual(chunk.model, "test-model")
 
   def test_stream_with_usage(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=True,
-      stream_options={"include_usage": True}
-    )
-
-    chunks = list(stream)
+    def generate(ids, **kwargs):
+      for token in (300, 301, 999):
+        ids.append(token)
+        yield token
+    with patch.object(self.mock_model, "generate", side_effect=generate):
+      chunks = list(self.client.chat.completions.create(
+        model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, stream_options={"include_usage": True}))
     last_chunk = chunks[-1]
 
-    self.assertIsNotNone(last_chunk.usage)
-    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
-    self.assertIsNotNone(last_chunk.usage.completion_tokens)
-    self.assertIsNotNone(last_chunk.usage.total_tokens)
+    self.assertEqual(last_chunk.usage.prompt_tokens, 3)
+    self.assertEqual(last_chunk.usage.completion_tokens, 2)
+    self.assertEqual(last_chunk.usage.total_tokens, 5)
 
   def test_multi_turn_conversation(self):
     stream = self.client.chat.completions.create(

test/null/test_schedule.py

@@ -600,7 +600,8 @@ class TestSchedule(unittest.TestCase):
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
     # TODO: this should be 3 if fix store hazard worked correctly
-    check_schedule(p, 4)
+    # NOTE: pad now always has an explicit fill mask (internal PAD is Invalid-filled), which materializes here
+    check_schedule(p, 5)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
     old_default_float, dtypes.default_float = dtypes.default_float, dtype

test/null/test_uop_symbolic.py

@@ -5,6 +5,7 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
+from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3
 
@@ -986,7 +987,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")
 
   def test_bitcast_chain(self):
-    a = Variable("a", 0, 3)
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
     self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)
 
   def test_negation_in_where(self):
@@ -1330,7 +1331,8 @@ class TestInvalidIndex(unittest.TestCase):
   def test_invalid_times_0(self):
     ridx = Variable("ridx", 0, 10)
     idx = (ridx<5).where(ridx, UOp.invalid())*0
-    self.assertIs(idx.simplify(), (ridx<5).where(0, UOp.invalid()), "multiplying an index by 0 should preserve the invalid")
+    self.assertIs(idx.simplify(), (ridx<5).where(UOp.const(dtypes.weakint, 0), UOp.invalid()),
+                  "multiplying an index by 0 should preserve the invalid")
 
   def test_alu_moves_inside_invalid(self):
     ridx = Variable("ridx", 0, 10)
@@ -1384,11 +1386,7 @@ class TestMoveWhereOnLoad(unittest.TestCase):
     idx = buf.index(a.valid(valid))
     expr = cond.where(idx, idx.const_like(0))
     out = graph_rewrite(expr, pm_move_where_on_load)
-    # any WHERE in the rewritten graph must have matched-dtype branches
-    for u in out.toposort():
-      if u.op is Ops.WHERE:
-        self.assertEqual(u.dtype, u.src[1].dtype, f"WHERE branch 1 dtype mismatch: {u}")
-        self.assertEqual(u.dtype, u.src[2].dtype, f"WHERE branch 2 dtype mismatch: {u}")
+    type_verify(out, spec_shared)  # Invalid matches any dtype
 
 class TestSymbolicRealWorld(unittest.TestCase):
   def test_resnet_half(self):

test/null/test_uop_vmin_vmax.py

@@ -168,8 +168,7 @@ class TestVminVmaxProperties(unittest.TestCase):
 
   def test_vmin_vmax_invalid_vconst(self):
     x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
-    self.assertLess(x.vmin, 0)
-    self.assertGreater(x.vmax, 4)
+    self.assertEqual((x.vmin, x.vmax), (0, 4))
 
 class TestVminVmaxDivMod(unittest.TestCase):
   def test_vmin_vmax_division_positive(self):

test/null/test_uops.py

@@ -7,14 +7,14 @@ from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
-from tinygrad.uop.symbolic import sym
+from tinygrad.uop.symbolic import sym, pm_remove_invalid
+from tinygrad.mixin.movement import MovementMixin
 from test.helpers import eval_uop, to_uops_list
 
 class TestDTypeFromUOp(unittest.TestCase):
   def test_broadcastable_promotion(self):
     self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
     self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32)
-    with self.assertRaises(KeyError): dtype_from_uop(Ops.ADD, (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.int8, 1)), None)
 
   def test_same_dtype_fast_path(self):
     src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2))
@@ -45,6 +45,28 @@ class TestDTypeFromUOp(unittest.TestCase):
       with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
       type_verify(UOp.const(concrete, value).sink(), spec_program)
 
+  def test_invalid_dtype_and_consumers(self):
+    invalid = UOp.invalid()
+    self.assertIs(invalid.dtype, dtypes.bool)
+    self.assertIs(UOp.const(dtypes.float32, Invalid), invalid)
+    self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved)
+    scratch = Tensor.invalids(4, dtype=dtypes.float32)
+    self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
+      if u.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
+    invalid, value = UOp.invalid(), UOp.const(dtypes.float32, 1)
+    for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
+              UOp.const(dtypes.bool, True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
+              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
+    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
+    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(dtypes.float, 0))
+    type_verify(out.sink(), spec_program)
+
+  def test_remove_invalid_stack_lanes(self):
+    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(dtypes.half, 1), UOp.invalid()))
+    out = graph_rewrite(stack, pm_remove_invalid)
+    self.assertEqual(out.src, (UOp.const(dtypes.half, 1), UOp.const(dtypes.half, 0)))
+    type_verify(out.sink(), spec_program)
+
 class TestLowerIndexDtype(unittest.TestCase):
   def test_gated_shrink_lowers_to_selected_width(self):
     # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
@@ -436,7 +458,7 @@ class TestContiguousViewOffset(unittest.TestCase):
   def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
   def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
   def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
-  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
+  def test_shrink_invalid(self): self._check(MovementMixin.pad(UOp.empty(4), ((2,2),))[0], None)
   def test_strided(self): self._check(UOp.empty(4)[::2], None)
 
 if __name__ == '__main__':

test/opt/test_tensor_cores.py

@@ -26,8 +26,14 @@ def run_program(prg:UOp, bufs:list[Buffer]):
   for u,b in zip(buf_uops, bufs): buffers[u] = b
   run_linear(UOp(Ops.LINEAR, src=(prg.call(*buf_uops),)))
 
+def _skip_unsupported_tc_dtypes(dtype_in:DType, dtype_out:DType):
+  supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
+  if unsupported := [f"{name}={dtype}" for name,dtype in (("dtype_in", dtype_in), ("dtype_out", dtype_out)) if dtype not in supported_dtypes]:
+    raise unittest.SkipTest(f"tensor core requires unsupported renderer dtype: {', '.join(unsupported)}")
+
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                          ensure_triggered:bool=True):
+  _skip_unsupported_tc_dtypes(dtype_in, dtype_out)
   a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
   r = a.matmul(b, dtype=dtype_out)
   sched = r.schedule_linear()
@@ -47,6 +53,7 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
     except KernelOptError: pass
 
 def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
+  _skip_unsupported_tc_dtypes(dtype_in, dtype_out)
   a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
   np_a, np_b = a.numpy(), b.numpy()
   r = a.matmul(b, dtype=dtype_out)

test/unit/test_dtype_spec.py

@@ -86,6 +86,8 @@ class TestTypeSpec(unittest.TestCase):
     _assert_eq(Tensor(2.34), dtypes.default_float, 2.34)
     _assert_eq(Tensor([]), dtypes.default_float, [])
     _assert_eq(Tensor([1]), dtypes.default_int, [1])
+    # list elements are python scalars; a numpy scalar in a list has no inferred dtype (use np.array or state a dtype)
+    with self.assertRaises(RuntimeError): Tensor([np.int32(1)])
     _assert_eq(Tensor([1.1]), dtypes.default_float, [1.1])
 
     _assert_eq(Tensor.eye(0), dtypes.default_float, np.eye(0))

test/unit/test_dtype_weak.py

@@ -1,4 +1,4 @@
-import tempfile, unittest
+import tempfile, unittest, math
 
 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import Context
@@ -137,5 +137,19 @@ class TestWeakMaterializationEntries(unittest.TestCase):
       self.assertEqual(empty.tolist(), [])
 
 
+class TestSignedUint64Weakfloat(unittest.TestCase):
+  # int64 and uint64 have no common integer supertype (JAX JEP), so the join defers to weakfloat instead of wrapping
+  def test_no_wrap(self):
+    r = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
+    self.assertEqual((r.dtype, r.item()), (dtypes.weakfloat, 0.0))
+
+  def test_weakfloat_lowers(self):
+    i64, u64 = Tensor([-1], dtype=dtypes.int64, device="CPU"), Tensor([3], dtype=dtypes.uint64, device="CPU")
+    r = i64 + u64 + Tensor([2], dtype=dtypes.float16, device="CPU")  # a concrete consumer takes the join
+    self.assertEqual((r.dtype, r.cast(dtypes.float32).item()), (dtypes.half, 4.0))
+    self.assertEqual((i64 < u64).item(), True)  # comparison meets at float
+    self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5)  # Unary lowers before transcendental
+
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_jit_footguns.py

@@ -137,11 +137,11 @@ class TestJitFootguns(unittest.TestCase):
     from tinygrad import Variable
     a = Tensor.rand(3, 10).realize()
 
-    # broken: pad is a view, BIND values frozen at capture (i=2)
+    # fixed: pad now has an explicit fill mask (internal PAD is Invalid-filled), which recomputes from the symbolic shape
     @TinyJit
     def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
     for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
-    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2)  # should be 4!
+    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
 
     # workaround: contiguous fuses pad into kernel
     @TinyJit

tinygrad/callify.py

@@ -142,11 +142,11 @@ pm_early_transform_tensor_graph = PatternMatcher([
   (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), 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, allow_any_len=True, name="copy"), lambda x,copy:
-   copy.replace(src=(x,)+copy.src[1:], tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+  (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 isinstance(x.device, str) and x.device.startswith("DISK") 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],)+copy.src[1:], tag=None),)+x.src[1:]) \
+   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) \
    if isinstance(x.device, str) and x.device.startswith("DISK") else None),
 
   # add CONTIGUOUS to tagged UOps
@@ -179,10 +179,9 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
+  if b.op is Ops.BIND: return b.param_like(len(ctx.replacements)-1)
   return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   b._min_max if b.op is Ops.BIND else None, name=b.src[0].expr if b.op is Ops.BIND else None,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL,
-                   multiple_of=b.src[0].arg.multiple_of if b.op is Ops.BIND else None)
+                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
 
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.AFTER, name="x"), finalize_after),

tinygrad/codegen/__init__.py

@@ -8,7 +8,7 @@ from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -22,11 +22,11 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
 from tinygrad.schedule.rangeify import pm_mops
-from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize, ranges_to_loops
+from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
 from tinygrad.helpers import all_same, flatten, argsort, partition
-from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element
+from tinygrad.uop.ops import _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts
 
 def do_number_param(ctx:list[int], x:UOp):
@@ -88,25 +88,20 @@ expander2 = PatternMatcher([
   (UPat(Ops.WMMA, name="u"), expand_wmma),
 ])+pm_flatten_range+mop_cleanup
 
-def broadcast_binary(x:UOp):
+def expand_broadcast(x:UOp):
   shapes = [u._shape for u in x.src]
   if any(s is None for s in shapes) or all_same(shapes): return None
-  shaped_aligned = _align_left(*shapes)
-  broadcasted = _broadcast_shape(*shapes)
-  src_reshaped = [u.reshape(shp).expand(broadcasted) for u,shp in zip(x.src, shaped_aligned)]
-  return x.replace(src=tuple(src_reshaped))
+  shape = _broadcast_shape(*shapes)
+  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
 
 def broadcast_and_devec_wmma(b:UOp):
   shapes = [u.shape[:-1] for u in b.src]
   if all_same(shapes): return None
-  shaped_aligned = _align_left(*shapes)
-  broadcasted = _broadcast_shape(*shapes)
-  src_reshaped = [u.reshape(shp+(u.shape[-1],)).expand(broadcasted+(u.shape[-1],))
-                  for u,shp in zip(b.src, shaped_aligned)]
+  shape = _broadcast_shape(*shapes)
+  src_expanded = tuple([u.expand(shape+(u.shape[-1],)) for u in b.src])
   src = []
   for idx in itertools.product(*[range(i) for i in b.shape[:-1]]):
-    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
-    src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
+    src.append(b.replace(src=tuple([x.index(*idx) for x in src_expanded])))
   return UOp.stack(*src).reshape(b.shape)
 
 pm_wmma_add = PatternMatcher([
@@ -119,19 +114,19 @@ pm_wmma_add = PatternMatcher([
     lambda wmma,reshape,permute,add: (wmma + add.permute(argsort(permute.arg)).reshape(wmma.shape)).reshape(reshape.shape).permute(permute.arg)),
 ])
 
-unbroadcast = pm_wmma_add+PatternMatcher([
-  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), broadcast_binary),
+pm_expand_broadcast = pm_wmma_add+PatternMatcher([
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
   (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])
 
 def do_devectorize(b:UOp):
   if b.shape == (): return None
-  # broadcasting needs to be already unpacked
-  if not all_same([x.shape for x in b.src]): return None
+  # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
+  if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None
   src = []
   for idx in itertools.product(*[range(x) for x in b.shape]):
     idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
-    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+    src.append(b.replace(src=tuple(x.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src)))
   return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
 
 def do_stack_wmma(u:UOp):
@@ -302,7 +297,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # **** optimizations are done, now we lower to actual code ****
 
-  sink = graph_rewrite(sink, symbolic_simple+unbroadcast+pm_add_loads, name="*** unbroadcast / add loads")
+  sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
   sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
@@ -349,12 +344,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
-  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC: type_verify(sink, spec_program)
-
-  # rewrite bounded ranges to loops for renderers without range support, after validation like instruction selection
-  if not ren.supports_ranges: sink = ranges_to_loops(sink)
-
   # this was the linearizer
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
@@ -362,6 +351,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
+  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
+  if SPEC: type_verify(sink, spec_program)
+
   # return the rewritten sink
   return sink
 
@@ -411,8 +403,7 @@ def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
 
 def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
   src = ctx.render(list(lin.src))
-  new_arg = replace(prg.arg, aux=tuple(ctx.aux(list(lin.src)))) if ctx.has_aux else prg.arg
-  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=new_arg)
+  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),))
 
 def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
   if DEBUG >= 4: print(source.arg)
@@ -445,14 +436,14 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   elif ast.op is Ops.SINK:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
-    prog_info = ProgramInfo.from_sink(full_sink)
+    prog_info = ProgramInfo.from_sink(full_sink, renderer.target)
     # instruction selection
     if isinstance(renderer, ISARenderer):
       full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
       full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
     prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info)
   else: raise RuntimeError(f"can't call to_program on {ast.op}")
-  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
+  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0], renderer.target))
   prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
   if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
   return prg

tinygrad/codegen/decomp/dtype.py

@@ -140,7 +140,7 @@ pm_long_decomp = PatternMatcher([
    l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
                                        if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag] if x.tag is not None else None),
   (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
-   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag).replace(dtype=l2i_dt[x.dtype]),))),
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag).replace(dtype=l2i_dt[x.dtype]),)) if x.tag is not None else None),
   (UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
    UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
 ])

tinygrad/codegen/decomp/transcendental.py

@@ -226,7 +226,7 @@ def xlog2(d:UOp) -> UOp:
   denormal_exp = 10 if d.dtype == dtypes.float16 else 64
   FLT_MIN = d.const_like({dtypes.float16: 6.1e-5, dtypes.float32: 1e-4, dtypes.float64: 1e-4}[d.dtype])
   is_denormal = d<FLT_MIN
-  a = is_denormal.where(d * (2 ** denormal_exp), d)
+  a = is_denormal.where(d * (2.0 ** denormal_exp), d)
 
   e = ilogb2k(a * (1.0 / 0.75)).cast(a.dtype)
   m = ldexp3k(a, -e)

tinygrad/codegen/late/gater.py

@@ -2,6 +2,10 @@
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes
 
+def move_where_load(gate, l, a, w):
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.arg is Invalid else
+    a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)
+
 pm_move_gates_from_index = PatternMatcher([
   # for image idx (must be first)
   (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
@@ -18,8 +22,8 @@ pm_move_gates_from_index = PatternMatcher([
    .store(UPat.var("data")), lambda mop,gate,idx,data: mop.replace(src=(mop.src[0],idx)+mop.src[2:]).store(data, gate)),
 
   # Where after gated load becomes alt value
-  (UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")), lambda gate,l,a:
-   l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
-  (UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()), lambda gate,l,a:
-   l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
+  (UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")).named("w"),
+   move_where_load),
+  (UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()).named("w"),
+   move_where_load),
 ])

tinygrad/codegen/late/linearizer.py

@@ -1,7 +1,7 @@
 import heapq
 from typing import Any
 from collections import defaultdict
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str, ParamArg, AxisType
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 
@@ -93,42 +93,3 @@ pm_split_ends = PatternMatcher([
   # split the ends
   (UPat(Ops.END, name="e"), do_split_ends),
 ])
-
-def ranges_to_loops(sink:UOp) -> UOp:
-  # rewrite bounded ranges to bound-less loops with a register counter: i = 0; loop { body; i += 1; loop again while i < bound }
-  slot = max((u.arg.slot for u in sink.toposort() if u.op is Ops.BUFFER and u.addrspace == AddrSpace.REG), default=-1) + 1
-  ends = [u for u in sink.toposort() if u.op is Ops.END and any(x.op is Ops.RANGE and x.dtype is not dtypes.void for x in u.src[1:])]
-  # e.ranges over-approximates nesting (it flows ranges through ordering deps), so compute true nesting from the body slices
-  # NOTE: uop identity is not stable (the uop cache is weak), all lookups are by uop key
-  end_for_range = {r.key: e for e in ends for r in e.src[1:] if r.op is Ops.RANGE and r.dtype is not dtypes.void}
-  body_ends = {e.key: {u.key for u in e.src[0].toposort()} for e in ends}
-  repl: dict[UOp, UOp] = {}
-  range_to_loop: dict[bytes, UOp] = {}
-  for e in ends:
-    # the counter init is placed after the enclosing loops so it resets every outer iteration, the loop header depends on it so it runs first
-    enclosing = tuple(r for r in e.ranges if (er:=end_for_range.get(r.key)) is not None and e.key in body_ends[er.key])
-    e = e.substitute(repl)
-    assert len(e.src) == 2, f"expected a split END with one range, got {len(e.src)-1} ranges"
-    r = e.src[1]
-    i = UOp(Ops.BUFFER, src=(UOp.const(dtypes.int, 1),), arg=ParamArg(slot, r.dtype, addrspace=AddrSpace.REG))
-    slot += 1
-    z = UOp.const(dtypes.int, 0)
-    init = i.after(*enclosing).index(z).store(UOp.const(r.dtype, 0))
-    i = i.after(init)
-    # a do-while can't skip its first iteration, so a range with a possibly zero bound gets a one-time entry guard on the loop header
-    guard = () if r.src[0].vmin >= 1 else (UOp.const(r.dtype, 0) < r.src[0],)
-    l = range_to_loop[r.key] = UOp(Ops.RANGE, dtypes.void, src=(init,)+guard, arg=(r.arg[0], AxisType.LOOP))
-    iv = i.after(l).index(z).load()
-    inc = iv + UOp.const(r.dtype, 1)
-    body = e.src[0].substitute({r: iv})
-    # the counter store is part of the loop body, an AFTER body can't be in a GROUP so sequence it with a dep instead
-    ret = body.after(i.index(z).store(inc)) if body.op is Ops.AFTER else UOp.group(body, i.index(z).store(inc))
-    repl[e] = ret.end(l, inc < r.src[0])
-    # keep the tracked loop headers up to date: their init deps on enclosing ranges get rewritten by the same substitution
-    for k in range_to_loop: range_to_loop[k] = range_to_loop[k].substitute({r: iv})
-  if not len(repl): return sink
-  out = sink.substitute(repl)
-  # ordering deps on the old ranges (scope AFTERs outside the loop bodies) point at the loop headers
-  fix = {a: a.replace(src=(a.src[0],) + tuple(range_to_loop[s.key] if s.key in range_to_loop else s for s in a.src[1:]))
-         for a in out.toposort() if a.op is Ops.AFTER and any(s.key in range_to_loop for s in a.src[1:])}
-  return out.substitute(fix) if len(fix) else out

tinygrad/codegen/simplify.py

@@ -95,11 +95,12 @@ pm_reduce_unparented = PatternMatcher([
 ])
 
 pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
-  # lift x+y out of reduce on lt
-  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on lt. only fire if x still has the range: with both sides range-free it just shuffles constants
+  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"),
+   lambda x,y,c: (x < (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
   ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
+    lambda x,y,c: (x < ((c+y-1) // y)) if not no_range(x) and no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
   # sum over r in [0,N) of [lower<=r<upper]*val -> clamp(min(upper,N) - max(lower,0), 0, N) * val
   (UPat.any(
     (UPat(Ops.RANGE, name="r") < UPat.var("upper")).where(UPat.var("val"), 0),
@@ -119,8 +120,9 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
 ])+symbolic
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
-  # lift x+y out of reduce on ne
-  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on ne (same range guard as the lt version)
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"),
+   lambda x,y,c: (x != (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
   # reduce on gated load becomes can substitute the range and remove the reduce
   ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),

tinygrad/device.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
+from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap
 from tinygrad.dtype import DType, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -283,12 +283,25 @@ class Compiler:
     return lib
   def disassemble(self, lib:bytes): pass
 
+@dataclass
+class TinyELF:
+  lib: bytes
+  name: str
+  target: Target
+  # tuple of (name, slot, dtype, shape)
+  signature: tuple[tuple[str|None, int, DType, tuple], ...]
+
+class Program(Generic[DeviceType]):
+  def __init__(self, dev:DeviceType, obj:TinyELF): pass
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(),
+               wait=False) -> float|None: pass
+
 class Compiled:
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
-  def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime, graph=None, arch=None):
+  def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
-    self.device, self.allocator, self.runtime, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
+    self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
     self.arch = arch
     self.cached_renderer:dict[Any, Renderer] = {}
 
@@ -300,6 +313,8 @@ class Compiled:
     if (ret:=self.renderer.compiler) is None: raise RuntimeError(f"no compiler for {self.device}")
     return ret
 
+  def runtime(self, obj:TinyELF) -> Program[Self]: return unwrap(self.runtime_t)(self, obj)
+
   def _renderer_name(self, r:type[Renderer]) -> str:
     return r.__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 

tinygrad/dtype.py

@@ -111,7 +111,7 @@ class dtypes:
     return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
             dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
   void: Final[DType] = DType.new(-1, 0, "void", None)
-  weakint: Final[DType] = DType.new(0, 800, "weakint", None)  # NOTE: not in the promo lattice: index math never mixes dtypes
+  weakint: Final[DType] = DType.new(0, 800, "weakint", None)  # the weak int position in the promo lattice
   bool: Final[DType] = DType.new(0, 1, "bool", '?')
   int8: Final[DType] = DType.new(1, 8, "signed char", 'b')
   uint8: Final[DType] = DType.new(2, 8, "unsigned char", 'B')
@@ -129,7 +129,6 @@ class dtypes:
   fp8e4m3fnuz: Final[DType] = DType.new(10, 8, "float8_e4m3fnuz", None)
   fp8e5m2fnuz: Final[DType] = DType.new(11, 8, "float8_e5m2fnuz", None)
   float16: Final[DType] = DType.new(12, 16, "half", 'e')
-  # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
   bfloat16: Final[DType] = DType.new(13, 16, "__bf16", None)
   float32: Final[DType] = DType.new(14, 32, "float", 'f')
   float64: Final[DType] = DType.new(15, 64, "double", 'd')
@@ -153,7 +152,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
-  weaks = (weakfloat,)
+  weaks = (weakint, weakfloat)
   all = floats + ints + (bool,) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
@@ -163,14 +162,13 @@ if (env_default_float := getenv("DEFAULT_FLOAT", "")):
 DTypeLike = str|DType
 def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
 def strong_dtype(dtype:DType) -> DType:
-  # TODO: weakint
-  return dtypes.default_float if dtype == dtypes.weakfloat else dtype
+  return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype)
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
 # we don't support complex type
-# TODO: weakint
-promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
-  dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
+promo_lattice = { dtypes.bool: [dtypes.weakint], dtypes.weakint: [dtypes.int8, dtypes.uint8],
+  dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
+  dtypes.int64: [dtypes.weakfloat], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
   dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.weakfloat],
   dtypes.weakfloat: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
   dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16],

tinygrad/engine/realize.py

@@ -4,7 +4,7 @@ import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
@@ -112,9 +112,8 @@ def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
 
 runtime_cache: dict[tuple[bytes, str], Any] = {}
 def get_runtime(device:str, ast:UOp, cache=True):
-  assert ast.op is Ops.PROGRAM and isinstance(ast.arg, ProgramInfo), "get_runtime should only be called with a PROGRAM ast"
   if (runtime:=runtime_cache.get(key:=(ast.key, device))) is None:
-    runtime = Device[device].runtime(ast.arg.function_name, ast.src[3].arg, *ast.arg.aux, runtimevars=ast.arg.runtimevars, prg=ast)
+    runtime = Device[device].runtime(ast.to_elf())
     if cache: runtime_cache[key] = runtime
   return runtime
 

tinygrad/llm/gguf.py

@@ -38,7 +38,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   def q_to_uint8(t: Tensor, b: int) -> Tensor:
     # TODO: rewrite with arange?
     shift_tensor, bitmask = Tensor.const(t.dtype, tuple(2**(i*b) for i in range(8//b))), 0xff >> (8 - b)
-    return t.unsqueeze(-1).expand((*t.shape,8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
+    return t.unsqueeze(-1).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
@@ -68,7 +68,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     if ggml_type == 14:
       xl, xh = q_to_uint8(blocks[:,:128].reshape((-1, 2, 64)), 4), q_to_uint8(blocks[:,128:192].reshape((-1, 2, 32)), 2).lshift(4)
       scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))
-      d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
+      d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32)
       return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
     if ggml_type == 18:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))

tinygrad/llm/serve.py

@@ -68,6 +68,7 @@ class Handler(HTTPRequestHandler):
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
   def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0):
     model, tok = self.server.model, self.server.tok
+    prompt_tokens = len(ids)
     cache_start_pos = model.get_start_pos(ids)
     stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
     tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
@@ -79,7 +80,7 @@ class Handler(HTTPRequestHandler):
     dec = tok.stream_decoder()
     router = StreamRouter()
     for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
       if tok.is_end(next_id): break
       out.append(next_id)
       for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
@@ -101,7 +102,8 @@ class Handler(HTTPRequestHandler):
       if finish_reason == "stop": finish_reason = "tool_calls"
     yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
     if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
+                                      "total_tokens": prompt_tokens + len(out)}, **tmpl}
     et = time.perf_counter()
     stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
                f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")

tinygrad/mixin/creation.py

@@ -1,5 +1,5 @@
 from typing import TYPE_CHECKING, Callable, Self
-from tinygrad.dtype import ConstType, DTypeLike, Invalid, dtypes, to_dtype
+from tinygrad.dtype import ConstType, DTypeLike, Invalid, dtypes, to_dtype, strong_dtype
 from tinygrad.helpers import argfix, prod
 from tinygrad.mixin.dtype import DTypeMixin
 from tinygrad.mixin.movement import MovementMixin
@@ -78,9 +78,13 @@ class CreationMixin(DTypeMixin, MovementMixin):
     from tinygrad.uop.ops import UOp
     new_shape = argfix(shape)
     dt = to_dtype(dtype) if dtype is not None else fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value)
+    # materializing commits an inferred weak width
+    if dtype is None and buffer: dt = strong_dtype(dt)
     val = cls.const(dt, fill_value)
     val = val.reshape((1,)*len(new_shape)).expand(new_shape)
-    return val.clone(device=device) if buffer else val
+    if not buffer or val._uop.base.arg is not Invalid or val.dtype == dt: return val.clone(device=device) if buffer else val
+    ret = val.empty_like(dt, device)
+    return cls._wrap_uop(ret._uop.after(ret._uop.store(val._uop)))
 
   def full_like(self, fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, buffer=True) -> Self:
     """

tinygrad/mixin/dtype.py

@@ -1,5 +1,5 @@
 from typing import TYPE_CHECKING, Self
-from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype
+from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype, Invalid
 from tinygrad.uop import Ops
 
 if TYPE_CHECKING:
@@ -30,7 +30,7 @@ class DTypeMixin:
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.arg is Invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
 
   def bitcast(self, dtype:DTypeLike) -> Self:
     """

tinygrad/mixin/elementwise.py

@@ -200,7 +200,7 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
-    if dtypes.is_int(a.dtype): return a.alu(Ops.FLOORMOD, b)
+    if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): return a.alu(Ops.FLOORMOD, b)
     return a - a.div(b, rounding_mode="floor") * b
 
   def fmod(self, x: Self | ConstType) -> Self:
@@ -213,7 +213,7 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x)
-    if dtypes.is_int(a.dtype): return a.alu(Ops.CMOD, b)
+    if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): return a.alu(Ops.CMOD, b)
     return a - a.div(b, rounding_mode="trunc") * b
 
   def div(self, x: Self | ConstType, reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self:
@@ -237,7 +237,7 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
-    if dtypes.is_int(a.dtype):
+    if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype):
       if rounding_mode == "trunc": return a.alu(Ops.CDIV, b)
       if rounding_mode == "floor": return a.alu(Ops.FLOORDIV, b)
     d = a * b.reciprocal()

tinygrad/mixin/movement.py

@@ -265,7 +265,8 @@ class MovementMixin:
     return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   def pad_to(self, shape, *args) -> Self:
-    return self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))
+    # NOTE: this calls the overridden pad (OpMixin.pad when available) so the fill is an explicit 0, not Invalid
+    return self.pad(tuple((0, 0) if ns is None else (0, ns-s) for s, ns in zip(self.shape, argfix(shape, *args), strict=True)))
 
   def view(self, shape, *args) -> Self:
     """`.view` is an alias for `.reshape`."""

tinygrad/mixin/op.py

@@ -286,8 +286,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     X = self.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))) if has_neg else self
     pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
     base = MovementMixin.pad(X, pads)
-    if value == 0: return base
-    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, value)
+    if base is X: return X  # no padding, nothing to fill
+    # the fill is always explicit: internal PAD fills with Invalid, so the mask is required for every value (incl. 0)
+    # for 0 use a literal that loses every promotion: a Python 0.0 would promote int/bool tensors to float, int 0 bool to int
+    fill = (False if X.dtype == dtypes.bool else 0) if value == 0 else value
+    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, fill)
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
@@ -440,7 +443,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     sz = merge_dicts([dict(zip(s, x.shape)) for s, x in zip(inputs, xs)])
     alpha = sorted(sz)
     # align all tensors to alphabet, multiply, sum non-output, permute to output order
-    xs = [x.permute(*[s.index(c) for c in sorted(s)]).reshape([sz[c] if c in s else 1 for c in alpha]).expand([sz[c] for c in alpha]) if s else x
+    xs = [x.permute(*[s.index(c) for c in sorted(s)]).reshape([sz[c] if c in s else 1 for c in alpha]) if s else x
           for s, x in zip(inputs, xs)]
     return xs[0].uprod(*xs[1:]).sum([i for i,c in enumerate(alpha) if c not in rhs], dtype=dtype).permute(argsort(argsort(list(rhs))))
 
@@ -738,11 +741,9 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     SPLIT = 256
     value = identity_element(op, self.dtype)
     if not isinstance(s:=self.shape[axis], int) or s <= SPLIT*2: return self._cumalu(axis, op)
-    ret = self.transpose(axis,-1)._pad_constant((None,)*(self.ndim-1)+((round_up(s,SPLIT)-s,0),), value).unflatten(-1,(-1,SPLIT))._cumalu(-1, op)
-    base = ret[..., -1]._cumalu(-1, op)._pad_constant((None,)*(ret.ndim-2) + ((1, -1),), value)
-    base = base.unsqueeze(-1).expand(*base.shape, ret.shape[-1])
-    def fix(x: Self) -> Self: return x.flatten(start_dim=-2)[..., -s:].transpose(axis,-1)
-    return fix(ret).alu(op, fix(base))
+    chunks = self.transpose(axis,-1)._pad_constant((None,)*(self.ndim-1)+((round_up(s,SPLIT)-s,0),), value).unflatten(-1,(-1,SPLIT))._cumalu(-1, op)
+    base = chunks[..., -1]._cumalu(-1, op)._pad_constant((None,)*(chunks.ndim-2) + ((1, -1),), value)
+    return chunks.alu(op, base.unsqueeze(-1)).flatten(start_dim=-2)[..., -s:].transpose(axis,-1)
 
   def cumsum(self, axis:int=0) -> Self:
     """
@@ -831,7 +832,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if self.ndim == 0: return self
     x = self.transpose(axis, -1)
     last_dim_size = x.shape[-1]
-    x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
+    x_unsqueezed = x.unsqueeze(-2)
     x_cummax = x.cummax(-1)[0].detach()
     mask = type(self).ones(last_dim_size, last_dim_size, buffer=False, dtype=dtypes.bool).tril()
     ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax

tinygrad/mixin/rand.py

@@ -269,7 +269,7 @@ class RandMixin(OpMixin):
     if replacement or num_samples == 1:
       cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
       unif_samples = type(self).rand(num_samples, cdf.shape[0], 1).to(self.device)  # type: ignore[attr-defined]
-      indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+      indices = (unif_samples >= cdf).sum(2).permute((1, 0))
     else:
       # Efraimidis-Spirakis
       indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]

tinygrad/nn/__init__.py

@@ -40,7 +40,7 @@ class BatchNorm:
 
   def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]:
     shape_mask: list[int] = [1, -1, *([1]*(x.ndim-2))]
-    if self.track_running_stats and not TRAINING: return self.running_mean, self.running_var.reshape(shape=shape_mask).expand(x.shape)
+    if self.track_running_stats and not TRAINING: return self.running_mean, self.running_var.reshape(shape=shape_mask)
     # This requires two full memory accesses to x
     # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
     # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm

tinygrad/renderer/__init__.py

@@ -64,7 +64,6 @@ class Renderer:
   has_local: bool = True
   has_threads: bool = False
   has_shared: bool = True
-  has_aux: bool = False # additional program info, eg. image shapes
   # NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims
   global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
   local_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
@@ -73,8 +72,6 @@ class Renderer:
   tensor_cores: list[TensorCore] = []
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
-  # renderers without range support get all bounded ranges rewritten to loops in codegen
-  supports_ranges: bool = True
 
   compiler: Compiler = Compiler()
 
@@ -82,7 +79,6 @@ class Renderer:
   def __reduce__(self): return self.__class__, (self.target,)
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")
   def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler")
-  def aux(self, uops:list[UOp]) -> dict: raise NotImplementedError("needs aux")
   def supported_dtypes(self) -> set[DType]:
     # double can't be bitcast to anything without long support
     return set(dtypes.all) - ({dtypes.double} if dtypes.long in EMULATED_DTYPES.tolist(dtypes) else set())

tinygrad/renderer/cstyle.py

@@ -76,9 +76,9 @@ base_rewrite = PatternMatcher([
 
 def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
   patterns = PatternMatcher([
-    (UPat(Ops.WHERE, src=(UPat.var("b"), UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
-     lambda b,x,y: UOp(Ops.WHERE, src=(b,x.cast(dtypes.float),y.cast(dtypes.float))).cast(x.dtype)),
-    (UPat(GroupOp.ALU, dtype=dts, name="x"),
+    (UPat(Ops.WHERE, dtype=dts, src=(UPat.var("b"), UPat.var("x"), UPat.var("y")), name="w"),
+     lambda w,b,x,y: b.where(x.cast(dtypes.float), y.cast(dtypes.float)).cast(w.dtype)),
+    (UPat(GroupOp.ALU-{Ops.WHERE}, dtype=dts, name="x"),
      lambda x: UOp(x.op, src=tuple(vv.cast(dtypes.float) for vv in x.src), arg=x.arg).cast(x.dtype)),
     (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
      lambda alu,x,y: UOp(alu.op, src=(x.cast(dtypes.float), y.cast(dtypes.float)), arg=alu.arg))])
@@ -303,8 +303,6 @@ class ClangRenderer(CStyleLanguage):
     self.compiler = ClangCompiler(target.arch.split(","))
 
 class OpenCLRenderer(CStyleLanguage):
-  has_aux = True
-
   # language options
   kernel_typedef = "__kernel void"
   buffer_prefix = "__global "
@@ -336,13 +334,6 @@ class OpenCLRenderer(CStyleLanguage):
     if any(uop.dtype == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or []))
     return super().render_kernel(function_name, kernel, bufs, uops, prefix)
 
-  def aux(self, uops:list[UOp]):
-    arg_dtypes:list[list[tuple[int, DType, tuple|None]]] = []
-    for i,u in enumerate(u for u in uops if u.op is Ops.PARAM):
-      while len(arg_dtypes) <= u.arg.slot: arg_dtypes.append([])
-      arg_dtypes[u.arg.slot].append((i, u.dtype, u._shape))
-    return tuple(tuple(a) for a in arg_dtypes),
-
   def supported_dtypes(self): return {d for d in super().supported_dtypes()
                                       if (d != dtypes.half or "cl_khr_fp16" in self.target.arch) and
                                       (d != dtypes.double or "cl_khr_fp64" in self.target.arch) and d not in dtypes.fp8s}

tinygrad/renderer/isa/x86.py

@@ -13,7 +13,7 @@ from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target
 class X86Ops(FastEnum):
   # NOTE: X86Ops with i suffix are variants that take an immediate, m suffix are variants that can write to memory instead of read from
   # these aren't real instructions, DEFINE is a register placeholder that defines a register without emitting an instruction
-  FRAME_INDEX = auto(); LABEL = auto(); DEFINE = auto()
+  FRAME_INDEX = auto(); LABEL = auto(); DEFINE = auto(); LOOP_CMP = auto()
   # index
   LEA = auto()
   # register / memory / immediate moves
@@ -328,6 +328,7 @@ def _xmm_sz_m(x: UOp) -> X86Ops:
 def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
   # register placeholders with real registers
   if x.arg is X86Ops.DEFINE and x.tag is not None: return None
+  if x.arg is X86Ops.LOOP_CMP: return None
   # this is an immediate
   if x.arg is X86Ops.FRAME_INDEX: return None
   # no register definition
@@ -353,6 +354,9 @@ isel_matcher = PatternMatcher([
   # range is lowered to acc, cmp, jmp after regalloc
   (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])),
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
+  # really all a backedge END is is an IF with a tag referencing the RANGE start label
+  (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
+    lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])),
   # **** Op -> X86Op ****
   # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc
   # so regalloc builds the prologue/epilogue naturally
@@ -562,22 +566,37 @@ pre_regalloc_matcher = PatternMatcher([
 # TODO: control flow should be overhauled so that this isn't necessary
 def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
   loop_label = "_".join(str(i) for i in x.arg[:-1])
-  acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
   label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}")
-  cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
-  jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
-  ctx.loop_label[acc] = loop_label
-  return (acc, [acc, label, cmp, jump_out])
+  # loop, cmp on backedge all we need is a jmp tag
+  if x.dtype is dtypes.void: return (label, [label])
+  else:
+    acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
+    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
+    jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
+    ctx.loop_label[acc] = loop_label
+    return (acc, [acc, label, cmp, jump_out])
+
+def lower_end(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
+  end_label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")
+  jmp = UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}")
+  inc = x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),))
+  return (inc, [inc, jmp, end_label])
+
+def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
+  cond = x.replace(op=x.tag, src=x.src[:2])
+  jmp = isel_matcher.rewrite(UOp(Ops.IF, src=(cond,)))
+  return (jmp.src[0], [jmp.src[0], jmp.replace(tag=x.src[3].tag)])
 
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
   # rewrite FRAME_INDEX to IMM now that the stack size is known
   (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])),
+  # expand the cmp here so we can preserve rng src edge to get label from ctx
+  (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
-  (UPat(Ops.RANGE, name="x"), lambda ctx,x: lower_range(ctx, x)),
+  (UPat(Ops.RANGE, name="x"), lower_range),
   # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target
-  (UPat(Ops.END, name="x"), lambda ctx,x: (jmp:=UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}"),
-   [x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)), jmp, UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")])),
+  (UPat(Ops.END, name="x"), lower_end),
   # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
   (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
    [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg in X86GroupOp.TwoAddress else None),
@@ -845,6 +864,7 @@ class X86Renderer(ISARenderer):
     binary = bytearray()
     for u in uops:
       if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue
+      if u.arg is X86Ops.LOOP_CMP: continue
       if u.arg is X86Ops.LABEL:
         targets[u.tag] = len(binary)
         continue

tinygrad/renderer/llvmir.py

@@ -3,7 +3,7 @@ from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast
 from tinygrad.codegen.decomp.transcendental import xexp2, xlog2
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
 from tinygrad.dtype import dtypes, float_to_fp8, DType, truncate, AddrSpace
 from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX
 
@@ -101,13 +101,28 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.WHERE, name="x"), lambda ctx,x:
    f"  {ctx[x]} = select {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[2].dtype)} {ctx[x.src[2]]}"),
 
-  # loop (ranges are rewritten to loops in codegen), a bool src is a one-time entry guard for possibly zero trip counts
-  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx,l:
-    f"  br i1 {ctx[g]}, label %loop_{ctx[l][1:]}, label %loop_exit_{ctx[l][1:]}\nloop_{ctx[l][1:]}:" \
-    if (g:=next((s for s in l.src if s.dtype is dtypes.bool), None)) is not None else f"  br label %loop_{ctx[l][1:]}\nloop_{ctx[l][1:]}:"),
+  # loop (a RANGE with no src is an unbounded loop header)
+  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx,l: f"  br label %loop_{ctx[l][1:]}\nloop_{ctx[l][1:]}:"),
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx,l,c:
     f"  br i1 {ctx[c]}, label %loop_{ctx[l][1:]}, label %loop_exit_{ctx[l][1:]}\nloop_exit_{ctx[l][1:]}:"),
 
+  # range
+  (UPat(Ops.RANGE, name="r"), lambda ctx,r:
+   f"  br label %loop_entry_{range_str(r)}\n"
+   f"loop_entry_{range_str(r)}:\n"
+   f"  br label %loop_latch_{range_str(r)}\n"
+   f"loop_latch_{range_str(r)}:\n"
+   f"  {ctx[r]} = phi {ldt(r.dtype)} [ 0, %loop_entry_{range_str(r)} ], [ {ctx[r]}phi, %loop_footer_{range_str(r)} ]\n"
+   f"  {ctx[r]}phi = add {ldt(r.dtype)} {ctx[r]}, 1\n"
+   f"  {ctx[r]}cmp = icmp ult {ldt(r.dtype)} {ctx[r]}, {ctx[r.src[0]]}\n"
+   f"  br i1 {ctx[r]}cmp, label %loop_body_{range_str(r)}, label %loop_exit_{range_str(r)}\n"
+   f"loop_body_{range_str(r)}:"),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda r:
+   f"  br label %loop_footer_{range_str(r)}\n"
+   f"loop_footer_{range_str(r)}:\n"
+   f"  br label %loop_latch_{range_str(r)}\n"
+   f"loop_exit_{range_str(r)}:"),
+
   # if
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"  br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
   (UPat(Ops.ENDIF, name="x"), lambda ctx,x: f"  br label %ifskip_{ctx[x.src[0]][1:]}\nifskip_{ctx[x.src[0]][1:]}:"),
@@ -117,7 +132,6 @@ base_rewrite = PatternMatcher([
 
 class LLVMRenderer(Renderer):
   supports_float4 = True
-  supports_ranges = False
   abi: str | None
   string_rewrite: PatternMatcher
   code_for_op = {k:lambda:None for v in lop.values() for k in v.keys()}

tinygrad/renderer/nir.py

@@ -2,8 +2,8 @@ from typing import Callable, Any
 from tinygrad.dtype import AddrSpace, DType, dtypes, truncate
 from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape
 from tinygrad.renderer import Renderer
-from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
+from tinygrad.renderer.cstyle import CUDARenderer
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.runtime.autogen import mesa, libc
 from tinygrad.runtime.support.c import POINTER
 import base64, ctypes, struct, functools, inspect, itertools
@@ -117,7 +117,6 @@ def nidx(b:mesa.nir_builder, buf, off, space, itemsize, gate=None) -> mesa.nir_d
 class NIRRenderer(Renderer):
   suffix = "NIR"
   nir_options: bytes
-  supports_ranges = False
   global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
   code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
 
@@ -188,7 +187,8 @@ class NIRRenderer(Renderer):
     self.prerender(uops)
     for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg
     self.r: dict[UOp, Any] = {}
-    self.param_idx, loop_ifs = 0, []
+    self.param_idx = 0
+    ranges: list[mesa.nir_def|None] = []
 
     for u in uops:
       if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass
@@ -204,17 +204,29 @@ class NIRRenderer(Renderer):
         self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
         self.b.shader.contents.info.shared_size += u.max_numel()*u.dtype.itemsize
       elif u.op == Ops.RANGE:
-        # ranges are rewritten to loops in codegen: just open the loop, the END adds the conditional backedge
-        # a bool src is a one-time entry guard for possibly zero trip counts
-        assert u.dtype == dtypes.void, "NIRRenderer does not support ranges"
-        guard = next((s for s in u.src if s.dtype is dtypes.bool), None)
-        loop_ifs.append(mesa.nir_push_if(self.b, self.r[guard]) if guard is not None else None)
-        mesa.nir_push_loop(self.b)
+        if u.dtype == dtypes.void:
+          # a RANGE with no bound is a loop header: just open the loop, the END adds the conditional backedge
+          ranges.append(None)
+          mesa.nir_push_loop(self.b)
+        else:
+          ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{range_str(u)}".encode()).contents))
+          nstore(self.b, AddrSpace.REG, i, nimm(self.b, 0, u.dtype))
+          mesa.nir_push_loop(self.b)
+          self.r[u] = nload(self.b, AddrSpace.REG, i, u)
+          nif(self.b, nalu(self.b, "ilt", self.r[u], self.r[u.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
       elif u.op == Ops.END:
-        # loop again while the condition is true
-        nif(self.b, self.r[u.src[2]], lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
-        mesa.nir_pop_loop(self.b, None)
-        if (nif_ref:=loop_ifs.pop()) is not None: mesa.nir_pop_if(self.b, nif_ref)
+        r = u.src[1]
+        if r.dtype == dtypes.void:
+          # loop again while the condition is true
+          nif(self.b, self.r[u.src[2]], lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
+          ranges.pop()
+          mesa.nir_pop_loop(self.b, None)
+        else:
+          next_i = nalu(self.b, "iadd", self.r[r], nimm(self.b, 1, r.dtype))
+          # TODO: this nif should be removable ... but TestMultiTensor.test_double_matmul_shard_W_0 segfaults with it gone
+          nif(self.b, nalu(self.b, "ilt", next_i, self.r[r.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
+          nstore(self.b, AddrSpace.REG, ranges.pop(), next_i),
+          mesa.nir_pop_loop(self.b, None)
       else:
         d: mesa.nir_def|None = self.def_rewrite.rewrite(u, ctx=self)
         if d is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
@@ -270,9 +282,7 @@ _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2
   srcs=lambda b,img,idx_y,idx_x:[nsrc(x) for x in [img, tovec(b, idx_y, idx_x), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
       lambda b,img,idx_y,idx_x,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
 
-class IR3Renderer(NIRRenderer, OpenCLRenderer):
-  has_aux = True
-
+class IR3Renderer(NIRRenderer):
   def nload_img(ctx,img,idx_y,idx_x):
     ctx.texs.add(img)
     return _nload_img(ctx.b, ctx.r[img], ctx.r[idx_y], ctx.r[idx_x], img.dtype)

tinygrad/renderer/ptx.py

@@ -116,13 +116,18 @@ string_rewrite = PatternMatcher([
   # simple
   (UPat(Ops.BUFFER, name="x"), lambda ctx, x: [] if x.addrspace == AddrSpace.REG else [
     f".shared .align 16 .b8 local{x.arg.slot}[{x.max_numel()*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg.slot}[0];"]),
-  # loop (ranges are rewritten to loops in codegen), a bool src is a one-time entry guard for possibly zero trip counts
-  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx, l:
-    [f"@!{ctx.r[g]} bra WAITLOOP_EXIT_{ctx.uops.index(l)};", f"WAITLOOP_{ctx.uops.index(l)}:"] \
-    if (g:=next((s for s in l.src if s.dtype is dtypes.bool), None)) is not None else f"WAITLOOP_{ctx.uops.index(l)}:"),
+  (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx, l: f"WAITLOOP_{ctx.uops.index(l)}:"),
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx, l, c:
-    [f"@{ctx.r[c]} bra WAITLOOP_{ctx.uops.index(l)};"] +
-    ([f"WAITLOOP_EXIT_{ctx.uops.index(l)}:"] if any(s.dtype is dtypes.bool for s in l.src) else [])),
+    f"@{ctx.r[c]} bra WAITLOOP_{ctx.uops.index(l)};"),
+  (UPat(Ops.RANGE, name="r"), lambda ctx, r: [
+    f"mov.u32 {ctx.r[r]}, -1;",
+    f"bra END_{ctx.r[r][1:]};",
+    "LOOP_" + f"{ctx.r[r][1:]}:"]),
+  (UPat(Ops.END, name="x", src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda ctx, x, r: [
+    "END_" + f"{ctx.r[r][1:]}:",
+    ctx.code_for_op[Ops.ADD](ctx.r[r], ctx.r[r], "1", dtypes.int, ctx.types[dtypes.int]),
+    ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[r], ctx.r[r.src[0]], dtypes.int, ctx.types[dtypes.int]),
+    f"@{ctx.r[x]} bra LOOP_{ctx.r[r][1:]};"]),
   (UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"),
   (UPat(Ops.ENDIF, name="x"), lambda ctx, x: f"IF_{ctx.r[x.src[0].src[0]][1:]}_{ctx.uops.index(x.src[0])}:"),
   (UPat(Ops.WMMA, name="x"), lambda ctx, x: list(render_wmma(ctx, x))),
@@ -131,7 +136,6 @@ string_rewrite = PatternMatcher([
 
 class PTXRenderer(Renderer):
   suffix = "PTX"
-  supports_ranges = False
   global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
   tc_sm80 = [x for x in tc.cuda_sm80 if x.dtype_in in [dtypes.half, dtypes.float]]
   code_for_op = asm_for_op

tinygrad/runtime/ops_amd.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices, hcq_profile
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec
+from tinygrad.device import Compiled, BufferSpec, TinyELF
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
 from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
@@ -557,10 +557,10 @@ class AMDCopyQueue(HWQueue):
 
     sdma_queue.signal_doorbell(dev)
 
-class AMDProgram(HCQProgram):
-  def __init__(self, dev:AMDDevice, name:str, lib:bytes, **kwargs):
+class AMDProgram(HCQProgram['AMDDevice']):
+  def __init__(self, dev:AMDDevice, obj:TinyELF):
     # TODO; this API needs the type signature of the function and global_size/local_size
-    self.dev, self.name, self.lib = dev, name, lib
+    self.dev, self.name, self.lib = dev, obj.name, obj.lib
 
     image, sections, relocs = elf_loader(self.lib)
 
@@ -608,17 +608,17 @@ class AMDProgram(HCQProgram):
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
                wait=False, timeout:int|None=None):
-    if self.dev.sqtt_enabled: cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
+    if self.dev.sqtt_enabled: cast(AMDComputeQueue, unwrap(self.dev.hw_compute_queue_t)()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
     res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
     if self.dev.pmc_enabled:
-      cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).pmc_read(self.dev.pmc_buffer, self.dev.pmc_sched) \
-                                                          .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+      cast(AMDComputeQueue, unwrap(self.dev.hw_compute_queue_t)()).pmc_read(self.dev.pmc_buffer, self.dev.pmc_sched) \
+                                                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
       self.dev.allocator._copyout(pmc_buf:=memoryview(bytearray(self.dev.pmc_buffer.size)), self.dev.pmc_buffer)
       Compiled.profile_events += [ProfilePMCEvent(self.dev.device, self.prof_prg_counter, self.dev.pmc_sched, bytes(pmc_buf),
                                                   self.dev.prof_exec_counter)]
     if self.dev.sqtt_enabled:
-      cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_stop(self.dev.sqtt_wptrs) \
-                                                          .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+      cast(AMDComputeQueue, unwrap(self.dev.hw_compute_queue_t)()).sqtt_stop(self.dev.sqtt_wptrs) \
+                                                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
       self.dev.synchronize()
 
       for se, buf in enumerate(self.dev.sqtt_buffers):
@@ -991,7 +991,7 @@ class AMDDevice(HCQCompiled):
     self.sdma_queues:dict = {}
     self.has_sdma_queue = self.sdma_queue(0) is not None
 
-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], functools.partial(AMDProgram, self), AMDSignal,
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], AMDProgram, AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
                      functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size) if self.has_sdma_queue else None,
                      kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000,

tinygrad/runtime/ops_cl.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 from typing import cast
-import ctypes, functools, hashlib
+import ctypes, hashlib
 from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.runtime.support import c
 from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing, unwrap, round_up, is_image_shape
 from tinygrad.renderer.cstyle import OpenCLRenderer
-from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, TinyELF, Program
 
 CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
 BP_CB = c.CFUNCTYPE[None, [cl.cl_program, c.POINTER[None]]]
@@ -36,15 +36,15 @@ class CLCompiler(Compiler):
     check(cl.clReleaseProgram(program))
     return bytes(binary)
 
-class CLProgram:
-  def __init__(self, device:CLDevice, name:str, lib:bytes, arg_dtypes=[], **kwargs):
-    self.dev, self.name, self.lib, self.arg_dtypes = device, name, device.cl_compiler.compile_cached(lib.decode()), arg_dtypes
+class CLProgram(Program['CLDevice']):
+  def __init__(self, device:CLDevice, obj:TinyELF):
+    self.dev, self.lib, self.signature = device, device.cl_compiler.compile_cached(obj.lib.decode()), obj.signature
     self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(self.lib)),
                                                         to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                         errcode_ret := ctypes.c_int32()), errcode_ret)
     check(binary_status.value)
     check(cl.clBuildProgram(self.program, 1, device.device_id, None, BP_CB(), None)) # NOTE: OSX requires this
-    self.kernel = checked(cl.clCreateKernel(self.program, name.encode(), status := ctypes.c_int32()), status)
+    self.kernel = checked(cl.clCreateKernel(self.program, obj.name.encode(), status := ctypes.c_int32()), status)
 
   def __del__(self):
     try: check(cl.clReleaseKernel(self.kernel))
@@ -54,17 +54,15 @@ class CLProgram:
 
   def __call__(self, *bufs:cl.cl_mem, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, vals:tuple[int, ...]=(),
                wait=False, **kw) -> float|None:
-    i = 0
-    for i,b in enumerate(bufs):
-      for real_i, dt, shape in self.arg_dtypes[i]:
-        if is_image_shape(shape):
-          pitch = (round_up(shape[1], 256) if OSX else shape[1]) * 4 * dt.itemsize
-          fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])
-          desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, shape[1], shape[0], image_row_pitch=pitch, buffer=b)
-          img = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status)
-          check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(img), ctypes.byref(img)))
-        else: check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(b), ctypes.byref(b)))
-    for i,v in enumerate(vals,start=i+1): check(cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v))))
+    for i, (_, slot, dt, shape) in enumerate(self.signature):
+      b = bufs[slot] if slot < len(bufs) else ctypes.c_int32(vals[slot-len(bufs)])
+      if is_image_shape(shape):
+        pitch = (round_up(shape[1], 256) if OSX else shape[1]) * 4 * dt.itemsize
+        fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])
+        desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, shape[1], shape[0], image_row_pitch=pitch, buffer=b)
+        img = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status)
+        check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(img), ctypes.byref(img)))
+      else: check(cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b)))
     if local_size is not None: global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size)))
     event = cl.cl_event() if wait else None
     check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size),
@@ -123,7 +121,7 @@ class CLDevice(Compiled):
     if "cl_khr_image2d_from_buffer" in self.device_exts:
       check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
       arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}"
-    super().__init__(device, CLAllocator(self), [OpenCLRenderer], functools.partial(CLProgram, self), arch=arch)
+    super().__init__(device, CLAllocator(self), [OpenCLRenderer], CLProgram, arch=arch)
 
   def count(self) -> int: return len(unwrap(self.device_ids))
 

tinygrad/runtime/ops_cpu.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 import platform, sys, os, ctypes, functools, mmap, threading, array
+from dataclasses import replace
+from typing import cast
 from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le
-from tinygrad.device import BufferSpec
+from tinygrad.device import Buffer, BufferSpec, TinyELF
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
 from tinygrad.renderer.cstyle import ClangRenderer
@@ -64,18 +66,20 @@ class CPUComputeQueue(HWQueue):
     args:list[sint|None] = [args_state.buf.va_addr] if lvp else [*[x.va_addr for x in args_state.bufs], *args_state.vals]
     assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
     for tid in range(1 if lvp else (global_size or (1,))[0]):
-      if not lvp and 'core_id' in prg.runtimevars: args[len(args_state.bufs)+prg.runtimevars['core_id']] = tid
+      if not lvp and 'core_id' in prg.runtimevars: args[prg.runtimevars['core_id']] = tid
       self.q(prg, *[unwrap(x) for x in args], *([0] * (MAX_ARGS - len(args))))
     return self
   def wait(self, signal, value=0): return self._cmd(wait_prog, (signal.base_buf,), (value,))
-  def timestamp(self, signal): return self._cmd(timestamp_prog, (signal.base_buf.offset(8, 8), self.dev.func_table.offset(0, 8)))
+  def timestamp(self, signal): return self._cmd(timestamp_prog, (signal.base_buf.offset(8, 8), self.dev.func_table._buf.offset(0, 8)))
   def signal(self, signal, value:sint=0): return self._cmd(signal_prog, (signal.base_buf,), (value,))
   def _submit(self, dev):
+    dev.ensure_worker()
+    ring_view = dev.ring.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
     for off in range(0, len(self._q), CMD_SIZE):
       entry = [self._q[off].addr, *self._q[off+1:off+CMD_SIZE]]
-      dev.ring_view[(base:=(dev.ring_pos % RING_SLOTS) * CMD_SIZE):base+CMD_SIZE] = array.array('Q', (int(x) & ((1<<64)-1) for x in entry))
+      ring_view[(base:=(dev.ring_pos % RING_SLOTS) * CMD_SIZE):base+CMD_SIZE] = array.array('Q', (int(x) & ((1<<64)-1) for x in entry))
       dev.ring_pos += 1
-      if WIN: dev.sys_view[0] = dev.ring_pos
+      if WIN: dev.sys.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = dev.ring_pos
       else: assert libc.sem_post(dev.sem) == 0
 
 class LVPArgsState(CLikeArgsState):
@@ -84,32 +88,33 @@ class LVPArgsState(CLikeArgsState):
 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800
 
-class CPUProgram(HCQProgram):
+class CPUProgram(HCQProgram['CPUDevice']):
   rt_lib = None
   try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
   except OSError: pass
 
-  def __init__(self, dev, name:str, lib:bytes, runtimevars:dict[str, int]|None=None, native=False, **kwargs):
-    self.runtimevars = runtimevars or {}
+  def __init__(self, dev:CPUDevice, obj:TinyELF):
+    self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
 
-    LVP = isinstance(dev.renderer, LVPRenderer) and not native
+    LVP = obj.target.renderer == "LVP"
     if sys.platform == "win32": # mypy doesn't understand when WIN is used here
       PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
       ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
-      self.addr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(lib)), MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE)
-      ctypes.memmove(self.addr, lib, len(lib))
+      self.addr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(obj.lib)), MEM_COMMIT | MEM_RESERVE,
+                                                      PAGE_EXECUTE_READWRITE)
+      ctypes.memmove(self.addr, obj.lib, len(obj.lib))
       ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
       proc = ctypes.windll.kernel32.GetCurrentProcess()
-      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(lib)))
+      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(obj.lib)))
       self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
     else:
       # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
       # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
-      self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
+      self.mem = mmap.mmap(-1, len(obj.lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
       self.addr = mv_address(self.mem)
 
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
-      if LVP: lib = jit_loader(lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m'])
+      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if LVP else obj.lib
       self.mem.write(lib)
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
@@ -124,7 +129,7 @@ class CPUProgram(HCQProgram):
 
       self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
 
-    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, name, kernargs_alloc_size=12+256 if LVP else 0)
+    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj.name, kernargs_alloc_size=12+256 if LVP else 0)
 
   @suppress_finalizing
   def __del__(self):
@@ -145,38 +150,46 @@ class CPUAllocator(HCQAllocator):
 
 class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
-    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], functools.partial(CPUProgram, self),
-      HCQSignal, functools.partial(CPUComputeQueue, self), arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram, HCQSignal,
+      functools.partial(CPUComputeQueue, self), arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
 
-    self.ring = self.allocator.alloc(RING_SLOTS * CMD_SIZE * 8, BufferSpec())
-    self.ring_view, self.ring_pos = self.ring.cpu_view().view(fmt='Q'), 0
+    self.ring_pos = 0
 
     # posix uses sem to put cpus into sleep
-    if WIN:
-      self.sys = self.allocator.alloc(8, BufferSpec())
-      self.sys_view, sem_addr = self.sys.cpu_view().view(fmt='Q'), 0
-    else:
+    self.sem_addr = 0
+    if not WIN:
       self.sem = libc.sem_open(sem_name:=f"/tinygrad-{os.getpid()}-{id(self):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg]
-      if (sem_addr:=unwrap(ctypes.cast(self.sem, ctypes.c_void_p).value)) == ctypes.c_void_p(-1).value or libc.sem_unlink(sem_name):
-        raise OSError(ctypes.get_errno(), "semaphore")
-      self.sem_buf = HCQBuffer(sem_addr, 1, owner=self)
-
-    # TODO: move to hcq2 infra
-    self.func_table = self.allocator.alloc(32, BufferSpec())
-    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0] if WIN else  # type: ignore[attr-defined]
-           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close])
-    self.func_table.cpu_view().view(fmt='Q')[:] = array.array('Q', [unwrap(ctypes.cast(f, ctypes.c_void_p).value) if f else 0 for f in fns])
+      self.sem_addr = unwrap(ctypes.cast(self.sem, ctypes.c_void_p).value)
+      if self.sem_addr == ctypes.c_void_p(-1).value or libc.sem_unlink(sem_name): raise OSError(ctypes.get_errno(), "semaphore")
 
     # TODO: move to hcq2
     with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
       prgs = {f: f().sink(arg=KernelInfo(f.__name__), tag=1) for f in (signal_prog, wait_prog, timestamp_prog, quit_prog, worker_prog)}
-      self.prgs = {f: self.runtime(f.__name__, do_to_program(v, ClangRenderer(self.renderer.target)).src[3].arg, native=True) for f,v in prgs.items()}
+      self.prgs = {f: self.runtime(do_to_program(v, ClangRenderer(replace(self.renderer.target, renderer="CLANG"))).to_elf()) for f,v in prgs.items()}
 
-    self.worker:threading.Thread|None = threading.Thread(target=self.prgs[worker_prog].fxn, args=(ctypes.c_uint64(self.ring.va_addr),
-      ctypes.c_uint64(self.sys.va_addr if WIN else self.func_table.va_addr+16), ctypes.c_uint64(sem_addr)), daemon=True)
-    self.worker.start()
+  @functools.cached_property
+  def ring(self) -> Buffer: return Buffer(self.device, RING_SLOTS * CMD_SIZE, dtypes.uint64, preallocate=True)
+  @functools.cached_property
+  def sys(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint64, preallocate=True)
+  @functools.cached_property
+  def sem_buf(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint8, options=BufferSpec(external_ptr=self.sem_addr), preallocate=True)
+
+  # TODO: move to hcq2 infra
+  @functools.cached_property
+  def func_table(self) -> Buffer:
+    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0] if WIN else  # type: ignore[attr-defined]
+           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close])
+    addrs = array.array('Q', [unwrap(ctypes.cast(f, ctypes.c_void_p).value) if f else 0 for f in fns])
+    (ft:=Buffer(self.device, len(fns), dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[:] = addrs
+    return ft
+
+  @functools.cache
+  def ensure_worker(self):
+    threading.Thread(target=cast(CPUProgram, self.prgs[worker_prog]).fxn, daemon=True, args=[ctypes.c_uint64(x) for x in
+      [self.ring._buf.va_addr, self.sys._buf.va_addr if WIN else self.func_table._buf.va_addr+16, self.sem_addr]]).start()
 
   def finalize(self):
-    if self.worker is None: return
-    CPUComputeQueue(self)._cmd(quit_prog, (self.func_table.offset(8, 8),) if WIN else (self.func_table.offset(8, 24), self.sem_buf)).submit(self)
-    self.worker = None
+    if self.ring_pos == 0: return # the worker starts with the first submit
+    ft = self.func_table._buf
+    CPUComputeQueue(self)._cmd(quit_prog, (ft.offset(8, 8),) if WIN else (ft.offset(8, 24), self.sem_buf._buf)).submit(self)
+    self.ring_pos = 0

tinygrad/runtime/ops_cuda.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
-import ctypes, functools
+import ctypes
 from tinygrad.helpers import DEBUG, DEV, getenv, mv_address, suppress_finalizing
-from tinygrad.device import Compiled, BufferSpec, LRUAllocator
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator, Program, TinyELF
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
@@ -33,18 +33,18 @@ def cu_time_execution(cb, enable=False) -> float|None:
   for ev in evs: cuda.cuEventDestroy_v2(ev)
   return ret.value * 1e-3
 
-class CUDAProgram:
-  def __init__(self, dev:CUDADevice, name:str, lib:bytes, smem:int=0, **kwargs):
-    self.dev, self.name, self.lib, self.smem = dev, name, lib, smem
-    if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(lib.decode('utf-8')).split("\n"))]))
+class CUDAProgram(Program['CUDADevice']):
+  def __init__(self, dev:CUDADevice, obj:TinyELF, smem:int=0):
+    self.dev, self.name, self.lib, self.smem = dev, obj.name, obj.lib, smem
+    if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(obj.lib.decode('utf-8')).split("\n"))]))
 
     check(cuda.cuCtxSetCurrent(self.dev.context))
     self.module = cuda.CUmodule()
-    status = cuda.cuModuleLoadData(ctypes.byref(self.module), lib)
+    status = cuda.cuModuleLoadData(ctypes.byref(self.module), obj.lib)
     if status != 0:
       del self.module
       raise RuntimeError(f"module load failed with status code {status}: {cuda.enum_cudaError_enum.get(status)}")
-    check(cuda.cuModuleGetFunction(ctypes.byref(prg := cuda.CUfunction()), self.module, name.encode("utf-8")))
+    check(cuda.cuModuleGetFunction(ctypes.byref(prg := cuda.CUfunction()), self.module, obj.name.encode("utf-8")))
     self.prg = prg
     if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem))
 
@@ -117,8 +117,8 @@ class CUDADevice(Compiled):
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
-    super().__init__(device, CUDAAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer], functools.partial(CUDAProgram, self),
-                     None if MOCKGPU else CUDAGraph, arch=f"sm_{major.value}{minor.value}")
+    super().__init__(device, CUDAAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer], CUDAProgram, None if MOCKGPU else CUDAGraph,
+                     arch=f"sm_{major.value}{minor.value}")
 
   def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(cuda.cuDeviceGetCount(ctypes.byref(x)))).value
 

tinygrad/runtime/ops_dsp.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
-import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
+import ctypes, os, mmap, tempfile, pathlib, array, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target
@@ -74,9 +74,9 @@ def rpc_prep_args(ins=None, outs=None, in_fds=None):
   for i, mv in enumerate(ins + outs): pra[i].buf.pv, pra[i].buf.len = ctypes.c_void_p(mv_address(mv) if mv.nbytes > 0 else 0), mv.nbytes
   return pra, fds, attrs, (ins, outs)
 
-class DSPProgram:
-  def __init__(self, dev:DSPDevice, name:str, lib:bytes, **kwargs):
-    self.dev, self.lib = dev, lib
+class DSPProgram(Program['DSPDevice']):
+  def __init__(self, dev:DSPDevice, obj:TinyELF):
+    self.dev, self.lib = dev, obj.lib
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
@@ -144,7 +144,7 @@ class DSPDevice(Compiled):
     if getenv("MOCKDSP"): super().__init__(device, DSPAllocator(self), [MockDSPRenderer], MockDSPProgram)
     else:
       self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
-      super().__init__(device, DSPAllocator(self), [DSPRenderer], functools.partial(DSPProgram, self))
+      super().__init__(device, DSPAllocator(self), [DSPRenderer], DSPProgram)
       fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
       self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
       ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
@@ -287,8 +287,8 @@ class MockDSPRenderer(DSPRenderer):
     msrc.append('exit(0); }')
     return '\n'.join(msrc)
 
-class MockDSPProgram:
-  def __init__(self, name:str, lib:bytes, **kwargs): self.lib = lib
+class MockDSPProgram(Program[DSPDevice]):
+  def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib = obj.lib
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
       dsp_lib.write(self.lib)

tinygrad/runtime/ops_hip.py

@@ -1,6 +1,6 @@
-import ctypes, functools
+import ctypes
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import Compiled, LRUAllocator, BufferSpec
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec, Program, TinyELF
 from tinygrad.runtime.autogen import hip
 from tinygrad.renderer.cstyle import HIPRenderer
 from tinygrad.runtime.support.c import init_c_var, init_c_struct_t
@@ -15,7 +15,7 @@ class HIPDevice(Compiled):
     self.arch = init_c_var(hip.hipDeviceProp_t, lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
     self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t, lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
 
-    super().__init__(device, HIPAllocator(self), [HIPRenderer], functools.partial(HIPProgram, self), arch=self.arch)
+    super().__init__(device, HIPAllocator(self), [HIPRenderer], HIPProgram, arch=self.arch)
 
   def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(hip.hipGetDeviceCount(x))).value
 
@@ -23,12 +23,12 @@ class HIPDevice(Compiled):
     check(hip.hipSetDevice(self.device_id))
     check(hip.hipDeviceSynchronize())
 
-class HIPProgram:
-  def __init__(self, dev:HIPDevice, name:str, lib:bytes, **kwargs):
-    self.dev, self.name, self.lib = dev, name, lib
+class HIPProgram(Program[HIPDevice]):
+  def __init__(self, dev:HIPDevice, obj:TinyELF):
+    self.dev, self.name, self.lib = dev, obj.name, obj.lib
     check(hip.hipSetDevice(self.dev.device_id))
-    self.module = init_c_var(hip.hipModule_t, lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), lib)))
-    self.prg = init_c_var(hip.hipFunction_t, lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, name.encode("utf-8"))))
+    self.module = init_c_var(hip.hipModule_t, lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), obj.lib)))
+    self.prg = init_c_var(hip.hipFunction_t, lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, obj.name.encode("utf-8"))))
 
   @suppress_finalizing
   def __del__(self):

tinygrad/runtime/ops_metal.py

@@ -1,7 +1,7 @@
 import subprocess, pathlib, struct, ctypes, tempfile, functools, decimal, platform
 from tinygrad.helpers import prod, to_mv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
 import tinygrad.runtime.support.objc as objc
-from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent
+from tinygrad.device import Compiled, Compiler, CompileError, Program, TinyELF, LRUAllocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support.c import DLL
@@ -45,9 +45,9 @@ class MetalDevice(Compiled):
     from tinygrad.runtime.graph.metal import MetalGraph
     # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
     # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), [MetalRenderer],
-      functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None,
-      arch=metal.enum_MTLGPUFamily[check_family("Apple") or check_family("Mac")][12:])
+    super().__init__(device, MetalAllocator(self), [MetalRenderer], MetalProgram,
+                     MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None,
+                     arch=metal.enum_MTLGPUFamily[check_family("Apple") or check_family("Mac")][12:])
 
   def synchronize(self):
     for cbuf in self.mtl_buffers_in_flight:
@@ -111,13 +111,13 @@ class MetalCompiler(Compiler):
       ret = proc.wait()
       if ret: print("Disassembler Error: Make sure you have https://github.com/dougallj/applegpu cloned to tinygrad/extra/disassemblers/applegpu")
 
-class MetalProgram:
-  def __init__(self, dev:MetalDevice, name:str, lib:bytes, **kwargs):
-    self.dev, self.name, self.lib = dev, name, lib
-    data = objc.dispatch_data_create(lib, len(lib), None, None)
+class MetalProgram(Program[MetalDevice]):
+  def __init__(self, dev:MetalDevice, obj:TinyELF):
+    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+    data = objc.dispatch_data_create(obj.lib, len(obj.lib), None, None)
     self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
     error_check(error_lib)
-    self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
+    self.fxn = self.library.newFunctionWithName(to_ns_str(obj.name)).retained()
     descriptor = metal.MTLComputePipelineDescriptor.new()
     descriptor.setComputeFunction(self.fxn)
     descriptor.setSupportIndirectCommandBuffers(True)

tinygrad/runtime/ops_null.py

@@ -1,5 +1,5 @@
-import inspect, functools, math
-from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent
+import inspect, math
+from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl
 from tinygrad.renderer.cstyle import CStyleLanguage
@@ -16,8 +16,8 @@ class NullRenderer(CStyleLanguage):
     from tinygrad.renderer.amd.elf import assemble_linear
     return assemble_linear(prg, lin, self.target.arch)
 
-class NullProgram:
-  def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, name
+class NullProgram(Program['NullDevice']):
+  def __init__(self, dev:'NullDevice', obj:TinyELF): self.device, self.name = dev.device, obj.name
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     with cpu_profile(self.name, self.device): return 1e-3
 
@@ -54,4 +54,4 @@ class NullDevice(Compiled):
       "EMULATE is deprecated, use DEV=NULL:HIP:"+{"AMD":"gfx1100", "AMD_RDNA4":"gfx1201", "AMD_CDNA4":"gfx950"}.get(emu, "<arch>")
     renderers = [NullRenderer] + [r for m in [cstyle, nir, ptx, llvmir, wgsl] for r in m.__dict__.values()
                                   if inspect.isclass(r) and issubclass(r, Renderer)]
-    super().__init__(device, NullAllocator(self), dedup(renderers), functools.partial(NullProgram, device), NullGraph)
+    super().__init__(device, NullAllocator(self), dedup(renderers), NullProgram, NullGraph)

tinygrad/runtime/ops_nv.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, hcq_filter_visible_devices, hcq_profile
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec
+from tinygrad.device import Compiled, BufferSpec, TinyELF
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, hi32, lo32, PROFILE, ContextVar, VIZ, ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
@@ -244,14 +244,14 @@ class NVArgsState(CLikeArgsState):
     super().__init__(buf, prg, bufs, vals=vals, prefix=prg.cbuf_0 or None)
 
 class NVProgram(HCQProgram['NVDevice']):
-  def __init__(self, dev:NVDevice, name:str, lib:bytes, **kwargs):
-    self.dev, self.name, self.lib = dev, name, lib
+  def __init__(self, dev:NVDevice, obj:TinyELF):
+    self.dev, self.name, self.lib = dev, obj.name, obj.lib
     self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]]
 
     if (NAK:=isinstance(dev.renderer, NAKRenderer)):
-      image, self.cbuf_0 = memoryview(bytearray(lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(lib)):])), []
+      image, self.cbuf_0 = memoryview(bytearray(obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):])), []
       self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
-    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(lib) + b'\x00' * (4 - len(lib)%4)).cast("I"), [], [] # type: ignore
+    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(obj.lib) + b'\x00' * (4 - len(obj.lib)%4)).cast("I"), [], [] # type: ignore
     else: image, sections, relocs = elf_loader(self.lib, force_section_align=128)
     # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
     self.lib_gpu = self.dev.allocator.alloc(round_up((prog_sz:=image.nbytes), 0x1000) + 0x1000, buf_spec:=BufferSpec(nolru=True))
@@ -266,7 +266,7 @@ class NVProgram(HCQProgram['NVDevice']):
           self.constbufs[int(m.group(1))] = (self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size)
         elif sh.name.startswith(".nv.info"):
           for typ, param, data in self._parse_elf_info(sh):
-            if sh.name == f".nv.info.{name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
+            if sh.name == f".nv.info.{obj.name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
             elif sh.name == ".nv.info" and param == 0x12: self.lcmem_usage = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
             elif sh.name == ".nv.info" and param == 0x2f: self.regs_usage = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
 
@@ -630,8 +630,8 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], functools.partial(NVProgram, self), NVSignal,
-                     NVComputeQueue, NVCopyQueue, arch=self.arch)
+    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], NVProgram, NVSignal, NVComputeQueue,
+                     NVCopyQueue, arch=self.arch)
 
     self.pma_enabled = PMA.value > 0 and PROFILE >= 1
     if self.pma_enabled: self._prof_init()

tinygrad/runtime/ops_python.py

@@ -7,7 +7,7 @@ import pickle, base64, itertools, time, sys, functools
 from dataclasses import replace
 from tinygrad.dtype import DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
-from tinygrad.device import Buffer, Compiled, Compiler, Allocator
+from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast
 from tinygrad.renderer import Renderer
@@ -39,9 +39,9 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_
         out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
   return out
 
-class PythonProgram:
-  def __init__(self, name:str, lib:bytes, **kwargs):
-    self.uops: list[UOp] = pickle.loads(lib)
+class PythonProgram(Program['PythonDevice']):
+  def __init__(self, dev:'PythonDevice', obj:TinyELF):
+    self.uops: list[UOp] = pickle.loads(obj.lib)
     self.uop_to_index: dict[UOp, int] = {u:i for i,u in enumerate(self.uops)}
     self.loop_ends: dict[UOp, int] = {u.src[1]:i for i, u in enumerate(self.uops) if u.op == Ops.END}
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):

tinygrad/runtime/ops_qcom.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
 from typing import Any, cast
-from tinygrad.device import BufferSpec, Device
+from tinygrad.device import BufferSpec, Device, TinyELF
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa
@@ -25,7 +25,7 @@ def dcache_flush():
   flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
   sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"))
   prg = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer)
-  return Device["CPU"].runtime(prg.arg.function_name, prg.src[3].arg)
+  return Device["CPU"].runtime(prg.to_elf())
 
 #Parse C-style defines: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:
 # qreg.<regname>(<field_x>=..., <field_y>=..., ..., <field_n>=...)
@@ -200,8 +200,8 @@ class QCOMArgsState(HCQArgsState):
     super().__init__(buf, prg, bufs, vals=vals)
     ctypes.memset(int(self.buf.va_addr), 0, prg.kernargs_alloc_size)
 
-    ubos = [b for i,b in enumerate(bufs) for _,dt,shape in prg.buf_dtypes[i] if not is_image_shape(shape)]
-    uavs = [(dt,shape,b) for i,b in enumerate(bufs) for _,dt,shape in prg.buf_dtypes[i] if is_image_shape(shape)]
+    ubos = [bufs[slot] for _,slot,_,shape in prg.signature if slot < len(bufs) and not is_image_shape(shape)]
+    uavs = [(dt,shape,bufs[slot]) for _,slot,dt,shape in prg.signature if slot < len(bufs) and is_image_shape(shape)]
     # NIR can reorder images to different texture slots
     ibos, texs = uavs[:prg.ibo_cnt], [uavs[prg.ibo_cnt + (prg.tex_to_image[i] if prg.NIR else i)] for i in range(prg.tex_cnt)]
     for cnst_val,cnst_off,cnst_sz in prg.consts_info:
@@ -227,14 +227,14 @@ class QCOMArgsState(HCQArgsState):
     self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off)
     self.bind_sints_to_buf(*flatten(map(functools.partial(_tex, ibo=True), ibos)), buf=self.buf, fmt='I', offset=prg.ibo_off)
 
-class QCOMProgram(HCQProgram):
-  def __init__(self, dev: QCOMDevice, name: str, lib: bytes, buf_dtypes=[], **kwargs):
+class QCOMProgram(HCQProgram['QCOMDevice']):
+  def __init__(self, dev: QCOMDevice, obj: TinyELF):
     self.dev: QCOMDevice = dev
-    self.buf_dtypes, self.name, self.NIR = buf_dtypes, name, isinstance(dev.renderer, IR3Renderer)
+    self.signature, self.name, self.NIR = obj.signature, obj.name, isinstance(dev.renderer, IR3Renderer)
 
     if self.NIR:
       from tinygrad.runtime.support.compiler_mesa import IR3Compiler
-      v, cs, imm_vals, self.image = IR3Compiler.unpack_lib(lib)
+      v, cs, imm_vals, self.image = IR3Compiler.unpack_lib(obj.lib)
       self.prg_offset, self.brnchstck, self.image_size, self.pvtmem, self.shmem = 0, v.branchstack, v.info.size, v.pvtmem_size, v.shared_size
       self.wgsz = alloc.offset_vec4 * 4 + 8 if (alloc:=cs.allocs.consts[mesa.IR3_CONST_ALLOC_DRIVER_PARAMS]).size_vec4 else 0xfc
 
@@ -252,7 +252,7 @@ class QCOMProgram(HCQProgram):
 
       self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
       self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
-    else: self._parse_lib(lib)
+    else: self._parse_lib(obj.lib)
 
     self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
     to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
@@ -369,8 +369,8 @@ class QCOMDevice(HCQCompiled):
     if PROFILE and self.gpu_id[:2] < (7, 3):
       System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
 
-    super().__init__(device, QCOMAllocator(self), [QCOMCLRenderer, IR3Renderer], functools.partial(QCOMProgram, self), QCOMSignal,
-                     functools.partial(QCOMComputeQueue, self), arch=("a%d%d%d" + (",IMAGE_PITCH_ALIGNMENT=64" if IMAGE else "")) % self.gpu_id)
+    super().__init__(device, QCOMAllocator(self), [QCOMCLRenderer, IR3Renderer], QCOMProgram, QCOMSignal, functools.partial(QCOMComputeQueue, self),
+                     arch=("a%d%d%d" + (",IMAGE_PITCH_ALIGNMENT=64" if IMAGE else "")) % self.gpu_id)
 
   def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
     flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP

tinygrad/runtime/ops_webgpu.py

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import Compiled, Allocator, BufferSpec
+from tinygrad.device import Compiled, Allocator, BufferSpec, Program, TinyELF
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing, getenv, to_mv
 from tinygrad.runtime.autogen import webgpu
@@ -49,12 +49,12 @@ InstanceRequestAdapter = synchronous(webgpu.enum_WGPURequestAdapterStatus, True)
 AdapterRequestDevice = synchronous(webgpu.enum_WGPURequestDeviceStatus, True)(webgpu.wgpuAdapterRequestDevice2)
 QueueOnSubmittedWorkDone = synchronous(webgpu.enum_WGPUQueueWorkDoneStatus)(webgpu.wgpuQueueOnSubmittedWorkDone2)
 
-class WebGPUProgram:
-  def __init__(self, dev:'WebGpuDevice', name:str, lib:bytes, **kwargs):
-    self.dev, self.name = dev, to_wgpu_str(name)
+class WebGPUProgram(Program['WebGpuDevice']):
+  def __init__(self, dev:'WebGpuDevice', obj:TinyELF):
+    self.dev, self.name = dev, to_wgpu_str(obj.name)
 
     # Creating shader module
-    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(lib.decode()),
+    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(obj.lib.decode()),
                                                    chain=webgpu.WGPUChainedStruct(sType=webgpu.WGPUSType_ShaderSourceWGSL))
     module = webgpu.WGPUShaderModuleDescriptor(nextInChain=ctypes.cast(ctypes.pointer(shader), ctypes.POINTER(webgpu.struct_WGPUChainedStruct)))
 
@@ -186,7 +186,7 @@ class WebGpuDevice(Compiled):
 
     webgpu.wgpuAdapterRelease(adapter_res)
 
-    super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], functools.partial(WebGPUProgram, self),
+    super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], WebGPUProgram,
                      arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in self.features))
 
   def synchronize(self): QueueOnSubmittedWorkDone(self.queue)

tinygrad/runtime/support/hcq.py

@@ -6,7 +6,7 @@ try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap
 from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
-from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
+from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.memory import BumpAllocator
@@ -329,7 +329,7 @@ class CLikeArgsState(HCQArgsState[ProgramType]):
     assert None not in vals
     self.bind_sints_to_buf(*cast(tuple[sint, ...], vals), buf=self.buf, fmt='I', offset=len(prefix or []) * 4 + len(bufs) * 8)
 
-class HCQProgram(Generic[HCQDeviceType]):
+class HCQProgram(Program[HCQDeviceType]):
   def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, name:str, kernargs_alloc_size:int, lib:bytes|None=None, base:int|None=None):
     self.args_state_t, self.dev, self.name, self.kernargs_alloc_size = args_state_t, dev, name, kernargs_alloc_size
     self.prof_prg_counter = next(self.dev.prof_prg_counter)
@@ -389,9 +389,9 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
   cpu_devices: list[HCQCompiled] = []
 
-  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime, signal_t:Type[SignalType]|None=None,
-               comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20),
-               sigalloc_size=0x1000, can_recover:bool=False, arch=None):
+  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
+               signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
+               kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph

tinygrad/schedule/__init__.py

@@ -2,7 +2,7 @@ import time, inspect
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
 from tinygrad.uop.spec import type_verify, spec_tensor
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup
 
 # **** schedule linearizer
 
@@ -134,6 +134,39 @@ pm_schedule = PatternMatcher([
   (UPat(Ops.SINK, name="function"), lower_sink_to_linear),
 ])
 
+def assert_all_same_devices(ast:UOp):
+  devices = dedup([x.device for x in ast.toposort() if x.op is Ops.PARAM and x.device is not None])
+  if len(devices) >= 2: raise RuntimeError(f"all buffers must be on the same device: {devices}")
+
+def copy_kernel_to_copy_uop(call:UOp, dst:UOp, src:UOp, r:UOp|None=None):
+  if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
+  return call.replace(src=(UOp(Ops.COPY, dtype=src.dtype, src=(src,), arg=dst.device),) + call.src[1:])
+
+def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp):
+  # NOTE: this is a codegen for SDMA devices
+  if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
+  from tinygrad.codegen.simplify import pm_flatten_range, pm_simplify_ranges
+  from tinygrad.schedule.rangeify import pm_mops
+  from tinygrad.uop.symbolic import sym
+  sink = graph_rewrite(ast, sym+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy")
+  return call.replace(src=(sink,) + call.src[1:])
+
+pm_copy_from_store = PatternMatcher([
+  # simplify copy kernels
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"), UPat.var("dst"), UPat.var("src")), name="call"), simplify_copy_kernel),
+
+  # replace this with a copy if it's a copy
+  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.CONST, arg=0))
+                .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.CONST, arg=0))).sink(),),
+                name="call", allow_any_len=True), copy_kernel_to_copy_uop),
+  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.RANGE, name="r"))
+                .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.RANGE, name="r"))).end(UPat(Ops.RANGE, name="r")).sink(),),
+                name="call", allow_any_len=True), copy_kernel_to_copy_uop),
+
+  # if it wasn't copy, it currently can't be cross device
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices),
+])
+
 @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
 def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
   # big_sink srcs are all the Tensors
@@ -142,6 +175,9 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
   # this recursively resolves the linear_call and allocates buffers
   linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
 
+  # create copies
+  linear = graph_rewrite(linear, pm_copy_from_store, name="create COPY kernels for SDMA")
+
   # vars used in the schedule
   used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
   # get var_vals

tinygrad/schedule/allreduce.py

@@ -1,7 +1,6 @@
 import functools, itertools
 from tinygrad.helpers import all_int, prod, DEBUG, RING, ALL2ALL, getenv
 from tinygrad.uop.ops import UOp
-from tinygrad.dtype import Invalid
 
 # *** allreduce implementation ***
 def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
@@ -57,7 +56,7 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
 
 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
-  if output is None: output = UOp.const(red.dtype, Invalid, shape=red.shape).clone(device=red.device)
+  if output is None: output = UOp.invalids(red.shape, dtype=red.dtype, device=red.device)
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(*red.arg)

tinygrad/schedule/indexing.py

@@ -1,13 +1,13 @@
 from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
-from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
+from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
 
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE,
                       Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                       Ops.LOAD, Ops.CALL, Ops.FUNCTION}
 
@@ -18,8 +18,8 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
 def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
-  # don't realize COPY/SLICE when they are the direct source of STORE+AFTER — the target buffer is the output
-  if src.op in {Ops.COPY, Ops.SLICE} and src in ctx \
+  # don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output
+  if src.op is Ops.SLICE and src in ctx \
      and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
     del ctx[src]
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
@@ -27,9 +27,9 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
 
 pm_generate_realize_map = PatternMatcher([
   # always realize
-  (UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
+  (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
   # realize srcs of these
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
+  (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # sometimes we need to realize the src of STORE if there's a self-access
   (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src),
 ])
@@ -60,14 +60,21 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
   baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape)
   return tuple(r.const_like(0) if j in baxes else r for j,r in enumerate(rngs) if j >= nleft)
 
+# TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
+def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
+  if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
+  if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
+  return src
+
 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
   new_srcs = []
+  # shape/bound/index args that are not data src should not be indexed
+  data_src_count = len(data_srcs(x.op, x.src))
   for i, s in enumerate(x.src):
     new_src = s
     src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
-    # shape args of movement ops are at src[1:] and should not be indexed
     if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
-      if x in ctx.range_map and not (x.op in GroupOp.Movement and i > 0): new_src = new_src.index(*src_rngs)
+      if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]
       assert isinstance(realized_ranges, list), "realize map must contain range list"
@@ -77,8 +84,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
         new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
         del ctx.realize_map[s]
       else:
-        # the Bufferize before a COPY is not removable unless it's a view. there should be a better way to do this
-        removable = (x.op is not Ops.COPY or s.has_buffer_identity()) and s.op not in ALWAYS_CONTIGUOUS
+        removable = s.op not in ALWAYS_CONTIGUOUS
         # LOCAL: None in the device assigns it a number later
         opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
                BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
@@ -95,7 +101,8 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
   if x not in ctx.range_map: return None
   bx = create_bufferize_and_index_based_on_ranges(ctx, x)
   valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
-  return valid.where(bx.src[0], UOp.const(x.dtype, 0))
+  # internal PAD fills with Invalid. bool keeps 0-fill: False is the bool reduce identity and external pad masks need it
+  return valid.where(bx.src[0], UOp.const(x.dtype, 0 if x.dtype == dtypes.bool else Invalid))
 
 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   if x.arg[1] == 0: return None
@@ -180,7 +187,11 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
   # get the consumer map
   with cpu_profile("consumer map in rangeify", "TINY"):
-    consumer_map = consumer_map_from_toposort(tsink_toposort:=tsink.toposort(gate_kernel_sink))
+    tsink_toposort = tsink.toposort(gate_kernel_sink)
+    consumer_map: dict[UOp, dict[UOp, None]] = {x:{} for x in tsink_toposort}
+    for c in tsink_toposort:
+      for x in data_srcs(c.op, c.src):
+        if x in consumer_map: consumer_map[x][c] = None
 
   # explicit rangeify
   ending_ranges: dict[UOp, list[UOp]] = {}
@@ -194,7 +205,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # treat MSTACK/MSELECT like SINK
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 
-    if x.dtype == dtypes.weakint: continue  # TODO: why do I need this?
     ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
     # ranges the consumers iterate that this node broadcasts over
     ended = [rctx.range_map[c][0][i] for c in consumer_map[x] if c in rctx.range_map and c.op in GroupOp.Broadcastable

tinygrad/schedule/multi.py

@@ -1,6 +1,6 @@
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, broadcast_axes, _broadcast_shape
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, Invalid
 from tinygrad.schedule.allreduce import handle_allreduce
 
 # ***** multi rewrite MSELECT/MSTACK *****
@@ -18,10 +18,14 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
       ret.append(apply_shrink(x, i).contiguous())
   return ms.replace(src=tuple(ret))
 
+def lower_broadcast_copy(c:UOp, x:UOp):
+  if not (isinstance(c.device, tuple) and isinstance(x.device, str)): return None
+  if (sx:=x.simplify()).device is None and sx.base.op is Ops.CONST: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
+  return UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device))
+
 replace_allreduce = PatternMatcher([
   # BROADCAST: explicitly expand broadcast copies and combine with MSTACK
-  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x:
-    UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device)) if isinstance(c.device, tuple) and isinstance(x.device, str) else None),
+  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lower_broadcast_copy),
   # COPY_TO_ONE: if copying from multidevice to one, MSELECT the first (TODO: a little from each?)
   (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x:
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
@@ -39,6 +43,34 @@ _early_allreduce = PatternMatcher([
 ])
 if not getenv("LATE_ALLREDUCE", 1): replace_allreduce = _early_allreduce + replace_allreduce
 
+# ***** symbolic multi rewrite (SYMBOLIC_MULTI) *****
+# replaces MULTI/MSELECT/MSTACK with the symbolic _device_num representation:
+#   MULTI(x, axis) -> x._unshard(axis): PAD with _device_num-dependent bounds back to full shape, other shards Invalid
+#   MSELECT(x, i) -> dnum==i ? x : Invalid
+#   MSTACK(srcs)  -> STACK(srcs).index(dnum), lowered to nested dnum==k ? src_k : Invalid
+# the per-device specialization binds _device_num at exec time (unwrap_multi in engine/realize.py)
+
+def _dnum(ndev:int) -> UOp: return UOp.variable("_device_num", 0, ndev-1)
+
+def mselect_to_where(ms:UOp) -> UOp:
+  return _dnum(len(ms.src[0].device)).eq(ms.arg).where(ms.src[0], ms.src[0].const_like(Invalid))
+
+def mstack_to_stack_index(ms:UOp) -> UOp:
+  return UOp(Ops.STACK, src=ms.src).index(_dnum(len(ms.src)))
+
+def index_stack_to_where(stack:UOp, var:UOp) -> UOp:
+  ret = stack.src[0].const_like(Invalid)
+  for k in range(len(stack.src)-1, -1, -1): ret = var.eq(k).where(stack.src[k], ret)
+  return ret
+
+symbolic_multi_pm = PatternMatcher([
+  (UPat(Ops.MULTI, src=(UPat(),), name="multi"), lambda multi: multi.src[0]._unshard(multi.arg)),
+  (UPat(Ops.MSELECT, src=(UPat(),), name="ms"), mselect_to_where),
+  (UPat(Ops.MSTACK, name="ms"), mstack_to_stack_index),
+  # lower INDEX into a value-STACK to nested selects (must fire before the pointer-INDEX spec in codegen)
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="stack"), UPat.var("var"))), index_stack_to_where),
+])
+
 # ***** multi functions *****
 
 def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
@@ -124,7 +156,7 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...]):
   if isinstance(device, str):
     pieces = [multi.src[0].mselect(i).copy_to_device(device) for i in range(len(multi.device))]
     return pieces[0].cat(*pieces[1:], dim=multi.axis)
-  return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device)
+  return multi.src[0]._unshard_fill(multi.axis).allreduce(Ops.ADD, device)
 
 def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).multi(src.axis)
 

tinygrad/schedule/rangeify.py

@@ -4,14 +4,14 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, pm_invalid_reduce_identity
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
+from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
-from tinygrad.schedule.multi import multi_pm
+from tinygrad.schedule.multi import multi_pm, symbolic_multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
 # creation can recurse a lot
@@ -149,11 +149,16 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"),), name="c"),
    lambda c,r: c.replace(src=(r.contiguous(),)) if resolve(r.numel() != r.base.numel(), False) or r.contiguous_view_offset() is None else None),
 
-  # copying mselect to same device is just mselect (no NOOP kernel)
-  (UPat(Ops.COPY, src=(UPat(Ops.MSELECT, name="ms"),), name="copy"), lambda ms,copy: ms if ms.device == copy.device else None),
+  # copy to same device is a no-op
+  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
 
-  # copy only to different device
-  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x.f(Ops.NOOP) if x.device == copy.device else None),
+  # copy on reshape is reshape on copy
+  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
+
+  # reshaping on STORE can be a NOOP
+  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
+                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
+   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
 
   # ** store rules **
 
@@ -184,7 +189,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.NOOP}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.NOOP}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -305,8 +310,6 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   # indexing an after with all fully invalid stores is invalid
   (UPat(Ops.INDEX, src=(UPat(Ops.AFTER, name="after"),), allow_any_len=True, name="idx"),
    lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
-  # copy on CONST is CONST
-  (UPat(Ops.COPY, src=(UPat.cvar("x"),), name="copy"), lambda copy,x: copy.const_like(x.arg)),
   # hack if a noop turned to a const
   (UPat(Ops.NOOP, src=(UPat.cvar("c"),)), lambda c: c),
   # mstack on CONST is CONST
@@ -514,32 +517,47 @@ def split_store(x:UOp) -> UOp|None:
   lctx = LocalAddBufferContext()
   ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
 
-  # SINK requires all buffers on the same device, but COPY is cross-device
-  if ret.op is Ops.STORE: stored = ret.src[1]
-  elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
-  else: raise RuntimeError(f"unknown kernel type {ret.op}")
-  if stored.op is Ops.COPY: ret = stored.replace(src=stored.src + ret.ended_ranges)
-  else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
-
-  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
-  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
-    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
-  return kernel
+  # create the Kernel. NOTE: buffers can be on different devices here now, they are compiled to SDMA copies later by schedule
+  return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values(), *lctx.vars.keys())
 
 split_kernels = PatternMatcher([
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
 ])
 
+def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
+  input_src = copy.src[0]
+  if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
+  input_src = input_src.flatten()
+  if existing_buf is not None:
+    # if the existing buffer is not a full buffer, we can't use it
+    if not existing_buf.has_buffer_identity(after_ok=True): return None
+    # if there's already a buffer, we just use it
+    return existing_buf.flatten().store(input_src)
+  # create the output buffer
+  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  # reshape back to input
+  return buf.after(buf.store(input_src)).reshape(copy.shape)
+
+pm_copy_to_store = PatternMatcher([
+  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
+  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+])
+
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  # SYMBOLIC_MULTI replaces the per-shard multi_pm rules with the symbolic _device_num representation (Stage 1)
+  tsink = graph_rewrite(sink, symbolic_multi_pm if getenv("SYMBOLIC_MULTI") else multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
+  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
+
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf")
+  # Invalid (from internal PAD) in reduce inputs contributes the reduce identity, must run after gate lifting
+  tsink = graph_rewrite(tsink, pm_invalid_reduce_identity, name="reduce invalid to identity")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")

tinygrad/tensor.py

@@ -4,9 +4,9 @@ import time, functools, sys, inspect, pathlib, hashlib, weakref
 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, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
-from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, TracingKey
+from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
-from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, ConstLike
+from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
 from tinygrad.device import Buffer, canonicalize_device
@@ -70,23 +70,17 @@ class Tensor(RandMixin):
     self.is_param:bool = True
 
     # create a UOp from the different types of inputs
-    if isinstance(data, UOp):
-      # if data is dtype.weakint that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
-      if data.dtype == dtypes.weakint: data = _index_to_concrete_int(data)
-    elif data is None:
+    if data is None:
       data = UOp.const(_dtype or dtypes.default_float, 0)
     elif isinstance(data, get_args(ConstType)):
       data = UOp.const(_dtype or dtypes.from_py(data), data)
     elif is_numpy_ndarray(data) and data.shape == ():
       data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
-    else:
+    elif not isinstance(data, UOp):
       if _dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {_dtype}")
       if isinstance(data, bytes): data = UOp._frompy(data, _dtype or dtypes.uint8, _device)
       elif isinstance(data, (list, tuple)):
-        if _dtype is None:
-          if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
-          else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
-        data = UOp._frompy(data, _dtype, _device)
+        data = UOp._frompy(data, _dtype or dtypes.from_py(data), _device)
       elif is_numpy_ndarray(data):
         data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)
       elif isinstance(data, pathlib.Path):
@@ -179,7 +173,9 @@ class Tensor(RandMixin):
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
     """Creates the LINEAR UOp needed to realize these Tensor(s), with Variables."""
-    if any(t.dtype in dtypes.weaks for t in (self,)+lst): raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
+    # 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]))
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)

tinygrad/uop/ops.py

@@ -4,13 +4,13 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
 from dataclasses import dataclass, replace
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
-from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype
 from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
+from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey
-from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY, DISALLOW_BROADCAST, get_shape, fully_flatten
-from tinygrad.helpers import colored, ansilen, printable
+from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY, DISALLOW_BROADCAST, get_shape, fully_flatten, to_tuple
+from tinygrad.helpers import colored, ansilen, printable, Target
 if TYPE_CHECKING:
   from tinygrad.renderer import Estimates
 
@@ -45,7 +45,7 @@ axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisTy
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
-               Ops.COPY: 1, Ops.SLICE: 2, Ops.LINEAR: 0}
+               Ops.SLICE: 2, Ops.LINEAR: 0}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op])
@@ -94,6 +94,8 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
   return ret
 
 def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp:
+  for x in arg:
+    if isinstance(x, UOp) and not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}")
   if len(arg) == 0: return UOp(Ops.STACK)
   elif len(arg) == 1: return UOp.const(dtypes.weakint, arg[0])
   else: return UOp(Ops.STACK, src=tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))
@@ -169,7 +171,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
       return arg
     case Ops.CONST:
       # derived from the value. order matters: bool is an int subclass, ConstFloat is a float subclass
-      if isinstance(arg, InvalidType): return dtypes.bool  # Invalid is the lattice bottom, typed by its consumer
+      if isinstance(arg, InvalidType): return dtypes.bool  # Invalid is always bool, the promo lattice bottom
       if isinstance(arg, bool): return dtypes.bool
       if isinstance(arg, int): return None
       if isinstance(arg, float): return dtypes.weakfloat
@@ -184,9 +186,11 @@ class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
   def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
                metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
+    if op is Ops.CONST and arg is Invalid: dtype = dtypes.bool
     if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
     # CONST derives its dtype by value only when the constructor omits one; an explicit (strong) const dtype is legal until the field is removed
-    if SPEC == 2 and op is not Ops.CONST and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
+    if SPEC == 2 and op is not Ops.CONST and not any(s.base.arg is Invalid for s in src) and \
+       (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
       raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
     if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
     UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
@@ -224,6 +228,7 @@ class recursive_property(property):
 
 # we import this late so we can use resolve/smax in mixins
 from tinygrad.mixin.op import OpMixin
+from tinygrad.mixin.movement import MovementMixin
 from tinygrad.mixin.rand import RandMixin
 
 # NOTE: this should be frozen, but frozen is slower
@@ -475,8 +480,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def ended_ranges(self) -> tuple[UOp, ...]:
-    # an END only ends ranges, the loop backedge condition is not an ended range
-    if self.op is Ops.END: return tuple(x for x in self.src[1:] if x.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:]]))
     return ()
@@ -591,6 +594,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return UOp(Ops.STORE, src=srcs, **kwargs)
   def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, src=(self,)+src, **kwargs) if len(src) else self
+  @property
+  def without_after(self) -> UOp: return self.src[0] if self.op is Ops.AFTER else self
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
   def contract(self, *rngs:UOp):
@@ -603,7 +608,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if isinstance(b, UOp): return b.cast(dtype)
     # NOTE: it always has to be STACK now, even if they are all the same
     if isinstance(b, tuple):
-      stk = [UOp(Ops.CONST, dtype, arg=dtype.const(c), src=()) for c in b]
+      stk = [UOp.const(dtype, c) for c in b]
       ret = UOp.stack(*stk)
     else:
       ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
@@ -630,15 +635,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     ret = UOp(Ops.REDUCE, src=(self.permute(perm),), arg=(op, len(reduce_axis)))
     return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret
   @staticmethod
-  def invalid(): return UOp.const(dtypes.weakint, Invalid)
+  def invalid(): return UOp.const(dtypes.bool, Invalid)
   def valid(self, cond):
     return cond.where(self, self.const_like(Invalid))
   def get_idx(self) -> UOp:
-    assert dtypes.is_int(self.dtype), "Can only call get_idx on index dtype"
     if self.op is Ops.STACK: return UOp.stack(*(x.get_idx() for x in self.src))
     return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self
   def get_valid(self) -> UOp:
-    assert dtypes.is_int(self.dtype), "Can only call get_valid on index dtype"
     if self.op is Ops.STACK: return UOp.stack(*(x.get_valid() for x in self.src))
     return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
   def reduce(self, *src:UOp, **kwargs):
@@ -702,6 +705,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return src_axis
 
   def _unshard(self, axis:int) -> UOp:
+    bsz, dcount = self.shape[axis], len(self.device)
+    dnum = UOp.variable("_device_num", 0, dcount-1)
+    # raw PAD with _device_num-dependent bounds: the other shards' regions are Invalid ("no data"), never a fill value
+    return MovementMixin.pad(self, tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
+
+  def _unshard_fill(self, axis:int) -> UOp:
+    # 0-filled variant of _unshard for ALU allreduce: materialized pad regions must be explicit 0, not gated-stale Invalid
     bsz, dcount = self.shape[axis], len(self.device)
     dnum = UOp.variable("_device_num", 0, dcount-1)
     return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
@@ -722,7 +732,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}")
     return UOp(Ops.COPY, src=(inp,), arg=device)
   def mselect(self, arg:int) -> UOp: return UOp(Ops.MSELECT, src=(self,), arg=arg)
-  def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs)
+  def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs) if len(srcs) else self
   @property
   def metadata(self) -> tuple[Metadata, ...]|None: return all_metadata.get(self, None)
 
@@ -772,6 +782,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   unique_num = itertools.count(0)
 
+  def getaddr(self, device=None) -> UOp:
+    if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
+    return UOp(Ops.GETADDR, dtypes.uint64, src=(self,), arg=device or to_tuple(self.device)[0])
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}")
@@ -860,9 +873,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
     return out.arg if out.op is Ops.CONST and isinstance(out.arg, int) else None
 
-  def has_buffer_identity(self):
+  def has_buffer_identity(self, after_ok=False):
     """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain)."""
-    if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity()
+    if self.op in {Ops.RESHAPE, Ops.MULTI, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
+    if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
     return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
 
   def _base_buffer_is_realized(self) -> bool:
@@ -1086,13 +1100,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop high level syntactic sugar ***
 
   @staticmethod
-  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, device=None):
     if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
+      ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device))
     else:
       assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
-      buf_shape = (prod(shape),)
-      ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, dtype, addrspace=addrspace))
+      assert device is None, "LOCAL and REG placeholders cannot have a device"
+      ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
   def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL):
@@ -1137,6 +1151,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
     return [s.after(kernel) for s in contig_srcs]
 
+  def to_elf(self) -> TinyELF:
+    assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
+    sig = tuple((u.arg.name, u.arg.slot, u.dtype, u._shape)
+                for u in tuple(filter(lambda u: u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU, self.src[1].src)) + self.arg.vars)
+    return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig)
+
 @dataclass(frozen=True)
 class KernelInfo:
   name: str = "test"            # name of the kernel
@@ -1158,7 +1178,7 @@ class ProgramInfo:
   globals: tuple[int, ...] = ()
   outs: tuple[int, ...] = ()
   ins: tuple[int, ...] = ()
-  aux: tuple = ()
+  target: Target = Target()
 
   @property
   def function_name(self): return to_function_name(self.name)
@@ -1176,7 +1196,7 @@ class ProgramInfo:
     except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None
 
   @staticmethod
-  def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo:
+  def from_sink(sink:UOp, target:Target=Target()) -> ProgramInfo:
     _vars: list[UOp] = []
     _globals: list[int] = []
     outs: list[int] = []
@@ -1196,7 +1216,7 @@ class ProgramInfo:
       if u.op is Ops.PARAM and u in _vars and u.expr == 'core_id': global_size[0] = int(u.vmax) + 1
     return ProgramInfo(sink.arg.name if isinstance(sink.arg, KernelInfo) else "test", tuple(global_size),
                        tuple(local_size) if local_size is not None else None, tuple(sorted(dedup(_vars), key=lambda v: v.arg.slot)),
-                       tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), aux)
+                       tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), target)
 
 @dataclass(frozen=True)
 class CallInfo:
@@ -1316,6 +1336,8 @@ class UPat(OpMixin):
   def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs)
 
   # copied from UOp
+  def sink(*srcs:UPat|None, **kwargs):  # pylint: disable=no-self-argument
+    return UPat(Ops.SINK, src=tuple([x for x in srcs if x is not None]), **kwargs)
   def index(self, *srcs:UPat|None, **kwargs):
     return UPat(Ops.INDEX, src=(self,)+tuple(x for x in srcs if x is not None), **kwargs)
   def cast(self, dtype=None, **kwargs):
@@ -1690,28 +1712,23 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
   rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
   return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
 
-def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
+def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)
 
 def select_dtype(u:UOp):
+  if u.dtype is dtypes.weakfloat: return dtypes.default_float
   return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
-def lower_alu_dtype(u:UOp, x:UOp, y:UOp, dt:DType) -> UOp:
-  src = u.src[:-2]+(x.cast(dt), y.cast(dt))
-  return src[0].alu(u.op, *src[1:]).cast(u.dtype)
-pm_lower_weakint = PatternMatcher([
-  # There are no Unary ops at this point in symbolic, those are introduced later
-  (UPat(Ops.CONST, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None),
-  # Binary can widen the dtype, WHERE cannot
-  (UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))),
-   lambda u,x,y: lower_alu_dtype(u, x, y, least_upper_dtype(select_dtype(u), x.dtype, y.dtype))),
-  (UPat(Ops.WHERE, dtypes.weakint, src=(UPat(), UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint)), name="u"),
-   lambda u,x,y: lower_alu_dtype(u, x, y, least_upper_dtype(x.dtype, y.dtype))),
-  # in a weakint WHERE, an Invalid branch takes the dtype of the other branch
-  (UPat.var("gate").where(UPat.var("idx", dtypes.ints).cast(dtypes.weakint), UPat(Ops.CONST, arg=Invalid)),
-   lambda gate,idx: idx.valid(gate).cast(dtypes.weakint)),
-  (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)),
-  (UPat(Ops.STACK, src=UPat().cast(dtypes.weakint), name="v"),
-    lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt) for s in v.src)).cast(dtypes.weakint)),
+def lower_weak_node(u:UOp) -> UOp|None:
+  start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
+  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
+  dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
+                    else unwrap(dtype_from_uop(u.op, src, u.arg)))
+  return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype)
+pm_lower_weak = PatternMatcher([
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype)),
+  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
+  # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
+  (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK}, name="u"), lower_weak_node),
   # special can only be int32
   (UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.weakint),), name="u"),
     lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.weakint)),
@@ -1724,23 +1741,22 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
   if ctx is None: ctx = {}
   def lower(s:UOp) -> UOp:
     if (r:=ctx.get(s)) is None:
-      r = graph_rewrite(s, pm_lower_weakint)
+      r = graph_rewrite(s, pm_lower_weak)
       # the consumer absorbs the cast on its own edge
-      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype == dtypes.weakint else r
+      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
     return r
   # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands
-  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype == dtypes.weakint else s for s in u.src))
+  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
   return None if ret is u else ret
 pm_lower_index_dtype = PatternMatcher([
   (UPat(GroupOp.All, name="u"),
-   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype != dtypes.weakint and any(s.dtype == dtypes.weakint for s in u.src) else None),
+   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
   # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)
   # TODO: more generic
   (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
         allow_any_len=True, name="u"),
    lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
 ])
-def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])])

tinygrad/uop/spec.py

@@ -44,7 +44,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
         raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
 
 # ***** new specs *****
-
+def matches_dtype(x:UOp, dtype:DType) -> bool: return x.dtype == dtype or x.base.arg is Invalid  # Invalid matches any dtype
 # these ops can be used in the tensor graph and programs
 spec_shared = PatternMatcher([
   # NOTE: for testing, we let sinks be anything
@@ -59,27 +59,27 @@ spec_shared = PatternMatcher([
   # STACK is everywhere too
   (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True),
   (UPat(Ops.STACK, src=(UPat(),), allow_any_len=True, name="s"),
-   lambda s: all_same([x.shape for x in s.src]) and all(x.dtype == s.dtype for x in s.src)),
+   lambda s: all_same([x.shape for x in s.src]) and all(matches_dtype(x, s.dtype) for x in s.src)),
 
   # ALUs: operands match the result dtype, except comparisons/WHERE; renderer-lowered shifts may use a uint32 count
   # a weak dtype matches any dtype (TODO: make python scalars weak consts)
   (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat(), UPat())),
-   lambda w: all(s.dtype == w.dtype or s.dtype in dtypes.weaks for s in w.src[1:])),
+   lambda w: all(matches_dtype(s, w.dtype) or s.dtype in dtypes.weaks for s in w.src[1:])),
   (UPat(GroupOp.Comparison, dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))),
-   lambda x,y: x.dtype == y.dtype or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
+   lambda x,y: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
   (UPat((Ops.AND, Ops.OR, Ops.XOR, Ops.SHL, Ops.SHR), name="x"), lambda x: False if any(dtypes.is_float(s.dtype) for s in x.src) else None),
-  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat(dtype=dtypes.uint)), name="a"), lambda a,x: a.dtype == x.dtype or None),
+  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat(dtype=dtypes.uint)), name="a"), lambda a,x: matches_dtype(x, a.dtype) or None),
   (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
-  (UPat(GroupOp.ALU, name="x"), lambda x: all(y.dtype == x.dtype or y.dtype in dtypes.weaks for y in x.src)),
+  (UPat(GroupOp.ALU, name="x"), lambda x: all(matches_dtype(y, x.dtype) or y.dtype in dtypes.weaks for y in x.src)),
 
   # CAST
   (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: isinstance(x.arg, DType)),
 
   # RANGE can be in the big graph now. a void RANGE is a bound-less loop header, the arg is an axis id like RANGE
   (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x:
-    rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
+    matches_dtype(x, rng.dtype) and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
       all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
-  (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(dtypes.is_int(y.dtype) for y in x.src[1:]) or None),
+  (UPat(Ops.INDEX, name="x"), lambda x: len(x.src)>0 and all(dtypes.is_int(y.dtype) or y.base.arg is Invalid for y in x.src[1:]) or None),
   # END closes RANGEs
   (UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:]) or None),
   # a loop-ended END requires a trailing bool condition for the backedge (loop again while true)
@@ -96,14 +96,14 @@ spec_shared = PatternMatcher([
   # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.INDEX,
                                                      Ops.AFTER, Ops.MULTI, Ops.BITCAST, Ops.INS})),),
-        allow_any_len=True), lambda: True),
+        allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
 
   # CUSTOM (inline and non inline)
   (UPat((Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
 
   # CALL of an external function
   (UPat(Ops.CALL, src=(UPat(),), allow_any_len=True, name="x"),
-   lambda x: x.src[0].dtype is dtypes.uint64 if x.src[0].dtype is not dtypes.void else None),
+   lambda x: matches_dtype(x.src[0], dtypes.uint64) if x.src[0].dtype is not dtypes.void else None),
 
   # pattern compiler IR ops (not in tensor/program graphs, but spec-compliant)
   (UPat(Ops.PYLITERAL), lambda: True),
@@ -117,7 +117,7 @@ spec_shared = PatternMatcher([
   # LOAD(idx) / STORE(idx, val) with gates on the LOAD/STORE
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(), validate_index),
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().load(UPat.var("alt"), UPat.var("gate", dtype=dtypes.bool), name="load"),
-   lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == load.dtype else False),
+   lambda uidx,gate,alt,load: validate_index(uidx, gate) if matches_dtype(alt, load.dtype) else False),
   (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),
 
@@ -130,8 +130,7 @@ spec_shared = PatternMatcher([
 
 def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d))
 
-def valid_gettuple(g:UOp, t:UOp):
-  return isinstance(g.arg, int) and 0 <= g.arg < len(t.src) and g.dtype == t.src[g.arg].dtype
+def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < len(t.src) and matches_dtype(t.src[g.arg], g.dtype)
 
 # these ops can exist in tensor but not programs. example: movement
 spec_tensor = PatternMatcher([
@@ -139,7 +138,7 @@ spec_tensor = PatternMatcher([
 
   # BUFFER
   (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
-   (isinstance(buf.dtype, DType) and buf.src[0].dtype == dtypes.weakint and is_device(buf.arg.device))
+   (isinstance(buf.dtype, DType) and matches_dtype(buf.src[0], dtypes.weakint) and is_device(buf.arg.device))
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # Tensor variable bindings
@@ -158,10 +157,7 @@ spec_tensor = PatternMatcher([
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), valid_gettuple),
 
   # SPECIAL is index before index lowering. custom_kernel currently has this
-  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.weakint),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
-
-  # inputs to movement ops
-  (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True),
+  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.weakint),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
 
   # movement ops
   (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True),
@@ -173,19 +169,19 @@ spec_tensor = PatternMatcher([
    lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in GroupOp.Reduce
    and isinstance(x.arg[1], int) and all(y.dtype in (dtypes.weakint, dtypes.int) for y in x.src[1:])),
 
-  # COPY. TODO: this should not have allow_any_len, but something is adding ranges
-  (UPat(Ops.COPY, name="copy", src=(UPat.var("x"),), allow_any_len=True), lambda copy,x: copy.dtype == x.dtype and is_device(copy.arg)),
-  (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"),)), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, tuple) and
+  # COPY
+  (UPat(Ops.COPY, name="copy", src=(UPat.var("x"),)), lambda copy,x: matches_dtype(x, copy.dtype) and is_device(copy.arg)),
+  (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"),)), lambda red,x: matches_dtype(x, red.dtype) and isinstance(red.arg, tuple) and
    len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])),
 
   # MULTI/MSELECT/MSTACK
-  (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
+  (UPat(Ops.MULTI, name="multi"), lambda multi: all(matches_dtype(x, multi.dtype) for x in multi.src) and isinstance(multi.arg, int)),
   (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
   (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None),
-   lambda root,x: root.dtype == x.dtype),
+   lambda root,x: matches_dtype(x, root.dtype)),
 
   # TODO: this should not be here. STAGE is transformed to BUFFER later
   (UPat(Ops.STAGE, src=(UPat(),), allow_any_len=True), lambda: True),
@@ -222,7 +218,7 @@ spec_program = PatternMatcher([
   (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
 
   # SPECIAL is int32 after index lowering
-  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
+  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
 ])+spec_shared
 
 spec_hcq = PatternMatcher([

tinygrad/uop/symbolic.py

@@ -1,7 +1,7 @@
 # all of symbolic lives here now
 import math, struct
 from collections import defaultdict
-from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu, identity_element
 from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid
 from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
 from tinygrad.uop.divandmod import div_and_mod_symbolic
@@ -69,10 +69,11 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 pm_data_invalid = PatternMatcher([
+  (invalid_pat.broadcast(), lambda i: i),
   (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,), name="op"), lambda i,op: i.cast(op.dtype)),
   (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_gate,), name="op"),
    lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i.cast(op.dtype))),
-  # binary ops move inside the gate, with Invalid cast to the result dtype (bool for comparisons)
+  # binary ops move inside the gate, with Invalid in the false branch
   (UPat(GroupOp.Binary, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i.cast(alu.dtype))),
   (UPat(GroupOp.Binary, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i.cast(alu.dtype))),
   (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
@@ -92,9 +93,37 @@ pm_data_invalid = PatternMatcher([
 ])
 
 pm_remove_invalid = PatternMatcher([
-  (invalid_pat, lambda i: i.const_like(0)),
+  (invalid_gate.named("w"), lambda cond,x,i,w: w.replace(src=(cond,x,w.const_like(0)))),
+  (UPat(Ops.STACK, name="s"), lambda s: s.replace(src=tuple(UOp.const(s.dtype, 0) if x.arg is Invalid else x for x in s.src))
+   if any(x.arg is Invalid for x in s.src) else None),
 ])
 
+# Invalid in a reduce input means "no data": those lanes must contribute the reduce identity, not poison the accumulator.
+# this must run after pm_data_invalid gate lifting (gates float to the top of the reduce input) and before codegen builds
+# the accumulator loop, otherwise acc+where(c,x,Invalid) gate-lifts and one invalid lane poisons the whole reduction.
+# only WHERE alt gates whose condition involves a reduce range are rewritten: those are "this element has no data".
+# an Invalid behind a reduce-range-independent gate (e.g. gather with an Invalid index) poisons the whole lane: keep it.
+def _invalid_to_identity(u:UOp, ident:UOp, red_ranges:frozenset[UOp]) -> UOp|None:
+  if u.op is Ops.WHERE:
+    if u.src[2].base.op is Ops.CONST and u.src[2].base.arg is Invalid and u.src[2].dtype == ident.dtype:
+      alt = ident if not red_ranges.isdisjoint(u.src[0].ranges) else None
+    else: alt = _invalid_to_identity(u.src[2], ident, red_ranges)
+    then = _invalid_to_identity(u.src[1], ident, red_ranges)
+    if alt is None and then is None: return None
+    return u.replace(src=(u.src[0], u.src[1] if then is None else then, u.src[2] if alt is None else alt))
+  if u.op in GroupOp.Elementwise-{Ops.WHERE}:
+    new_srcs = tuple(_invalid_to_identity(s, ident, red_ranges) for s in u.src)
+    if all(n is None for n in new_srcs): return None
+    return u.replace(src=tuple(s if n is None else n for s,n in zip(u.src, new_srcs)))
+  return None
+
+def reduce_invalid_identity(r:UOp) -> UOp|None:
+  red_ranges = frozenset(x for x in r.src[1:] if x.op is Ops.RANGE)
+  new_src = _invalid_to_identity(r.src[0], r.const_like(identity_element(r.arg[0], r.dtype)), red_ranges)
+  return r.replace(src=(new_src,)+r.src[1:]) if new_src is not None else None
+
+pm_invalid_reduce_identity = PatternMatcher([(UPat(Ops.REDUCE, name="r"), reduce_invalid_identity)])
+
 symbolic_simple = pm_data_invalid + PatternMatcher([
   # ** self folding **
   (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
@@ -177,6 +206,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
+  # nested where with the same condition in the then position: c ? (c ? t : f) : f2 -> c ? t : f2
+  (UPat.var("c").where(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("f2")), lambda c,t,f,f2: c.where(t, f2)),
   # a.where(c, b.where(c, d)) -> (a | b).where(c, d)
   (UPat.var("a").where(UPat.var("c"), UPat.var("b").where(UPat.var("c"), UPat.var("d"))), lambda a,b,c,d: (a|b).where(c,d)),
 ])+mop_cleanup

tinygrad/viz/README

@@ -14,7 +14,7 @@ This can:
 VIZ pkls can be viewed in two ways:
 
 1. Web browser: python -m tinygrad.viz.serve
-2. Command line: python -m tinygrad.viz.cli
+2. Command line: python -m tinygrad.viz.cli (add --json for srcipting)
 
 By default, VIZ UIs automatically load the latest files.
 

tinygrad/viz/cli.py

@@ -218,7 +218,8 @@ def main(args) -> None:
     for k in (produce_top_kernels if args.t else produce_all_kernels)(): render_event(k)
 
 def get_arg_parser() -> argparse.ArgumentParser:
-  parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli")
+  parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli", epilog="DEBUG modes (cumulative): 3=base AST, 4=generated source, "
+                                   "5=rewrite steps and kernel graph, 6=all UOp graphs, 7=all rewrites")
   parser.add_argument("-s", "--src", nargs="+", default=[], metavar="NAME", help="Select a data source (default: all)")
   parser.add_argument("--list", "--ls", dest="list", action="store_true", help="List sources")
   parser.add_argument("--interval", nargs="+", metavar=("START", "END"), help="Optional start and end marker")

tinygrad/viz/serve.py

@@ -146,7 +146,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
       if u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
       if u.op in {Ops.CALL, Ops.FUNCTION}:
-        label += f"\n{u.src[0].key.hex()[:8]}"
+        label += f"\n{u.src[0].key.hex()[:8]}\n{u.src[0].op}"
       if u.op in {Ops.INDEX, Ops.STAGE}:
         label += f"\n{u.render()}" if sum(len(s.toposort()) for s in u.src[1:]) < 30 else "\nINDEX TOO LARGE"
         ranges: list[UOp] = []