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

viz: support axis colors in UOp nodes (#11009)

Browse Source 
* work

* javascript

* optional defaultColor

* fine
This commit is contained in:
qazal
2025-06-27 23:02:55 +03:00
committed by GitHub
parent 44257f25e4
commit 382fa6a325
2 changed files with 17 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -50,8 +50,17 @@ async function renderDag(graph, additions, recenter=false) {
const x = (d.width-d.padding*2)/2;
const y = (d.height-d.padding*2)/2+STROKE_WIDTH;
return `translate(-${x}, -${y})`;
}).selectAll("text").data(d => [d.label.split("\n")]).join("text").selectAll("tspan").data(d => d).join("tspan").text(d => d).attr("x", "0")
.attr("dy", 14).attr("xml:space", "preserve");
}).selectAll("text").data(d => {
const ret = [[]];
for (const { st, color } of parseColors(d.label, defaultColor="initial")) {
for (const [i, l] of st.split("\n").entries()) {
if (i > 0) ret.push([]);
ret.at(-1).push({ st:l, color });
}
}
return [ret];
}).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
.attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve");
const tags = 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})`);
tags.selectAll("circle").data(d => [d]).join("circle");
@@ -218,8 +227,8 @@ function renderMemoryGraph(graph) {
}
const ANSI_COLORS = ["#b3b3b3", "#ff6666", "#66b366", "#ffff66", "#6666ff", "#ff66ff", "#66ffff", "#ffffff"];
const parseColors = (name) => [...name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g)].map(([_, code, colored_st, st]) =>
({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : "#ffffff" }));
const parseColors = (name, defaultColor="#ffffff") => [...name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g)]
.map(([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : defaultColor }));
// ** profiler graph

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

@@ -11,11 +11,12 @@ onmessage = (e) => {
if (additions.length !== 0) g.setNode("addition", {label:"", style:"fill: rgba(26, 27, 38, 0.5);", padding:0});
for (let [k, {label, src, ref, ...rest }] of Object.entries(graph)) {
const idx = ref ? ctxs.findIndex(k => k.ref === ref) : -1;
// replace colors in label
// replace JSON.parse string literal with real ESC
label = label.replace(/\\x1b\r?\n*\[/g, "\u001B[");
if (idx != -1) label += `\ncodegen@${ctxs[idx].function_name}`;
// adjust node dims by label size + add padding
// adjust node dims by label size (excluding escape codes) + add padding
let [width, height] = [0, 0];
for (line of label.split("\n")) {
for (line of label.replace(/\u001B\[(?:K|.*?m)/g, "").split("\n")) {
width = Math.max(width, ctx.measureText(line).width);
height += LINE_HEIGHT;
}