From 40633ab34db6a469ea808a1513ffd699cd4ff45b Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:51:36 +0800 Subject: [PATCH] list buffer args to kernel in profiler (#12826) * list buffer args to kernel in profiler * stable order * back button works * deselect also works --- tinygrad/viz/js/index.js | 51 +++++++++++++++++++++++++++++----------- 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index dd803bda5b..42fc837349 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -173,10 +173,16 @@ function tabulate(rows) { return root; } -var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity; +var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity, shapeMetadata = new Map(); +function focusShape(shape) { + saveToHistory({ shape:focusedShape }); + focusedShape = shape?.key; d3.select("#timeline").call(canvasZoom.transform, zoomLevel); + return document.querySelector(".metadata").replaceChildren(shapeMetadata.get(focusedShape) ?? ""); +} + async function renderProfiler() { displayGraph("profiler"); - d3.select(".metadata").node().replaceChildren(focusedShape?.html ?? ""); + d3.select(".metadata").node().replaceChildren(shapeMetadata.get(focusedShape) ?? ""); // layout once! if (data != null) return updateProgress({ start:false }); const profiler = d3.select(".profiler").html(""); @@ -242,6 +248,7 @@ async function renderProfiler() { } const html = document.createElement("div"); html.appendChild(tabulate([["Name", colored(e.name)], ["Duration", formatTime(e.dur)], ["Start Time", formatTime(e.st)]]).node()); + const argsDiv = document.createElement("div"); argsDiv.id = "args"; html.appendChild(document.createElement("br")); html.appendChild(argsDiv); if (e.info != null) html.appendChild(document.createElement("p")).innerText = "\n"+e.info; if (shapeRef != null) { const a = html.appendChild(document.createElement("a")); @@ -250,7 +257,8 @@ async function renderProfiler() { } // tiny device events go straight to the rewrite rule const key = k.startsWith("TINY") ? null : `${k}-${j}`; - const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), html, key, ...shapeRef }; + if (key != null) shapeMetadata.set(key, html); + const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), key, ...shapeRef }; if (e.key != null) shapeMap.set(e.key, arg); // offset y by depth shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor }); @@ -294,17 +302,30 @@ async function renderProfiler() { const rows = [["DType", dtype], ["Len", formatUnit(sz)], ["Size", formatUnit(nbytes, "B")], ["Lifetime", formatTime(dur)]]; if (users != null) rows.push(["Users", users.length]); const info = html.appendChild(tabulate(rows).node()); + const arg = {tooltipText:info.outerHTML, key:`${k}-${num}`}; for (let u=0; u focusShape(shape); + const args = shapeMetadata.get(shape.key).querySelector("#args"); + const bufArg = d3.create("p").text(`${bufInfo} ${rows[2][1]}`).style("cursor", "pointer").style("margin-top", "4px").on("click", () => { + const device = document.getElementById(k); + if (!isExpanded(device)) device.click(); + focusShape(arg); + }).node(); + bufArg.dataset.num = num; + let before = null; + for (const c of args.children) { if (+c.dataset.num > num) { before = c; break; } } + args.insertBefore(bufArg, before); } } - const arg = {tooltipText:info.outerHTML, html, key:`${k}-${num}`}; + shapeMetadata.set(arg.key, html) shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) }); } // generic polygon merger @@ -337,6 +358,7 @@ async function renderProfiler() { else if (tid === focusedDevice) { track.shapes = track.views[0]; offset += rescaleTrack(track, tid, 1/track.scaleFactor); } } data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:"B" } : null; + toggleCls(document.getElementById(focusedDevice), document.getElementById(newFocus), "expanded"); focusedDevice = newFocus; return resize(); }); @@ -395,7 +417,7 @@ async function renderProfiler() { lw += e.label[li].width; } } - if (focusedShape?.key && e.arg?.key === focusedShape.key) { paths.push([p, pcolor]); } + if (focusedShape != null && e.arg?.key === focusedShape) { paths.push([p, pcolor]); } } } // draw axes @@ -462,15 +484,11 @@ async function renderProfiler() { } } - function focusShape(shape) { - focusedShape = shape; render(zoomLevel); - return document.querySelector(".metadata").replaceChildren(shape?.html ?? ""); - } canvas.addEventListener("click", e => { e.preventDefault(); const foundRect = findRectAtPosition(e.clientX, e.clientY); if (foundRect?.step != null && foundRect?.key == null) { return setCtxWithHistory(foundRect.ctx, foundRect.step); } - if (foundRect?.key != focusedShape?.key) { focusShape(foundRect); } + if (foundRect?.key != focusedShape) { focusShape(foundRect); } }); canvas.addEventListener("mousemove", e => { @@ -583,15 +601,20 @@ function setState(ns) { const getSubrewrites = (ul) => ul.querySelectorAll(":scope > ul"); +function saveToHistory(ns) { + // NOTE: browser does a structured clone, passing a mutable object is safe. + history.replaceState(ns, ""); + history.pushState(ns, ""); +} + // set a new context and keep the old one in browser history function setCtxWithHistory(newCtx, step=0) { - // NOTE: browser does a structured clone, passing a mutable object is safe. - history.replaceState(state, ""); - history.pushState(state, ""); + saveToHistory(state); setState({ expandSteps:true, currentCtx:newCtx+1, currentStep:step, currentRewrite:0 }); } window.addEventListener("popstate", (e) => { + if (e.state?.shape != null) return focusShape({ key:e.state?.shape }); if (e.state != null) setState(e.state); });