sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

viz: collapse PROGRAM nodes like CALL (codex) (#17438)

Browse Source 
Co-authored-by: qazal <[email protected]>
This commit is contained in:
George Hotz
2026-08-07 11:44:34 +09:00
committed by GitHub
co-authored by qazal
parent f253c4469d
commit 1858f1fd9a
2 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/viz/js/index.js
+3 -3
View File
@@ -120,9 +120,9 @@ const drawGraph = (data) => {
.attr("transform", d => `translate(${d.width/2-8}, ${-d.height/2+8})`).datum(e => ({ rect:true, width:10, height:10, fill:e.addrspace, stroke:"none" })));
const CALL_TAG_WIDTH = 14;
addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag clickable ${d.collapsed ? 'collapsed' : 'expanded'}`)
.attr("transform", d => d.callNode ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
.datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.callNode ? null : d.color,
...(d.callNode && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
.attr("transform", d => d.collapsePorts != null ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
.datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.collapsePorts != null ? null : d.color,
...(d.collapsePorts != null && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
e.stopPropagation();
const t = d3.zoomTransform(document.getElementById("graph-svg"));
const [x, y] = t.apply([d.x, d.y]);
tinygrad/viz/js/worker.js
+8 -6
View File
@@ -54,15 +54,17 @@ const layoutUOp = (g, { graph, change }, opts) => {
width = Math.max(width, ctx.measureText(line).width);
height += lineHeight;
}
const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
const op = label.split("\n", 1)[0];
const callNode = op === "CALL" || op === "FUNCTION", programNode = op === "PROGRAM";
const collapsePorts = callNode ? [0] : programNode ? [0, 1] : null;
if (callNode) callCount++;
g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace,
g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, collapsePorts, exclude, addrspace,
className:label.startsWith("REWRITE_ERROR") ? "err" : null});
// add edges
const edgeCounts = {};
for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
...(callNode && port === 0 && {color:"#a0a1b8"})});
...(collapsePorts?.includes(port) && {color:"#a0a1b8"})});
if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
}
// optionally hide nodes from the layout
@@ -87,11 +89,11 @@ const layoutUOp = (g, { graph, change }, opts) => {
const consumer = g.node(consumerId);
// add +- toggle if this consumer has collapsible sources
const edge = g.edge(n, consumerId);
const collapsible = consumer.callNode ? edge?.label?.text === 0 : node.exclude;
const collapsible = consumer.collapsePorts != null ? consumer.collapsePorts.includes(edge?.label?.text) : node.exclude;
if (!collapsible) continue;
consumer.collapsible = true;
// increase width of call/function nodes to make space for a toggle
if (consumer.callNode) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
// increase width of call/function/program nodes to make space for a toggle
if (consumer.collapsePorts != null) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
// make sources invisible if UI has toggled it off
const collapsed = consumer.callNode ? opts.showCallSrc === opts.callSrcMask.has(consumerId) : !opts.expandedNodes.has(consumerId);
if (!collapsed) continue;