schedule: realize custom kernel inputs that don't resolve to a buffer state

rangeify assigns ranges backward from consumers and CALL contributes none,
so the subgraph above a custom kernel input gets no ranges unless something
in it is realized, and reduce conversion crashes with a KeyError. realize
call inputs that don't resolve to a buffer state.

only view-only movement ops preserve the underlying buffer: anything
computed (ALU, REDUCE, ...) must be realized even if one of its sources
resolves to a buffer, since the whole subgraph above the call has no
ranges. unwrapping src[0] unconditionally missed const branches hanging
off non-src[0] children and silently resolved REDUCEs to their source
buffer. includes regression tests for pure const, mixed buffer+const, and
view-over-buffer inputs

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")
 

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)

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}

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