Merge branch 'master' into rewrite_rangeify2

tinygrad/function.py

@@ -46,7 +46,7 @@ class _function(Generic[ReturnType]):
     params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values()
 
     # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None])
+    call_uops: list[UOp] = dedup([u for t in params if (u:=t._uop).device is not None or u.is_bound_var])
 
     # disable realize/schedule while this is running
     # run it and do surgery later
@@ -64,7 +64,10 @@ class _function(Generic[ReturnType]):
       raise RuntimeError(f"function return type {type(ret)} not supported")
 
     # replace the known inputs with params (using deduplicated slots)
-    subs = {x:x.param_like(i) for i,x in enumerate(call_uops)}
+    def make_param(x:UOp, i:int) -> UOp:
+      p = x.param_like(i)
+      return p.replace(arg=replace(p.arg, name=f"p{i}")) if x.is_bound_var else p
+    subs = {x:make_param(x, i) for i,x in enumerate(call_uops)}
     uret = uret.substitute(subs)
 
     # the BUFFERs that are left are the implicit inputs

tinygrad/mixin/op.py

@@ -549,6 +549,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     squares = (self - self.mean(axis=axis, keepdim=True)).square()
     n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
     numerator = squares.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
+    if resolve(n == 1, False) and correction >= 1: return self.sum(axis=axis, keepdim=keepdim).cast(output_dtype) * math.nan
     return numerator.div(smax(n - correction, 0)).cast(output_dtype)
 
   def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:

tinygrad/schedule/__init__.py

@@ -80,8 +80,9 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
+#from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.schedule.rangeify2 import get_kernel_graph
 from tinygrad.schedule.prepare import prepare_rangeify
-from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace

tinygrad/schedule/allreduce.py

@@ -58,8 +58,11 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
 
 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
-  if output is None: output = UOp.invalids(red.shape, dtype=red.dtype, device=red.device)
+  if output is None:
+    call_output = UOp.invalids(red.max_shape, dtype=red.dtype, device=red.device)
+    output = call_output.shrink_to(red.shape)
+  else: call_output = output
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(*red.arg)
-  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))
+  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(call_output, buf.contiguous(), name="allreduce", precompile=True))

tinygrad/schedule/multi.py

@@ -33,6 +33,8 @@ replace_allreduce = PatternMatcher([
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
+  # Identical device-less values do not need a multi wrapper.
+  (UPat(Ops.MSTACK, src=(UPat.var("x"),), allow_any_len=True, name="m"), lambda m,x: x if x.device is None and all_same(m.src) else None),
   # move shrink before MSTACK
   (UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), allow_any_len=True, name="shrink"), mstack_early_shrink),
   # move MSELECT before movement/ALU ops

tinygrad/schedule/prepare.py

@@ -12,6 +12,12 @@ def walk_mop(u:UOp):
   if u.op in GroupOp.Movement or u.op in {Ops.INDEX, Ops.UNSHARD}: return walk_mop(u.src[0])
   return u
 
+def has_buffer_view(u:UOp) -> bool:
+  # CALL argument lowering currently passes the base allocation, so only an
+  # offset-zero contiguous view backed by a real buffer can avoid a copy.
+  if u.has_buffer_identity(after_ok=True): return True
+  return (cv:=u.contiguous_view()) is not None and cv[1] == 0 and cv[0].has_buffer_identity(after_ok=True)
+
 def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
   if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
   while True:
@@ -150,9 +156,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
 
   # reshaping on STORE can be a NOOP
-  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
-                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
-   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
+  #(UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
+  #                      UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
+  # lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
 
   # ** store rules **
 
@@ -179,6 +185,19 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # handle size 0
   (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
 
+  # ** new prepare **
+
+  # CALL inputs need buffer identity (and to be flat)
+  (UPat(Ops.CALL, name="c"),
+   lambda c: c.replace(src=c.src[0:1]+tuple(x if has_buffer_view(x) else x.contiguous() for x in c.src[1:]))),
+
+  # MSTACK inputs need buffer identity
+  (UPat(Ops.MSTACK, name="c"),
+   lambda c: c.replace(src=tuple(x.contiguous() if not x.has_buffer_identity(after_ok=True) else x for x in c.src))),
+
+  # STORE to () is reshaped to (1,)
+  (UPat(Ops.STORE, name="s"), lambda s: s.src[0].reshape((1,)).store(s.src[1].reshape((1,))) if s.shape == () else None),
+
   # remove movement ops from SINK/AFTER. TODO: should be generic
   (UPat(Ops.SINK, name="s"), lambda s: s.replace(src=tuple(walk_mop(u) for u in s.src if u.op is not Ops.NOOP))),
   (UPat(Ops.AFTER, name="s"), lambda s: s.replace(src=(s.src[0],)+tuple(walk_mop(u) for u in s.src[1:] if u.op is not Ops.NOOP))),
@@ -196,12 +215,13 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
     return existing_buf.flatten().store(input_src)
   # create the output buffer
   buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  buf = buf.shrink_to(input_src.shape)
   # reshape back to input
   return buf.after(buf.store(input_src)).reshape(copy.shape)
 
 pm_copy_to_store = PatternMatcher([
   (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+  (UPat((Ops.COPY, Ops.CONTIGUOUS), name="copy"), convert_copy_to_store),
 ])
 
 @rewrite_group(new_ctx=False)

tinygrad/schedule/rangeify2.py

@@ -0,0 +1,426 @@
+from dataclasses import dataclass, field
+import itertools
+from tinygrad.dtype import AddrSpace, Invalid, strong_dtype
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, KernelInfo
+from tinygrad.uop.ops import graph_rewrite, AxisType, rewrite_group, remove_all_tags, resolve, shape_to_shape_arg
+from tinygrad.helpers import all_int, prod, VIZ, SPEC, Context, panic
+from tinygrad.schedule.indexing import BufferizeOpts, apply_movement_op
+from tinygrad.schedule.prepare import has_buffer_view
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen.simplify import pm_reduce_simplify
+
+# *** preparation ***
+
+fix_mselect_mstack = PatternMatcher([
+  # move RESHAPEs through MSELECT/MSTACK
+  (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
+   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
+])
+
+from tinygrad.helpers import all_same
+from tinygrad.uop.ops import _broadcast_shape
+
+def expand_broadcast(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shape = _broadcast_shape(*shapes)
+  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
+
+pm_lil_prepare_graph = PatternMatcher([
+  # expand broadcasts first
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
+])+fix_mselect_mstack
+
+# *** RANGE creation ***
+
+def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
+  if red.arg[1] == 0: return None
+  if idx is None and len(red.shape) > 0: return None
+  # TODO: is AxisType.REDUCE a real thing?
+  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
+  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
+
+def rangeify_on_store(ctx, x:UOp):
+  if x.shape == (): return None
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
+
+def rangeify_on_stage(ctx, x:UOp):
+  if x.src[0].shape == (): return None
+  # size 1 dims don't get ranges, they are reshaped out and back in
+  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
+    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
+  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
+  return x.replace(src=(x.src[0].index(*rngs), *rngs))
+
+pm_range_creation = PatternMatcher([
+  # reduce/store are what creates ranges
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
+  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
+  (UPat(Ops.STORE, name="x"), rangeify_on_store),
+  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
+])
+
+# *** RANGE migration ***
+
+# movement op on INDEX as a PatternMatcher
+def _mop_index(r:UOp, idx:UOp):
+  idxs = idx.src[1:]
+  if len(idxs) == len(r.shape):
+    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
+    if r.op is Ops.PAD:
+      # NOTE: neither 0 or ret.const_like(0) is correct here.
+      # const_like breaks because it adds casts, and 0 is wrong if ret is a bool
+      invalid_value = UOp.const(ret.dtype.const(0))
+      # insert invalid_value for PAD with where
+      a = UOp.const(True)
+      for s in UOp.sink(*ret.src[1:]).simplify().src:
+        if s.is_invalid: return invalid_value
+        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: a = a & s.src[0]
+      ret = a.where(ret, invalid_value)
+    return ret
+  if r.op is Ops.RESHAPE:
+    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
+    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
+      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
+      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
+      return ret if ret.shape == idx.shape else None
+
+# TODO: this should be in _mop_index
+def index_on_stack(stack:UOp, idx:UOp):
+  srcs = [s.index(*idx.src[2:]) for s in stack.src]
+  r0 = idx.src[1]
+  ret = srcs[-1]
+  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
+  return ret
+
+pm_range_migration = PatternMatcher([
+  # STAGE on shape () is nothing
+  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
+  # reshape of a single element shaped value to scalar is an index
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
+  # handle movement ops on INDEX
+  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
+  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
+  # move movement ops and INDEX after AFTER
+  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
+   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
+  # block bitcast that changes shape
+  (UPat(Ops.BITCAST, name="b").index(allow_any_len=True),
+   lambda b: panic(RuntimeError, "shape changing bitcast not allowed in rangeify") if b.src[0].shape != b.shape else None),
+  # pass index through elementwise
+  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
+   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
+  # INDEX without src is nothing (must be at the bottom)
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+])
+
+# *** split into kernels ***
+
+@dataclass
+class SplitCtx:
+  call_args:list[UOp] = field(default_factory=list)
+  buffers:dict[UOp, int] = field(default_factory=dict)
+  range_number:int = -1
+  addrspace:AddrSpace = AddrSpace.GLOBAL
+
+def _split_graph(ctx:SplitCtx, u:UOp) -> UOp|None:
+  if u.tag is not None: return None
+  if u.addrspace != ctx.addrspace: return None
+  if u.addrspace == AddrSpace.ALU: return u.param_like(-1).rtag().reshape(u.shape)
+
+  # A kernel takes each underlying buffer state once. In particular, AFTER and its buffer must use the same slot, with AFTER kept as the call
+  # argument so its dependencies are preserved.
+  key = u.buf_uop if u.op is Ops.AFTER else u
+  if (slot:=ctx.buffers.get(key)) is None:
+    slot = ctx.buffers[key] = len(ctx.call_args)
+    ctx.call_args.append(u)
+  elif u.op is Ops.AFTER:
+    ctx.call_args[slot] = u
+
+  # Parameters describe the max-sized physical allocation. A symbolic logical shape is a view of that allocation, not part of the PARAM itself.
+  param = u.param_like(slot).rtag().replace(src=(shape_to_shape_arg((u.max_numel(),)),))
+  return param.reshape(u.max_shape).shrink_to(u.shape)
+
+def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp|None:
+  if u.tag is not None: return None
+  ctx.range_number += 1
+  return u.replace(arg=(ctx.range_number, u.arg[-1])).rtag()
+
+pm_split_graph = pm_range_migration+PatternMatcher([
+  (UPat((Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK), name="u"), _split_graph),
+  (UPat(Ops.RANGE, name="u"), _renumber_range),
+])
+
+def _is_fully_invalid_state(x:UOp) -> bool:
+  while x.op in GroupOp.Movement|{Ops.INDEX}: x = x.src[0]
+  if x.op is not Ops.AFTER or len(x.src) != 2: return False
+  end = x.src[1]
+  st = end.src[0] if end.op is Ops.END else end
+  if st.op is not Ops.STORE or not st.src[1].base.is_invalid: return False
+  if end.op is not Ops.END: return st.src[0].max_numel() == x.max_numel()
+  covered = 1
+  for r in end.src[1:]:
+    if r.op is not Ops.RANGE or r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int): return False
+    covered *= r.src[0].val
+  return covered == x.max_numel()
+
+def split_store(x:UOp) -> UOp|None:
+  st = x.src[0] if x.op is Ops.END else x
+  if st.op is Ops.STORE and st.src[0].is_variable: return None
+  if st.op is Ops.STORE and st.src[0] is st.src[1]: return UOp(Ops.NOOP)
+  # A directly-invalid value makes this store a no-op. An AFTER carrying an
+  # invalid partial store is still a valid buffer state: uncovered elements
+  # must continue to read from the previous state.
+  if st.op is Ops.STORE and (st.src[1].base.is_invalid or _is_fully_invalid_state(st.src[1])): return UOp(Ops.NOOP)
+  ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True)
+  # TODO: params and args should be able to be in any order
+  ctx.addrspace = AddrSpace.ALU
+  ret = graph_rewrite(ret, pm_split_graph, ctx, name="split kernel (vars)", bottom_up=True)
+  ret = graph_rewrite(ret, remove_all_tags, name="remove split tags", bottom_up=True)
+  return ret.sink(arg=KernelInfo()).call(*ctx.call_args)
+
+split_kernels = PatternMatcher([
+  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
+])
+
+# cleanups
+
+def strip_zero_offset_shrink(x:UOp) -> UOp:
+  return x.src[0] if x.op is Ops.SHRINK and all(resolve(start == 0, False) for start,_ in x.marg) else x
+
+def no_indexing_calls(u:UOp):
+  new_srcs = []
+  for x in u.src:
+    if x.op is Ops.INDEX:
+      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
+      # TODO: we should add safety checks here for contiguous
+      new_srcs.append(x.src[0])
+    elif x.op is Ops.SHRINK:
+      # SHRINK with offset 0 is fine
+      new_srcs.append(strip_zero_offset_shrink(x))
+    elif x.op is Ops.MSTACK:
+      new_srcs.append(x.replace(src=tuple(strip_zero_offset_shrink(s) for s in x.src)))
+    else:
+      # everything else we pass through
+      new_srcs.append(x)
+  return u.replace(src=tuple(new_srcs))
+
+pm_no_indexing_calls = PatternMatcher([
+  (UPat(Ops.CALL, name="u"), no_indexing_calls),
+  (UPat(Ops.AFTER, name="u"), lambda u: u.replace(src=tuple(s for s in u.src if s.op is not Ops.NOOP))),
+])
+
+# *** main rangeify ***
+
+debug_tag_factor = PatternMatcher([
+  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[0][x] if x not in ctx[1] else 'REAL') if x.tag is None else None),
+])
+
+def remove_stage(ctx, x:UOp) -> UOp:
+  dtype = strong_dtype(x.dtype)
+  buf = UOp.new_buffer(x.arg.device, x.max_numel(), dtype, num=next(ctx))
+  val = x.src[0] if x.src[0].dtype == dtype else x.src[0].cast(dtype)
+  return buf.after(buf.reshape(x.shape).index(*x.src[1:]).store(val).end(*x.src[1:])).reshape(x.shape)
+
+pm_remove_stage = PatternMatcher([
+  (UPat(Ops.STAGE, name="x"), remove_stage),
+])+fix_mselect_mstack
+
+def remove_selected_stage(ctx:set[UOp], stage:UOp, idx:UOp) -> UOp|None:
+  return stage.src[0] if stage in ctx and stage.src[1:] == idx.src[1:] else None
+
+pm_remove_selected_stage = PatternMatcher([
+  (UPat(Ops.STAGE, name="stage").index(name="idx", allow_any_len=True), remove_selected_stage),
+])
+
+def inline_stage_index(stage:UOp, idx:UOp) -> UOp:
+  replacements, cache = dict(zip(stage.src[1:], idx.src[1:])), {}
+  def replace(x:UOp) -> UOp:
+    if x in replacements: return replacements[x]
+    if x.has_buffer_identity(after_ok=True) or x.op is Ops.STAGE: return x
+    if x not in cache: cache[x] = x.replace(src=tuple(replace(s) for s in x.src))
+    return cache[x]
+  return replace(stage.src[0])
+
+MAX_RECOMPUTE = 8
+MAX_SCALAR_RECOMPUTE = 64
+
+def recompute_cost(x:UOp, seen:set[UOp]|None=None) -> int|None:
+  if seen is None: seen = set()
+  if x in seen or x.op is Ops.STAGE or x.has_buffer_identity(after_ok=True): return 0
+  seen.add(x)
+  if x.op is Ops.REDUCE: return None
+  costs = [recompute_cost(s, seen) for s in (x.src[:1] if x.op is Ops.INDEX else x.src)]
+  return None if any(c is None for c in costs) else sum(c for c in costs if c is not None) + (x.op in GroupOp.Elementwise)
+
+def materialize_call_args(c:UOp) -> UOp:
+  srcs:list[UOp] = []
+  for x in c.src[1:]:
+    device = x.device or c.device
+    srcs.append(x if x.op is Ops.STAGE or x.is_bound_var or has_buffer_view(x) or x.shape == () or device is None
+                else x.bufferize(arg=BufferizeOpts(device=device)))
+  return c.replace(src=(c.src[0], *srcs))
+
+def materialize_mselect(m:UOp, x:UOp) -> UOp|None:
+  if x.device is None or x.op is Ops.STAGE or (x.op not in GroupOp.ALU and x.has_buffer_identity(after_ok=True)): return None
+  return m.replace(src=(x.bufferize(arg=BufferizeOpts(device=x.device)),))
+
+pm_materialize_call_args = PatternMatcher([
+  (UPat(Ops.CALL, name="c"), materialize_call_args),
+  (UPat(Ops.MSELECT, src=(UPat(name="x"),), name="m"), materialize_mselect),
+])
+
+@rewrite_group(new_ctx=False)
+def get_kernel_graph(sink:UOp) -> UOp:
+  tsink = graph_rewrite(sink, pm_lil_prepare_graph, bottom_up=True, name="prepare graph")
+  # Calls can only receive buffer states. Materialize lazy constants/computations instead of silently unwrapping them to a nonexistent base buffer.
+  tsink = graph_rewrite(tsink, pm_materialize_call_args, name="materialize call args")
+
+  read_cache:dict[UOp, set[UOp]] = {}
+  def read_buffers(x:UOp) -> set[UOp]:
+    if x not in read_cache:
+      read_cache[x] = {x.buf_uop} if x.has_buffer_identity(after_ok=True) else set().union(*(read_buffers(s) for s in x.src))
+    return read_cache[x]
+  stores = [u for u in tsink.toposort() if u.op is Ops.STORE and not u.src[0].is_variable]
+  dests = [u.src[0].buf_uop for u in stores]
+  reads = [read_buffers(u.src[1]) for u in stores]
+  force_stage:set[UOp] = set()
+  param_writes = [i for i,dest in enumerate(dests) if dest.op is Ops.PARAM]
+  if len(param_writes) == 2:
+    i, j = param_writes
+    if stores[j].src[1].op_in_backward_slice_with_self(Ops.REDUCE) and dests[i] is not dests[j] and \
+       dests[i] in reads[j] and dests[j] in reads[i]: force_stage.add(stores[j].src[1])
+
+  # add safe STAGEs to never duplicate compute
+  # we compute the number of times a buffer is consumed. if > 1, we realize
+  realize = {}
+  consumes = {tsink:0}
+  for u in reversed(tsink.toposort()):
+    assert u in consumes, f"{u.op} not in consumes"
+    if u in force_stage:
+      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device, removable=False))
+      consumes[u] = 1
+    elif (u.op in GroupOp.ALU or u.op is Ops.REDUCE) and consumes[u] > 1 and u.device is not None:
+      # TODO: rename to stage
+      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device))
+      consumes[u] = 1
+    if u.op is Ops.STORE: consumes[u] = 1
+    if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel()
+    for i,s in enumerate(u.src):
+      if s not in consumes: consumes[s] = 0
+      if u.op is not Ops.STORE or i > 0:
+        consumes[s] += consumes[u]
+  if VIZ:
+    with Context(TRACK_MATCH_STATS=0): ctags = graph_rewrite(tsink, debug_tag_factor, ctx=(consumes, realize), bottom_up=True)
+    graph_rewrite(ctags, PatternMatcher([]), name="View Consumes")
+
+  # add stages
+  tsink = graph_rewrite(tsink.substitute(realize), remove_all_tags, name="untag")
+
+  # simple rangeify
+  tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
+
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="pre-fusion reduce simplify")
+
+  # remove stage boundaries when this doesn't duplicate expensive compute or nest reductions
+  while 1:
+    staged:dict[UOp, list[UOp]] = {}
+    children:dict[UOp, list[UOp]] = {}
+    for u in tsink.toposort():
+      for s in u.src: children.setdefault(s, []).append(u)
+      if u.op is Ops.INDEX and u.src[0].op is Ops.STAGE: staged.setdefault(u.src[0], []).append(u)
+
+    boundary_cache:dict[UOp, set[UOp]] = {}
+    def boundaries(x:UOp) -> set[UOp]:
+      if x not in boundary_cache:
+        boundary_cache[x] = set().union(*({c} if c.op in {Ops.STAGE, Ops.STORE, Ops.CALL} else boundaries(c)
+                                          for c in children.get(x, [])))
+      return boundary_cache[x]
+
+    reduce_cache:dict[UOp, bool] = {}
+    def feeds_reduce(x:UOp) -> bool:
+      if x not in reduce_cache:
+        reduce_cache[x] = any(c.op is Ops.REDUCE or (c.op not in {Ops.STAGE, Ops.STORE, Ops.CALL} and feeds_reduce(c))
+                              for c in children.get(x, []))
+      return reduce_cache[x]
+
+    def boundary_work(boundary:UOp) -> int|None:
+      value = boundary.src[1] if boundary.op is Ops.STORE else boundary.src[0]
+      if any(r.src[0].op is not Ops.CONST or not isinstance(r.src[0].val, int) for r in value.ranges): return None
+      return prod(r.src[0].val for r in value.ranges)
+
+    def recompute_work(idxs:list[UOp]) -> int|None:
+      works = [boundary_work(next(iter(bs))) for idx in idxs if len(bs:=boundaries(idx)) == 1]
+      return sum(x for x in works if x is not None) if len(works) == len(idxs) and all(x is not None for x in works) else None
+
+    replacements:dict[UOp, UOp] = {}
+    range_replacements:dict[UOp, UOp] = {}
+    selected_stages:set[UOp] = set()
+    for stage,idxs in staged.items():
+      if not stage.arg.removable or children.get(stage) != idxs: continue
+      inlinable = stage.src[0].op in GroupOp.ALU or stage.src[0].op is Ops.REDUCE
+      cost = recompute_cost(stage.src[0])
+
+      # passthrough stages don't duplicate compute when indexing is unchanged
+      if not inlinable:
+        if len(idxs) == 1 and stage.src[1:] == idxs[0].src[1:]:
+          replacements[idxs[0]] = stage.src[0]
+          break
+        continue
+
+      # duplicate cheap elementwise stages; reductions require identical indexing into one output boundary
+      if len(idxs) > 1:
+        work = recompute_work(idxs)
+        if cost is not None and cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel() * len(idxs):
+          replacements.update((idx, inline_stage_index(stage, idx)) for idx in idxs)
+        elif cost is None and all(idx.src[1:] == idxs[0].src[1:] for idx in idxs) and not any(feeds_reduce(idx) for idx in idxs):
+          stage_boundaries = set().union(*(boundaries(idx) for idx in idxs))
+          if len(stage_boundaries) == 1 and boundary_work(next(iter(stage_boundaries))) == stage.max_numel():
+            range_replacements.update(zip(stage.src[1:], idxs[0].src[1:]))
+            selected_stages.add(stage)
+        if replacements or range_replacements: break
+        continue
+
+      idx = idxs[0]
+      stage_boundaries = boundaries(idx)
+      work = boundary_work(next(iter(stage_boundaries))) if len(stage_boundaries) == 1 else None
+      scalar = stage.max_numel() == 1 and cost is not None and cost <= MAX_SCALAR_RECOMPUTE and all(r.op is Ops.CONST for r in idx.src[1:])
+      small = cost is not None and cost <= MAX_SCALAR_RECOMPUTE and stage.max_numel() <= 8 and idx.max_numel() <= 8
+      if cost is not None:
+        if stage.src[0].op_in_backward_slice_with_self(Ops.THREEFRY) or scalar or small or \
+           (cost <= MAX_RECOMPUTE and work is not None and work <= stage.max_numel()):
+          replacements[idx] = inline_stage_index(stage, idx)
+      elif len(stage_boundaries) == 1 and not feeds_reduce(idx) and work == stage.max_numel():
+        if all(r.op is Ops.RANGE for r in idx.src[1:]):
+          range_replacements.update(zip(stage.src[1:], idx.src[1:]))
+          selected_stages.add(stage)
+        else: replacements[idx] = inline_stage_index(stage, idx)
+      if replacements or range_replacements: break
+
+    if not replacements and not range_replacements: break
+    tsink = tsink.substitute(replacements).substitute(range_replacements)
+    if selected_stages:
+      selected_stages = {stage.substitute(range_replacements) for stage in selected_stages}
+      tsink = graph_rewrite(tsink, pm_remove_selected_stage, ctx=selected_stages, name="remove selected stage")
+
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify, name="reduce simplify")
+
+  # ***** MERGING AND SPLITTING (should be totally optional) *****
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Merged Rangeify")
+
+  next_buffer_num = itertools.count(1000)
+  tsink = graph_rewrite(tsink, symbolic+pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
+  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
+  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if SPEC:
+    # validate the kernel graph
+    from tinygrad.uop.spec import type_verify, spec_kernel_graph
+    type_verify(tsink, spec_kernel_graph, enter_calls=False)
+  return tsink
+

tinygrad/tensor.py

@@ -114,7 +114,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
   # add the outputs to the call
   srcs = c.src[0].src
   resolved = [c.gettuple(i) for i in range(len(srcs))]
-  outs = tuple(r.empty_like() for r in resolved)
+  # CALL outputs are max-sized physical buffers. Keep symbolic shapes as views so writable arguments and returned buffers stay identical.
+  outs = tuple(r.pad_to(r.max_shape).empty_like() for r in resolved)
   targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
 
   subs:dict[UOp, UOp] = {}
@@ -133,11 +134,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
 
   # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
   new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg)
-  rets = tuple(o.after(new_call) for o in outs)
-
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
-  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
+  # NOTE: use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes.
+  rets = tuple(o.after(new_call).shrink_to(rs.shape) for o,rs in zip(outs, resolved))
 
   return UOp.maketuple(*rets)
 
@@ -220,6 +218,32 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.AFTER, name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.is_bound_var else None),
 ])
 
+def _check_state_cycles(sink:UOp):
+  # Track only the first storage state on each path. Combining a raw read with an assigned state is stale, while distinct AFTERs can be valid
+  # independent snapshots (for example RNG state). Memoizing the two state sets avoids repeatedly walking large training graphs.
+  states:dict[UOp, tuple[frozenset[UOp], frozenset[UOp]]] = {}
+  state:tuple[frozenset[UOp], frozenset[UOp]]
+  for u in sink.toposort(enter_calls=False):
+    if u.op is Ops.BUFFER and u.addrspace == AddrSpace.GLOBAL: state = (frozenset((u,)), frozenset())
+    elif u.op is Ops.AFTER and u.addrspace == AddrSpace.GLOBAL:
+      key = u.buf_uop
+      stores = [x for x in u.src[1:] if x.op is Ops.STORE and x.src[0].buf_uop is key]
+      # Ordering dependencies can STORE to another buffer without changing this state. Self-dependent updates continue the existing state lineage;
+      # only a write independent of the old value creates a conflicting state.
+      self_update = any(key in states[x.src[1]][0] or key in states[x.src[1]][1] for x in stores)
+      state = states[u.src[0]] if not stores or self_update else (frozenset(), frozenset((key,)))
+    else:
+      srcs = u.src[1:] if u.op in {Ops.CALL, Ops.FUNCTION} else u.src
+      raw = frozenset().union(*(states[x][0] for x in srcs))
+      assigned = frozenset().union(*(states[x][1] for x in srcs))
+      # CONTIGUOUS snapshots stale raw reads before a pending assignment, but cannot make a post-assignment read happen earlier.
+      state = (frozenset() if u.op is Ops.CONTIGUOUS else raw, assigned)
+    states[u] = state
+    if u.op in GroupOp.ALU:
+      branches = [states[x] for x in u.src]
+      if any((a[0] & b[1]) or (a[1] & b[0]) for i,a in enumerate(branches) for b in branches[i+1:]):
+        raise RuntimeError("cycle detected while combining buffer states")
+
 @rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
@@ -406,7 +430,9 @@ class Tensor(RandMixin):
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+    _check_state_cycles(big_sink)
+    big_sink, becomes_map = transform_to_call(big_sink)
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
 

tinygrad/uop/ops.py

@@ -1173,12 +1173,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
     return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
   def param_like(self, slot:int):
-    # Variables become ALU params in the call body; the stored value (if bound) stays in the call args
-    if self.is_bound_var or self.is_variable:
-      b = self.src[0] if self.op is Ops.AFTER else self
-      return UOp(Ops.PARAM, src=b.src, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
-    addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
+    # if it's a PARAM or BUFFER, we just replace the slot
+    buf = self
+    while buf.op is Ops.AFTER: buf = buf.src[0]
+    if buf.op in {Ops.PARAM, Ops.BUFFER}: return UOp(Ops.PARAM, src=buf.src, arg=replace(buf.arg, slot=slot))
+    # otherwise we create a new param
+    addrspace = buf.addrspace if buf.addrspace is not None else AddrSpace.GLOBAL
+    return UOp.param(slot, buf.dtype, buf.shard_shape if buf.axis is not None else buf._shape, self.device, addrspace=addrspace, axis=buf.axis)
 
   @staticmethod
   def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

tinygrad/uop/spec.py

@@ -257,12 +257,18 @@ spec_kernel_graph = PatternMatcher([
   # const + stack to make vconsts and shape args. a 0-size/bound reduce keeps its const casted
   (UPat(Ops.CONST, src=()), lambda: True),
   (UPat(Ops.CAST, src=(UPat(Ops.CONST, src=()),)), lambda: True),
+  # symbolic shape expressions can remain outside kernels (for example, flattening (n, 4) produces n*4)
+  (UPat(GroupOp.ALU, name="x"), lambda x: x.dtype in dtypes.ints+(dtypes.weakint,) or None),
   (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
   # linear for more kernels (TODO: we should enter non sink calls)
   #(UPat(Ops.LINEAR), lambda: True),
   # param is outside buffer, buffer is local buffer
   (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
   (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.GLOBAL, AddrSpace.ALU)),
+  # indexing/movement views on kernel buffers/call results are allowed to carry symbolic shape expressions
+  (UPat({Ops.INDEX}|GroupOp.Movement,
+        src=(UPat({Ops.INDEX}|GroupOp.Movement|{Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST}),),
+        allow_any_len=True), lambda: True),
   # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
   (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
   # mstack/mselect