From 414995d1f5cf8ff030afb6c40f81379798d6cbbf Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Thu, 9 Jul 2026 21:57:42 -0400 Subject: [PATCH] Revert "remove late buffer view" (#16960) --- tinygrad/callify.py | 31 +++++++++++-------------------- tinygrad/schedule/rangeify.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 42 insertions(+), 21 deletions(-) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 176894000b..2896cc7d5e 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 not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag) + if assigned_to.op is not Ops.BUFFER: 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. Otherwise None.""" + """If movement ops on src collapse to a contiguous range, return SLICE.reshape(src.shape). Otherwise None.""" if (offset := src.contiguous_view_offset()) is None: return None buf = src.base if buf.op is Ops.SLICE: @@ -61,26 +61,25 @@ 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()) + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.index, offset)), src.numel()).reshape(src.shape) def contiguous_mops_to_view(c:UOp, src:UOp): - """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" + """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(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} and c.op is not Ops.BITCAST: return None - if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None + if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE}: return None # no symbolic shape if not all_int(c.shape): return None # check if view is supported from tinygrad.device import Device - devs = (src.device,) if isinstance(src.device, str) else src.device + devs = (c.device,) if isinstance(c.device, str) else c.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: - 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 + return view.contiguous(tag=c.tag) # 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): @@ -88,7 +87,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.reshape(resolved.src[0].shape).multi(resolved.arg).contiguous(tag=c.tag) + return view.multi(resolved.arg).contiguous(tag=c.tag) return None @@ -143,16 +142,8 @@ 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]), - # 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), + # 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), # 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 bf4a5e6e49..4dd91066f0 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -140,6 +140,14 @@ 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))), @@ -323,6 +331,28 @@ 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 @@ -411,7 +441,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+PatternMatcher([ +pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+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