try pad Invalid

AGENTS.md

@@ -3,4 +3,3 @@
 - 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, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW
 
   INITMLPERF = getenv("INITMLPERF")
   RUNMLPERF = getenv("RUNMLPERF")
@@ -1482,8 +1482,7 @@ def train_llama3():
 
   @TinyJit
   def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
     for g in grads: g.assign(0)
@@ -1668,7 +1667,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, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1776,8 +1775,7 @@ def train_gptoss():
 
   @TinyJit
   def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
     for g in grads: g.assign(0)

examples/mlperf/optim.py

@@ -21,12 +21,6 @@ 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)
@@ -50,21 +44,38 @@ 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 fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
-    updates, extra = self._step([], grads)
+  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)
     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')]
-    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
 
-  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))
+    Tensor.realize(*to_realize)
+    return extra[-1]
 
   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
@@ -77,7 +88,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
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
 
   def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
     w = master if master is not None else t
@@ -98,7 +109,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(t.shape)
+        ret = w_q.reshape(new_w.shape)
         return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
       if IMMEDIATE_SCALE:

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, GroupOp
+from tinygrad.uop.ops import Ops
 from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
@@ -18,16 +18,12 @@ 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, u = 0, x.uop
-  while u.op in alias_ops:
-    if u.op is Ops.SHRINK:
+  offset = 0
+  for u in x.uop.toposort():
+    if u.op == 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
@@ -224,7 +220,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):
-  if storage_offset is None: storage_offset = tensor.storage_offset()
+  storage_offset = storage_offset or tensor.storage_offset()
   return _as_strided(tensor, size, stride, storage_offset)
 
 @torch.library.impl("aten::_reshape_alias", "privateuseone")
@@ -232,16 +228,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=None, layout=None, device=None, pin_memory=False):
+def empty_strided(size, stride, dtype, 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 or torch.get_default_dtype()), device=_from_torch_device(device))
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype), device=_from_torch_device(device)).contiguous()
   # 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))
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device)).contiguous()
   return wrap(ret)
 
 @torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
@@ -656,6 +652,35 @@ 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:
@@ -669,7 +694,7 @@ def wrap_fxn(k,f):
     else: raise RuntimeError(f"unknown output type {type(out)}")
   return nf
 
-def wrap_inplace(f):
+def wrap_inplace(k,f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -677,7 +702,7 @@ def wrap_inplace(f):
     return orig
   return nf
 
-def wrap_inplace_view_op(f):
+def wrap_inplace_view_op(k,f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -710,17 +735,11 @@ def wrap_inplace_view_op(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():
-  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)
+  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))
 
 @torch.library.impl("aten::equal", "privateuseone")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -756,17 +775,21 @@ 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)
-  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)
+  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)
 
 # _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
 # we need torch.autograd.Function with explicit AutogradPrivateUse1 registration

extra/torch_backend/test.py

@@ -71,33 +71,6 @@ 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)
@@ -343,21 +316,6 @@ 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)
@@ -784,20 +742,6 @@ 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,21 +124,16 @@ 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) {
-  std::shared_ptr<c10::SafePyObject> tiny = tiny_impl(tensor)->opaque_handle();
+  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();
   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 = tiny_impl(tensor);
+  auto* impl = tensor.unsafeGetTensorImpl();
   impl->set_allow_tensor_metadata_change(true);
   impl->set_sizes_and_strides(sizes, strides, storage_offset);
 }

multi_migration.md

@@ -1,86 +0,0 @@
-# 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,10 +169,11 @@ 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 the AMD runtime."""
-  from tinygrad.device import Device, TinyELF
+  """Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import Target, flat_mv
+  from tinygrad.helpers import flat_mv
 
   dev = Device["AMD"]
   compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -222,7 +223,7 @@ amdhsa.kernels:
 """
 
   lib = compiler.compile(asm_src)
-  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
+  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
 
   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 Target, flat_mv
+from tinygrad.helpers import 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,7 +42,8 @@ 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, TinyELF
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -85,7 +86,7 @@ amdhsa.kernels:
 ...
 .end_amdgpu_metadata
 """
-  prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
+  prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
   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 Target, flat_mv
+from tinygrad.helpers import flat_mv
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
@@ -37,7 +37,8 @@ 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, TinyELF
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -83,7 +84,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
+  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
   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 Target, flat_mv
+from tinygrad.helpers import 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,7 +36,8 @@ 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, TinyELF
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device['AMD']
@@ -84,7 +85,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
+  prg = AMDProgram(dev, 'test', lib)
   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.int, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
+  var = UOp.param(2, dtypes.weakint, 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,16 +7,15 @@ 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())
-compiled_prg = linear.src[-1].src[0]
-lib = bytearray(compiled_prg.src[3].arg)
+lib = bytearray(linear.src[-1].src[0].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
@@ -28,7 +27,7 @@ for i in range(0, len(lib) - 4, 4):
     break
 assert found, "s_endpgm not found"
 
-patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
+patched_prg = AMDProgram(dev, "patched", 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(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.multi(0), device=devs)
+    c = Tensor(UOp.const(dtypes.float, Invalid, shape=(2, 4)).clone(device=devs).multi(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     c.realize()

test/backend/test_tensor_variable.py

@@ -18,16 +18,13 @@ 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):
-    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
+    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
 
   def test_variable_tensor_dtype_arg(self):
     vv = Variable("a", 1, 10).bind(2)

test/device/test_metal.py

@@ -1,8 +1,7 @@
 import unittest
-from tinygrad.device import CompileError, Device, BufferSpec, TinyELF
-from tinygrad.helpers import Target
+from tinygrad.device import CompileError, Device, BufferSpec
 if Device.DEFAULT=="METAL":
-  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler
+  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
 class TestMetal(unittest.TestCase):
   def test_alloc_oom(self):
@@ -49,7 +48,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
-      device.runtime(TinyELF(compiled, "r_5", Target("METAL"), ()))
+      MetalProgram(device, "r_5", compiled)
 
   def test_free(self):
     size = 2**16

test/device/test_ocl.py

@@ -1,20 +1,18 @@
 import unittest
 from unittest.mock import patch
 from tinygrad import Device
-from tinygrad.device import Buffer, TinyELF
+from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import Target
-from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler
+from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
 
 @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; }"
-    obj = TinyELF(src.encode(), "cached_test", Target("CL"), ())
-    device.runtime(obj)
+    CLProgram(device, name="cached_test", lib=src.encode())
     with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
-      device.runtime(obj)
+      CLProgram(device, name="cached_test", lib=src.encode())
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
@@ -29,7 +27,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      device.runtime(TinyELF(b"__kernel void test(__global int* a) { a[0] = 1; }", "", Target("CL"), ()))
+      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
     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, TinyELF; from tinygrad.helpers import Target; '
+  return _run('from tinygrad.device import Device; from tinygrad.runtime.ops_amd import AMDProgram; '
               'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
-              f'dev.runtime(TinyELF(HIPCompiler(dev.arch).compile("""{asm_src}"""), "test", Target("AMD", arch=dev.arch), ()))('
+              f'AMDProgram(dev, "test", HIPCompiler(dev.arch).compile("""{asm_src}"""))('
               '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,8 +5,7 @@ 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, TinyELF
-from tinygrad.helpers import Target
+from tinygrad.device import Device
 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",
@@ -43,8 +42,8 @@ class TestGPUCrash(unittest.TestCase):
       self.fail("Device not working before test")
 
   def _run(self, code: str):
-    prg = self.dev.runtime(TinyELF(self.compiler.compile(assemble(code, is_cdna=self.is_cdna)), "test",
-                                   Target("AMD", arch=self.dev.arch), ()))
+    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.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 PythonRenderer
+from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
 
 def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp:
   if ren is None: ren = Renderer(Target())
@@ -83,15 +83,14 @@ 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, ...]=()):
-  dev = Device['PYTHON']
-  allocator = dev.allocator
+  allocator = Device['PYTHON'].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 = dev.runtime(prg.to_elf())
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src[2].arg))
   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(adt, i)).load().cast(dtypes.uint64) << _u64(i * 8))
+        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, 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(adt, 1)).load().cast(dtypes.uint32)
+        hi = mindex(idx + _const(dtypes.int, 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(adt, i)).load().cast(dtypes.uint32) << _u32(i * 8))
+        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, 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,8 +205,7 @@ 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
-    # 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.int64, dtypes.uint64) == dtypes.uint64
     assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
     assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
 
@@ -225,9 +224,8 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
 
   def test_weakint_promo(self):
-    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
+    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.weakint)
+    with self.assertRaises(KeyError): least_upper_dtype(dtypes.weakint, dtypes.int8)
 
   def test_weakfloat_promo(self):
     # weakfloat is a float, but is not one of dtypes.floats

test/null/test_schedule.py

@@ -600,8 +600,7 @@ class TestSchedule(unittest.TestCase):
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
     # TODO: this should be 3 if fix store hazard worked correctly
-    # NOTE: pad now always has an explicit fill mask (internal PAD is Invalid-filled), which materializes here
-    check_schedule(p, 5)
+    check_schedule(p, 4)
 
   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,7 +5,6 @@ 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
 
@@ -987,7 +986,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")
 
   def test_bitcast_chain(self):
-    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
+    a = Variable("a", 0, 3)
     self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)
 
   def test_negation_in_where(self):
@@ -1331,8 +1330,7 @@ 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(UOp.const(dtypes.weakint, 0), UOp.invalid()),
-                  "multiplying an index by 0 should preserve the invalid")
+    self.assertIs(idx.simplify(), (ridx<5).where(0, UOp.invalid()), "multiplying an index by 0 should preserve the invalid")
 
   def test_alu_moves_inside_invalid(self):
     ridx = Variable("ridx", 0, 10)
@@ -1386,7 +1384,11 @@ 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)
-    type_verify(out, spec_shared)  # Invalid matches any dtype
+    # 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}")
 
 class TestSymbolicRealWorld(unittest.TestCase):
   def test_resnet_half(self):

test/null/test_uop_vmin_vmax.py

@@ -168,7 +168,8 @@ class TestVminVmaxProperties(unittest.TestCase):
 
   def test_vmin_vmax_invalid_vconst(self):
     x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
-    self.assertEqual((x.vmin, x.vmax), (0, 4))
+    self.assertLess(x.vmin, 0)
+    self.assertGreater(x.vmax, 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, pm_remove_invalid
-from tinygrad.mixin.movement import MovementMixin
+from tinygrad.uop.symbolic import sym
 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,28 +45,6 @@ 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
@@ -458,7 +436,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(MovementMixin.pad(UOp.empty(4), ((2,2),))[0], None)
+  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
   def test_strided(self): self._check(UOp.empty(4)[::2], None)
 
 if __name__ == '__main__':

test/test_tiny.py

@@ -31,9 +31,13 @@ class TestTiny(unittest.TestCase):
     out = Tensor.ones(16).contiguous() + Tensor.ones(16).contiguous()
     self.assertListEqual(out.tolist(), [2]*16)
 
+  def test_stack(self):
+    out = Tensor.stack(Tensor.ones(8).contiguous(), Tensor.zeros(8).contiguous())
+    self.assertListEqual(out.flatten().tolist(), [1]*8+[0]*8)
+
   def test_cat(self):
-    out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.zeros(8).contiguous())
-    self.assertListEqual(out.tolist(), [1]*8+[0]*8)
+    out = Tensor.cat(Tensor.ones(8).contiguous(), Tensor.zeros(5).contiguous())
+    self.assertListEqual(out.tolist(), [1]*8+[0]*5)
 
   def test_sum(self, N=getenv("SUM_N", 256)):
     out = Tensor.ones(N).contiguous().sum()

test/unit/test_dtype_weak.py

@@ -1,4 +1,4 @@
-import tempfile, unittest, math
+import tempfile, unittest
 
 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import Context
@@ -137,19 +137,5 @@ 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()
 
-    # fixed: pad now has an explicit fill mask (internal PAD is Invalid-filled), which recomputes from the symbolic shape
+    # broken: pad is a view, BIND values frozen at capture (i=2)
     @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()), 4)
+    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2)  # should be 4!
 
     # workaround: contiguous fuses pad into kernel
     @TinyJit

tinygrad/callify.py

@@ -179,9 +179,10 @@ 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,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+                   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)
 
 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, Invalid
+from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -121,12 +121,12 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([
 
 def do_devectorize(b:UOp):
   if b.shape == (): 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
+  # broadcasting needs to be already unpacked
+  if not all_same([x.shape 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.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src)))
+    src.append(b.replace(src=tuple([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):
@@ -403,7 +403,8 @@ 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))
-  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=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)
 
 def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
   if DEBUG >= 4: print(source.arg)
@@ -436,14 +437,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, renderer.target)
+    prog_info = ProgramInfo.from_sink(full_sink)
     # 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], renderer.target))
+  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
   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/late/gater.py

@@ -2,10 +2,6 @@
 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)),
@@ -22,8 +18,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")).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),
+  (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)),
 ])

tinygrad/codegen/simplify.py

@@ -95,12 +95,11 @@ pm_reduce_unparented = PatternMatcher([
 ])
 
 pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
-  # 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 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
   ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-    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),
+   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),
   # 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),
@@ -120,9 +119,8 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
 ])+symbolic
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
-  # 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),
+  # 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),
   # 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, Self, TYPE_CHECKING
+from typing import Any, Generic, TypeVar, Iterator, Generator, 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, Target, unwrap
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize
 from tinygrad.dtype import DType, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -283,25 +283,12 @@ 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:type[Program[Self]]|None, graph=None, arch=None):
+  def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime, graph=None, arch=None):
     from tinygrad.renderer import Renderer
-    self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
+    self.device, self.allocator, self.runtime, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
     self.arch = arch
     self.cached_renderer:dict[Any, Renderer] = {}
 
@@ -313,8 +300,6 @@ 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)  # the weak int position in the promo lattice
+  weakint: Final[DType] = DType.new(0, 800, "weakint", None)  # NOTE: not in the promo lattice: index math never mixes dtypes
   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,6 +129,7 @@ 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')
@@ -152,7 +153,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
-  weaks = (weakint, weakfloat)
+  weaks = (weakfloat,)
   all = floats + ints + (bool,) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
@@ -162,13 +163,14 @@ 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:
-  return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype)
+  # TODO: weakint
+  return dtypes.default_float if dtype == dtypes.weakfloat else dtype
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
 # we don't support complex type
-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],
+# 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],
   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
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
@@ -112,8 +112,9 @@ 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.to_elf())
+    runtime = Device[device].runtime(ast.arg.function_name, ast.src[3].arg, *ast.arg.aux, runtimevars=ast.arg.runtimevars, prg=ast)
     if cache: runtime_cache[key] = runtime
   return runtime
 

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, strong_dtype
+from tinygrad.dtype import ConstType, DTypeLike, Invalid, dtypes, to_dtype
 from tinygrad.helpers import argfix, prod
 from tinygrad.mixin.dtype import DTypeMixin
 from tinygrad.mixin.movement import MovementMixin
@@ -78,13 +78,9 @@ 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)
-    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)))
+    return val.clone(device=device) if buffer else val
 
   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, Invalid
+from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype
 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)) or self._uop.base.arg is Invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
 
   def bitcast(self, dtype:DTypeLike) -> Self:
     """

tinygrad/mixin/movement.py

@@ -265,8 +265,7 @@ 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:
-    # 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)))
+    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)))
 
   def view(self, shape, *args) -> Self:
     """`.view` is an alias for `.reshape`."""

tinygrad/mixin/op.py

@@ -286,11 +286,8 @@ 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 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)
+    if value == Invalid: return base
+    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, value)
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder

tinygrad/renderer/__init__.py

@@ -64,6 +64,7 @@ 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
@@ -79,6 +80,7 @@ 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

@@ -303,6 +303,8 @@ class ClangRenderer(CStyleLanguage):
     self.compiler = ClangCompiler(target.arch.split(","))
 
 class OpenCLRenderer(CStyleLanguage):
+  has_aux = True
+
   # language options
   kernel_typedef = "__kernel void"
   buffer_prefix = "__global "
@@ -334,6 +336,13 @@ 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/nir.py

@@ -2,7 +2,7 @@ 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
+from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer
 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
@@ -282,7 +282,9 @@ _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):
+class IR3Renderer(NIRRenderer, OpenCLRenderer):
+  has_aux = True
+
   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/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, TinyELF
+from tinygrad.device import Compiled, BufferSpec
 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['AMDDevice']):
-  def __init__(self, dev:AMDDevice, obj:TinyELF):
+class AMDProgram(HCQProgram):
+  def __init__(self, dev:AMDDevice, name:str, lib:bytes, **kwargs):
     # TODO; this API needs the type signature of the function and global_size/local_size
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+    self.dev, self.name, self.lib = dev, name, lib
 
     image, sections, relocs = elf_loader(self.lib)
 
@@ -608,17 +608,17 @@ class AMDProgram(HCQProgram['AMDDevice']):
 
   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, unwrap(self.dev.hw_compute_queue_t)()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
+    if self.dev.sqtt_enabled: cast(AMDComputeQueue, 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, 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)
+      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)
       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, 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)
+      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)
       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], AMDProgram, AMDSignal,
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], functools.partial(AMDProgram, self), 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, hashlib
+import ctypes, functools, 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, TinyELF, Program
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError
 
 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(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
+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
     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, obj.name.encode(), status := ctypes.c_int32()), status)
+    self.kernel = checked(cl.clCreateKernel(self.program, name.encode(), status := ctypes.c_int32()), status)
 
   def __del__(self):
     try: check(cl.clReleaseKernel(self.kernel))
@@ -54,15 +54,17 @@ class CLProgram(Program['CLDevice']):
 
   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:
-    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)))
+    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))))
     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),
@@ -121,7 +123,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], CLProgram, arch=arch)
+    super().__init__(device, CLAllocator(self), [OpenCLRenderer], functools.partial(CLProgram, self), arch=arch)
 
   def count(self) -> int: return len(unwrap(self.device_ids))
 

tinygrad/runtime/ops_cpu.py

@@ -1,9 +1,7 @@
 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 Buffer, BufferSpec, TinyELF
+from tinygrad.device import Buffer, BufferSpec
 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
@@ -66,7 +64,7 @@ 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[prg.runtimevars['core_id']] = tid
+      if not lvp and 'core_id' in prg.runtimevars: args[len(args_state.bufs)+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,))
@@ -88,33 +86,32 @@ class LVPArgsState(CLikeArgsState):
 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800
 
-class CPUProgram(HCQProgram['CPUDevice']):
+class CPUProgram(HCQProgram):
   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:CPUDevice, obj:TinyELF):
-    self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+  def __init__(self, dev, name:str, lib:bytes, runtimevars:dict[str, int]|None=None, native=False, **kwargs):
+    self.runtimevars = runtimevars or {}
 
-    LVP = obj.target.renderer == "LVP"
+    LVP = isinstance(dev.renderer, LVPRenderer) and not native
     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(obj.lib)), MEM_COMMIT | MEM_RESERVE,
-                                                      PAGE_EXECUTE_READWRITE)
-      ctypes.memmove(self.addr, obj.lib, len(obj.lib))
+      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))
       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(obj.lib)))
+      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(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(obj.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(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)
-      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if LVP else obj.lib
+      if LVP: lib = jit_loader(lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m'])
       self.mem.write(lib)
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
@@ -129,7 +126,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
 
       self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
 
-    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj.name, kernargs_alloc_size=12+256 if LVP else 0)
+    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, name, kernargs_alloc_size=12+256 if LVP else 0)
 
   @suppress_finalizing
   def __del__(self):
@@ -150,8 +147,8 @@ class CPUAllocator(HCQAllocator):
 
 class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
-    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")
+    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")
 
     self.ring_pos = 0
 
@@ -165,7 +162,7 @@ class CPUDevice(HCQCompiled):
     # 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(do_to_program(v, ClangRenderer(replace(self.renderer.target, renderer="CLANG"))).to_elf()) for f,v in prgs.items()}
+      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()}
 
   @functools.cached_property
   def ring(self) -> Buffer: return Buffer(self.device, RING_SLOTS * CMD_SIZE, dtypes.uint64, preallocate=True)
@@ -185,7 +182,7 @@ class CPUDevice(HCQCompiled):
 
   @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
+    threading.Thread(target=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):

tinygrad/runtime/ops_cuda.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
-import ctypes
+import ctypes, functools
 from tinygrad.helpers import DEBUG, DEV, getenv, mv_address, suppress_finalizing
-from tinygrad.device import Compiled, BufferSpec, LRUAllocator, Program, TinyELF
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator
 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(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"))]))
+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"))]))
 
     check(cuda.cuCtxSetCurrent(self.dev.context))
     self.module = cuda.CUmodule()
-    status = cuda.cuModuleLoadData(ctypes.byref(self.module), obj.lib)
+    status = cuda.cuModuleLoadData(ctypes.byref(self.module), 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, obj.name.encode("utf-8")))
+    check(cuda.cuModuleGetFunction(ctypes.byref(prg := cuda.CUfunction()), self.module, 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], CUDAProgram, None if MOCKGPU else CUDAGraph,
-                     arch=f"sm_{major.value}{minor.value}")
+    super().__init__(device, CUDAAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer], functools.partial(CUDAProgram, self),
+                     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, threading, contextlib, sys, subprocess, struct
+import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler
 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(Program['DSPDevice']):
-  def __init__(self, dev:DSPDevice, obj:TinyELF):
-    self.dev, self.lib = dev, obj.lib
+class DSPProgram:
+  def __init__(self, dev:DSPDevice, name:str, lib:bytes, **kwargs):
+    self.dev, self.lib = dev, 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], DSPProgram)
+      super().__init__(device, DSPAllocator(self), [DSPRenderer], functools.partial(DSPProgram, self))
       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(Program[DSPDevice]):
-  def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib = obj.lib
+class MockDSPProgram:
+  def __init__(self, name:str, lib:bytes, **kwargs): self.lib = 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
+import ctypes, functools
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import Compiled, LRUAllocator, BufferSpec, Program, TinyELF
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec
 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], HIPProgram, arch=self.arch)
+    super().__init__(device, HIPAllocator(self), [HIPRenderer], functools.partial(HIPProgram, self), 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(Program[HIPDevice]):
-  def __init__(self, dev:HIPDevice, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+class HIPProgram:
+  def __init__(self, dev:HIPDevice, name:str, lib:bytes, **kwargs):
+    self.dev, self.name, self.lib = dev, name, lib
     check(hip.hipSetDevice(self.dev.device_id))
-    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"))))
+    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"))))
 
   @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, Program, TinyELF, LRUAllocator, ProfileDeviceEvent
+from tinygrad.device import Compiled, Compiler, CompileError, 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], 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:])
+    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:])
 
   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(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)
+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)
     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(obj.name)).retained()
+    self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
     descriptor = metal.MTLComputePipelineDescriptor.new()
     descriptor.setComputeFunction(self.fxn)
     descriptor.setSupportIndirectCommandBuffers(True)

tinygrad/runtime/ops_null.py

@@ -1,5 +1,5 @@
-import inspect, math
-from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
+import inspect, functools, math
+from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent
 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(Program['NullDevice']):
-  def __init__(self, dev:'NullDevice', obj:TinyELF): self.device, self.name = dev.device, obj.name
+class NullProgram:
+  def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, 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), NullProgram, NullGraph)
+    super().__init__(device, NullAllocator(self), dedup(renderers), functools.partial(NullProgram, device), 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, TinyELF
+from tinygrad.device import Compiled, BufferSpec
 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, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+  def __init__(self, dev:NVDevice, name:str, lib:bytes, **kwargs):
+    self.dev, self.name, self.lib = dev, name, 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(obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):])), []
+      image, self.cbuf_0 = memoryview(bytearray(lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(lib)):])), []
       self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
-    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(obj.lib) + b'\x00' * (4 - len(obj.lib)%4)).cast("I"), [], [] # type: ignore
+    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(lib) + b'\x00' * (4 - len(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.{obj.name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
+            if sh.name == f".nv.info.{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], NVProgram, NVSignal, NVComputeQueue,
-                     NVCopyQueue, arch=self.arch)
+    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], functools.partial(NVProgram, self), 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, Program, TinyELF
+from tinygrad.device import Buffer, Compiled, Compiler, Allocator
 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(Program['PythonDevice']):
-  def __init__(self, dev:'PythonDevice', obj:TinyELF):
-    self.uops: list[UOp] = pickle.loads(obj.lib)
+class PythonProgram:
+  def __init__(self, name:str, lib:bytes, **kwargs):
+    self.uops: list[UOp] = pickle.loads(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, TinyELF
+from tinygrad.device import BufferSpec, Device
 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.to_elf())
+  return Device["CPU"].runtime(prg.arg.function_name, prg.src[3].arg)
 
 #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 = [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)]
+    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)]
     # 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['QCOMDevice']):
-  def __init__(self, dev: QCOMDevice, obj: TinyELF):
+class QCOMProgram(HCQProgram):
+  def __init__(self, dev: QCOMDevice, name: str, lib: bytes, buf_dtypes=[], **kwargs):
     self.dev: QCOMDevice = dev
-    self.signature, self.name, self.NIR = obj.signature, obj.name, isinstance(dev.renderer, IR3Renderer)
+    self.buf_dtypes, self.name, self.NIR = buf_dtypes, 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(obj.lib)
+      v, cs, imm_vals, self.image = IR3Compiler.unpack_lib(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['QCOMDevice']):
 
       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(obj.lib)
+    else: self._parse_lib(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], QCOMProgram, 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], functools.partial(QCOMProgram, self), 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, Program, TinyELF
+from tinygrad.device import Compiled, Allocator, BufferSpec
 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(Program['WebGpuDevice']):
-  def __init__(self, dev:'WebGpuDevice', obj:TinyELF):
-    self.dev, self.name = dev, to_wgpu_str(obj.name)
+class WebGPUProgram:
+  def __init__(self, dev:'WebGpuDevice', name:str, lib:bytes, **kwargs):
+    self.dev, self.name = dev, to_wgpu_str(name)
 
     # Creating shader module
-    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(obj.lib.decode()),
+    shader = webgpu.WGPUShaderModuleWGSLDescriptor(code=to_wgpu_str(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], WebGPUProgram,
+    super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], functools.partial(WebGPUProgram, self),
                      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, Program
+from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent
 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(Program[HCQDeviceType]):
+class HCQProgram(Generic[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: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):
+  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):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph

tinygrad/schedule/allreduce.py

@@ -1,6 +1,7 @@
 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:
@@ -56,7 +57,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.invalids(red.shape, dtype=red.dtype, device=red.device)
+  if output is None: output = UOp.const(red.dtype, Invalid, shape=red.shape).clone(device=red.device)
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(*red.arg)

tinygrad/schedule/indexing.py

@@ -1,9 +1,9 @@
 from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
-from tinygrad.dtype import dtypes, AddrSpace, Invalid
+from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
-from tinygrad.uop.ops import gate_kernel_sink
+from tinygrad.uop.ops import gate_kernel_sink, Invalid
 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
 
@@ -101,8 +101,7 @@ 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]])
-  # 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))
+  return valid.where(bx.src[0], UOp.const(x.dtype, Invalid))
 
 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   if x.arg[1] == 0: return None

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, Invalid
+from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
 
 # ***** multi rewrite MSELECT/MSTACK *****
@@ -43,34 +43,6 @@ _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]:
@@ -156,7 +128,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_fill(multi.axis).allreduce(Ops.ADD, device)
+  return multi.src[0]._unshard(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, pm_invalid_reduce_identity
+from tinygrad.uop.symbolic import symbolic
 from tinygrad.uop.movement import mop_cleanup
 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, symbolic_multi_pm
+from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
 # creation can recurse a lot
@@ -545,8 +545,7 @@ pm_copy_to_store = PatternMatcher([
 
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  # 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")
+  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
@@ -556,8 +555,6 @@ def get_kernel_graph(sink:UOp) -> UOp:
   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

@@ -6,7 +6,7 @@ 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, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
-from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike
+from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, ConstLike
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
 from tinygrad.device import Buffer, canonicalize_device
@@ -70,13 +70,16 @@ class Tensor(RandMixin):
     self.is_param:bool = True
 
     # create a UOp from the different types of inputs
-    if data is None:
+    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:
       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())
-    elif not isinstance(data, UOp):
+    else:
       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)):
@@ -173,9 +176,7 @@ 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."""
-    # 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")
+    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")
     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, strong_dtype
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace
 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, TinyELF
+from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
 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, to_tuple
-from tinygrad.helpers import colored, ansilen, printable, Target
+from tinygrad.helpers import colored, ansilen, printable
 if TYPE_CHECKING:
   from tinygrad.renderer import Estimates
 
@@ -94,8 +94,6 @@ 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))
@@ -171,7 +169,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 always bool, the promo lattice bottom
+      if isinstance(arg, InvalidType): return dtypes.bool  # Invalid is the lattice bottom, typed by its consumer
       if isinstance(arg, bool): return dtypes.bool
       if isinstance(arg, int): return None
       if isinstance(arg, float): return dtypes.weakfloat
@@ -186,11 +184,9 @@ 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 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:
+    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:
       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))
@@ -228,7 +224,6 @@ 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
@@ -608,7 +603,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.const(dtype, c) for c in b]
+      stk = [UOp(Ops.CONST, dtype, arg=dtype.const(c), src=()) for c in b]
       ret = UOp.stack(*stk)
     else:
       ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
@@ -635,13 +630,15 @@ 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.bool, Invalid)
+  def invalid(): return UOp.const(dtypes.weakint, 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):
@@ -705,13 +702,6 @@ 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))))
@@ -1151,12 +1141,6 @@ 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
@@ -1178,7 +1162,7 @@ class ProgramInfo:
   globals: tuple[int, ...] = ()
   outs: tuple[int, ...] = ()
   ins: tuple[int, ...] = ()
-  target: Target = Target()
+  aux: tuple = ()
 
   @property
   def function_name(self): return to_function_name(self.name)
@@ -1196,7 +1180,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, target:Target=Target()) -> ProgramInfo:
+  def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo:
     _vars: list[UOp] = []
     _globals: list[int] = []
     outs: list[int] = []
@@ -1216,7 +1200,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))), target)
+                       tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), aux)
 
 @dataclass(frozen=True)
 class CallInfo:
@@ -1716,19 +1700,24 @@ 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_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),
+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)),
   # 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)),
@@ -1741,22 +1730,23 @@ 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_weak)
+      r = graph_rewrite(s, pm_lower_weakint)
       # the consumer absorbs the cast on its own edge
-      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
+      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype == dtypes.weakint 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 in dtypes.weaks 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 == dtypes.weakint 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 not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
+   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),
   # 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(matches_dtype(x, s.dtype) for x in s.src)),
+   lambda s: all_same([x.shape for x in s.src]) and all(x.dtype == 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(matches_dtype(s, w.dtype) or s.dtype in dtypes.weaks for s in w.src[1:])),
+   lambda w: all(s.dtype == 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: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
+   lambda x,y: x.dtype == y.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: matches_dtype(x, a.dtype) or 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.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(matches_dtype(y, x.dtype) or y.dtype in dtypes.weaks for y in x.src)),
+  (UPat(GroupOp.ALU, name="x"), lambda x: all(y.dtype == 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:
-    matches_dtype(x, rng.dtype) and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
+    rng.dtype == x.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, 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),
+  (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),
   # 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, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
+        allow_any_len=True), lambda: True),
 
   # 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: matches_dtype(x.src[0], dtypes.uint64) if x.src[0].dtype is not dtypes.void else None),
+   lambda x: x.src[0].dtype is 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 matches_dtype(alt, load.dtype) else False),
+   lambda uidx,gate,alt,load: validate_index(uidx, gate) if alt.dtype == 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,7 +130,8 @@ 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 matches_dtype(t.src[g.arg], g.dtype)
+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
 
 # these ops can exist in tensor but not programs. example: movement
 spec_tensor = PatternMatcher([
@@ -138,7 +139,7 @@ spec_tensor = PatternMatcher([
 
   # BUFFER
   (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
-   (isinstance(buf.dtype, DType) and matches_dtype(buf.src[0], dtypes.weakint) and is_device(buf.arg.device))
+   (isinstance(buf.dtype, DType) and buf.src[0].dtype == dtypes.weakint and is_device(buf.arg.device))
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # Tensor variable bindings
@@ -157,7 +158,10 @@ 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: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
+  (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),
 
   # movement ops
   (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True),
@@ -170,18 +174,18 @@ spec_tensor = PatternMatcher([
    and isinstance(x.arg[1], int) and all(y.dtype in (dtypes.weakint, dtypes.int) for y in x.src[1:])),
 
   # 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
+  (UPat(Ops.COPY, name="copy", src=(UPat.var("x"),)), 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
    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(matches_dtype(x, multi.dtype) for x in multi.src) and isinstance(multi.arg, int)),
+  (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == 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: matches_dtype(x, root.dtype)),
+   lambda root,x: root.dtype == x.dtype),
 
   # TODO: this should not be here. STAGE is transformed to BUFFER later
   (UPat(Ops.STAGE, src=(UPat(),), allow_any_len=True), lambda: True),
@@ -218,7 +222,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: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
+  (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: s.dtype == x.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, identity_element
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
 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,11 +69,10 @@ 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 in the false branch
+  # binary ops move inside the gate, with Invalid cast to the result dtype (bool for comparisons)
   (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),
@@ -93,37 +92,9 @@ pm_data_invalid = PatternMatcher([
 ])
 
 pm_remove_invalid = PatternMatcher([
-  (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_pat, lambda i: i.const_like(0)),
 ])
 
-# 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
@@ -206,8 +177,6 @@ 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 (add --json for srcipting)
+2. Command line: python -m tinygrad.viz.cli
 
 By default, VIZ UIs automatically load the latest files.
 

tinygrad/viz/cli.py

@@ -218,8 +218,7 @@ 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", epilog="DEBUG modes (cumulative): 3=base AST, 4=generated source, "
-                                   "5=rewrite steps and kernel graph, 6=all UOp graphs, 7=all rewrites")
+  parser = argparse.ArgumentParser(prog="python -m tinygrad.viz.cli")
   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")