geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

remove late buffer view (#16941)

Browse Source
This commit is contained in:
sirhcm
2026-07-09 17:32:35 -04:00
committed by GitHub
parent 93338df753
commit 851e5727d2
2 changed files with 21 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/callify.py
+20 -11
View File
@@ -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"),
tinygrad/schedule/rangeify.py
+1 -31
View File
@@ -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