commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

reduce scheduler process replay overhead [pr] (#6981)

This commit is contained in:
qazal 2024-10-10 20:03:38 +03:00 committed by GitHub
parent 859d6d0407
commit 4f60252210
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tinygrad/engine/schedule.py
View File

@ -126,12 +126,18 @@ reduceop_fusor = PatternMatcher([
enumerate_bufs = PatternMatcher([(UPat(UOps.BUFFER, name="x"), lambda ctx,x: UOp(UOps.DEFINE_GLOBAL, x.dtype, (), ctx.bufs.index(x.arg[0])))])
PROCESS_REPLAY_CAPTURE: List[Tuple[UOp, ScheduleItemContext, UOp]] = []
if getenv("RUN_PROCESS_REPLAY"):
@atexit.register
def save_process_replay():
for base_sink,ctx,ret in PROCESS_REPLAY_CAPTURE: diskcache_put("schedule_process_replay", str(base_sink.key), (base_sink, ctx, ret))
@track_rewrites
def full_ast_rewrite(base_sink:UOp, ctx:ScheduleItemContext) -> UOp:
if not AST_REWRITE: return base_sink
sink = graph_rewrite(base_sink, reduceop_fusor)
ret = graph_rewrite(sink, enumerate_bufs, ctx)
if getenv("RUN_PROCESS_REPLAY"): diskcache_put("schedule_process_replay", str(base_sink.key), (base_sink, ctx, ret))
PROCESS_REPLAY_CAPTURE.append((base_sink, ctx, ret))
return ret
# *** List[LazyBuffer] lowering to ScheduleItem ***