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

remove apply_after [PR] (#16825)

Browse Source 
tighten buffer_map spec and becomes_map is much smaller now, all intermediate afters are not needed
This commit is contained in:
chenyu
2026-07-02 09:19:35 -04:00
committed by GitHub
parent 149fd91e22
commit b61285efa4
1 changed files with 4 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/callify.py
+4 -10
View File
@@ -24,20 +24,13 @@ def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
from_creation = isinstance(u.src[0].device, str) and any(u.src[0].device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"])
if from_creation: return tag_uop(ctx, u)
def apply_after(ctx:AllocCtx, u:UOp):
base = u.src[0]
while base.op is Ops.AFTER: base = base.src[0]
ctx.buffer_map[u] = base
# CONTIGUOUS and AFTER+STORE + parents are the only nodes that get updated
# CONTIGUOUS and AFTER + parents are the only nodes that get updated
add_tags = PatternMatcher([
(UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
# no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), name="x"), tag_uop),
(UPat(Ops.AFTER, name="u"), apply_after),
(UPat(Ops.CONTIGUOUS, name="x"), tag_uop),
(UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop),
(UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
])
@@ -216,8 +209,9 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
# here we can break the tensor graph. this is the only place you need to maintain numbered tags
big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
# here we construct the final buffer_map. this is everything that will go into the tensor map
# here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
return ret, ctx.buffer_map