From c8ef4b60f6d1942cab37f7c86875a3ee6dbf319c Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sun, 19 Oct 2025 14:30:07 +0800 Subject: [PATCH] viz: share match tracing and TINY device profiler (#12783) * set a default name for the traces * set profile_matches + renames * profile_matches test * traces 4 steps total --- test/unit/test_viz.py | 24 +++++++++++++++++++++++- tinygrad/schedule/indexing.py | 8 ++++---- tinygrad/uop/ops.py | 22 ++++++++++++---------- tinygrad/viz/js/index.js | 2 +- 4 files changed, 40 insertions(+), 16 deletions(-) diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index 067af1ee7c..a1d9ba05a2 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -2,7 +2,7 @@ import unittest, decimal, json, struct 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 +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS, 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 @@ -117,6 +117,28 @@ class TestViz(BaseTestViz): # NOTE: names from TracingKey do not get deduped self.assertEqual(lst[0]["name"], "custom_name") + def test_profile_matches(self): + @profile_matches + def nested_function(u:UOp): + for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}") + + @track_rewrites() + def main_rewrite(u:UOp): + graph_rewrite(u, PatternMatcher([]), name="init") + nested_function(u) + + main_rewrite(UOp.variable("a", 1, 10)+UOp.variable("b", 1, 10)) + steps = get_viz_list()[0]["steps"] + self.assertEqual(steps[0]["name"], "init") + self.assertEqual(steps[1]["name"], "nested_function") + self.assertEqual(len(steps), 4) + + def test_profile_matches_invalid_arg(self): + @profile_matches + def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([])) + with self.assertRaisesRegex(AssertionError, "invalid match tracing input"): + invalid_fxn("test") + def test_colored_label(self): # NOTE: dataclass repr prints literal escape codes instead of unicode chars @dataclass(frozen=True) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 45bf24e083..d4b048a823 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -2,9 +2,9 @@ from typing import Iterator import functools, operator, itertools from dataclasses import dataclass, field from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses -from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey, PCONTIG, colored +from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW, Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL, @@ -139,7 +139,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO case _: raise RuntimeError(f"{op} is not a MovementOp") return rngs -@cpu_profile(TracingKey("run_rangeify"), "TINY") +@profile_matches def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: rctx = IndexingContext() @@ -147,7 +147,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize") # get the traversal order - with cpu_profile(TracingKey("reverse toposort"), "TINY"): + with cpu_profile("reverse toposort", "TINY"): tsink_reverse_toposort = tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map()) # explicit rangeify diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index c58e64f8c1..55ee4a69c1 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -930,7 +930,7 @@ class TrackedGraphRewrite: loc:tuple[str, int] # location that called graph_rewrite sink:int # the sink input to graph_rewrite matches:list[tuple[int, int, tuple, float]] # before/after UOp, UPat location and time - name:str|None # optional name of the rewrite + name:str # name of the rewrite depth:int # depth if it's a subrewrite bottom_up:bool @@ -975,19 +975,21 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal return _decorator active_rewrites:list[TrackedGraphRewrite] = [] -def track_matches(func): - def _track_func(*args, **kwargs): +def profile_matches(fxn:Callable): + def wrap(*args, **kwargs): + name = str(kwargs.get("name", None) or fxn.__name__) + assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}" if tracking:=(TRACK_MATCH_STATS >= 2): loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno) depth = len(active_rewrites) - if not tracked_ctxs: add_trace_group(TracingKey(f"default {func.__name__}")) - tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], kwargs.get("name", None), depth, kwargs.get("bottom_up", False))) + if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}")) + tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False))) active_rewrites.append(ctx) - with cpu_profile(kwargs.get("name", ""), "TINY", display=tracking): - ret = func(*args, **kwargs) + with cpu_profile(name, "TINY", display=tracking): + ret = fxn(*args, **kwargs) if tracking: active_rewrites.pop() return ret - return _track_func + return wrap class TrackedPatternMatcher(PatternMatcher): def rewrite(self, uop:UOp, ctx=None) -> UOp|None: @@ -1127,12 +1129,12 @@ class RewriteContext: self.replace[n] = replaced_new_n return self.replace[root] -@track_matches +@profile_matches def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None) -> UOp: rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx) return rewrite_ctx.unified_rewrite(sink) -@track_matches +@profile_matches def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, input_map:dict[UOp, UOp]|None=None, ) -> dict[UOp, UOp]: rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index f5e9472609..90ad10b6be 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -608,7 +608,7 @@ async function main() { for (const [j,u] of steps.entries()) { const inner = ul.appendChild(document.createElement("ul")); inner.id = `step-${i}-${j}`; - inner.innerText = `${u.name ?? u.loc[0].replaceAll("\\", "/").split("/").pop()+':'+u.loc[1]}`+(u.match_count ? ` - ${u.match_count}` : ''); + inner.innerText = `${u.name}`+(u.match_count ? ` - ${u.match_count}` : ''); inner.style.marginLeft = `${8*u.depth}px`; inner.onclick = (e) => { e.stopPropagation();