github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-23 13:58:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

viz: profile RewriteSteps in TINY device (#11125)

Browse Source 
* viz: profile RewriteSteps in TINY device

* use TracingKey with category

* split by whitespace

* add tracing.py

* work

* tracing_key

* TRACK_MATCH_STATS=3, can this be in defaults?

* fallback name

* work

* javascript

* measure text is slow

* checkout

* profile graph_rewrite/graph_rewrite_map

* change that

* no as

* finally

* work

* linking works
This commit is contained in:
qazal
2025-07-10 17:45:57 +03:00
committed by GitHub
parent 32117402dd
commit dcc9704b6b
2 changed files with 21 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/uop/ops.py
View File

@@ -815,7 +815,8 @@ def track_matches(func):
depth = len(active_rewrites)
tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, track_uop(args[0]), [], kwargs.get("name", None), depth, kwargs.get("bottom_up", False)))
active_rewrites.append(ctx)
ret = func(*args, **kwargs)
with cpu_profile(kwargs.get("name", "<unnamed>"), "TINY", display=tracking):
ret = func(*args, **kwargs)
if tracking: active_rewrites.pop()
return ret
return _track_func

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

@@ -108,10 +108,12 @@ function formatTime(ts, dur=ts) {
}
const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
const devColors = {"TINY":["#1B5745", "#1D2E62"],
"DEFAULT":["#1D1F2A", "#2A2D3D", "#373B4F", "#444862", "#12131A", "#2F3244", "#3B3F54", "#4A4E65", "#181A23", "#232532", "#313548", "#404459"],}
const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "rgb(70 172 194)", "rgb(29, 46, 98)"],
"DEFAULT":["rgb(29,31,42)","rgb(42,45,61)","rgb(55,59,79)","rgb(68,72,98)","rgb(18,19,26)","rgb(47,50,68)","rgb(59,63,84)","rgb(74,78,101)","rgb(24,26,35)","rgb(35,37,50)","rgb(49,53,72)","rgb(64,68,89)"],}
const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"];
const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step)));
var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
async function renderProfiler() {
displayGraph("profiler");
@@ -141,15 +143,23 @@ async function renderProfiler() {
const { y:baseY, height:baseHeight } = rect(div);
const levelHeight = baseHeight-padding;
const offsetY = baseY-canvasTop+padding/2;
for (const [i,e] of timeline.shapes.entries()) {
const colorKey = e.cat ?? e.name;
let colorKey, currentRef;
for (const e of timeline.shapes) {
if (e.depth === 0) colorKey = e.cat ?? e.name;
if (!colorMap.has(colorKey)) {
const colors = devColors[k] ?? devColors.DEFAULT;
colorMap.set(colorKey, colors[i%colors.length]);
colorMap.set(colorKey, colors[colorMap.size%colors.length]);
}
const fillColor = lighten(colorMap.get(colorKey), e.depth);
const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
if (e.ref != null) currentRef = e.ref;
let ref = {ctx:currentRef, step:0};
if (e.ref == null && currentRef != null) {
const stepIdx = ctxs[currentRef+1].steps.findIndex((s) => s.name == e.name);
if (stepIdx !== -1) ref.step = stepIdx;
}
// offset y by depth
data.shapes.push({x:e.st-st, dur:e.dur, height:levelHeight, y:offsetY+levelHeight*e.depth, ref:e.ref, label, fillColor:colorMap.get(colorKey)});
data.shapes.push({x:e.st-st, dur:e.dur, height:levelHeight, y:offsetY+levelHeight*e.depth, ref, label, fillColor });
}
// position shapes on the canvas and scale to fit fixed area
const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2;
@@ -293,7 +303,7 @@ async function renderProfiler() {
canvas.addEventListener("click", e => {
e.preventDefault();
const foundRect = findRectAtPosition(e.clientX, e.clientY);
if (foundRect?.ref != null) return setCtxWithHistory(foundRect.ref);
if (foundRect?.ref != null) return setCtxWithHistory(foundRect.ref.ctx, foundRect.ref.step);
});
const tooltip = document.body.appendChild(document.createElement("div"));
@@ -402,12 +412,12 @@ function setState(ns) {
}
// set a new context and keep the old one in browser history
function setCtxWithHistory(newCtx) {
function setCtxWithHistory(newCtx, step=0) {
if (newCtx == null) return;
// NOTE: browser does a structured clone, passing a mutable object is safe.
history.replaceState(state, "");
history.pushState(state, "");
setState({ expandSteps:true, currentCtx:newCtx+1, currentStep:0, currentRewrite:0 });
setState({ expandSteps:true, currentCtx:newCtx+1, currentStep:step, currentRewrite:0 });
}
window.addEventListener("popstate", (e) => {