From 581bfdd94f0526e802639e321e6d41cefdc6fece Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:41:38 -0700 Subject: [PATCH] merge track_rewrites and profile_matches into rewrite_group [PR] (#17420) * merge track_rewrites and profile_matches into rewrite_group * bug * flip ctx polarity --- test/null/test_graph_rewrite.py | 4 +- test/null/test_upat_compile.py | 4 +- test/null/test_viz.py | 36 +++++++++--------- tinygrad/codegen/__init__.py | 4 +- tinygrad/engine/jit.py | 4 +- tinygrad/runtime/support/hcq2.py | 6 +-- tinygrad/schedule/__init__.py | 4 +- tinygrad/schedule/indexing.py | 4 +- tinygrad/schedule/rangeify.py | 4 +- tinygrad/tensor.py | 4 +- tinygrad/uop/ops.py | 63 +++++++++++++++----------------- 11 files changed, 67 insertions(+), 70 deletions(-) diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index a10d13ea8b..9194cc01bd 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -208,7 +208,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): import inspect -from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites +from tinygrad.uop.ops import graph_rewrite, _substitute, rewrite_group from tinygrad.uop.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): @@ -220,7 +220,7 @@ class TestBottomUpRewrite(unittest.TestCase): self.assertIs(gt, ret) # normally .substitute would be fine, but it's not tracked -@track_rewrites() +@rewrite_group() def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True) def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel) diff --git a/test/null/test_upat_compile.py b/test/null/test_upat_compile.py index b59f5f03c6..eeea86efad 100644 --- a/test/null/test_upat_compile.py +++ b/test/null/test_upat_compile.py @@ -1,11 +1,11 @@ import unittest from tinygrad.helpers import DEBUG, Context from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops +from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops from tinygrad.uop.upat import _get_code, upat_compile import dis -@track_rewrites() +@rewrite_group() def do_compile(up): print("\n***** COMPILE", up) match_code = _get_code(up, False) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 954a4f2252..c6cdb6213c 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -3,7 +3,7 @@ from pathlib import Path from dataclasses import dataclass from typing import Generator -from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches +from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, rewrite_group from tinygrad.uop.symbolic import sym from tinygrad.dtype import dtypes, AddrSpace from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker @@ -14,7 +14,7 @@ from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewr from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors from tinygrad.codegen import do_to_program -@track_rewrites(name=True) +@rewrite_group(name=True) def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp: for i,pm in enumerate(pm_lst): sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None) @@ -109,7 +109,7 @@ class TestViz(unittest.TestCase): def test_default_name(self): with save_viz() as viz: a = UOp.variable("a", 1, 10) - @track_rewrites() + @rewrite_group() def name_default(): return graph_rewrite(a, PatternMatcher([])) name_default() lst = viz.list_items() @@ -118,7 +118,7 @@ class TestViz(unittest.TestCase): # name can also come from a function that returns a string def test_dyn_name_fxn(self): with save_viz() as viz: - @track_rewrites(name=lambda *args,ret,**kwargs: ret.render()) + @rewrite_group(name=lambda *args,ret,**kwargs: ret.render()) def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([])) name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"]) lst = viz.list_items() @@ -128,18 +128,18 @@ class TestViz(unittest.TestCase): # name can also come from a function that returns a TracingKey def test_tracing_key(self): with save_viz() as viz: - @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,))) + @rewrite_group(name=lambda inp,ret: TracingKey("custom_name", (inp,))) def test(s:UOp): return graph_rewrite(s, PatternMatcher([])) test(UOp.variable("a", 1, 10)+1) lst = viz.list_items() # NOTE: names from TracingKey do not get deduped self.assertEqual(lst[0]["name"], "custom_name") - def test_nested_track_rewrites(self): + def test_nested_rewrite_group(self): with save_viz() as viz: - @track_rewrites(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,))) + @rewrite_group(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,))) def inner(x:UOp): return graph_rewrite(x, PatternMatcher([]), name="each") - @track_rewrites(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs") + @rewrite_group(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs") def outer(*xs:tuple[UOp, ...]): return graph_rewrite(UOp.sink(*[inner(x) for x in xs]), PatternMatcher([]), name="all") items = ["a", "b", "c"] outer(*[UOp.variable(x, 1, 10) for x in items]) @@ -156,13 +156,13 @@ class TestViz(unittest.TestCase): self.assertEqual(len(steps), 1) self.assertEqual(steps[0]["name"], "each") - def test_profile_matches(self): + def test_rewrite_group_nested(self): with save_viz() as viz: - @profile_matches + @rewrite_group(new_ctx=False) def nested_function(u:UOp): for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}") - @track_rewrites() + @rewrite_group() def main_rewrite(u:UOp): graph_rewrite(u, PatternMatcher([]), name="init") nested_function(u) @@ -173,9 +173,9 @@ class TestViz(unittest.TestCase): self.assertEqual(steps[1]["name"], "nested_function") self.assertEqual(len(steps), 4) - def test_profile_matches_invalid_arg(self): + def test_rewrite_group_invalid_arg(self): with save_viz(): - @profile_matches + @rewrite_group(new_ctx=False) def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([])) with self.assertRaisesRegex(AssertionError, "invalid match tracing input"): invalid_fxn("test") @@ -395,7 +395,7 @@ class TestVizIntegration(unittest.TestCase): graph = next(viz.get_details(0, 0))["graph"] self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1) - # tracing also works without a track_rewrites context + # tracing also works without a rewrite_group context # all graph_rewrites get put into the default group def test_default_tracing(self): with save_viz() as viz: @@ -407,11 +407,11 @@ class TestVizIntegration(unittest.TestCase): self.assertEqual(len(ls), 1) self.assertEqual(ls[0]["name"], "default graph_rewrite") - # using @track_rewrites organizes function calls into groups + # using @rewrite_group organizes function calls into groups # and nicely counts function calls. def test_group_traces(self): with save_viz() as viz: - @track_rewrites() + @rewrite_group() def test(root): return graph_rewrite(root, sym) test(c:=UOp.const(1)) @@ -420,11 +420,11 @@ class TestVizIntegration(unittest.TestCase): self.assertEqual(len(ls), 2) for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}") - # @track_rewrites always starts a new group. + # @rewrite_group always starts a new group. def test_group_combined(self): with save_viz() as viz: def default_test(root): return graph_rewrite(root, sym) - tracked_test = track_rewrites()(default_test) + tracked_test = rewrite_group()(default_test) c = UOp.const(1) default_test(c+1) # goes to the default group tracked_test(c) # all rewrites after this go inside the second group. diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 58fceb0882..abdaa3cf49 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -2,7 +2,7 @@ from dataclasses import replace, dataclass import itertools, functools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic -from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program @@ -448,7 +448,7 @@ pm_to_program = PatternMatcher([ (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile), ]) -@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True) +@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True) @Context(ALLOW_DEVICE_USAGE=0) def do_to_program(ast:UOp, renderer:Renderer) -> UOp: """ diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 97074388ff..fe11805c17 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor, all_tensors from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker from tinygrad.dtype import DType -from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite +from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite from tinygrad.renderer import Estimates from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins @@ -64,7 +64,7 @@ def _copy_input(u:UOp) -> UOp: run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),))) return new -@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}") +@rewrite_group(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}") def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp: if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear") diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 2933ddfba0..be32457624 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -5,7 +5,7 @@ from dataclasses import replace, dataclass from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker -from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp +from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp from tinygrad.uop.symbolic import symbolic from tinygrad.dtype import dtypes, truncate from tinygrad.runtime.support.hcq import MMIOInterface @@ -393,7 +393,7 @@ pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=( hcq_compile_cache:dict[bytes, UOp] = {} -@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") +@rewrite_group(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp: if input_uops is not None: slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))} @@ -477,7 +477,7 @@ def link_buf_key(a:UOp): return a.key, to_tuple(a.device) link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {} link_linear_cache:dict[bytes, UOp] = {} -@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") +@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}") def hcq_link(linear:UOp, cache=True) -> UOp: if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 99d228945f..82763ca74e 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -1,6 +1,6 @@ import time, inspect from collections import deque -from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo +from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo from tinygrad.uop.spec import type_verify, spec_tensor from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup @@ -167,7 +167,7 @@ pm_copy_from_store = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices), ]) -@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}") +@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}") def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: # big_sink srcs are all the Tensors linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index b3b6d6da68..8c02c389d1 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -2,7 +2,7 @@ from typing import Iterator import functools, itertools from dataclasses import dataclass, field, replace from tinygrad.dtype import dtypes, AddrSpace -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes from tinygrad.uop.ops import gate_kernel_sink from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC @@ -176,7 +176,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 -@profile_matches +@rewrite_group(new_ctx=False) def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: if debug: print("**************************") rctx = IndexingContext() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 3c4e3750aa..1a45579a20 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -3,7 +3,7 @@ from typing import cast import itertools from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg -from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element +from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element from tinygrad.uop.symbolic import symbolic from tinygrad.uop.movement import mop_cleanup from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS @@ -551,7 +551,7 @@ pm_copy_to_store = PatternMatcher([ (UPat(Ops.COPY, name="copy"), convert_copy_to_store), ]) -@profile_matches +@rewrite_group(new_ctx=False) def get_kernel_graph(sink:UOp) -> UOp: tsink = graph_rewrite(sink, multi_pm, name="multi_pm") if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters") diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 57a9fea6e3..33db031cfd 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -8,7 +8,7 @@ from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtyp _from_np_dtype, _to_np_dtype, PyConst, AddrSpace from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize -from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, track_rewrites +from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, rewrite_group from tinygrad.mixin.rand import RandMixin from tinygrad.schedule import create_linear_with_vars from tinygrad.device import Buffer, canonicalize_device @@ -211,7 +211,7 @@ pm_replace_buf = PatternMatcher([ (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer), ]) -@track_rewrites(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}") +@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}") def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph") # uop list is a list in the original_sink graph and we can map to the tags later diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 1d06c85a46..491d810377 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1515,55 +1515,52 @@ def add_trace_group(kt:TracingKey) -> None: tracked_ctxs.append([]) active_group:list[int] = [] -def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False): +active_rewrites:list[TrackedGraphRewrite] = [] +def rewrite_group(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False, new_ctx:bool=True): + if not new_ctx: assert not callable(name) and not replay, "name fxn and replay are only supported for new_ctx groups" def _decorator(func): def __wrapper(*args, **kwargs): + # without tracking, we just call the function (unless top-level, which always profiles) + if TRACK_MATCH_STATS < 2 and not new_ctx: return func(*args, **kwargs) fn = key = func.__name__ idx = -1 if TRACK_MATCH_STATS >= 2: - add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,))) - active_group.append(idx:=len(tracked_keys)-1) + if new_ctx: + add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,))) + active_group.append(idx:=len(tracked_keys)-1) + else: + rewrite_name = str(kwargs.get("name", None) or fn) + assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {rewrite_name} with {args}" + 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 {fn}")) + dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1 + tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], rewrite_name, depth, kwargs.get("bottom_up", False), + kwargs.get("walk", False), kwargs.get("enter_calls", False))) + active_rewrites.append(ctx) + key = rewrite_name # profile spans are named after the rewrite step with cpu_profile(key, "TINY") as e: ret = func(*args, **kwargs) - if TRACK_MATCH_STATS >= 2: active_group.pop() - if TRACK_MATCH_STATS >= 2 and callable(name): - name_ret = name(*args, **kwargs, ret=ret) - assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}" - tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret - e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys) + if TRACK_MATCH_STATS >= 2: + if new_ctx: active_group.pop() + else: active_rewrites.pop() + if callable(name): + name_ret = name(*args, **kwargs, ret=ret) + assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}" + tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret + e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys) if CAPTURE_PROCESS_REPLAY and replay: # find the unittest frame we're capturing in frm = sys._getframe(1) while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back - loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}" + replay_loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}" # capture global context vars and all the args passed in inputs = (fn, args, kwargs, ContextVar._cache) - replay_capture.append(pickle.dumps(inputs+(loc, ret))) + replay_capture.append(pickle.dumps(inputs+(replay_loc, ret))) return ret return __wrapper return _decorator -active_rewrites:list[TrackedGraphRewrite] = [] -def profile_matches(fxn:Callable): - def wrap_profile_matches(*args, **kwargs): - if TRACK_MATCH_STATS >= 2: - 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}" - 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 {fxn.__name__}")) - dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1 - tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False), - kwargs.get("walk", False), kwargs.get("enter_calls", False))) - active_rewrites.append(ctx) - with cpu_profile(name, "TINY"): - ret = fxn(*args, **kwargs) - active_rewrites.pop() - return ret - # without tracking, we just call the function - return fxn(*args, **kwargs) - return wrap_profile_matches - class TrackedPatternMatcher(PatternMatcher): def rewrite(self, uop:UOp, ctx=None): if len(pats:=self.pdict.get(uop.op, [])): @@ -1742,7 +1739,7 @@ class RewriteContext: if n in waitlist: stack.extend(waitlist.pop(n)) return self.replace[root] -@profile_matches +@rewrite_group(new_ctx=False) def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp: rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)