add simple profiler test (#18001)

* add simple profiler test

* dev cpu err

examples/gpt2.py

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

extra/models/llama.py

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

test/backend/test_assign.py

@@ -45,16 +45,6 @@ class TestAssign(unittest.TestCase):
     c.realize()
     assert_kernel_count(2 if is_hcq2_device() else 1)
 
-  def test_assign_copy_retained_uses(self):
-    for use in (lambda x: x.reshape(1, 3), lambda x: x + 1):
-      with self.subTest(use=use):
-        x = Tensor([1., 2, 3], device="PYTHON").to(None)
-        retained = use(x)
-        dest = Tensor.empty(3).assign(x)
-        del x
-        dest.realize().assign(0).realize()
-        self.assertEqual(retained.tolist(), [[1., 2, 3]] if retained.ndim == 2 else [2., 3, 4])
-
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
     xs = X[2:4]
@@ -1024,10 +1014,10 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(c.tolist(), [[0,0],[0,0]])
 
-  def test_clone(self):
+  def test_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    c = t.permute(1,0).clone()
-    self.assertIs(c.uop.base.op, Ops.AFTER)
+    c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
+    self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(c.tolist(), [[1,1],[2,1]])
 
@@ -1042,16 +1032,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(cb.tolist(), [[1,2],[3,4]])
 
-  def test_detach_buffer_assignment(self):
-    for realized in (False, True):
-      with self.subTest(realized=realized):
-        base = Tensor([1., 2., 3.])
-        if realized: base.realize()
-        detached = base.detach()
-        detached.assign(detached + 1).realize()
-        self.assertEqual(detached.tolist(), [2., 3., 4.])
-        self.assertEqual(base.tolist(), [2., 3., 4.])
-
   def test_detach_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
     d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
@@ -1063,10 +1043,10 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now
       self.assertEqual(d.tolist(), [[0,0],[0,0]])
 
-  def test_detach_clone(self):
+  def test_detach_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
-    d = t.permute(1,0).clone().detach()
-    self.assertIs(d.uop.base.op, Ops.AFTER)
+    d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
+    self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
     d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(d.tolist(), [[1,1],[2,1]])
 

test/backend/test_const_folding.py

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

test/backend/test_ops.py

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

test/backend/test_profiler.py

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

test/backend/test_schedule.py

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

test/backend/test_setitem.py

@@ -75,11 +75,6 @@ class TestSetitem(unittest.TestCase):
     t.detach()[1, 2] = 5
     self.assertEqual(t[1, 2].item(), 5.0)
 
-  def test_setitem_detach_whole(self):
-    t = Tensor.zeros((3, 3)).realize()
-    t.detach()[:] = 5
-    np.testing.assert_equal(t.numpy(), np.full((3, 3), 5.))
-
   def test_setitem_permute(self):
     # setitem on permuted tensor should modify original
     t = Tensor.zeros((2, 3)).contiguous().realize()

test/null/test_uops.py

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

test/null/test_validate_oob.py

@@ -123,6 +123,15 @@ class TestValidateOOB(unittest.TestCase):
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
       to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
 
+  def test_float_cast_in_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int, 1)
+      r = UOp.range(20, 0)
+      unknown = r.cast(dtypes.float).cast(dtypes.bool)  # a bool from a float is unconstrained
+      to_uops_list([buf.index(r.valid((r < 1) & unknown)).load()])
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid(unknown)).load()])
+
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
       buf = UOp.param(0, dtypes.int, 1)

test/null/test_viz.py

@@ -454,7 +454,7 @@ class TestVizIntegration(unittest.TestCase):
   def test_jit(self):
     with save_viz():
       @TinyJit
-      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).clone().assign(a.to(c.device)), b.assign(c.to(b.device))
+      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
       a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
       for _ in range(3): Tensor.realize(*f(a, b, c))
     out = load_profile(cpu_events)

test/unit/test_call.py

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

test/unit/test_callify.py

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

tinygrad/llm/kernels/amd.py

@@ -75,7 +75,7 @@ class Linear(nn.Linear):
       nbytes, nblocks = raw.max_numel(), raw.max_numel() // Q6_BYTES
       byte_view = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(nbytes, dtypes.uint8, raw_offset)))
       padded = byte_view.reshape((nblocks, Q6_BYTES)).pad_to((nblocks, Q6_PADDED)).bitcast(dtypes.uint32)
-      self.weight = padded.clone().reshape(nblocks * Q6_WORDS)
+      self.weight = padded.contiguous().reshape(nblocks * Q6_WORDS)
     else:
       self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
         .view(raw.max_numel() * raw.dtype.itemsize // dtypes.uint32.itemsize, dtypes.uint32, raw_offset)))

tinygrad/mixin/elementwise.py

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

tinygrad/mixin/rand.py

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

tinygrad/schedule/__init__.py

@@ -80,7 +80,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
-from tinygrad.schedule.prepare import prepare_rangeify, prepare_call_views
+from tinygrad.schedule.prepare import prepare_rangeify
 from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
@@ -122,8 +122,6 @@ def lower_sink_to_linear(call:UOp) -> UOp|None:
   if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None
   # value calls (with unbound outputs) are inlined positionally during prepare: their bodies are not programs to schedule
   if call.has_unbound_outputs: return None
-  call = prepare_call_views(call)
-  function = call.src[0]
   st = time.perf_counter()
   cache_key = function.key
   if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:

tinygrad/schedule/prepare.py

@@ -8,109 +8,6 @@ from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.schedule.allreduce import create_allreduce_function
 from tinygrad.schedule.multi import multi_pm
 
-def on_disk(u:UOp): return isinstance(u.device, str) and u.device.startswith("DISK")
-
-def contiguous_mops_to_view(ctx:list[UOp]|None, c:UOp, src:UOp):
-  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
-  # A list holds CALL arguments; None rewrites views in the live Tensor graph.
-  # Ordinary copies keep their source graph so JIT can substitute its input buffer.
-  if ctx is None and c.op is Ops.COPY and not on_disk(src): return None
-  buf = src.base
-  while buf.op is Ops.BITCAST: buf = buf.src[0].base
-  # no symbolic shape
-  if buf.op not in {Ops.BUFFER, Ops.PARAM, Ops.UNSHARD} or not all_int(c.shape): return None
-
-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then view the resolved shard
-  unshard = None
-  if buf.op is Ops.UNSHARD:
-    if isinstance(c.device, str): return None
-    if (unshard := graph_rewrite(src, multi_pm, name="multi_buffer_view")).op is not Ops.UNSHARD: return None
-    src = unshard.src[0]
-
-  # offset the base buffer by the collapsed movement ops and view it
-  if (cv := src.contiguous_view()) is None or (buf := cv[0]).op not in {Ops.BUFFER, Ops.PARAM}: return None
-  view = buf[cv[1]:cv[1] + src.max_numel() * src.element_size() // buf.element_size()].bitcast(src.dtype)
-  if ctx is not None and view.op in {Ops.SHRINK, Ops.BITCAST}:
-    arg = view.substitute({u: ctx[u.arg.slot] for u in view.toposort() if u.op is Ops.PARAM and u.arg.slot >= 0})
-    if arg not in ctx: ctx.append(arg)
-    view = view.param_like(ctx.index(arg))
-  elif on_disk(buf) and buf.op is Ops.BUFFER and not buf.is_unbound: view = UOp.from_buffer(view.buffer, device=buf.device)
-  view = view.reshape(src.shape).unshard(unshard.arg, unshard.src[1:]) if unshard is not None else view.reshape(c.shape)
-  return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
-
-# Fold contiguous movement operations into buffer views.
-pm_mops_to_view = PatternMatcher([
-  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
-  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
-
-  # remove contiguous on movement ops before a copy on disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
-   copy.replace(src=(x,), tag=None) if on_disk(x) else None),
-  # push copy past movement ops on disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
-   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) if on_disk(x) else None),
-])
-
-def transform_precompiled_call(c:UOp) -> UOp|None:
-  if c.arg is None or not c.arg.precompile or not c.has_unbound_outputs: return None
-  assert c.src[0].op is Ops.SINK, "precompiled call bodies are SINKs of stores into the output PARAMs"
-  # Bind output storage at the existing argument positions.
-  outs = {p: a.empty_like() for p,a in enumerate(c.src[1:]) if a.unsharded_base.is_unbound}
-  placed:dict[UOp, UOp] = {}
-  items = []
-  for st in c.src[0].src:
-    value = st.src[1]
-    while value.op is Ops.AFTER: value = value.src[0]
-    # A custom kernel's output buffer can be the call output directly. Rebind each buffer only once.
-    if value.op in {Ops.BUFFER, Ops.UNSHARD} and value.has_buffer_identity() and value not in placed:
-      placed[value] = st.src[0]
-      items.append(st.src[1])
-    else: items.append(st.src[0].after(st))
-  body = UOp.sink(*items).substitute(placed)
-  call = c.replace(src=(body, *(outs.get(i, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
-                                   for i, a in enumerate(c.src[1:]))))
-  return UOp.sink(*(c.src[1+p].store(o.after(call).shrink_to(c.src[1+p].shape)) for p,o in outs.items()))
-
-pm_resolve_call_outputs = PatternMatcher([
-  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
-  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
-])
-
-def buffer_view_subs(sink:UOp) -> dict[UOp, UOp]:
-  # Include intermediate nodes so every Tensor sharing a pending write receives the same view rewrite.
-  nodes = list(sink.toposort(enter_calls=False))
-  rewritten = graph_rewrite(UOp.sink(*nodes), pm_mops_to_view, bottom_up=True, name="fold buffer views")
-  return {u: v for u, v in zip(nodes, rewritten.src) if u is not v}
-
-def prepare_call_views(call:UOp) -> UOp:
-  # Lift contiguous views into call arguments, preserving their buffer/offset graph for JIT input substitution.
-  args = list(call.src[1:])
-  body = graph_rewrite(call.src[0], pm_mops_to_view, ctx=args, bottom_up=True, name="prepare call views")
-  return call.replace(src=(body, *args))
-
-def prepare_to_call(sink:UOp, tensor_roots:tuple[UOp, ...]) -> UOp:
-  # A copy used only to initialize another buffer can write directly into that destination.
-  # Include live Tensor graphs so retained copies and aliases keep their independent storage.
-  users:dict[UOp, set[UOp]] = {}
-  for u in UOp.sink(sink, *tensor_roots).toposort(enter_calls=False):
-    for src in u.src: users.setdefault(src, set()).add(u)
-  subs = {}
-  for store in sink.toposort(enter_calls=False):
-    if store.op is not Ops.STORE: continue
-    value = store.src[1]
-    if value.op is not Ops.AFTER or len(value.src) != 2: continue
-    buf, init = value.src
-    if init.op is not Ops.STORE or len(init.src) != 2 or init.src[0] is not buf or init.src[1].op is not Ops.COPY: continue
-    # Only this assignment may consume the copy, and only the initialization may use its storage.
-    if users.get(value) != {store} or users.get(buf) != {value, init}: continue
-    while buf.op is Ops.RESHAPE and users.get(buf.src[0]) == {buf}: buf = buf.src[0]
-    if buf.op is not Ops.BUFFER or buf.is_unbound or buf.buffer.is_allocated(): continue
-    subs[value] = init.src[1]
-  sink = sink.substitute(subs, walk=True)
-  sink = graph_rewrite(sink, pm_resolve_call_outputs, bottom_up=True, name="resolve call outputs")
-  return UOp.sink(*[u for u in sink.toposort(enter_calls=False)
-                   if u.op is Ops.AFTER and not u.is_bound_var and not u.src[0].unsharded_base.is_unbound])
-
 def walk_mop(u:UOp):
   if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD, Ops.BITCAST}: return walk_mop(u.src[0])
   if u.op is Ops.AFTER and (b:=walk_mop(u.src[0])) is not u.src[0]: return b.after(*u.src[1:])
@@ -127,8 +24,6 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
   ctx[x] = after
 
 # *** fold moved AFTERs (hack for openpilot) ***
-# These temporary stores exist only in the schedule; they do not persist Tensor intermediates.
-pm_contiguous_to_store = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="c"), lambda c: c.clone())])
 pm_fold_moved_after = PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
   # replace ALU sources with AFTER versions found above
@@ -234,10 +129,13 @@ def expand_bitcast(bc:UOp) -> UOp|None:
   parts = [tmp>>8*i*ns for i in range(os//ns)]
   return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
 
-earliest_rewrites = mop_cleanup+pm_resolve_call_outputs+PatternMatcher([
-  # Inline calls with unbound outputs.
+earliest_rewrites = mop_cleanup+PatternMatcher([
+  # resolve calls with RETURNED inputs (inline the body)
   (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.has_unbound_outputs else None),
 
+  # resolve AFTER on RETURNED (call outputs)
+  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
+
   # resolve allreduce (must be bottom up)
   (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
 
@@ -317,9 +215,7 @@ pm_copy_to_store = PatternMatcher([
 def prepare_rangeify(sink:UOp) -> UOp:
   # prepare for rangeify
   tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
-  if OPENPILOT_HACKS:
-    tsink = graph_rewrite(tsink, pm_contiguous_to_store, bottom_up=True, name="materialize contiguous")
-    tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
+  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
   tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
   return tsink

tinygrad/tensor.py

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

tinygrad/uop/ops.py

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

tinygrad/uop/spec.py

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

tinygrad/uop/validate.py

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

tinygrad/uop/weak.py

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