From 2a791f7924e227c23cf0ee3da535de066122aa41 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 2 Aug 2024 22:28:15 +0800
Subject: [PATCH] fuzz uops is simpler with List[UOp] [run_process_replay]
 (#5875)

* remove from fuzz_uops

* update fuzz_uops.py

* add to realize.py
---
 test/external/fuzz_uops.py   | 15 ++++++++-------
 tinygrad/codegen/uopgraph.py |  5 -----
 tinygrad/engine/realize.py   |  2 +-
 3 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/test/external/fuzz_uops.py b/test/external/fuzz_uops.py
index 35897858..4daeb6cc 100644
--- a/test/external/fuzz_uops.py
+++ b/test/external/fuzz_uops.py
@@ -12,7 +12,7 @@ from tinygrad.shape.symbolic import Variable
 from tinygrad.tensor import _to_np_dtype
 from test.external.fuzz_schedule import FUZZ_SCHEDULE_MAX_PATHS, find_all_toposorts
 
-def fuzz_uops(uops:UOpGraph) -> List[Tuple[UOp, ...]]:
+def fuzz_uops(uops:List[UOp]) -> List[Tuple[UOp, ...]]:
   blocks: List[List[UOp]] = [[]]
   for u in uops:
     if u.op in END_FOR_UOP: blocks.append([u])
@@ -38,24 +38,25 @@ def fuzz_uops(uops:UOpGraph) -> List[Tuple[UOp, ...]]:
 
 class UOpsFuzzerRunner(CompiledRunner):
   def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False):
-    assert self.p.uops is not None and len(self.p.uops._fuzz_paths) >= 1
+    assert self.p.uops is not None
+    fuzz_paths = fuzz_uops(self.p.uops)
     init_rawbufs, init_name = {x:x.as_buffer() for x in rawbufs}, self.p.function_name
-    init_globals = {i[0]:buf for i, buf in zip(self.p.globals, rawbufs)}
-    if DEBUG >= 1: print(colored(f"fuzzing {len(self.p.uops._fuzz_paths)} uop permutations for {init_name}", "yellow"))
+    init_globals = dict(zip(self.p.globals, rawbufs))
+    if DEBUG >= 1: print(colored(f"fuzzing {len(fuzz_paths)} uop permutations for {init_name}", "yellow"))
 
     super().__call__(rawbufs, var_vals, wait)
     ground_truth = {x:np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in rawbufs}
 
-    for i, path in enumerate(self.p.uops._fuzz_paths):
+    for i, path in enumerate(fuzz_paths):
       # setup prg
       uops = UOpGraph([])
       uops._uops = list(path)
       if DEBUG >= 5: uops.print()
-      self.p = replace(self.p, name=(name:=f"{init_name}fuzz{i}"), src=Device[self.p.dname].renderer.render(name, uops), uops=uops)
+      self.p = replace(self.p, name=(name:=f"{init_name}fuzz{i}"), src=Device[self.p.dname].renderer.render(name, uops.uops), uops=uops.uops)
       if DEBUG >= 4: print(self.p.src)
       self.lib = Device[self.p.dname].compiler.compile_cached(self.p.src)
       self.clprg = Device[self.p.dname].runtime(name, self.lib)
-      for x in (rawbufs:=[init_globals[i[0]] for i in self.p.globals]): x.copyin(init_rawbufs[x])
+      for x in (rawbufs:=[init_globals[i] for i in self.p.globals]): x.copyin(init_rawbufs[x])
       # verify
       super().__call__(rawbufs, var_vals, wait)
       for i, x in enumerate(rawbufs):
diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py
index e20fa1b8..da1379d5 100644
--- a/tinygrad/codegen/uopgraph.py
+++ b/tinygrad/codegen/uopgraph.py
@@ -567,9 +567,4 @@ class UOpGraph:
 
     # strip the SINK
     self._uops = self._uops[:-1]
-
-    if getenv("FUZZ_UOPS"):
-      from test.external.fuzz_uops import fuzz_uops
-      self._fuzz_paths = fuzz_uops(self)
-
     return self
diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py
index f6b1663b..ad65e70e 100644
--- a/tinygrad/engine/realize.py
+++ b/tinygrad/engine/realize.py
@@ -158,7 +158,7 @@ def get_runner(dname:str, ast:LazyOp) -> CompiledRunner:
     method_cache[ckey] = ret = CompiledRunner(replace(bret.p, dname=dname), bret.lib)
   else:
     prg: Program = get_kernel(Device[dname].renderer, ast).to_program()
-    if hasattr(prg.uops, "_fuzz_paths"):
+    if getenv("FUZZ_UOPS"):
       from test.external.fuzz_uops import UOpsFuzzerRunner
       return UOpsFuzzerRunner(replace(prg, dname=dname))
     method_cache[ckey] = method_cache[bkey] = ret = CompiledRunner(replace(prg, dname=dname))