github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-23 05:48:08 -05:00
Code Issues Packages Projects Releases Wiki Activity

save lines in the scheduler [run_process_replay] (#5890)

Browse Source
This commit is contained in:
qazal
2024-08-03 19:20:11 +08:00
committed by GitHub
parent 56ef9e453e
commit 37cc87ea75
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tinygrad/engine/schedule.py
View File

@@ -298,14 +298,12 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]):
top_reduce = reduceop.base.srcs[0].base
if len(children[top_reduce]) == 1: del realizes[top_reduce]
def _can_fold_reduce(r:LazyBuffer, group:Dict[LazyBuffer, None]) -> bool:
if DEBUG_ARANGE:=(getenv("DEBUG_ARANGE")): print(f"checking {r} {group=}")
if any(tr.forced_realize or tr in outs for tr in group): return False
if DEBUG_ARANGE: print(colored(f"folding {r}", "green"))
return True
for r in reduce_of_const:
if _can_fold_reduce(r, group:={tr:None for tr,rop in reduce_for_op.items() if rop is r}):
for tr in group: del realizes[tr]
group = {tr:None for tr,rop in reduce_for_op.items() if rop is r}
if DEBUG_ARANGE:=(getenv("DEBUG_ARANGE")): print(f"checking {r} {group=}")
if any(tr.forced_realize or tr in outs for tr in group): continue
if DEBUG_ARANGE: print(colored(f"folding {r}", "green"))
for tr in group: del realizes[tr]
output_groups: DefaultDict[LazyBuffer, List[LazyBuffer]] = defaultdict(list)
for buf in realizes: