From aaea6b97adcb98b811323c95c3b4c8cb2a99d81e Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sat, 23 Aug 2025 17:34:07 +0300
Subject: [PATCH] viz memory: compute nbytes (#11795)

* viz memory: compute nbytes

* local map
---
 tinygrad/device.py       |  2 +-
 tinygrad/viz/js/index.js |  7 ++++---
 tinygrad/viz/serve.py    | 16 +++++++++-------
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/tinygrad/device.py b/tinygrad/device.py
index 162532260a..c7e5cabdf6 100644
--- a/tinygrad/device.py
+++ b/tinygrad/device.py
@@ -139,7 +139,7 @@ class Buffer:
       if PROFILE:
         self._prof_num = num = len(Buffer.profile_events)
         ts = decimal.Decimal(time.perf_counter_ns())/1000
-        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":str(self.dtype),"sz":self.size,"nbytes":self.nbytes}))
+        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":self.dtype, "sz":self.size}))
     return self
   def deallocate(self):
     assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 21a94ce3d2..6f90736966 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -151,7 +151,7 @@ async function renderProfiler() {
   // layout once!
   if (data != null) return;
   const profiler = d3.select(".profiler").html("");
-  const { layout, dur, peak } = await (await fetch("/get_profile")).json();
+  const { layout, dur, peak, dtypes } = await (await fetch("/get_profile")).json();
   // place devices on the y axis and set vertical positions
   const [tickSize, padding] = [10, 8];
   const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
@@ -196,8 +196,9 @@ async function renderProfiler() {
       const shapes = [];
       for (const [i,e] of v.shapes.entries()) {
         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) });
+        const nbytes = dtypes[e.arg.dtype]*e.arg.sz;
+        const arg = {tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(nbytes, "B")}`};
+        shapes.push({ x, y0:e.y.map(yscale), y1:e.y.map(y => yscale(y+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
       }
       data.tracks.set(k, { shapes, offsetY, height, peak:v.peak, scaleFactor:maxheight*4/height });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index e179b07638..f392951981 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -146,7 +146,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int)
     shapes.append({"name":name, "ref":ref, "st":st-start_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]], start_ts:int, end_ts:int, peaks:list[int]) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtypes_map:dict[str, int]) -> dict:
   step, peak, mem = 0, 0, 0
   shps:dict[int, dict] = {}
   temp:dict[int, dict] = {}
@@ -154,21 +154,22 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_
   for st,_,_,e in events:
     if not isinstance(e, ProfilePointEvent): continue
     if e.name == "alloc":
-      shps[e.key] = temp[e.key] = {"x":[step], "y":[mem], "arg":e.arg}
+      shps[e.key] = temp[e.key] = {"x":[step], "y":[mem], "arg":{"dtype":e.arg["dtype"].name, "sz":e.arg["sz"]}}
+      dtypes_map.setdefault(e.arg["dtype"].name, e.arg["dtype"].itemsize)
       timestamps.append(int(e.ts)-start_ts)
       step += 1
-      mem += e.arg["nbytes"]
+      mem += e.arg["sz"]*e.arg["dtype"].itemsize
       if mem > peak: peak = mem
     if e.name == "free":
       timestamps.append(int(e.ts)-start_ts)
       step += 1
-      mem -= (removed:=temp.pop(e.key))["arg"]["nbytes"]
+      mem -= (free_nbytes:=(removed:=temp.pop(e.key))["arg"]["sz"]*dtypes_map[removed["arg"]["dtype"]])
       removed["x"].append(step)
       removed["y"].append(removed["y"][-1])
       for k,v in temp.items():
         if k > e.key:
           v["x"] += [step, step]
-          v["y"] += [v["y"][-1], v["y"][-1]-removed["arg"]["nbytes"]]
+          v["y"] += [v["y"][-1], v["y"][-1]-free_nbytes]
   for v in temp.values():
     v["x"].append(step)
     v["y"].append(v["y"][-1])
@@ -192,11 +193,12 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
   # return layout of per device events
   layout:dict[str, dict] = {}
   peaks:list[int] = []
+  dtypes_map:dict[str, int] = {}
   for k,v in dev_events.items():
     v.sort(key=lambda e:e[0])
     layout[k] = timeline_layout(v, start_ts)
-    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks)
-  return json.dumps({"layout":layout, "dur":unwrap(end_ts)-start_ts, "peak":max(peaks, default=0)}).encode("utf-8")
+    layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks, dtypes_map)
+  return json.dumps({"layout":layout, "dur":unwrap(end_ts)-start_ts, "peak":max(peaks, default=0), "dtypes":dtypes_map}).encode("utf-8")
 
 def get_runtime_stats(key) -> list[dict]:
   ret:list[dict] = []