github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

remove collapse_nested_assign [pr] (#14775)

Browse Source 
the else branch is dead code, and we can check directly with UPat
This commit is contained in:
chenyu
2026-02-15 18:04:47 -05:00
committed by GitHub
parent 17db43ab46
commit 1ded250bbe
1 changed files with 1 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tinygrad/schedule/rangeify.py
View File

@@ -33,9 +33,6 @@ pm_mops = PatternMatcher([
# *****************
# 0. do some cleanup rewrites, mostly copied from the old stuff
def collapse_nested_assign(assign:UOp, target:UOp, src:UOp):
"""nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__): collapse the redundant outer ASSIGN"""
if src.src[0].base is target.base: return src if src.src[0] is target else assign.replace(src=(target, src.src[1]))
def assign_to_contiguous(assign:UOp, target:UOp, src:UOp):
if (t := target.base).op is Ops.PARAM or (t.op is Ops.MSTACK and all(s.op is Ops.PARAM for s in t.src)): return None
@@ -148,7 +145,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
# ** assign rules **
# collapse nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__)
(UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, name="src")), name="assign"), collapse_nested_assign),
(UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat()), name="src"))), lambda target, src: src),
# move bitcast from assign target to source: a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
(UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src")), name="assign"),