tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 22:46:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

viz: amdgpu disassembly register highlighting UI (#14059)

Browse Source 
* viz: amdgpu disassembly register highlighting

* minor details

* details from IDA

* more details from IDA

* refactor token colors

* move tokenizer to python

* simplify

* minimal tokenizer for registers

* all the operand types
This commit is contained in:
qazal
2026-01-09 11:27:09 +09:00
committed by GitHub
parent b878f9d5a4
commit a071adffc0
4 changed files with 49 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/viz/index.html
+9 -2
View File
@@ -127,10 +127,17 @@
fill: none;
stroke-width: 1.4px;
}
g.node.highlight rect, .edgePath.highlight, g.port circle {
g.label rect.bg {
fill: transparent;
stroke: none;
}
g.label rect.bg.highlight {
fill: #5f0059;
}
g.node.highlight rect.node, .edgePath.highlight, g.port circle {
stroke: #89C9A2;
}
g.highlight.child rect, .edgePath.highlight.child {
g.highlight.child rect.node, .edgePath.highlight.child {
stroke: #C888B0;
}
#edge-labels g.port.highlight {
tinygrad/viz/js/index.js
+18 -5
View File
@@ -70,17 +70,19 @@ const drawGraph = (data) => {
if (parents == null && children == null) return;
const src = [...parents, ...children, d.id];
nodes.classed("highlight", n => src.includes(n.id)).classed("child", n => children.includes(n.id));
if (!e.target.classList.contains("token")) labels.selectAll("rect.bg").classed("highlight", false);
const matchEdge = (v, w) => (v===d.id && children.includes(w)) ? "highlight child " : (parents.includes(v) && w===d.id) ? "highlight " : "";
d3.select("#edges").selectAll("path.edgePath").attr("class", e => matchEdge(e.v, e.w)+"edgePath");
d3.select("#edge-labels").selectAll("g.port").attr("class", (_, i, n) => matchEdge(...n[i].id.split("-"))+"port");
e.stopPropagation();
});
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);
const STROKE_WIDTH = 1.4;
.attr("x", d => -d.width/2).attr("y", d => -d.height/2).classed("node", true);
const STROKE_WIDTH = 1.4, textSpace = g.graph().textSpace;
const labels = nodes.selectAll("g.label").data(d => [d]).join("g").attr("class", "label");
labels.attr("transform", d => `translate(-${d.labelWidth/2}, -${d.labelHeight/2+STROKE_WIDTH*2})`);
labels.selectAll("text").data(d => {
const rectGroup = labels.selectAll("g.rect-group").data(d => [d]).join("g").attr("class", "rect-group");
const tokens = labels.selectAll("g.text-group").data(d => [d]).join("g").attr("class", "text-group").selectAll("text").data(d => {
if (Array.isArray(d.label)) return [d.label];
const ret = [[]];
for (const s of parseColors(d.label, defaultColor="initial")) {
@@ -91,8 +93,19 @@ 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("fill", d => typeof d.color === "string" ? d.color : colorScale(d.color))
.text(d => d.st).attr("xml:space", "preserve");
.selectAll("tspan").data(d => d).join("tspan").attr("dx", (d, i) => i > 0 ? 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));
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) => {
if (!d.keys?.length) return;
const b = nodes[i].getBBox(); ret.push({ keys:d.keys, x:b.x, y:b.y, width:b.width, height:b.height });
});
return ret;
}).join("rect").attr("class", "bg").attr("x", d => d.x).attr("y", d => d.y).attr("width", d => d.width).attr("height", d => d.height);
tokens.on("click", (e, { keys }) => {
tokensBg.classed("highlight", (d, i, nodes) => !nodes[i].classList.contains("highlight") && d.keys.some(k => keys?.includes(k)));
});
addTags(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})`).datum(e => e.tag));
// draw edges
tinygrad/viz/js/worker.js
+9 -11
View File
@@ -12,22 +12,20 @@ onmessage = (e) => {
self.close();
}
const layoutCfg = (g, { blocks, paths, pc_table, counters, colors }) => {
const lineHeight = 16;
g.setGraph({ rankdir:"TD", font:"monospace", lh:lineHeight });
const layoutCfg = (g, { blocks, paths, pc_tokens, counters, colors }) => {
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;
let maxColor = 0, 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 text = pc_table[m][0];
if (counters != null) {
const num = counters[m]?.hit_count || 0;
if (num > maxColor) maxColor = num;
label.push([{st:text, color:num}]);
} else { const [inst, ...operands] = text.split(" "); label.push([{st:inst+" ", color:"#7aa2f7"}, {st:operands.join(" "), color:"#9aa5ce"}]); }
width = Math.max(width, ctx.measureText(text).width);
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]})));
width = Math.max(width, ctx.measureText(tokens.join(" ")).width);
height += lineHeight;
}
g.setNode(lead, { ...rectDims(width, height), label, id:lead, color:"#1a1b26" });
tinygrad/viz/serve.py
+13 -1
View File
@@ -393,6 +393,15 @@ def parse_branch(asm:str) -> int|None:
return (x - 0x10000 if x & 0x8000 else x)*4
return None
def amdgpu_tokenize(st:str) -> list[str]:
try:
from extra.assembly.amd.dsl import s, v, Reg, VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
from extra.assembly.amd.asm import _op2dsl
dsl = eval(_op2dsl(st), {'s':s, 'v':v, 'VCC_LO':VCC_LO, 'VCC_HI':VCC_HI, 'VCC':VCC, 'EXEC_LO':EXEC_LO, 'EXEC_HI':EXEC_HI, 'EXEC':EXEC,
'SCC':SCC, 'M0':M0, 'NULL':NULL, 'OFF':OFF})
return [f"{type(dsl).__name__[0].lower()}{dsl.idx + i}" for i in range(dsl.count)] if isinstance(dsl, Reg) else [st]
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:
@@ -423,7 +432,10 @@ def amdgpu_cfg(lib:bytes, target:int) -> dict:
if asm.startswith("s_branch"): paths[curr][nx+offset] = UNCOND
else: paths[curr].update([(nx+offset, COND_TAKEN), (nx, COND_NOT_TAKEN)])
elif nx in leaders: paths[curr][nx] = UNCOND
return {"data":{"blocks":blocks, "paths":paths, "pc_table":pc_table, "colors":cfg_colors}, "src":"\n".join(lines)}
pc_tokens:dict[int, list[dict]] = {}
for pc, (text, _) in pc_table.items():
pc_tokens[pc] = [{"st":s, "keys":amdgpu_tokenize(s.replace(",", "")) if i>0 else [s], "kind":int(i>0)} for i,s in enumerate(text.split(" "))]
return {"data":{"blocks":blocks, "paths":paths, "colors":cfg_colors, "pc_tokens":pc_tokens}, "src":"\n".join(lines)}
# ** Main render function to get the complete details about a trace event