diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 68e866a35f..55924ad47b 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -127,6 +127,41 @@ 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/test/backend/test_symbolic_ops.py b/test/backend/test_symbolic_ops.py index 8e7020dc2c..344bd54493 100644 --- a/test/backend/test_symbolic_ops.py +++ b/test/backend/test_symbolic_ops.py @@ -6,6 +6,15 @@ from examples.gpt2 import Attention import numpy as np class TestSymbolicOps(unittest.TestCase): + def test_negative_slice(self): + a = Tensor.rand(3, 10, 4) + for i in range(3, 10): + vi = Variable("i", 1, 10).bind(i) + # negative int bounds against a symbolic dim must resolve against the size, like slice.indices + np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6) + np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6) + def test_plus1(self): def f(a): return (a+1).realize() a = Tensor.rand(3, 10) diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 14632c3804..380045497e 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -90,8 +90,11 @@ class MovementMixin: if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step") start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop step = 1 if index.step is None else index.step + # resolve negative int bounds against the (possibly symbolic) size, like slice.indices + if isinstance(start, int) and start < 0: start = start + size + if isinstance(stop, int) and stop < 0: stop = stop + size if all_int((start, stop, step)): - # handle int slicing (resolve negative bounds, clamp, stride) + # handle int slicing (clamp, stride) *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size) bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound) return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False} diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 8c02c389d1..22f89362c6 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -25,7 +25,21 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp): # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce if dest.base in src.backward_slice_with_self: ctx[src] = None +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