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

ignaciosica · Jul 2, 2024 · a1044e6 · a1044e6
1 parent dfbee4f
commit a1044e6
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
@@ -361,8 +361,7 @@ def get_recursive_children(x:UOp, end:UOps, include_self=False) -> Set[UOp]:
 
     # 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):