The smallest promotion

tinygrad_repo/tinygrad/function.py

@@ -1,7 +1,7 @@
 import functools, time
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
-from tinygrad.uop.ops import UOp, Ops, ProgramInfo, graph_rewrite, PatternMatcher, UPat
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_state_dict
 
@@ -17,13 +17,11 @@ pm_ctx = PatternMatcher([
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
 ])
 
-def write_only_outputs(uret:UOp) -> set[UOp]:
-  ret: set[UOp] = set()
-  for call in uret.backward_slice_with_self:
-    if call.op is Ops.CALL and call.src[0].op is Ops.SINK:
-      info = ProgramInfo.from_sink(call.src[0])
-      ret.update(call.src[1+i].buf_uop for i in set(info.outs)-set(info.ins))
-  return ret
+def invalid_outputs(uret:UOp) -> set[UOp]:
+  # invalids() returns fresh write-only scratch: a clone storing CONST(Invalid)
+  # don't capture it as an input; only skip fresh buffers, not realized ones
+  return {u.src[0].buf_uop for u in uret.backward_slice_with_self
+          if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
@@ -49,8 +47,10 @@ class _function(Generic[ReturnType]):
     # run it and do surgery later
     with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
       _function.depth += 1
-      ret = self.fxn(*args, **kwargs)
-      _function.depth -= 1
+      try:
+        ret = self.fxn(*args, **kwargs)
+      finally:
+        _function.depth -= 1
     if isinstance(ret, Tensor):
       uret = ret.uop
     elif isinstance(ret, tuple) and all(isinstance(x, Tensor) for x in ret):
@@ -59,36 +59,24 @@ class _function(Generic[ReturnType]):
       raise RuntimeError(f"function return type {type(ret)} not supported")
 
     # replace the known inputs with params (using deduplicated slots)
-    subs = {}
-    for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
+    subs = {x:x.param_like(i) for i,x in enumerate(call_uops)}
     uret = uret.substitute(subs)
 
-    # add contiguous to call_uops
-    #call_uops = [x.contiguous() for x in call_uops]
-
     # the BUFFERs that are left are the implicit inputs
     num_explicit = len(call_uops)
-    uret = graph_rewrite(uret, pm_ctx, (call_uops, write_only_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
+    uret = graph_rewrite(uret, pm_ctx, (call_uops, invalid_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
       if implicit_buffers:
-        buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.arg}, device={b.device}" for i,b in enumerate(implicit_buffers))
+        buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.max_numel()}, device={b.device}" for i,b in enumerate(implicit_buffers))
         raise RuntimeError(f"function {name} has {len(implicit_buffers)} implicit buffer(s), but allow_implicit=False\n  {buf_strs}")
 
-    # assign output
-    #pbuffer = uret.param_like(len(call_uops))
-    #assigned = pbuffer.assign(uret).sink()
-    #buffer = UOp.new_buffer(pbuffer.device, pbuffer.size, pbuffer.dtype).reshape(uret.shape)
-    #call = assigned.call(*call_uops, buffer, name=name)
-    #ret = buffer.after(call)
-
     fret = uret.call(*call_uops, grad_fxn=self.grad_fxn, name=name, precompile=self.precompile,
                      precompile_backward=self.precompile_backward)
 
     if DEBUG >= 2:
-      #signature = [(x._shape, x.dtype, x.device) for x in call_uops]
-      print("  "*_function.depth+f"function {uret.key.hex()[:8]} in {(time.perf_counter()-st)*1000:8.2f} ms: {name}") # with sig {signature}")
+      print("  "*_function.depth+f"function {uret.key.hex()[:8]} in {(time.perf_counter()-st)*1000:8.2f} ms: {name}")
 
     if isinstance(ret, tuple):
       return cast(ReturnType, tuple(Tensor(fret.gettuple(i)) for i in range(len(ret))))