diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 9e1eaf002c..0caa6dfdce 100644 --- a/tinygrad/schedule/rangeify.py +++ b/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 diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 4879a6daa6..096a540188 100644 --- a/tinygrad/uop/__init__.py +++ b/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 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 72bb277b1e..fbf9d618be 100644 --- a/tinygrad/uop/ops.py +++ b/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