From 4f60252210857523e9a946afe60e4e9fdf754999 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 10 Oct 2024 20:03:38 +0300
Subject: [PATCH] reduce scheduler process replay overhead [pr] (#6981)

---
 tinygrad/engine/schedule.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
index 04bf0ef8..7e866cf5 100644
--- a/tinygrad/engine/schedule.py
+++ b/tinygrad/engine/schedule.py
@@ -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 ***