github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

cleanup bits of viz [pr] (#10681)

Browse Source
This commit is contained in:
qazal
2025-06-07 09:18:12 +03:00
committed by GitHub
parent e3805171e2
commit 86a19e19e8
2 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tinygrad/viz/js/index.js
View File

@@ -39,7 +39,7 @@ 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})`);
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("style", d => `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px; ${d.style}`);
.attr("x", d => -d.width/2).attr("y", d => -d.height/2).attr("style", d => d.style ?? `stroke:#4a4b57; stroke-width:${STROKE_WIDTH}px;`);
nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label").attr("transform", d => {
const x = (d.width-d.padding*2)/2;
const y = (d.height-d.padding*2)/2+STROKE_WIDTH;
@@ -280,21 +280,21 @@ async function main() {
}
const kernelList = document.querySelector(".kernel-list");
kernelList.innerHTML = "";
for (const [i,k] of kernels.entries()) {
for (const [i,[name, steps]] of kernels.entries()) {
const ul = kernelList.appendChild(document.createElement("ul"));
if (i === currentKernel) {
ul.className = "active";
requestAnimationFrame(() => ul.scrollIntoView({ behavior: "auto", block: "nearest" }));
}
const p = ul.appendChild(document.createElement("p"));
p.innerHTML = k[0].replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => {
p.innerHTML = name.replace(/\u001b\[(\d+)m(.*?)\u001b\[0m/g, (_, code, st) => {
const colors = ['gray','red','green','yellow','blue','magenta','cyan','white'];
return `<span style="${`color: color-mix(in srgb, ${colors[(parseInt(code)-30+60)%60]} 60%, white)`}">${st}</span>`;
});
p.onclick = () => {
setState(i === currentKernel ? { expandKernel:!expandKernel } : { expandKernel:true, currentKernel:i, currentUOp:0, currentRewrite:0 });
}
for (const [j,u] of k[1].entries()) {
for (const [j,u] of steps.entries()) {
const inner = ul.appendChild(document.createElement("ul"));
if (i === currentKernel && j === currentUOp) {
inner.className = "active";

7
tinygrad/viz/js/worker.js
View File

@@ -8,15 +8,16 @@ 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:"", style:"fill: rgba(26, 27, 38, 0.5); stroke: none;", padding:0});
for (const [k, {label, src, color}] of Object.entries(graph)) {
if (additions.length !== 0) g.setNode("addition", {label:"", style:"fill: rgba(26, 27, 38, 0.5);", padding:0});
for (const [k, {label, src, ...rest }] of Object.entries(graph)) {
// adjust node dims by label size + add padding
let [width, height] = [0, 0];
for (line of label.split("\n")) {
width = Math.max(width, ctx.measureText(line).width);
height += LINE_HEIGHT;
}
g.setNode(k, {label, color, width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING});
g.setNode(k, {width:width+NODE_PADDING*2, height:height+NODE_PADDING*2, padding:NODE_PADDING, label, ...rest});
// add edges
const edgeCounts = {}
for (const s of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
for (const s of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? edgeCounts[s] : null });