viz: store relative timestamps (#11787)

* viz: store relative timestamps

* err

* update test

test/unit/test_viz.py

@@ -268,15 +268,20 @@ class TestVizProfiler(unittest.TestCase):
 
   def test_perfetto_copy_node(self):
     prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
+            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
+            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
+            ProfileDeviceEvent(device='NV:2', comp_tdiff=decimal.Decimal(-800), copy_tdiff=decimal.Decimal(-80))]
 
     j = json.loads(get_profile(prof))
 
     event = j['layout']['NV']['shapes'][0]
     self.assertEqual(event['name'], 'COPYxx')
-    self.assertEqual(event['st'], 900) # diff clock
+    self.assertEqual(event['st'], 0)   # first event
     self.assertEqual(event['dur'], 10)
 
+    event2 = j['layout']['NV:2']['shapes'][0]
+    self.assertEqual(event2['st'], 20) # second event, diff clock
+
   def test_perfetto_graph(self):
     prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
             ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),

tinygrad/viz/js/index.js

@@ -187,7 +187,7 @@ async function renderProfiler() {
         }
         const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
         // offset y by depth
-        shapes.push({x:e.st-st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
+        shapes.push({x:e.st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
       }
       div.style("height", levelHeight*v.maxDepth+padding+"px").style("pointerEvents", "none");
     } else {
@@ -195,7 +195,7 @@ async function renderProfiler() {
       const yscale = d3.scaleLinear().domain([0, v.peak]).range([height, 0]);
       const shapes = [];
       for (const [i,e] of v.shapes.entries()) {
-        const x = e.x.map(tsIdx => v.timestamps[tsIdx]-st);
+        const x = e.x.map(tsIdx => v.timestamps[tsIdx]);
         const arg = {tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}`};
         shapes.push({ x, y0:e.y.map(yscale), y1:e.y.map(y => yscale(y+e.arg.nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
       }

tinygrad/viz/serve.py

@@ -123,7 +123,7 @@ def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decim
       for i,ent in enumerate(e.ents): yield (cpu_ts[i*2], cpu_ts[i*2+1], ent)
 
 # timeline layout stacks events in a contiguous block. When a late starter finishes late, there is whitespace in the higher levels.
-def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
+def timeline_layout(events:list[tuple[int, int, float, DevEvent]], min_ts:int) -> dict:
   shapes:list[dict] = []
   levels:list[int] = []
   exec_points:dict[str, dict] = {}
@@ -143,10 +143,10 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
     elif isinstance(e.name, TracingKey):
       name, cat = e.name.display_name, e.name.cat
       ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
+    shapes.append({"name":name, "ref":ref, "st":st-min_ts, "dur":dur, "depth":depth, "cat":cat, "info":info})
   return {"shapes":shapes, "maxDepth":len(levels)}
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], min_ts:int, max_ts:int) -> dict:
   step, peak, mem = 0, 0, 0
   shps:dict[int, dict] = {}
   temp:dict[int, dict] = {}
@@ -155,12 +155,12 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dic
     if not isinstance(e, ProfilePointEvent): continue
     if e.name == "alloc":
       shps[e.key] = temp[e.key] = {"x":[step], "y":[mem], "arg":e.arg}
-      timestamps.append(int(e.ts))
+      timestamps.append(int(e.ts)-min_ts)
       step += 1
       mem += e.arg["nbytes"]
       if mem > peak: peak = mem
     if e.name == "free":
-      timestamps.append(int(e.ts))
+      timestamps.append(int(e.ts)-min_ts)
       step += 1
       mem -= (removed:=temp.pop(e.key))["arg"]["nbytes"]
       removed["x"].append(step)
@@ -172,7 +172,7 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dic
   for v in temp.values():
     v["x"].append(step)
     v["y"].append(v["y"][-1])
-  timestamps.append(max_ts)
+  timestamps.append(max_ts-min_ts)
   return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
@@ -192,8 +192,8 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   layout:dict[str, dict] = {}
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
-    layout[k] = timeline_layout(v)
-    layout[f"{k} Memory"] = mem_layout(v, unwrap(max_ts))
+    layout[k] = timeline_layout(v, min_ts)
+    layout[f"{k} Memory"] = mem_layout(v, min_ts, unwrap(max_ts))
   return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8")
 
 def get_runtime_stats(key) -> list[dict]: