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

iterate over scoped uops once [run_process_replay] (#5255)

Browse Source
This commit is contained in:
qazal
2024-07-02 09:21:09 +03:00
committed by GitHub
parent dfbee4f0f5
commit a1044e6063
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/codegen/uops.py
View File

@@ -361,8 +361,7 @@ class UOpGraph:
# scope children impact the toposort and END* insertion
end_for_uop = {UOps.IF:(UOps.STORE, UOps.ENDIF), UOps.RANGE:(UOps.PHI, UOps.ENDRANGE)}
loops, ifs = [x for x in in_degree if x.op is UOps.RANGE], [x for x in in_degree if x.op is UOps.IF]
scope_children = {p:get_recursive_children(p, end_for_uop[p.op][0]) for p in (loops+ifs)[::-1]}
scope_children = {p:get_recursive_children(p, end_for_uop[p.op][0]) for p in reversed(in_degree) if p.op in end_for_uop}
queue:List[Tuple[int, UOp]] = []
def push(u:UOp):