check Program from HEAD [run_process_replay] (#4996)

* use the same prg [run_process_replay]

* put var back

test/external/replay_codegen.py

@@ -2,7 +2,7 @@
 import difflib, pickle
 from tqdm import tqdm
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.helpers import colored, db_connection, VERSION
+from tinygrad.helpers import colored, db_connection, VERSION, to_function_name
 
 page_size = 100
 conn = db_connection()
@@ -11,12 +11,10 @@
 for offset in tqdm(range(0, row_count, page_size)):
   cur.execute(f"SELECT val FROM 'process_replay_{VERSION}' LIMIT ? OFFSET ?", (page_size, offset))
   for row in cur.fetchall():
-    compare_k: Linearizer = pickle.loads(row[0])
-    compare_src = compare_k.opts.render("test", compare_k.uops)
+    compare_k, compare_src = pickle.loads(row[0])
     k = Linearizer(*compare_k.ast, opts=compare_k.opts)
     for opt in compare_k.applied_opts: k.apply_opt(opt)
-    good_uops = k.linearize().uops
-    good_src = k.opts.render("test", good_uops)
+    good_src = k.opts.render(to_function_name(compare_k.name), k.linearize().uops)
     try: assert compare_src == good_src
     except AssertionError as e:
       print("PROCESS REPLAY FAILED")

tinygrad/codegen/linearizer.py

@@ -482,7 +482,7 @@ def to_program(self) -> Program:
     self.linearize()
     info = get_lazyop_info(self.ast[0])
     src = self.opts.render(to_function_name(self.name), self.uops)
-    if getenv("RUN_PROCESS_REPLAY"): diskcache_put("process_replay", id(self), self)
+    if getenv("RUN_PROCESS_REPLAY"): diskcache_put("process_replay", id(self), (self, src))
     ops, mem = self.uops.flops_mem()
     run_count = prod((self.global_size if self.global_size else []) + (self.local_size if self.local_size else []))
     # NOTE: we use min here to ignore the indexing FLOPS

tinygrad/codegen/uops.py

@@ -12,7 +12,7 @@
 # the order of these UOps controls the order of the toposort
 class UOps(Enum):
   # ops that aren't rendered
-  SINK = auto()
+  SINK = auto(); VAR = auto() # noqa: E702
   DEFINE_GLOBAL = auto(); DEFINE_VAR = auto(); DEFINE_LOCAL = auto(); DEFINE_ACC = auto() # noqa: E702
   CONST = auto(); SPECIAL = auto() # noqa: E702
   NOOP = auto(); UNMUL = auto(); GEP = auto() # noqa: E702
@@ -25,8 +25,6 @@ class UOps(Enum):
   BARRIER = auto(); IF = auto(); RANGE = auto() # noqa: E702
   # these two are not graph nodes
   ENDRANGE = auto(); ENDIF = auto() # noqa: E702
-  # has to be at the end because otherwise it will mess up descriminants of other enum variants and process replay will freak out
-  VAR = auto()
 
 def ufix(dtype: Optional[DType], x): return UOp.const(dtype, x) if not isinstance(x, UOp) else x
 @dataclass(eq=False)