diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index f1de08f22e..c8aff7c6e3 100644 --- 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%; + } @@ -204,6 +211,7 @@ Profiler
+

@@ -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 index 437ca31d1a..1e75338adc 100755 --- 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 @@ class Handler(BaseHTTPRequestHandler): 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"