diff --git a/test/null/test_viz.py b/test/null/test_viz.py index b02a6ef675..a8d3369351 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -222,6 +222,16 @@ class TestViz(unittest.TestCase): graph_rewrite(a + 4, TrackedPatternMatcher(_substitute.patterns), {a:a+1}, walk=True) list(viz.get_details(0, 0)) + def test_enter_calls_rewrite(self): + pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))]) + with save_viz() as viz: + inner = UOp.const(dtypes.int, 3) + func = UOp(Ops.FUNCTION, src=(UOp(Ops.SINK, src=(inner,)),)) + call = UOp(Ops.CALL, src=(func,)) + graph_rewrite(call, TrackedPatternMatcher(pm.patterns), enter_calls=True) + details = list(viz.get_details(0, 0)) + self.assertTrue(details[-1]["change"], "viz replay should detect change inside CALL") + def test_const_node_visibility(self): with save_viz() as viz: a = UOp.variable("a", 0, 10, dtype=dtypes.int) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index d6f38cbb1e..4ef821d952 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -440,12 +440,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def __bool__(self): return self._eval((dtypes.bool,), bool) def __int__(self): return self._eval(dtypes.ints, int) def __float__(self): return float(self._eval(dtypes.floats, float)) - def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False): + def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False, enter_calls:bool=False): dvars = {k:v for k,v in dvars.items() if k is not v} if len(dvars) == 0: return self with Context(TRACK_MATCH_STATS=(0 if name is None else TRACK_MATCH_STATS.value)): return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars, - bottom_up=True, walk=walk, name=name) + bottom_up=True, walk=walk, enter_calls=enter_calls, name=name) # NOTE: this is not called by Tensor slice (Tensor handles UOps directly), but satisfies SupportsIndex for type checking def __index__(self): return self.__int__() @@ -1397,6 +1397,7 @@ class TrackedGraphRewrite: depth:int # depth if it's a subrewrite bottom_up:bool walk:bool + enter_calls:bool tracked_keys:list[TracingKey] = [] tracked_ctxs:list[list[TrackedGraphRewrite]] = [] @@ -1454,7 +1455,7 @@ def profile_matches(fxn:Callable): 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("walk", False), kwargs.get("enter_calls", False))) active_rewrites.append(ctx) with cpu_profile(name, "TINY"): ret = fxn(*args, **kwargs) diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index a1438e56cb..a633cb3823 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -184,7 +184,7 @@ def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None) replaces: dict[UOp, UOp] = {} for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches, disable=not ctx.matches): replaces[u0:=_reconstruct(data, u0_num, depth=depth)] = u1 = _reconstruct(data, u1_num, depth=depth) - try: new_sink = next_sink.substitute(replaces, walk=ctx.walk) + try: new_sink = next_sink.substitute(replaces, walk=ctx.walk, enter_calls=ctx.enter_calls) except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e)) match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc) yield {"graph":(sink_json:=uop_to_json(data, new_sink)), "uop":pystr(new_sink), "change":[id(x) for x in u1.toposort() if id(x) in sink_json],