viz: light up srcs and paths on UOp click (#11844)

* viz: light up srcs and paths on UOp click

* safari doesn't have context-stroke

* safari also has a bug

* safari acceptance

tinygrad/viz/index.html

@@ -97,6 +97,15 @@
     fill: none;
     stroke-width: 1.4px;
   }
+  .highlight rect, .edgePath.highlight {
+    stroke: #FFC53D;
+  }
+  .highlight rect {
+    stroke-width: 3px;
+  }
+  #arrowhead {
+    fill: #4a4b57;
+  }
   .main-container {
     display: flex;
     width: 100%;
@@ -338,7 +347,7 @@
         </g>
         <defs>
           <marker id="arrowhead" viewBox="0 -5 10 10" refX="10" refY="0" markerWidth="6" markerHeight="6" orient="auto">
-            <path d="M0,-5L10,0L0,5" fill="#4a4b57"></path>
+            <path d="M0,-5L10,0L0,5" fill="context-stroke"></path>
           </marker>
         </defs>
       </svg>

tinygrad/viz/js/index.js

@@ -59,6 +59,9 @@ async function renderDag(graph, additions, recenter=false) {
     const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g")
       .attr("transform", d => `translate(${d.x},${d.y})`).classed("clickable", d => d.ref != null).on("click", (_,d) => {
         if (d.ref != null) return setCtxWithHistory(d.ref);
+        const src = g.predecessors(d.id) || [];
+        nodes.classed("highlight", n => src.includes(n.id));
+        d3.select("#edges").selectAll("path.edgePath").classed("highlight", e => src.includes(e.v) && e.w===d.id);
       });
     nodes.selectAll("rect").data(d => [d]).join("rect").attr("width", d => d.width).attr("height", d => d.height).attr("fill", d => d.color)
       .attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("class", d => d.className ?? "node");

tinygrad/viz/js/worker.js

@@ -16,7 +16,7 @@ onmessage = (e) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += LINE_HEIGHT;
     }
-    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ref, ...rest});
+    g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ref, id:k, ...rest});
     // add edges
     const edgeCounts = {}
     for (const s of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;