From 851e5727d2c6627ea842f22bca02b70878ea8d7c Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Thu, 9 Jul 2026 17:32:35 -0400 Subject: [PATCH] remove late buffer view (#16941) --- tinygrad/callify.py | 31 ++++++++++++++++++++----------- tinygrad/schedule/rangeify.py | 32 +------------------------------- 2 files changed, 21 insertions(+), 42 deletions(-) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index c0bc98f39d..563a8a3c47 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -50,10 +50,10 @@ def replace_contig_with_store_after(u:UOp): def replace_store_after_with_contig(u:UOp, src:UOp): assigned_to = u while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.MULTI}: assigned_to = assigned_to.src[0].base - if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag) + if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag) def _make_buffer_view(src:UOp) -> UOp|None: - """If movement ops on src collapse to a contiguous range, return SLICE.reshape(src.shape). Otherwise None.""" + """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None.""" if (offset := src.contiguous_view_offset()) is None: return None buf = src.base if buf.op is Ops.SLICE: @@ -61,25 +61,26 @@ def _make_buffer_view(src:UOp) -> UOp|None: buf = buf.src[0] if byte_offset % buf.dtype.itemsize != 0: return None offset = byte_offset // buf.dtype.itemsize - return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.index, offset)), src.numel()).reshape(src.shape) + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.index, offset)), src.numel()) def contiguous_mops_to_view(c:UOp, src:UOp): - """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(SLICE) when movement ops collapse to a contiguous range.""" + """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" buf = src.base if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.MULTI}: return None - if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE}: return None + if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None + if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None # no symbolic shape if not all_int(c.shape): return None # check if view is supported from tinygrad.device import Device - devs = (c.device,) if isinstance(c.device, str) else c.device + devs = (src.device,) if isinstance(src.device, str) else src.device if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None - # NOTE: this contiguous is removed because this SLICE/RESHAPE has_buffer_identity if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None: - return view.contiguous(tag=c.tag) + view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape) + return c.replace(src=(view,)) if c.op is Ops.COPY else view # for MULTI tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result if not isinstance(c.device, str): @@ -87,7 +88,7 @@ def contiguous_mops_to_view(c:UOp, src:UOp): resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") if resolved.op is not Ops.MULTI: return None if (view := _make_buffer_view(resolved.src[0])) is None: return None - return view.multi(resolved.arg).contiguous(tag=c.tag) + return view.reshape(resolved.src[0].shape).multi(resolved.arg).contiguous(tag=c.tag) return None @@ -142,8 +143,16 @@ pm_early_transform_tensor_graph = PatternMatcher([ # resolve TUPLE+GETTUPLE (for precompiled calls) (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), - # CONTIGUOUS(MOPS(BUFFER/SLICE)) → CONTIGUOUS(SLICE) when movement ops collapse to contiguous range - (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="c"), contiguous_mops_to_view), + # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range + (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view), + + # remove contiguous on movement ops before a copy on disk + (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"), lambda x,copy: + copy.replace(src=(x,)+copy.src[1:], tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None), + # push copy past movement ops to disk + (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy: + x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:], tag=None),)+x.src[1:]) \ + if isinstance(x.device, str) and x.device.startswith("DISK") else None), # add CONTIGUOUS to tagged UOps (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"), diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 3ce18cb402..b5b746fd34 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -140,14 +140,6 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ # remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations) (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]), - # remove contiguous on movement ops before a copy on disk - (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"), - lambda x,copy: copy.replace(src=(x,)+copy.src[1:]) if isinstance(x.device, str) and x.device.startswith("DISK") else None), - # push copy past movement ops to disk - (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, allow_any_len=True, name="copy"), - lambda x,copy: x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:]),)+x.src[1:]) \ - if isinstance(x.device, str) and x.device.startswith("DISK") else None), - # SINK only ever references the base (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.base for y in x.src))), @@ -331,28 +323,6 @@ pm_remove_bufferize = PatternMatcher([ (UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x), ]) -def late_buffer_view(t:UOp, b:UOp): - if not (isinstance(b.device, str) and b.device.startswith(("DISK", "TINYFS"))): return b - shape = b.shape - size = prod(shape) - - # walk up for the INDEX - x = t - while not any(u.op is Ops.INDEX for u in x.src): - assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise" - x = x.src[0] - x = next(u for u in x.src if u.op is Ops.INDEX) - assert x.op is Ops.INDEX, "must be INDEX" - - if len(shape) == 0: offset = x.src[1].arg - else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0) - - return b.replace(src=(UOp(Ops.SLICE, t.dtype, (x.src[0], UOp.const(dtypes.index, offset)), size),)) - -to_bufferview = PatternMatcher([ - (UPat(Ops.STAGE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view), -]) - DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device? def limit_bufs(ctx:IndexingContext, root:UOp): if (device:=root.device) is None: return None # no device, index related calculations @@ -441,7 +411,7 @@ def remove_noop_afters(x:UOp) -> UOp|None: if len(src) != len(x.src): return src[0] if len(src) == 1 else x.replace(src=src) return None -pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([ +pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([ (UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)), # move RESHAPEs through MSELECT/MSTACK