viz: make call toggle easier to click on (#16411)

* call tag is a rect

* details

* colors

* simplify, better comment

tinygrad/viz/index.html

@@ -102,18 +102,18 @@
     fill: #FFD700;
     stroke: #B8860B;
   }
-  g.tag.collapsed circle {
+  g.tag.collapsed circle, g.tag.collapsed rect {
     fill: #5CD68D;
     stroke: #4a4b57;
   }
-  g.tag.expanded circle {
+  g.tag.expanded circle, g.tag.expanded rect {
     fill: #9FDDE6;
     stroke: #4a4b57;
   }
   g.port circle {
     fill: #b3dcc2;
   }
-  g.tag circle, #edge-labels circle {
+  g.tag circle, g.tag rect, #edge-labels circle {
     stroke-width: 0.8;
   }
   g.tag text, #edge-labels text {

tinygrad/viz/js/index.js

@@ -57,7 +57,9 @@ function intersectRect(r1, r2) {
 }
 
 function addTags(root, path) {
-  root.selectAll("circle").data(d => [d]).join("circle").attr("r", 5).style("fill", d => d.fill ?? null);
+  root.selectAll("circle").data(d => d.rect ? [] : [d]).join("circle").attr("r", 5).style("fill", d => d.fill ?? null);
+  root.selectAll("rect").data(d => d.rect ? [d] : []).join("rect").attr("x", d => -d.width/2).attr("y", d => -d.height/2)
+    .attr("width", d => d.width).attr("height", d => d.height).style("fill", d => d.fill ?? null);
   if (path != null) root.selectAll("path").data(d => [d]).join("path").attr("d", path);
   else root.selectAll("text").data(d => [d]).join("text").text(d => d.text).attr("dy", "0.35em");
 }
@@ -85,7 +87,7 @@ const drawGraph = (data) => {
     .attr("x", d => -d.width/2).attr("y", d => -d.height/2).classed("node", true);
   const STROKE_WIDTH = 1.4, textSpace = g.graph().textSpace;
   const labels = nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label");
-  labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
+  labels.attr("transform", d => `translate(${d.labelX-d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
   const rectGroup = labels.selectAll("g.rect-group").data(d => [d]).join("g").attr("class", "rect-group");
   const tokens = labels.selectAll("g.text-group").data(d => [d]).join("g").attr("class", "text-group").selectAll("text").data(d => {
     if (Array.isArray(d.label)) return [d.label];
@@ -113,8 +115,11 @@ const drawGraph = (data) => {
   });
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => ({ text:e.tag })));
-  addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag ${d.collapsed ? 'collapsed' : 'expanded'}`)
-    .attr("transform", d => `translate(${-d.width/2}, ${0})`).datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.callNode ? null : d.color })).on("click", (e,d) => {
+  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) => {
       e.stopPropagation();
       const t = d3.zoomTransform(document.getElementById("graph-svg"));
       const [x, y] = t.apply([d.x, d.y]);

tinygrad/viz/js/worker.js

@@ -56,7 +56,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
     }
     const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
     if (callNode) callCount++;
-    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag, callNode, exclude});
+    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude});
     // add edges
     const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
@@ -79,6 +79,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
   }
   // optionally remove node srcs, track affected nodes
   const disconnected = new Set();
+  const CALL_TAG_WIDTH = 14;
   for (const n of g.nodes()) {
     const node = g.node(n);
     for (const consumerId of (g.successors(n) || [])) {
@@ -88,6 +89,8 @@ const layoutUOp = (g, { graph, change }, opts) => {
       const collapsible = consumer.callNode ? edge?.label?.text === 0 : 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; }
       // 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;