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

extra/benchmark_llm.py

@@ -1,31 +0,0 @@
-import argparse, time
-from tinygrad.llm.model import Transformer
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", required=True, help="path to gguf model")
-  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
-  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
-  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
-  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
-  args = parser.parse_args()
-
-  st = time.perf_counter()
-  model, _ = Transformer.from_gguf(args.model, args.max_context)
-  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
-
-  st = time.perf_counter()
-  model.warmup()
-  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
-
-  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
-  gen = model.generate(prompt, chunk_size=args.chunk_size)
-  st = time.perf_counter()
-  # first token is time-to-first-token; counted as part of prefill
-  output = [next(gen)]
-  pt = time.perf_counter()
-  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
-
-  for _ in range(args.decode_tokens): output.append(next(gen))
-  et = time.perf_counter()
-  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)

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

tinygrad/schedule/indexing.py

@@ -27,7 +27,19 @@ 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