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

show buffer info in memory viz (#9562)

Browse Source
This commit is contained in:
qazal
2025-03-24 22:12:30 +08:00
committed by GitHub
parent efaee75656
commit 1c40873962
1 changed files with 22 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tinygrad/viz/lib/graph.js
View File

@@ -76,12 +76,15 @@ function renderMemoryGraph(graph) {
// ** construct alloc/free traces
// we can map reads/writes from the kernel graph
const actions = [];
const children = new Map();
for (const [k,v] of Object.entries(graph)) {
if (!(v.label.startsWith("ASSIGN"))) continue;
actions.push({ op: "write", buffer: v.src[0] });
for (const s of graph[v.src[1]].src) {
const snode = graph[s];
const srcBuf = snode.label.startsWith("ASSIGN") ? snode.src[0] : s;
if (!children.has(srcBuf)) children.set(srcBuf, new Map());
children.get(srcBuf).set(v.src[1]);
if (srcBuf !== v.src[0]) actions.push({ op: "read", buffer: srcBuf });
}
}
@@ -108,7 +111,7 @@ function renderMemoryGraph(graph) {
let memUsed = 0; // y
for (const id of prealloc) {
const buf = getBuffer(graph[id]);
ret[id] = { x: [timestep], y: [memUsed], buf };
ret[id] = { x: [timestep], y: [memUsed], buf, id };
memUsed += buf.nbytes;
}
let peak = memUsed;
@@ -119,7 +122,7 @@ function renderMemoryGraph(graph) {
// alloc
if (idx === -1) {
liveBufs.push(t.buffer);
ret[t.buffer] = { x: [timestep], y: [memUsed], buf };
ret[t.buffer] = { x: [timestep], y: [memUsed], buf, id: t.buffer };
memUsed += buf.nbytes;
peak = Math.max(memUsed, peak);
timestep += 1;
@@ -150,7 +153,7 @@ function renderMemoryGraph(graph) {
const xscale = d3.scaleLinear().domain([0, timestep]).range([0, 1024]);
const xaxis = d3.axisBottom(xscale);
const axesGroup = render.append("g").attr("id", "axes");
axesGroup.append("g").call(d3.axisLeft(yscale).tickFormat(d3.format('.3s')));
axesGroup.append("g").call(d3.axisLeft(yscale).tickFormat(d3.format(".3~s")));
axesGroup.append("g").attr("transform", `translate(0, ${yscale.range()[0]})`).call(d3.axisBottom(xscale).tickFormat(() => ""));
const polygonGroup = render.append("g").attr("id", "polygons");
const colors = ["7aa2f7", "ff9e64", "f7768e", "2ac3de", "7dcfff", "1abc9c", "9ece6a", "e0af68", "bb9af7", "9d7cd8", "ff007c"];
@@ -161,7 +164,22 @@ function renderMemoryGraph(graph) {
const p0 = xs.map((x, i) => `${x},${y1[i]}`);
const p1 = xs.map((x, i) => `${x},${y2[i]}`).reverse();
return `${p0.join(' ')} ${p1.join(' ')}`;
}).attr("fill", d => `#${colors[d.buf.num % colors.length]}`);
}).attr("fill", d => `#${colors[d.buf.num % colors.length]}`).on("mouseover", (e, { id, buf, x }) => {
d3.select(e.currentTarget).attr("stroke", "rgba(26, 27, 38, 0.8)").attr("stroke-width", 0.8);
const metadata = document.querySelector(".container.metadata");
document.getElementById("current-buf")?.remove();
const { num, dtype, ...rest } = buf;
let label = `<BUFFER n${num} ${dtype}>\n${Object.entries(rest).map(([k, v]) => `${k}=${v}`).join('\n')}\nalive for ${x[x.length-1]-x[0]} timesteps`;
const buf_children = children.get(id);
if (buf_children) {
const n = buf_children.size;
label += `\n${n} `+(n === 1 ? "child" : "children")
}
metadata.appendChild(Object.assign(document.createElement("pre"), { innerText: label, id: "current-buf", className: "wrap" }));
}).on("mouseout", (e, _) => {
d3.select(e.currentTarget).attr("stroke", null).attr("stroke-width", null);
document.getElementById("current-buf")?.remove()
});
// TODO: add the toposort graph here
d3.select("#nodes").html("");
d3.select("#edges").html("");