From f7aeff6061ee281e9230e235219b5dec2bb3bec6 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Mon, 2 Mar 2026 12:24:38 +0200 Subject: [PATCH] viz: cli.py cleanups, do not require PYTHONPATH (#15085) * cleanup the print * sys.exit * equal check * cleanup unpacker * cli doesn't need PYTHONPATH * no semicolons * %s/PYTHONPATH=. //g --- .../tinybox_8xMI350X/profile.sh | 2 +- extra/viz/README | 6 +- extra/viz/cli.py | 67 +++++++++++++------ test/null/test_viz.py | 38 +---------- 4 files changed, 55 insertions(+), 58 deletions(-) diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh index dab82946a0..188aea51af 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh @@ -3,4 +3,4 @@ export BENCHMARK=5 export EVAL_BS=0 export VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh -PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20 +extra/viz/cli.py --profile --device "AMD" --top 20 diff --git a/extra/viz/README b/extra/viz/README index b1c76bad66..a7db4f9f66 100644 --- a/extra/viz/README +++ b/extra/viz/README @@ -1,17 +1,17 @@ A command line tool for exploring the VIZ trace. -After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files. +After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files. ## Inspect runtime profiling -Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices. +Use `extra/viz/cli.py --profile` to list all traced devices. List top slowest kernels on a device: `--profile --device "AMD"` List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3` ## Inspect codegen and PatternMatcher -Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels. +Use `extra/viz/cli.py --rewrites` to list all traced kernels. List all codegen steps for a kernel: `--rewrites --kernel E_3` Get source code: `--rewrites --kernel E_3 --select "View Source"` diff --git a/extra/viz/cli.py b/extra/viz/cli.py index a8f942dede..e6dfe04d99 100755 --- a/extra/viz/cli.py +++ b/extra/viz/cli.py @@ -1,44 +1,73 @@ #!/usr/bin/env python3 import os os.environ["VIZ"] = "0" -import argparse, pathlib +import argparse, pathlib, sys, struct, json from typing import Iterator from tinygrad.viz import serve as viz from tinygrad.uop.ops import RewriteTrace from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen -from test.null.test_viz import load_profile + +# ** generic helpers def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg def print_data(data:dict) -> None: if isinstance(data.get("value"), Iterator): for m in data["value"]: - if m.get("uop"): - print("Input UOp:") - print(m["uop"]) - if not m["diff"]: continue - print("Rewrites:") - fp = pathlib.Path(m["upat"][0][0]) - print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}") - print(m["upat"][1]) - for line in m["diff"]: - color = "red" if line.startswith("-") else "green" if line.startswith("+") else None - print(colored(line, color)) + if m.get("uop"): print(f"Input UOp:\n{m['uop']}") + if m.get("diff"): + loc = pathlib.Path(m["upat"][0][0]) + print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}") + for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None)) if data.get("src") is not None: print(data["src"]) +# ** Profiler trace decoder + +# 0 means None, otherwise it's an enum value +def option(i:int) -> int|None: return None if i == 0 else i-1 + +def decode_profile(data:bytes) -> dict: + ret, off = data, 0 + def u(fmt:str) -> tuple: + nonlocal off + vals = struct.unpack_from(fmt, ret, off) + off += struct.calcsize(fmt) + return vals + total_dur, global_peak, index_len, layout_len = u(" 0.01 else None) if et is not None else "" name = e["name"]+(" " * (46 - ansilen(e["name"]))) print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms "+e['fmt'].replace('\n', ' | ')+" ") @@ -81,7 +110,7 @@ if __name__ == "__main__": other_t = total-sum(t for _, (t, _) in sel) table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"]) print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github")) - exit(0) + sys.exit(0) for k in viz.ctxs: if not optional_eq(k, args.kernel): continue diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 8d8c74e6d9..b564d2eaae 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -1,4 +1,4 @@ -import unittest, decimal, json, struct, sys +import unittest, decimal, sys from dataclasses import dataclass from typing import Generator @@ -357,41 +357,9 @@ class TestVizIntegration(BaseTestViz): from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry from tinygrad.viz.serve import get_profile +from extra.viz.cli import decode_profile -class TinyUnpacker: - def __init__(self, buf): self.buf, self.offset = buf, 0 - def __call__(self, fmt:str) -> tuple: - ret = struct.unpack_from(fmt, self.buf, self.offset) - self.offset += struct.calcsize(fmt) - return ret - -# 0 means None, otherwise it's an enum value -def option(i:int) -> int|None: return None if i == 0 else i-1 - -def load_profile(lst:list[ProfileEvent]) -> dict: - ret = get_profile(lst) - u = TinyUnpacker(ret) - total_dur, global_peak, index_len, layout_len = u(" dict: return decode_profile(get_profile(lst)) class TestVizProfiler(BaseTestViz): def test_transfer_uses_copy_device(self):