diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 29328fe217..41ed9a62c3 100644 --- a/test/backend/test_custom_kernel.py +++ b/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.""" diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 85b1e11c6e..75d2f48137 100644 --- a/tinygrad/callify.py +++ b/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) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index f3f375fa0a..66b07307a6 100644 --- a/tinygrad/schedule/rangeify.py +++ b/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), ])