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

move process replay to grouper (#9830)

Browse Source 
* simpler

* sched
This commit is contained in:
qazal
2025-04-10 18:27:42 +08:00
committed by GitHub
parent c8f47c1d07
commit 16afe04f45
3 changed files with 24 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
test/external/process_replay/process_replay.py vendored
View File

@@ -2,11 +2,11 @@
# compare kernels created by HEAD against master
import os, multiprocessing, logging, pickle, sqlite3, difflib, functools, warnings
from typing import Callable, cast
from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, dedup
from tinygrad.engine.grouper import get_becomes_map
from tinygrad.codegen.kernel import Kernel, Opt
from tinygrad.renderer import Renderer
from tinygrad.ops import UOp
from tinygrad.ops import UOp, Ops
# *** process replay settings
@@ -34,8 +34,9 @@ class ProcessReplayWarning(Warning): pass
# *** recreators
def recreate_sched(big_sink:UOp) -> list[UOp]:
sched, _, __ = create_schedule_with_vars(big_sink)
return [x.ast for x in sched]
sched_sink = get_becomes_map(big_sink)[0][big_sink]
return dedup(u.src[1].arg.ast for u in sched_sink.toposort if u.op is Ops.ASSIGN)
def recreate_kernel(ast:UOp, opts:Renderer, applied_opts:list[Opt], name:str, _) -> str:
k = Kernel(ast, opts=opts)
for opt in applied_opts: k.apply_opt(opt)