From 159c0e92fa134eff2b5b9d07a1a3b07c1a3adf26 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 19 Dec 2025 14:08:19 +0900
Subject: [PATCH] viz: infrastructure for basic block graphs (#13751)

---
 tinygrad/viz/js/index.js | 10 +++++-----
 tinygrad/viz/serve.py    | 17 +++++++----------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 3a21984ec5..4674fc36f8 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -757,8 +757,8 @@ async function main() {
   if (ckey in cache) {
     ret = cache[ckey];
   }
-  // ** Disassembly view
-  if (!ckey.startsWith("/rewrites")) {
+  // ** Text view
+  if (!ckey.startsWith("/graph")) {
     if (!(ckey in cache)) cache[ckey] = ret = await fetchValue(ckey);
     if (ret.steps?.length > 0) {
       const el = select(state.currentCtx, state.currentStep);
@@ -822,8 +822,8 @@ async function main() {
     });
     return document.querySelector("#custom").replaceChildren(root.node());
   }
-  // ** UOp view (default)
-  // if we don't have a complete cache yet we start streaming rewrites in this step
+  // ** Graph view
+  // if we don't have a complete cache yet we start streaming graphs in this step
   if (!(ckey in cache) || (cache[ckey].length !== step.match_count+1 && activeSrc == null)) {
     ret = [];
     cache[ckey] = ret;
@@ -847,7 +847,7 @@ async function main() {
   toggle.onchange = (e) => render({ showIndexing:e.target.checked });
   // ** right sidebar code blocks
   const codeElement = codeBlock(ret[currentRewrite].uop, "python", { wrap:false });
-  metadata.replaceChildren(toggleLabel, codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeElement);
+  if (step.code_line != null) metadata.replaceChildren(toggleLabel, codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeElement);
   if (step.trace) {
     const trace = d3.create("pre").append("code").classed("hljs", true);
     for (let i=step.trace.length-1; i>=0; i--) {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index f29c25d02d..459273ae7b 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -37,8 +37,8 @@ ref_map:dict[Any, int] = {}
 def get_rewrites(t:RewriteTrace) -> list[dict]:
   ret = []
   for i,(k,v) in enumerate(zip(t.keys, t.rewrites)):
-    steps = [create_step(s.name, ("/rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc), trace=k.tb if j==0 else None,
-                         depth=s.depth) for j,s in enumerate(v)]
+    steps = [create_step(s.name, ("/graph-rewrites", i, j), loc=s.loc, match_count=len(s.matches), code_line=printable(s.loc),
+                         trace=k.tb if j==0 else None, depth=s.depth) for j,s in enumerate(v)]
     if isinstance(k.ret, ProgramSpec):
       steps.append(create_step("View UOp List", ("/uops", i, len(steps)), k.ret))
       steps.append(create_step("View Program", ("/code", i, len(steps)), k.ret))
@@ -380,6 +380,7 @@ def amd_readelf(lib:bytes) -> list[dict]:
 
 def get_render(i:int, j:int, fmt:str) -> dict:
   data = ctxs[i]["steps"][j]["data"]
+  if fmt == "graph-rewrites": return {"value":get_full_rewrite(trace.rewrites[i][j]), "content_type":"text/event-stream"}
   if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data.uops or [])), "lang":"txt"}
   if fmt == "code": return {"src":data.src, "lang":"cpp"}
   if fmt == "asm":
@@ -462,14 +463,10 @@ class Handler(HTTPRequestHandler):
         if url.path.endswith(".css"): content_type = "text/css"
       except FileNotFoundError: status_code = 404
     elif (query:=parse_qs(url.query)):
-      i, j = get_int(query, "ctx"), get_int(query, "step")
-      if (fmt:=url.path.lstrip("/")) == "rewrites":
-        try: return self.stream_json(get_full_rewrite(trace.rewrites[i][j]))
-        except (KeyError, IndexError): status_code = 404
-      else:
-        render_src = get_render(i, j, fmt)
-        if "content_type" in render_src: ret, content_type = render_src["value"], render_src["content_type"]
-        else: ret, content_type = json.dumps(render_src).encode(), "application/json"
+      render_src = get_render(get_int(query, "ctx"), get_int(query, "step"), url.path.lstrip("/"))
+      if "content_type" in render_src: ret, content_type = render_src["value"], render_src["content_type"]
+      else: ret, content_type = json.dumps(render_src).encode(), "application/json"
+      if content_type == "text/event-stream": return self.stream_json(render_src["value"])
     elif url.path == "/ctxs":
       lst = [{**c, "steps":[{k:v for k, v in s.items() if k != "data"} for s in c["steps"]]} for c in ctxs]
       ret, content_type = json.dumps(lst).encode(), "application/json"