github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

viz: cleanup cfg renderer (#14127)

Browse Source 
* remove colorDomains from sqtt

* colors in js

* work
This commit is contained in:
qazal
2026-01-13 11:10:42 -05:00
committed by GitHub
parent 1364449cab
commit 7fe91e5db9
3 changed files with 9 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -55,11 +55,8 @@ function addTags(root) {
root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em");
}
const colorScale = d3.scaleSequential(t => t > 0 ? d3.interpolateLab(colorScheme.ACTIVE[1], colorScheme.ACTIVE[2])(t) : colorScheme.ACTIVE[0]).clamp(true);
const drawGraph = (data) => {
const g = dagre.graphlib.json.read(data);
if (data.value.colorDomain != null) colorScale.domain(data.value.colorDomain);
// draw nodes
d3.select("#graph-svg").on("click", () => d3.selectAll(".highlight").classed("highlight", false));
const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g").attr("class", d => d.className ?? "node")
@@ -93,8 +90,8 @@ const drawGraph = (data) => {
}
return [ret];
}).join("text").style("font-family", g.graph().font).selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", g.graph().lh)
.selectAll("tspan").data(d => d).join("tspan").attr("dx", (d, i) => i > 0 && d.st !== "," ? textSpace: 0).text(d => d.st).attr("xml:space", "preserve")
.classed("token", true).attr("fill", d => typeof d.color === "string" ? d.color : colorScale(d.color));
.selectAll("tspan").data(d => d).join("tspan").attr("dx", (d, i) => i > 0 && d.st !== "," ? textSpace: 0).text(d => d.st).classed("token", true)
.attr("xml:space", "preserve").attr("fill", d => d.color);
const tokensBg = rectGroup.selectAll("rect.bg").data((d, i, nodes) => {
const ret = [];
d3.select(nodes[i].parentElement).select("g.text-group").selectAll("tspan.token").each((d, i, nodes) => {
@@ -174,8 +171,7 @@ const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
const colorScheme = {TINY:new Map([["Schedule","#1b5745"],["get_program","#1d2e62"],["compile","#63b0cd"],["DEFAULT","#354f52"]]),
DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SE:new Map([["OCC", "#101725"], ["INST", "#0A2042"]]),
ACTIVE:["#565f89", "#c8d3f5", "#7aa2f7"]}
BUFFER:["#342483", "#3E2E94", "#4938A4", "#5442B4", "#5E4CC2", "#674FCA"], SE:new Map([["OCC", "#101725"], ["INST", "#0A2042"]]),}
const cycleColors = (lst, i) => lst[i%lst.length];
const rescaleTrack = (source, tid, k) => {

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

@@ -12,28 +12,26 @@ onmessage = (e) => {
self.close();
}
const layoutCfg = (g, { blocks, paths, pc_tokens, counters, colors }) => {
const layoutCfg = (g, { blocks, paths, pc_tokens }) => {
const lineHeight = 18;
g.setGraph({ rankdir:"TD", font:"monospace", lh:lineHeight, textSpace:"1ch" });
ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
// basic blocks render the assembly in nodes
let maxColor = 0, tokenColors = {0:"#7aa2f7", 1:"#9aa5ce"};
const tokenColors = {0:"#7aa2f7", 1:"#9aa5ce"};
for (const [lead, members] of Object.entries(blocks)) {
let [width, height, label] = [0, 0, []];
for (const m of members) {
const tokens = pc_tokens[m];
const num = counters?.[m]?.hit_count ?? 0;
if (num > maxColor) maxColor = num;
label.push(tokens.map((t, i) => ({st:t.st, keys:t.keys, color:counters != null ? num : tokenColors[t.kind]})));
label.push(tokens.map((t, i) => ({st:t.st, keys:t.keys, color:tokenColors[t.kind]})));
width = Math.max(width, ctx.measureText(tokens.map((t) => t.st).join("")).width);
height += lineHeight;
}
g.setNode(lead, { ...rectDims(width, height), label, id:lead, color:"#1a1b26" });
}
g.graph().colorDomain = [0, maxColor];
// paths become edges between basic blocks
const pathColors = {0:"#3f7564", 1:"#7a4540", 2:"#3b5f7e"};
for (const [lead, value] of Object.entries(paths)) {
for (const [id, color] of Object.entries(value)) g.setEdge(lead, id, {label:{type:"port", text:""}, color:colors[color]});
for (const [id, color] of Object.entries(value)) g.setEdge(lead, id, {label:{type:"port", text:""}, color:pathColors[color]});
}
dagre.layout(g);
}

3
tinygrad/viz/serve.py
View File

@@ -403,7 +403,6 @@ def amdgpu_tokenize(st:str) -> list[str]:
except (ImportError, NameError, SyntaxError, TypeError): return []
COND_TAKEN, COND_NOT_TAKEN, UNCOND = range(3)
cfg_colors = {COND_TAKEN: "#3f7564", COND_NOT_TAKEN: "#7a4540", UNCOND: "#3b5f7e"}
def amdgpu_cfg(lib:bytes, target:int) -> dict:
# disassemble
pc_table = llvm_disasm(target, lib)
@@ -435,7 +434,7 @@ def amdgpu_cfg(lib:bytes, target:int) -> dict:
pc_tokens:dict[int, list[dict]] = {}
for pc, (text, _) in pc_table.items():
pc_tokens[pc] = [{"st":s, "keys":amdgpu_tokenize(s) if i>0 else [s], "kind":int(i>0)} for i,s in enumerate(text.replace(",", " , ").split(" "))]
return {"data":{"blocks":blocks, "paths":paths, "colors":cfg_colors, "pc_tokens":pc_tokens}, "src":"\n".join(lines)}
return {"data":{"blocks":blocks, "paths":paths, "pc_tokens":pc_tokens}, "src":"\n".join(lines)}
# ** Main render function to get the complete details about a trace event