two stage removal

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

@@ -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