From c12e218751a3e4a7edaa54b116bc11a7b50cfb46 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Wed, 26 Nov 2025 21:12:40 +0800
Subject: [PATCH] viz: double click on INST wave (#13458)

---
 tinygrad/viz/js/index.js | 15 ++++++++++++---
 tinygrad/viz/serve.py    |  7 ++++---
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 011a800d05..8630a56610 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -205,8 +205,10 @@ function getMetadata(shape) {
     html.append(() => tabulate([["Name", d3.create("p").html(e.arg.tooltipText.split("\n")[0]).node()],
                                 ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]).node());
     if (e.arg.ctx != null) {
-      html.append("a").text("View codegen rewrite").on("click", () => switchCtx(e.arg.ctx, e.arg.step));
-      html.append("a").text("View program").on("click", () => switchCtx(e.arg.ctx, ctxs[e.arg.ctx+1].steps.findIndex(s => s.name==="View Program")));
+      const i = e.arg.ctx; s = e.arg.step;
+      html.append("a").text(ctxs[i+1].steps[s].name).on("click", () => switchCtx(i, s));
+      const prgSrc = ctxs[i+1].steps.findIndex(s => s.name === "View Program");
+      if (prgSrc !== -1) html.append("a").text("View program").on("click", () => switchCtx(i, prgSrc));
     }
   }
   if (track.eventType === EventTypes.BUF) {
@@ -310,6 +312,13 @@ async function renderProfiler(path, unit, opts) {
           const start = ref.step>0 ? ref.step+1 : 0;
           const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
           if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
+        } else {
+          const steps = ctxs[state.currentCtx].steps;
+          for (let i=state.currentStep+1; i<steps.length; i++) {
+            const loc = steps[i].loc;
+            if (loc == null) break;
+            if (loc === e.name) { shapeRef = {ctx:state.currentCtx-1, step:i}; break; }
+          }
         }
         // tiny device events go straight to the rewrite rule
         const key = k.startsWith("TINY") ? null : `${k}-${j}`;
@@ -631,6 +640,7 @@ const evtSources = [];
 // context: collection of steps
 const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false};
 function setState(ns) {
+  saveToHistory(state);
   const { ctx:prevCtx, step:prevStep } = select(state.currentCtx, state.currentStep);
   const prevRewrite = state.currentRewrite;
   Object.assign(state, ns);
@@ -638,7 +648,6 @@ function setState(ns) {
   const { ctx, step } = select(state.currentCtx, state.currentStep);
   toggleCls(prevCtx, ctx, "expanded", state.expandSteps);
   if (ctx?.id !== prevCtx?.id) {
-    saveToHistory({ currentCtx:deselect(prevCtx).ctx, currentStep:deselect(prevStep).step || 0, currentRewrite:prevRewrite, expandSteps:true });
     toggleCls(prevCtx, ctx, "active");
   }
   if (ctx?.id !== prevCtx?.id || step?.id !== prevStep?.id) {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 279751c7ca..b7c6b97373 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -236,8 +236,8 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
       if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
       n = next(inst_units[u])
       if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
-      events.append(ProfileRangeEvent(w.simd_loc, f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
-      wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "run_number":n}
+      events.append(ProfileRangeEvent(w.simd_loc, loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
+      wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "run_number":n, "loc":loc}
     # occupancy events
     units:dict[str, itertools.count] = {}
     wave_start:dict[str, int] = {}
@@ -258,7 +258,8 @@ def load_sqtt(profile:list[ProfileEvent]) -> None:
       steps.append(create_step(cu, ("/counters", len(ctxs), len(steps)),
                                {"value":get_profile(events, sort_fn=row_tuple), "content_type":"application/octet-stream"}, depth=2))
       for k in sorted(wave_insts.get(cu, []), key=row_tuple):
-        steps.append(create_step(k.replace(cu, ""), ("/sqtt-insts", len(ctxs), len(steps)), wave_insts[cu][k], depth=3))
+        data = wave_insts[cu][k]
+        steps.append(create_step(k.replace(cu, ""), ("/sqtt-insts", len(ctxs), len(steps)), data, loc=data["loc"], depth=3))
   ctxs.append({"name":"Counters", "steps":steps})
 
 def device_sort_fn(k:str) -> tuple[int, str, int]: