tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-09-05 07:46:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix assign copy

Browse Source
This commit is contained in:
geohot
2026-02-18 12:14:18 +08:00
parent fdb0b48ce9
commit ca906371fa
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/schedule/indexing.py
+5
View File
@@ -167,6 +167,11 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
# get ops to realize
graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, bottom_up=True, name="get realize")
# don't realize COPY/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
for u in tsink.toposort():
if u.op is Ops.ASSIGN and u.src[1].op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC} and u.src[1] in rctx.realize_map \
and not u.src[0].op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
del rctx.realize_map[u.src[1]]
# get the consumer map
with cpu_profile("consumer map in rangeify", "TINY"):