From dcc9704b6bb30bb283feaa3b0386b1f6b03311f8 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 10 Jul 2025 17:45:57 +0300
Subject: [PATCH] viz: profile RewriteSteps in TINY device (#11125)

* 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
---
 tinygrad/uop/ops.py      |  3 ++-
 tinygrad/viz/js/index.js | 28 +++++++++++++++++++---------
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 0da3fd6449..2e186097d3 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -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
diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index fa38e4b202..7f7e88a3ef 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -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) => {