mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 08:28:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor scheduler parents search (#4402)

Browse Source
This commit is contained in:
qazal
2024-05-03 17:16:34 +03:00
committed by GitHub
parent 0627e26140
commit cf3ccb809f
1 changed files with 5 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tinygrad/engine/schedule.py
View File

@@ -167,16 +167,11 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> Tuple[Defaul
forced_realize = True
break
if len(realized_children) > 1:
for rc in realized_children:
rc_parents = deque(x.base for x in rc.srcs)
while rc_parents:
if (p:=rc_parents.pop()).realized or p.op is LoadOps.CONST: continue
if p is r: continue
# max one reduceop per kernel
if p.op in ReduceOps:
forced_realize = True
break
for x in p.srcs: rc_parents.append(x.base)
rc_parents = deque(realized_children)
while rc_parents and not forced_realize:
# max one reduceop per kernel
if (p:=rc_parents.pop()).op in ReduceOps: forced_realize = True
else: rc_parents.extend(x.base for x in p.srcs if x.base.realized is None and x.base is not r)
continue
for tr_next in children[tr].keys():
if not tr_next.realized: