fix _prepare_jit_inputs for weak [pr] (#17222)

test/unit/test_dtype_weak.py

@@ -1,10 +1,11 @@
 import tempfile, unittest, math
 
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
 from tinygrad.uop.ops import UOp, Ops, dtype_from_uop
 from tinygrad.uop.spec import spec_shared, type_verify
+from tinygrad.engine.jit import JitError
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -151,6 +152,18 @@ class TestWeakMaterializationEntries(unittest.TestCase):
       for entry in (lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
         with self.assertRaises(RuntimeError): entry(weak_val())
 
+  def test_weak_is_virtual(self):
+    # NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak
+    devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
+    for t in (Tensor.const(dtypes.weakfloat, 0.5), devful):
+      self.assertTrue(t.uop.is_virtual)
+      # realize is a no-op, so a weak input can never become the real buffer TinyJit needs
+      with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t)
+    # callify must not silently commit a weak CONTIGUOUS to storage
+    c = devful.alu(Ops.CONTIGUOUS)
+    c.callify()
+    self.assertIs(c.dtype, dtypes.weakfloat)
+
   def test_empty_reads_commit(self):
     for weak, strong in ((dtypes.weakfloat, dtypes.default_float),):
       empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))

tinygrad/callify.py

@@ -38,8 +38,8 @@ add_tags = PatternMatcher([
 ])
 
 def replace_contig_with_store_after(u:UOp):
-  # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs)
-  if u.device is None: return None
+  # 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/TINYFS tensors, they are left alone
@@ -203,7 +203,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   # uop list is a list in the original_sink graph and we can map to the tags later
   # same predicate as Tensor.realize
-  ctx = AllocCtx(bases={base for x in big_sink.src if (base:=x.base).device is not None and not base.has_buffer_identity()
+  ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity()
                         and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU})
 
   # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up

tinygrad/engine/jit.py

@@ -207,7 +207,7 @@ def _prepare_jit_inputs(args, kwargs):
     it = x if isinstance(x, (tuple,list)) else x.values() if isinstance(x, dict) else []
     tensors += [t for t in it if t.__class__ is Tensor and not any(t is y for y in tensors)]
   def get_input_uops() -> list[UOp]: return flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors])
-  if any(u.device is None for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()")
+  if any(u.is_virtual for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()")
   if len(unrealized_tensors := [x for x in tensors if not x.uop.is_realized]): Tensor.realize(*unrealized_tensors)
   input_uops = get_input_uops()
   # collect buffer UOps (including MultiBuffer)

tinygrad/tensor.py

@@ -189,7 +189,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 x.uop.device is not None and not x.uop.has_buffer_identity() and x.dtype not in dtypes.weaks]
+    to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
     if len(to_realize):
       run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
     return self

tinygrad/uop/ops.py

@@ -828,6 +828,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     for x in self.src:
       if x.device is not None: return x.device
     return None
+  @property
+  def is_virtual(self) -> bool:
+    # NOTE: no device means no place to store, weak means no width to store. neither can back a buffer as-is
+    # TODO: unify with has_buffer_identity
+    return self.device is None or self.dtype in dtypes.weaks
   @recursive_property
   def addrspace(self) -> AddrSpace|None:
     if self.op is Ops.PARAM: return self.arg.addrspace