nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

refactor: consolidate replace [run_process_replay] (#5403)

Browse Source
This commit is contained in:
Roelof van Dijk
2024-07-12 07:36:57 -07:00
committed by GitHub
parent 497274f663
commit b18aa00bba
1 changed files with 2 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/uopgraph.py
+2 -7
View File
@@ -143,13 +143,8 @@ def expand_nodes(parents:Set[UOp], expands:List[UOp], base:UOp) -> List[UOp]:
replace[cc] = super_replace
else:
for replace in replaces:
if cc in replace:
# NOTE: handle expands that are already replaced
tcc = replace[cc]
replace[cc] = UOp(tcc.op, tcc.dtype, tuple(replace.get(x, x) for x in tcc.src), tcc.arg)
else:
replace[cc] = UOp(cc.op, cc.dtype, tuple(replace.get(x, x) for x in cc.src), cc.arg)
tcc = replace.get(cc, cc) # NOTE: handle expands that are already replaced
replace[cc] = UOp(tcc.op, tcc.dtype, tuple(replace.get(x, x) for x in tcc.src), tcc.arg)
return [x.get(base, base) for x in replaces]
# ***** reduce+image+contract handling *****