mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

viz: store memory scale (#11794)

Browse source
This commit is contained in:
qazal 2025-08-23 16:19:44 +03:00 • committed by GitHub
commit 58653b5eae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tinygrad/viz/js/index.js
View file

@ -151,7 +151,7 @@ async function renderProfiler() {
// layout once!
if (data != null) return;
const profiler = d3.select(".profiler").html("");
const { layout, dur } = await (await fetch("/get_profile")).json();
const { layout, dur, peak } = 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");
@ -163,7 +163,7 @@ async function renderProfiler() {
// color by key (name/category/device)
const colorMap = new Map();
data = {tracks:new Map(), axes:{}};
const heightScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxheight=100]);
const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
for (const [k, v] of Object.entries(layout)) {
if (v.shapes.length === 0) continue;
const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px");

8
tinygrad/viz/serve.py
View file

@ -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) -> dict:
def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int]) -> dict:
step, peak, mem = 0, 0, 0
shps:dict[int, dict] = {}
temp:dict[int, dict] = {}
@ -173,6 +173,7 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_
v["x"].append(step)
v["y"].append(v["y"][-1])
timestamps.append(end_ts-start_ts)
peaks.append(peak)
return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}
def get_profile(profile:list[ProfileEvent]) -> bytes|None:
@ -190,11 +191,12 @@ def get_profile(profile:list[ProfileEvent]) -> bytes|None:
if start_ts is None: return None
# return layout of per device events
layout:dict[str, dict] = {}
peaks:list[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))
return json.dumps({"layout":layout, "dur":unwrap(end_ts)-start_ts}).encode("utf-8")
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")
def get_runtime_stats(key) -> list[dict]:
ret:list[dict] = []