Release 260111

tinygrad/viz/js/worker.js

@@ -0,0 +1,29 @@
+const NODE_PADDING = 10;
+const LINE_HEIGHT = 14;
+const canvas = new OffscreenCanvas(0, 0);
+const ctx = canvas.getContext("2d");
+ctx.font = `${LINE_HEIGHT}px sans-serif`;
+
+onmessage = (e) => {
+  const { graph, additions } = e.data;
+  const g = new dagre.graphlib.Graph({ compound: true });
+  g.setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
+  if (additions.length !== 0) g.setNode("addition", {label:"", className:"overlay", padding:0});
+  for (let [k, {label, src, ref, ...rest }] of Object.entries(graph)) {
+    // adjust node dims by label size (excluding escape codes) + add padding
+    let [width, height] = [0, 0];
+    for (line of label.replace(/\u001B\[(?:K|.*?m)/g, "").split("\n")) {
+      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, id:k, ...rest});
+    // 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}});
+    if (additions.includes(parseInt(k))) g.setParent(k, "addition");
+  }
+  dagre.layout(g);
+  postMessage(dagre.graphlib.json.write(g));
+  self.close();
+}