diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 10dd40c0d0..eb18d12012 100644 --- a/test/unit/test_dtype_weak.py +++ b/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),)) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 0fce987b0c..7280ab4aad 100644 --- a/tinygrad/callify.py +++ b/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 diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 221577f4ac..d0ab840bd0 100644 --- a/tinygrad/engine/jit.py +++ b/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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index b905717b4a..699f45e280 100644 --- a/tinygrad/tensor.py +++ b/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 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 6a006b8bf1..ae91af7939 100644 --- a/tinygrad/uop/ops.py +++ b/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