From b61285efa4276ea2332bad01fbe6fa2cd08f3b1f Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 2 Jul 2026 09:19:35 -0400 Subject: [PATCH] remove apply_after [PR] (#16825) tighten buffer_map spec and becomes_map is much smaller now, all intermediate afters are not needed --- tinygrad/callify.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 8dfad59647..d012c4b9c7 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -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