fix viz paginate + cleanups [pr] (tinygrad#8973)

* fix viz paginate [pr]

* cleanups

* remove the extra font definition

* more work

* none for the first graph

Author: Qazalin

tinygrad/viz/index.html

@@ -26,7 +26,7 @@
     padding: 0;
     width: 100%;
     height: 100%;
-    font-family: "Noto Sans", sans-serif;
+    font-family: sans-serif;
     font-optical-sizing: auto;
     font-weight: 400;
     font-style: normal;
@@ -356,25 +356,28 @@
       }
       else {
         const p = document.querySelector(".progress");
+        p.style.display = "none";
         const limit = 100;
         const pageCount = Math.ceil((kernels[currentKernel][1][currentUOp].match_count+1)/limit);
-        p.style.display = pageCount == 1 ? "none" : "flex";
+        if (pageCount > 1) {
+          p.style.display = "flex";
+          d3.select("#graph-svg g").selectAll("*").remove();
+        }
         for (let i=0; i < pageCount; i++) {
           p.innerText = `fetching data ${i+1}/${pageCount}`;
           const chunk = await (await fetch(`/kernels?kernel=${currentKernel}&idx=${currentUOp}&offset=${i*limit}&limit=${limit}`)).json();
           if (i === 0) ret = chunk
           else {
             // TODO: this shouldn't exist after the viz api refactor
             for (const [k,v] of Object.entries(chunk)) {
-              if (["uops", "graphs"].includes(k)) v.splice(0, 1);
               if (Array.isArray(v) && k !== "loc") ret[k].push(...v);
             }
           }
         }
         p.style.display = "none";
         cache[cacheKey] = ret;
       }
-      renderGraph(ret.graphs[currentRewrite], currentRewrite == 0 ? [] : ret.changed_nodes[currentRewrite-1]);
+      renderGraph(ret.graphs[currentRewrite], ret.changed_nodes[currentRewrite]);
     }
     // ***** RHS metadata
     const metadata = document.querySelector(".container.metadata");
@@ -402,8 +405,8 @@
         if (i === currentRewrite) {
           gUl.classList.add("active");
           if (i !== 0) {
-            const diff = ret.diffs[i-1];
-            const [loc, pattern] = ret.upats[i-1];
+            const diff = ret.diffs[i];
+            const [loc, pattern] = ret.upats[i];
             const parts = loc.join(":").split("/");
             const div = Object.assign(document.createElement("div"), { className: "rewrite-container" });
             const link = vsCodeOpener(parts);

tinygrad/viz/serve.py

@@ -2,7 +2,7 @@
 import multiprocessing, pickle, functools, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, decimal
 from http.server import HTTPServer, BaseHTTPRequestHandler
 from urllib.parse import parse_qs, urlparse
-from typing import Any, Callable, TypedDict
+from typing import Any, Callable, TypedDict, cast
 from tinygrad.helpers import colored, getenv, to_function_name, tqdm, unwrap, word_wrap
 from tinygrad.ops import TrackedGraphRewrite, UOp, Ops, lines, GroupOp
 from tinygrad.codegen.kernel import Kernel
@@ -29,7 +29,7 @@ class GraphRewriteDetails(GraphRewriteMetadata):
   changed_nodes: list[list[int]] # the changed UOp id + all its parents ids
   code_line: str                 # source code calling graph_rewrite
   kernel_code: str|None          # optionally render the final kernel code
-  upats: list[tuple[tuple[str, int], str]]
+  upats: list[tuple[tuple[str, int], str]|None]
 
 # NOTE: if any extra rendering in VIZ fails, we don't crash
 def pcall(fxn:Callable[..., str], *args, **kwargs) -> str:
@@ -69,9 +69,10 @@ def get_metadata(keys:list[Any], contexts:list[list[TrackedGraphRewrite]]) -> li
 def _prg(k:Kernel): return k.to_program().src
 def get_details(k:Any, ctx:TrackedGraphRewrite, metadata:GraphRewriteMetadata, offset=0, limit=200) -> GraphRewriteDetails:
   ret:GraphRewriteDetails = {"uops":[pcall(str, sink:=ctx.sink)], "graphs":[uop_to_json(sink)], "code_line":lines(ctx.loc[0])[ctx.loc[1]-1].strip(),
-                             "kernel_code":pcall(_prg, k) if isinstance(k, Kernel) else None, "diffs":[], "upats":[], "changed_nodes":[], **metadata}
+                             "kernel_code":pcall(_prg, k) if isinstance(k, Kernel) else None, **metadata,
+                             "diffs":[[]], "upats":[None], "changed_nodes":[[]]} # NOTE: the first graph just renders the input UOp
   replaces: dict[UOp, UOp] = {}
-  for i,(u0,u1,upat) in enumerate(tqdm(ctx.matches[offset:offset+limit])):
+  for i,(u0,u1,upat) in enumerate(tqdm(ctx.matches)):
     replaces[u0] = u1
     new_sink = sink.substitute(replaces)
     ret["graphs"].append(new_sink_js:=uop_to_json(new_sink))
@@ -80,7 +81,9 @@ def get_details(k:Any, ctx:TrackedGraphRewrite, metadata:GraphRewriteMetadata, o
     ret["upats"].append((upat.location, upat.printable()))
     # TODO: this is O(n^2)!
     ret["uops"].append(str(sink:=new_sink))
-  return ret
+  # if the client requested a chunk we only send that chunk
+  # TODO: is there a way to cache the replaces dict here?
+  return cast(GraphRewriteDetails, {k:v[offset:offset+limit] if isinstance(v,list) else v for k,v in ret.items()})
 
 # Profiler API
 devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}