github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

viz/cli: cleanups, add memory printer (#15762)

Browse Source 
* simple repro

* use context

* work

* memory printer

* rm

* memory printer

* pylint
This commit is contained in:
qazal
2026-04-16 16:44:47 +03:00
committed by GitHub
parent f57380cbc2
commit 126cda45f8
2 changed files with 21 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
extra/viz/cli.py
View File

@@ -1,5 +1,6 @@
#!/usr/bin/env python3
import argparse, pathlib, signal, sys, struct, json, itertools
import argparse, pathlib, signal, sys, struct, json, itertools, os
os.environ["VIZ"] = "0"
if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL)
from typing import Iterator
from tinygrad.viz import serve as viz
@@ -24,9 +25,8 @@ def decode_profile(data:bytes) -> dict:
klen = u("<B")[0]
k = ret[off:off+klen].decode()
off += klen
v:dict = {"events":[]}
layout[k] = v
event_type, event_count = u("<BI")
layout[k] = v = {"event_type":event_type, "events":[]}
if event_type == 0:
for _ in range(event_count):
name, ref, key, st, dur, fmt = u("<IIIIfI")
@@ -41,7 +41,8 @@ def decode_profile(data:bytes) -> dict:
else:
alloc, ts, key = u("<BII")
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, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[(k, strings[rep], num, mode) \
for k,rep,num,mode in [u("<IIIB") for _ in range(u("<I")[0])]]}})
return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
def get(data:dict, key:str):
@@ -114,6 +115,17 @@ def main(args) -> None:
print(tabulate(rows, headers=cols, tablefmt="github"))
return None
# ** Memory printer
if data["event_type"] == 1 and data.get("events", []):
print(f"Peak: {data['peak']}"+"\n"+f"{'TS':<10} {'Event':<6} {'Key':>8} Info")
modes = ("read","write","write+read")
for e in data["events"]:
info = str(e.get("arg", {}))
if e["event"] == "free":
info = ', '.join([f"{format_colored(kernel)} {['read','write','write+read'][mode]}@data{num}" for _,kernel,num,mode in e["arg"]["users"]])
print(f"{e['ts']:<10} {e['event']:<6} {e.get('key', ''):>8} {info}")
return None
# ** Profiler printer
agg:dict[str, tuple[float, int]] = {}
total = 0

22
test/null/test_viz.py
View File

@@ -2,11 +2,11 @@ import unittest, decimal, sys, json, contextlib
from dataclasses import dataclass
from typing import Generator
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS, profile_matches
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
from tinygrad.uop.symbolic import sym
from tinygrad.dtype import dtypes
from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
from tinygrad.helpers import VIZ, cpu_profile, ProfilePointEvent, unwrap
from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
from tinygrad.device import Buffer
from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
@@ -37,25 +37,13 @@ class VizTrace:
@contextlib.contextmanager
def save_viz():
# clear previous traces
for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
Buffer.profile_events.clear()
cpu_events.clear()
# set the context vars to enable VIZ
prev_viz = VIZ.value
VIZ.value = -1
prev_tms = TRACK_MATCH_STATS.value
TRACK_MATCH_STATS.value = 2
prev_profile = PROFILE.value
PROFILE.value = 1
viz = VizTrace()
try:
with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1):
yield viz
finally:
viz.set_data()
TRACK_MATCH_STATS.value = prev_tms
PROFILE.value = prev_profile
VIZ.value = prev_viz
viz.set_data()
class TestViz(unittest.TestCase):
def test_simple(self):