From 27c9ed5a844717e325fb22c65ad8aecf71a653ec Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 1 Sep 2025 14:16:47 +0300
Subject: [PATCH] viz: more consistent naming of events (#11948)

* s/shapes/events in test_viz

* s/bufs/events in the memory packer
---
 test/unit/test_viz.py | 26 +++++++++++++-------------
 tinygrad/viz/serve.py | 12 ++++++------
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py
index 9af29a5a54..53ed18c7d6 100644
--- a/test/unit/test_viz.py
+++ b/test/unit/test_viz.py
@@ -274,18 +274,18 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
     klen = u("<B")[0]
     k = ret[u.offset:u.offset+klen].decode()
     u.offset += klen
-    layout[k] = v = {"shapes":[]}
+    layout[k] = v = {"events":[]}
     event_type, event_count = u("<BI")
     if event_type == 0:
       for _ in range(event_count):
         name, ref, st, dur, _ = u("<IIIfI")
-        v["shapes"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
+        v["events"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
     else:
       v["peak"] = u("<Q")[0]
       for _ in range(event_count):
         alloc, ts, key = u("<BII")
-        if alloc: v["shapes"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["shapes"].append({"event":"free", "ts":ts, "key":key})
+        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key})
   return {"dur":dur, "peak":global_peak, "layout":layout}
 
 class TestVizProfiler(unittest.TestCase):
@@ -295,7 +295,7 @@ class TestVizProfiler(unittest.TestCase):
 
     j = load_profile(prof)
 
-    dev_events = j['layout']['NV']['shapes']
+    dev_events = j['layout']['NV']['events']
     self.assertEqual(len(dev_events), 1)
     event = dev_events[0]
     self.assertEqual(event['name'], 'E_2')
@@ -311,12 +311,12 @@ class TestVizProfiler(unittest.TestCase):
 
     j = load_profile(prof)
 
-    event = j['layout']['NV']['shapes'][0]
+    event = j['layout']['NV']['events'][0]
     self.assertEqual(event['name'], 'COPYxx')
     self.assertEqual(event['st'], 0)   # first event
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2']['shapes'][0]
+    event2 = j['layout']['NV:2']['events'][0]
     self.assertEqual(event2['st'], 20) # second event, diff clock
 
   def test_perfetto_graph(self):
@@ -334,18 +334,18 @@ class TestVizProfiler(unittest.TestCase):
     self.assertEqual(tracks[1], 'NV')
     self.assertEqual(tracks[2], 'NV:1')
 
-    nv_events = j['layout']['NV']['shapes']
+    nv_events = j['layout']['NV']['events']
     self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
     self.assertEqual(nv_events[0]['st'], 0)
     self.assertEqual(nv_events[0]['dur'], 2)
     #self.assertEqual(j['devEvents'][6]['pid'], j['devEvents'][0]['pid'])
 
-    nv1_events = j['layout']['NV:1']['shapes']
+    nv1_events = j['layout']['NV:1']['events']
     self.assertEqual(nv1_events[0]['name'], 'NV -> NV:1')
     self.assertEqual(nv1_events[0]['st'], 954)
     #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid'])
 
-    graph_events = j['layout']['NV Graph']['shapes']
+    graph_events = j['layout']['NV Graph']['events']
     self.assertEqual(graph_events[0]['st'], nv_events[0]['st'])
     self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
 
@@ -377,7 +377,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{a.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 2)
+    self.assertEqual(len(ret["events"]), 2)
 
   def test_del_once(self):
     a = _alloc(1)
@@ -386,7 +386,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{b.device} Memory"]
     self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["shapes"]), 3)
+    self.assertEqual(len(ret["events"]), 3)
 
   def test_alloc_free(self):
     a = _alloc(1)
@@ -396,7 +396,7 @@ class TestVizMemoryLayout(BaseTestViz):
     profile_ret = load_profile(Buffer.profile_events)
     ret = profile_ret["layout"][f"{c.device} Memory"]
     self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["shapes"]), 4)
+    self.assertEqual(len(ret["events"]), 4)
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index c504dd191b..0354f6c733 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -154,24 +154,24 @@ def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:
     events.append(struct.pack("<IIIfI", enum_str(name, scache), option(ref), st-start_ts, dur, enum_str(info or "", scache)))
   return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None
 
-def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
+def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
                scache:dict[str, int]) -> bytes|None:
   peak, mem = 0, 0
   temp:dict[int, int] = {}
-  bufs:list[bytes] = []
-  for st,_,_,e in events:
+  events:list[bytes] = []
+  for st,_,_,e in dev_events:
     if not isinstance(e, ProfilePointEvent): continue
     if e.name == "alloc":
-      bufs.append(struct.pack("<BIIIQ", 1, int(e.ts)-start_ts, e.key, enum_str(e.arg["dtype"].name, scache), e.arg["sz"]))
+      events.append(struct.pack("<BIIIQ", 1, int(e.ts)-start_ts, e.key, enum_str(e.arg["dtype"].name, scache), e.arg["sz"]))
       dtype_size.setdefault(e.arg["dtype"].name, e.arg["dtype"].itemsize)
       temp[e.key] = nbytes = e.arg["sz"]*e.arg["dtype"].itemsize
       mem += nbytes
       if mem > peak: peak = mem
     if e.name == "free":
-      bufs.append(struct.pack("<BII", 0, int(e.ts)-start_ts, e.key))
+      events.append(struct.pack("<BII", 0, int(e.ts)-start_ts, e.key))
       mem -= temp.pop(e.key)
   peaks.append(peak)
-  return struct.pack("<BIQ", 1, len(bufs), peak)+b"".join(bufs) if bufs else None
+  return struct.pack("<BIQ", 1, len(events), peak)+b"".join(events) if events else None
 
 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   # start by getting the time diffs