From 0877497bad466c5bfcec529393ac52467a9d533c Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 27 Mar 2025 11:52:12 +0800
Subject: [PATCH] hotfix: use captured uops in viz render [pr] (#9593)

* hotfix: use captured uops in viz render [pr]

* better error
---
 tinygrad/viz/serve.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 2986eed2c8..5e09abccb6 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -21,7 +21,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
 # NOTE: if any extra rendering in VIZ fails, we don't crash
 def pcall(fxn:Callable[..., str], *args, **kwargs) -> str:
   try: return fxn(*args, **kwargs)
-  except Exception as e: return f"ERROR: {e}"
+  except Exception as e: return f"ERROR in {fxn.__name__}: {e}"
 
 # ** Metadata for a track_rewrites scope
 
@@ -33,10 +33,10 @@ class GraphRewriteMetadata(TypedDict):
   name: str|None                         # optional name of the rewrite
 
 @functools.lru_cache(None)
-def _prg(k:Kernel): return k.to_program().src
+def render_program(k:Kernel): return k.opts.render(k.uops)
 def to_metadata(k:Any, v:TrackedGraphRewrite) -> GraphRewriteMetadata:
   return {"loc":v.loc, "match_count":len(v.matches), "code_line":lines(v.loc[0])[v.loc[1]-1].strip(),
-          "kernel_code":pcall(_prg, k) if isinstance(k, Kernel) else None, "name":v.name}
+          "kernel_code":pcall(render_program, k) if isinstance(k, Kernel) else None, "name":v.name}
 def get_metadata(keys:list[Any], contexts:list[list[TrackedGraphRewrite]]) -> list[tuple[str, list[GraphRewriteMetadata]]]:
   return [(k.name if isinstance(k, Kernel) else str(k), [to_metadata(k, v) for v in vals]) for k,vals in zip(keys, contexts)]