add pagination to viz [pr] (tinygrad#8794)

Qazalin · web-flow · commit 199a36d079eb · 2025-01-29T04:21:53.000+02:00
* add pagination to viz [pr]

* work

* lint
diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html
@@ -193,6 +193,13 @@
     border-radius: 8px;
     padding: 8px;
   }
+  .progress {
+    font-size: 20px;
+    z-index: 2;
+    align-self: center;
+    justify-content: center;
+    width: 100%;
+  }
   </style>
 </head>
 <body>
@@ -204,6 +211,7 @@
       <a class="btn nav-btn" href="/profiler">Profiler</a>
     </div>
     <div class="container kernel-list-parent"><div class="container kernel-list"></div></div>
+    <p class="progress"></p>
     <div class="graph">
       <svg id="graph-svg">
         <g id="render"></g>
@@ -347,7 +355,23 @@
         ret = cache[cacheKey];
       }
       else {
-        ret = await (await fetch(`/kernels?kernel=${currentKernel}&idx=${currentUOp}`)).json();
+        const p = document.querySelector(".progress");
+        const limit = 100;
+        const pageCount = Math.ceil((kernels[currentKernel][1][currentUOp].match_count+1)/limit);
+        p.style.display = pageCount == 1 ? "none" : "flex";
+        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]);
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
@@ -66,11 +66,11 @@ def get_metadata(keys:list[Any], contexts:list[list[TrackedGraphRewrite]]) -> li
 
 @functools.lru_cache(None)
 def _prg(k:Kernel): return k.to_program().src
-def get_details(k:Any, ctx:TrackedGraphRewrite, metadata:GraphRewriteMetadata) -> GraphRewriteDetails:
+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}
   replaces: dict[UOp, UOp] = {}
-  for i,(u0,u1,upat) in enumerate(tqdm(ctx.matches)):
+  for i,(u0,u1,upat) in enumerate(tqdm(ctx.matches[offset:offset+limit])):
     replaces[u0] = u1
     new_sink = sink.substitute(replaces)
     ret["graphs"].append(new_sink_js:=uop_to_json(new_sink))
@@ -127,10 +127,10 @@ def do_GET(self):
         if url.path.endswith(".css"): content_type = "text/css"
       except FileNotFoundError: status_code = 404
     elif url.path == "/kernels":
-      query = parse_qs(url.query)
-      if (qkernel:=query.get("kernel")) is not None:
-        kidx, ridx = int(qkernel[0]), int(query["idx"][0])
-        jret:Any = get_details(contexts[0][kidx], contexts[1][kidx][ridx], kernels[int(qkernel[0])][1][int(query["idx"][0])])
+      if "kernel" in (query:=parse_qs(url.query)):
+        def getarg(k:str,default=0): return int(query[k][0]) if k in query else default
+        kidx, ridx = getarg("kernel"), getarg("idx")
+        jret:Any = get_details(contexts[0][kidx], contexts[1][kidx][ridx], kernels[kidx][1][ridx], getarg("offset", 0), getarg("limit", 200))
       else: jret = kernels
       ret, content_type = json.dumps(jret).encode(), "application/json"
     elif url.path == "/get_profile" and perfetto_profile is not None: ret, content_type = perfetto_profile, "application/json"
