diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 55924ad47b..68e866a35f 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -127,41 +127,6 @@ class TestCustomKernel(unittest.TestCase): # https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3]) - def test_lazy_const_srcs_with_reduce(self): - # lazy const expressions above a custom kernel call don't resolve to a buffer state, they must be realized. - # without this, the rangeify doesn't assign ranges to the subgraph above the call and reduce conversion crashes - x = Tensor.linspace(-1.0, 1.0, 64) # Tensor.arange is cumsum-based, so this contains a REDUCE with no buffer anchor - out = Tensor.empty_like(x) - def copy_kernel(out:UOp, inp:UOp) -> UOp: - i = UOp.range(inp.numel(), 0) - return UOp.group(out[i].store(inp[i])).end(i).sink(arg=KernelInfo(name="copy")) - # forge the call like llm/kernels does: params and call args, no Tensor.custom_kernel contiguous - params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate((out.uop, x.uop))) - call = copy_kernel(*params).call(out.uop, x.uop) - np.testing.assert_allclose(Tensor(out.uop.after(call)).realize().numpy(), x.realize().numpy(), rtol=1e-6) - - def test_mixed_buffer_and_lazy_const_srcs(self): - # a computed input must be realized even if one of its sources resolves to a buffer: the CALL gives the whole - # subgraph no ranges, so the const branch still crashes reduce conversion if only the buffer branch is found - x = Tensor.linspace(-1.0, 1.0, 64) # lazy const expression with a REDUCE - y = Tensor.ones(64).contiguous().realize() - out = Tensor.empty_like(x) - def copy_kernel(out:UOp, inp:UOp) -> UOp: - i = UOp.range(inp.numel(), 0) - return UOp.group(out[i].store(inp[i])).end(i).sink(arg=KernelInfo(name="copy")) - for expr in (y + x, x + y, y * 2.0): # buffer on either side, and a scalar const over a buffer - params = tuple(UOp.placeholder_like(u, slot=i) for i,u in enumerate((out.uop, expr.uop))) - call = copy_kernel(*params).call(out.uop, expr.uop) - np.testing.assert_allclose(Tensor(out.uop.after(call)).realize().numpy(), expr.realize().numpy(), rtol=1e-6) - # view-only movement ops over a buffer resolve to the buffer state and must NOT be realized - expr = y.reshape(8, 8).reshape(64) - expected = expr.numpy() - params = tuple(UOp.placeholder_like(u, slot=i) for i,u in enumerate((out.uop, expr.uop))) - call = copy_kernel(*params).call(out.uop, expr.uop) - GlobalCounters.kernel_count = 0 - np.testing.assert_allclose(Tensor(out.uop.after(call)).realize().numpy(), expected, rtol=1e-6) - self.assertEqual(GlobalCounters.kernel_count, 1, "a view over a buffer should not add a realize kernel") - def test_simple_sharded(self): devs = ("CPU:0", "CPU:1") diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 22f89362c6..740c66ebe7 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -27,19 +27,7 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp): BUFFER_STATE_OPS: set[Ops] = {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND} -def realize_custom_kernel_srcs(ctx:dict[UOp, None], c:UOp) -> None: - # the inputs of a custom kernel must resolve to a buffer state. realize the ones that don't (e.g. lazy const - # expressions above the call), otherwise a reduce in that subgraph has no ranges and crashes in rangeify. - # NOTE: only view-only movement ops preserve the underlying buffer. anything computed (ALU, REDUCE, ...) must be - # realized even if one of its sources is a buffer, since the CALL gives the whole subgraph no ranges - for s in c.src[1:]: - t = s - while t.op in GroupOp.Movement and len(t.src): t = t.src[0] - if t.op not in BUFFER_STATE_OPS: ctx[s] = None - pm_generate_realize_map = PatternMatcher([ - # realize the inputs of custom kernel calls - (UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="c", allow_any_len=True), realize_custom_kernel_srcs), # always realize (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize), # realize srcs of these