fixups

tinygrad/renderer/cstyle.py

@@ -1,7 +1,7 @@
 from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
-from tinygrad.codegen.opt import tc
+from tinygrad.codegen.opt import tc, axis_letters
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
@@ -163,7 +163,7 @@ class CStyleLanguage(Renderer):
       # naming
       prefix = None
       if u.op is Ops.SPECIAL: r[u] = u.arg
-      elif u.op is Ops.RANGE: r[u] = "ridx"+range_str(u)
+      elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
       else:
         prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                   Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",

tinygrad/schedule/rangeify.py

@@ -211,6 +211,38 @@ pm_cleanups = pm_mops+PatternMatcher([
    lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
 ])
 
+def second_stage_removal(src:UOp, buf:UOp):
+  if buf.arg.addrspace != AddrSpace.GLOBAL: return None
+  # if it's user contiguous, we never remove it
+  if src.op in ALWAYS_RUN_OPS: return None
+  accessed_buffers: list[UOp] = []
+  def red_gate(x:UOp):
+    if x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL:
+      accessed_buffers.append(x)
+      return False
+    if x.op is Ops.BUFFER:
+      accessed_buffers.append(x)
+    return True
+  src.toposort(gate=red_gate)
+  del red_gate
+  accessed_buffers = dedup(accessed_buffers)
+  out_in_ratio = prod(buf.shape) / sum([x.size for x in accessed_buffers])
+  print(f"ratio {out_in_ratio:.2f} {buf.shape} from {[x.size for x in accessed_buffers]}")
+  if out_in_ratio > 10: return buf.replace(op=Ops.REMOVE, arg=None)
+
+pm_cleanups_2 = PatternMatcher([
+  (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf"), second_stage_removal),
+])
+
+def do_remove(rem:UOp, idx:UOp):
+  assert len(rem.src) == len(idx.src), f"remove on wrong bufferize, {len(rem.src)} != {len(idx.src)}"
+  return rem.src[0].substitute({k:v for k,v in zip(rem.src[1:], idx.src[1:]) if k.op is not Ops.CONST})
+
+pm_cleanup_remove = PatternMatcher([
+  # NOTE: this is better than substitute
+  (UPat(Ops.REMOVE, name="rem").f(Ops.INDEX, allow_any_len=True, name="idx"), do_remove),
+])
+
 def late_buffer_view(t:UOp, b:UOp):
   if isinstance(b.device, str) and (b.device.startswith("DISK") or b.device.startswith("TINYFS")):
     rngs = b.src[1:]
@@ -506,6 +538,8 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
   # TODO: can you substitute and remove costly buffers at the same time?
   tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
+  tsink = graph_rewrite(tsink, pm_cleanups_2, name="remove (more) costly buffers")
+  tsink = graph_rewrite(tsink, pm_cleanup_remove, name="actually remove")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph

tinygrad/uop/__init__.py

@@ -20,6 +20,7 @@ class Ops(FastEnum):
 
   # create buffer
   BUFFERIZE = auto()
+  REMOVE = auto()
   SUBSTITUTE = auto()
 
   # ops that adjust the behavior of the scheduler

tinygrad/uop/ops.py

@@ -17,7 +17,7 @@ class AxisType(Enum):
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
   THREAD = auto()
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
+range_start = {Ops.BUFFERIZE: 1, Ops.REMOVE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -186,7 +186,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
       case Ops.BUFFER: return (self.arg,)
       case Ops.BUFFER_VIEW: return (self.arg[0],)
-      case Ops.BUFFERIZE: return tuple([int(r.vmax+1) for r in self.src[1:]])
+      case Ops.BUFFERIZE | Ops.REMOVE: return tuple([int(r.vmax+1) for r in self.src[1:]])
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
 
       # passthrough ops