remove more invalid (#16227)

test/backend/test_custom_kernel.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp, GlobalCounters, Context
+from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 
@@ -284,6 +284,30 @@ class TestCustomKernel(unittest.TestCase):
     self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
     self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
 
+  def test_invalids_into_custom_kernel_no_empty_kernel(self):
+    from tinygrad.engine.realize import compile_linear
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+    out = Tensor.invalids(*a.shape, dtype=a.dtype)
+    out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
+    compiled = compile_linear(out.schedule_linear())
+    for call in compiled.src:
+      prg = call.src[0]
+      if prg.op is not Ops.PROGRAM: continue
+      self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
+  def test_invalids_into_custom_kernel_with_beam(self):
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
+      out = Tensor.invalids(*a.shape, dtype=a.dtype)
+      out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
+      result = out.flatten().tolist()
+    self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")
+
   @unittest.skip("what are anonymous buffers?")
   def test_anonymous_buffers_in_function(self):
     """Test that custom kernels with anonymous output buffers work inside @function."""

tinygrad/callify.py

@@ -107,13 +107,16 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   subs:dict[UOp, UOp] = {}
   items:list[UOp] = []
   for s, t in zip(srcs, targets):
-    while s.op is Ops.AFTER: s = s.src[0]
+    after_deps:list[UOp] = []
+    while s.op is Ops.AFTER:
+      after_deps.extend(s.src[1:])
+      s = s.src[0]
     base = s.base
     if base.op in {Ops.CONTIGUOUS, Ops.BUFFER} and base.shape == t.shape and base not in subs:
       subs[base] = t.after(t.store(base.src[0])) if base.op is Ops.CONTIGUOUS else t
-      items.append(s)
+      items.append(s.after(*after_deps) if after_deps else s)
     else:
-      items.append(t.after(t.store(s)))
+      items.append(t.after(t.store(s), *after_deps))
   fxn = UOp.sink(*(x.substitute(subs) for x in items))
 
   # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)

tinygrad/schedule/rangeify.py

@@ -442,6 +442,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
 
   # remove invalid writes
   (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)),
+  (UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONST, arg=Invalid))), lambda: UOp(Ops.NOOP)),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.NOOP, src=()))), lambda x: x),
   (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.END, src=(UPat(Ops.NOOP, src=()),), allow_any_len=True))), lambda x: x),
 ])