import unittest import decimal, sys, json, contextlib, tempfile, pickle, io, math, pathlib from dataclasses import dataclass from typing import Generator 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 from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap, VIZ, BEAM from tinygrad.device import Buffer from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace 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 @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) return sink # small container class for the viz server module class VizTrace: # loader init def __init__(self): self._data:VizData|None = None @property def data(self) -> VizData: return unwrap(self._data) def set_data(self) -> None: data = VizData(RewriteTrace(tracked_keys.copy(), tracked_ctxs.copy(), uop_fields.copy())) load_rewrites(data) self._data = data # the API def list_items(self) -> list[dict]: return self.data.ctxs def get_details(self, rewrite_idx:int, step:int) -> Generator[dict, None, None]: assert len(self.data.trace.rewrites) > rewrite_idx, f"only loaded {len(self.data.trace.rewrites)} traces, expecting at least {rewrite_idx}" return get_full_rewrite(self.data, self.data.trace.rewrites[rewrite_idx][step]) @contextlib.contextmanager def save_viz(): for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear() Buffer.profile_events.clear() cpu_events.clear() viz = VizTrace() with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1, PARALLEL=0): yield viz viz.set_data() needs_tracked_pm = unittest.skipUnless(VIZ, "using TrackedPatternMatcher requires global VIZ=1") class TestViz(unittest.TestCase): def test_simple(self): with save_viz() as viz: a = UOp.variable("a", 0, 10) exec_rewrite((a+0)*1, [sym]) lst = viz.list_items() # VIZ displays rewrites in groups of tracked functions self.assertEqual(len(lst), 1) # each group has a list of steps self.assertEqual(len(lst[0]["steps"]), 1) # each step has a list of matches self.assertEqual(lst[0]["steps"][0]["match_count"], 2) def test_rewrites(self): with save_viz() as viz: a = UOp.variable("a", 0, 10) exec_rewrite(a*1, [sym]) exec_rewrite(a*2, [sym]) lst = viz.list_items() self.assertEqual(len(lst), 2) # names dedup using a counter self.assertEqual(lst[0]["name"], "exec_rewrite n1") self.assertEqual(lst[1]["name"], "exec_rewrite n2") def test_steps(self): with save_viz() as viz: a = UOp.variable("a", 0, 10) exec_rewrite(a+1, [PatternMatcher([]), PatternMatcher([])], ["x", "y"]) steps = viz.list_items()[0]["steps"] # steps can optionally have a name self.assertEqual(steps[0]["name"], "x") self.assertEqual(steps[1]["name"], "y") def test_rewrite_location(self): def inner(sink): return graph_rewrite(sink, PatternMatcher([])) def outer(sink): return inner(sink) with save_viz() as viz: outer(UOp.variable("a", 1, 10)) lst = viz.list_items() # step location comes from inner rewrite fp, lineno = lst[0]["steps"][0]["loc"] self.assertEqual(fp, inner.__code__.co_filename) self.assertEqual(lineno, inner.__code__.co_firstlineno) def test_exceptions(self): # VIZ tracks rewrites up to and including the error def count_3(x:UOp): assert x.val <= 3 return UOp.const(x.val+1, x.dtype) err_pm = PatternMatcher([(UPat.cvar("x"), count_3),]) a = UOp.const(1) with save_viz() as viz: with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm]) lst = viz.list_items() err_step = lst[0]["steps"][0] self.assertEqual(err_step["match_count"], 4) # 3 successful rewrites + 1 err def test_default_name(self): with save_viz() as viz: a = UOp.variable("a", 1, 10) @rewrite_group() def name_default(): return graph_rewrite(a, PatternMatcher([])) name_default() lst = viz.list_items() self.assertEqual(lst[0]["name"], "name_default n1") # name can also come from a function that returns a string def test_dyn_name_fxn(self): with save_viz() as viz: @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() # name gets deduped by the function call counter self.assertEqual(lst[0]["name"], "(a+1) n1") # name can also come from a function that returns a TracingKey def test_tracing_key(self): with save_viz() as viz: @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_rewrite_group(self): with save_viz() as viz: @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") @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]) lst = viz.list_items() # inner calls fall outside the outer call self.assertEqual(len(lst), len(items)+1) self.assertEqual(lst[0]["name"], f"outer rewrite of {len(items)} inputs n1") steps = lst[0]["steps"] self.assertEqual(len(steps), 1) self.assertEqual(steps[0]["name"], "all") for i in range(len(items)): self.assertEqual(lst[i+1]["name"], f"inner fxn for {items[i]}") steps = lst[i+1]["steps"] self.assertEqual(len(steps), 1) self.assertEqual(steps[0]["name"], "each") def test_rewrite_group_nested(self): with save_viz() as viz: @rewrite_group(new_ctx=False) def nested_function(u:UOp): for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}") @rewrite_group() 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 = viz.list_items()[0]["steps"] self.assertEqual(steps[0]["name"], "init") self.assertEqual(steps[1]["name"], "nested_function") self.assertEqual(len(steps), 4) def test_rewrite_group_invalid_arg(self): with save_viz(): @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") def test_colored_label(self): # NOTE: dataclass repr prints literal escape codes instead of unicode chars @dataclass(frozen=True) class TestStruct: colored_field: str a = UOp(Ops.PYLITERAL, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue"))) a2 = uop_to_json(VizData(), a)[id(a)] self.assertEqual(ansistrip(a2["label"]), f"PYLITERAL\n{TestStruct.__qualname__}(colored_field='xyz12345')") def test_colored_label_multiline(self): with save_viz() as viz: arg = colored("x", "green")+"\n"+colored("y", "red")+colored("z", "yellow")+colored("ww\nw", "magenta") src = [Tensor.empty(1).uop for _ in range(10)] a = UOp(Ops.PYLITERAL, src=tuple(src), arg=arg) exec_rewrite(a, [PatternMatcher([])]) a2 = next(viz.get_details(0, 0))["graph"][id(a)] self.assertEqual(ansistrip(a2["label"]), "PYLITERAL\nx\nyzww\nw") def test_inf_loop(self): a = UOp.const(3) b = UOp.const(4) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with save_viz() as viz: # use smaller stack limit for faster test (default is 250000) with Context(REWRITE_STACK_LIMIT=100): self.assertRaises(RuntimeError, exec_rewrite, a, [pm]) graphs = flatten(x["graph"].values() for x in viz.get_details(0, 0)) self.assertEqual(graphs[0], uop_to_json(VizData(), a)[id(a)]) self.assertEqual(graphs[1], uop_to_json(VizData(), b)[id(b)]) # fallback to REWRITE_ERROR with the error message self.assertIn("REWRITE_ERROR\nTraceback", graphs[2]["label"]) # cut after the first error, instead of going through all REWRITE_STACK_LIMIT matches self.assertEqual(len(graphs), 3) def test_walk_rewrite(self): from tinygrad.uop.ops import _substitute with save_viz() as viz: a = UOp.variable("a", 0, 10) 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: UOp.const(4, x.dtype))]) with save_viz() as viz: inner = UOp.const(3) call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),)) func = UOp(Ops.CALL, src=(UOp(Ops.TUPLE, src=(call,)),)) graph_rewrite(func, 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) z = UOp.const(0) y = UOp.const(math.pi) alu = a*z ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym]) lst = viz.list_items() self.assertEqual(len(lst), 1) graphs = [x["graph"] for x in viz.get_details(0, 0)] # const is always in the graph, client side hides exclude=True nodes by default self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)]) self.assertTrue(graphs[0][id(z)]["exclude"]) self.assertTrue(graphs[0][id(y)]["exclude"]) self.assertFalse(graphs[0][id(alu)]["exclude"]) self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"]) self.assertEqual(list(graphs[1]), [id(u) for u in ret.toposort()]) # rewrite graph keys follow the rewritten sink's toposort def test_const_reshape_expand_folded(self): # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes c = UOp.const(1.0).expand((3,4)) # creates CONST->EXPAND chain a = UOp.variable("a", 0.0, 10.0, dtypes.float) alu = a + c with save_viz() as viz: graph_rewrite(alu, PatternMatcher([])) graph = [x["graph"] for x in viz.get_details(0, 0)][0] excluded_nodes = {v["label"].split("\n")[0] for v in graph.values() if v["exclude"]} self.assertIn("CONST", excluded_nodes) self.assertIn("STACK", excluded_nodes) self.assertIn("EXPAND", excluded_nodes) self.assertIn("CONST1 1", graph[id(alu)]["label"]) def test_stack_movement_not_folded_unless_all_const(self): a = UOp.variable("a", 0, 10, dtype=dtypes.int) c = UOp.const(1) stack = a.stack(c) reshaped = stack.reshape((1, 2)) graph = uop_to_json(VizData(), reshaped) self.assertFalse(graph[id(stack)]["exclude"]) const_stack = c.stack(UOp.const(2)) const_reshaped = const_stack.reshape((1, 2)) const_graph = uop_to_json(VizData(), const_reshaped) self.assertTrue(const_graph[id(const_stack)]["exclude"]) reshape_node = const_graph[id(const_reshaped)] self.assertFalse(reshape_node["exclude"]) self.assertIn("STACK0 {1,2} Ops.CONST", reshape_node["label"].split("\n")) # VIZ displays nested graph_rewrites in a tree view def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None leaf = TrackedPatternMatcher([(UPat(Ops.PARAM, name="x"), leaf_rewrite)]) def branch_rewrite(x:UOp, y:UOp): if x.tag is not None: return x2 = graph_rewrite(x, leaf, name="leaf_left") y2 = graph_rewrite(y, leaf, name="leaf_right") return x2 * y2 branch = TrackedPatternMatcher([(UPat.var("x")+UPat.var("y"), branch_rewrite)]) def root_rewrite(root:UOp): new_src = tuple(graph_rewrite(b, branch, name=f"branch_{i}") for i,b in enumerate(root.src)) return root.replace(src=new_src) root = TrackedPatternMatcher([(UPat(Ops.SINK, src=UPat(Ops.ADD), name="root"), root_rewrite),]) class TestVizTree(unittest.TestCase): def assertStepEqual(self, step:dict, want:dict): for k,v in want.items(): self.assertEqual(step[k], v, f"failed at '{k}': {v} != {step[k]}\n{step=}") def test_tree_view(self): with save_viz() as viz: a = UOp.variable("a",0,10,param=True) b = UOp.variable("b",0,10,param=True) c = UOp.variable("c",0,10,param=True) d = UOp.variable("d",0,10,param=True) sink = UOp.sink(a+b, c+d) def tree_rewrite(): return graph_rewrite(sink, root, name="root") tree_rewrite() lst = viz.list_items() steps = lst[0]["steps"] self.assertEqual(len(steps), 1+2+4) self.assertStepEqual(steps[0], {"name":"root", "depth":0, "match_count":1}) self.assertStepEqual(steps[1], {"name":"branch_0", "depth":1, "match_count":1}) self.assertStepEqual(steps[2], {"name":"leaf_left", "depth":2, "match_count":1}) self.assertStepEqual(steps[3], {"name":"leaf_right", "depth":2, "match_count":1}) self.assertStepEqual(steps[4], {"name":"branch_1", "depth":1, "match_count":1}) self.assertStepEqual(steps[5], {"name":"leaf_left", "depth":2, "match_count":1}) self.assertStepEqual(steps[6], {"name":"leaf_right", "depth":2, "match_count":1}) import gc def bufs_allocated() -> int: gc.collect() return sum([type(x).__name__ == "Buffer" and type(x).__module__ == "tinygrad.device" for x in gc.get_objects()]) class TestVizGC(unittest.TestCase): def test_gc(self): with save_viz() as viz: init = bufs_allocated() a = UOp.new_buffer("NULL", 10, dtypes.char) a.buffer.allocate() exec_rewrite(a, [PatternMatcher([])]) del a self.assertEqual(bufs_allocated()-init, 0) lst = viz.list_items() self.assertEqual(len(lst), 1) @unittest.skip("it's not generic enough to handle arbitrary UOps in arg") def test_gc_uop_in_arg(self): with save_viz() as viz: init = bufs_allocated() a = UOp.new_buffer("NULL", 10, dtypes.char) a.buffer.allocate() exec_rewrite(UOp(Ops.PYLITERAL, src=(a,), arg=a), [PatternMatcher([])]) del a self.assertEqual(bufs_allocated()-init, 0) lst = viz.list_items() self.assertEqual(len(lst), 1) # VIZ integrates with other parts of tinygrad from tinygrad import Tensor, Device, TinyJit, Variable, function class TestVizIntegration(unittest.TestCase): def test_link_sched_codegen(self): c1 = Tensor.empty(4, device="NULL") c2 = Tensor.empty(8, device="NULL") # uniquely named A = B + 1 kernel kernel_name = f"custom_add1_link_sched_codegen_{BEAM.value}" def custom_add1(A:UOp, B:UOp): return A[0].store(B[0]+1).sink(arg=KernelInfo(kernel_name)) with save_viz() as viz: c1 = Tensor.custom_kernel(c1, c2, fxn=custom_add1)[0] c1.realize() lst = viz.list_items() # schedule graph CALL nodes have a link to jump to codegen sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule")) viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph") graph = next(viz.get_details(sched_idx, viz_kernel))["graph"] call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")] for i,n in enumerate(call_nodes): assert n["ref"] is not None self.assertEqual(lst[n["ref"]]["name"], kernel_name) assert kernel_name[i] in n["label"], f"CALL must contain kernel name, got {n['label']}" # UOp addrspace is colored for u in graph.values(): if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL]) def test_link_sched_codegen_beam(self): with Context(BEAM=2): self.test_link_sched_codegen() @Context(TRACEMETA=2) def test_metadata_tracing(self): with save_viz() as viz: a = Tensor.empty(1) b = Tensor.empty(1) metadata = (alu:=a+b).uop.metadata alu.schedule_linear() 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 rewrite_group context # all graph_rewrites get put into the default group def test_default_tracing(self): with save_viz() as viz: def test(root): return graph_rewrite(root, sym) test(c:=UOp.const(1)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 1) self.assertEqual(ls[0]["name"], "default graph_rewrite") # using @rewrite_group organizes function calls into groups # and nicely counts function calls. def test_group_traces(self): with save_viz() as viz: @rewrite_group() def test(root): return graph_rewrite(root, sym) test(c:=UOp.const(1)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 2) for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}") # @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 = 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. default_test(c+2) ls = viz.list_items() self.assertEqual(len(ls), 2) graph = next(viz.get_details(0, 0))["graph"] # both operands of c+1 are the same bare weak CONST, so the graph has two nodes self.assertEqual(list(graph), [id(c), id(c+1)]) self.assertTrue(graph[id(c)]["exclude"]) self.assertFalse(graph[id(c+1)]["exclude"]) self.assertEqual(list(next(viz.get_details(1, 0))["graph"]), [id(c)]) graph = next(viz.get_details(1, 1))["graph"] self.assertEqual(list(graph), [id(c), id((c+2).src[1]), id(c+2)]) self.assertTrue(graph[id(c)]["exclude"]) self.assertTrue(graph[id((c+2).src[1])]["exclude"]) self.assertFalse(graph[id(c+2)]["exclude"]) def test_recurse(self): with save_viz() as viz: a = Tensor.empty(10) for _ in range(10_000): a += a graph_rewrite(a.uop, PatternMatcher([])) lst = viz.list_items() assert len(lst) == 1 def test_jit(self): with save_viz(): @TinyJit def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device)) a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1") for _ in range(3): Tensor.realize(*f(a, b, c)) out = load_profile(cpu_events) self.assertEqual(["NULL", "NULL Graph", "NULL:SDMA:0", "NULL:1", "NULL:1:SDMA:0"], [k for k in out["layout"] if k.startswith("NULL")]) self.assertEqual(len(out["layout"]["NULL"]["events"]), 2*3) self.assertEqual(len(out["layout"]["NULL:SDMA:0"]["events"]), 3) self.assertEqual(len(out["layout"]["NULL Graph"]["events"]), 2) for graph in out["layout"]["NULL Graph"]["events"]: graph_st, graph_et = graph["st"], graph["st"]+graph["dur"] for k in ["NULL", "NULL:1", "NULL:SDMA:0", "NULL:1:SDMA:0"]: events = [e for e in out["layout"][k]["events"] if graph_st <= e["st"] and e["st"]+e["dur"] <= graph_et] self.assertGreater(len(events), 0) self.assertEqual([e["st"] for e in events], [graph_st+i*events[0]["dur"] for i in range(len(events))]) @needs_tracked_pm def test_view_source(self): def custom_fn(X:UOp): X = X.flatten() i = UOp.range(X.numel(), 0) custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg=("{} + undeclared_name", X.dtype)) return X[i].store(custom_op).end(i).sink(arg=KernelInfo(name=f"custom_fn_{X.numel()}")) x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_fn)[0] with save_viz() as viz: with self.assertRaises(Exception) as e: x.realize() lst = viz.list_items() codegen_idx = len(lst)-1 steps = lst[codegen_idx]["steps"] lin_idx = next((i for i,s in enumerate(steps) if s["name"] == "View UOp List"), None) src_idx = next((i for i,s in enumerate(steps) if s["name"] == "View Source"), None) bin_idx = next((i for i,s in enumerate(steps) if s["name"] == "View Disassembly"), None) assert all(i is not None for i in [lin_idx, src_idx, bin_idx]), f"linear, source and disasm must be visible in {steps}" # Ops.LINEAR renders lin_render = get_render(viz.data, steps[lin_idx]["query"])["src"] self.assertIn("Ops.SINK", lin_render) self.assertIn("Ops.CUSTOMI", lin_render) # Ops.SOURCE renders src_render = get_render(viz.data, steps[src_idx]["query"])["src"] self.assertIn("undeclared_name", src_render) # Ops.BINARY shows the error message since compile failed bin_render = get_render(viz.data, steps[bin_idx]["query"])["src"] self.assertIn(type(e.exception).__name__, bin_render) def test_view_source_alt(self): src = "void E_3(float* data0_3) {}" binary = Device["CPU"].renderer.compiler.compile(src) def custom_binary(X:UOp): sink = UOp.sink(X, arg=KernelInfo("custom_binary")) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=sink.src+(sink,)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary))) x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_binary)[0] with save_viz() as viz: x.realize() lst = viz.list_items() codegen_idx = len(lst)-1 steps = lst[codegen_idx]["steps"] src_idx = next((i for i,s in enumerate(steps) if s["name"] == "View Source"), None) assert src_idx is not None, "must have source rendering in list" src_render = get_render(viz.data, steps[src_idx]["query"])["src"] self.assertEqual(src, src_render) def test_profiler_duplicate_name(self): kernel_name = "duplicate_name" def one(A:UOp): return A[0].store(UOp.const(1.0, dtypes.float)).sink(arg=KernelInfo(kernel_name)) def zero(A:UOp): return A[0].store(UOp.const(0.0, dtypes.float)).sink(arg=KernelInfo(kernel_name)) with save_viz() as viz: @TinyJit def f(a:Tensor, b:Tensor): return Tensor.custom_kernel(a, fxn=one)[0], Tensor.custom_kernel(b, fxn=zero)[0] a, b = Tensor.empty(4, device="NULL"), Tensor.empty(4, device="NULL") # warmup for _ in range(2): Tensor.realize(*f(a, b)) Tensor.realize(*f(a, b)) kernels = {i for i,c in enumerate(viz.list_items()) if c["name"] == kernel_name} profile = decode_profile(unwrap(get_profile(viz.data, cpu_events))) events = [e for e in profile["layout"]["NULL"]["events"] if e["name"] == kernel_name] self.assertEqual({e["ref"] for e in events}, kernels) from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry from tinygrad.viz.serve import get_profile from tinygrad.viz.cli import decode_profile def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(VizData(), lst)) class TestVizProfiler(unittest.TestCase): def test_transfer_uses_copy_device(self): with save_viz(): a = Tensor.ones(1, device="NULL").contiguous().realize() a.to("NULL:1").realize() range_events = [e for e in cpu_events if isinstance(e, ProfileRangeEvent)] compute_events = [e for e in range_events if e.device == "NULL"] copy_events = [e for e in range_events if e.device.endswith(":SDMA:0")] self.assertGreater(len(compute_events), 0, "expected compute events on base device") self.assertGreater(len(copy_events), 0, "transfer must produce events with ':SDMA' device suffix") def test_node(self): prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000))] j = load_profile(prof) dev_events = j['layout']['NV']['events'] self.assertEqual(len(dev_events), 1) event = dev_events[0] self.assertEqual(event['name'], 'E_2') self.assertEqual(event['st'], 0) self.assertEqual(event['dur'], 10) assert event['ref'] is None def test_copy_node(self): prof = [ProfileRangeEvent(device='NV:SDMA:0', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileRangeEvent(device='NV:2:SDMA:0', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileDeviceEvent(device='NV:SDMA:0', tdiff=decimal.Decimal(-100)), ProfileDeviceEvent(device='NV:2:SDMA:0', tdiff=decimal.Decimal(-80))] j = load_profile(prof) event = j['layout']['NV:SDMA:0']['events'][0] self.assertEqual(event['name'], 'COPYxx') self.assertEqual(event['st'], 0) # first event self.assertEqual(event['dur'], 10) event2 = j['layout']['NV:2:SDMA:0']['events'][0] self.assertEqual(event2['st'], 20) # second event, diff clock self.assertEqual(j["dur"], (event2["st"]+event2["dur"])-event["st"]) def test_copy_node_bandwidth(self): sz = 256*1024*1024 dur = 10_000 prof = [ProfileRangeEvent(device='NV:SDMA:0', name=TracingKey("NV -> NV:1", ret=sz), st=decimal.Decimal(1000), en=decimal.Decimal(1000+dur)), ProfileDeviceEvent(device='NV:SDMA:0', tdiff=decimal.Decimal(-1000))] j = load_profile(prof) event = j['layout']['NV:SDMA:0']['events'][0] self.assertEqual(event['fmt'], {"B/s": sz/(dur*1e-6), "B": sz}) def test_graph(self): prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)), ProfileDeviceEvent(device='NV:1:SDMA:0', tdiff=decimal.Decimal(-50)), ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV', name='E_25_4n2', st_id=0, en_id=1), ProfileGraphEntry(device='NV:1:SDMA:0', name='NV -> NV:1', st_id=2, en_id=3)], deps=[[], [0]], sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])] j = load_profile(prof) tracks = list(j['layout']) self.assertEqual(tracks[0], 'NV') self.assertEqual(tracks[1], 'NV Graph') self.assertEqual(tracks[2], 'NV:1:SDMA:0') nv_events = j['layout']['NV']['events'] self.assertEqual(nv_events[0]['name'], 'E_25_4n2') self.assertEqual(nv_events[0]['st'], 0) self.assertEqual(nv_events[0]['dur'], 2) sdma_events = j['layout']['NV:1:SDMA:0']['events'] self.assertEqual(sdma_events[0]['name'], 'NV -> NV:1') self.assertEqual(sdma_events[0]['st'], 954) graph_events = j['layout']['NV Graph']['events'] self.assertEqual(graph_events[0]['st'], nv_events[0]['st']) self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], sdma_events[0]['st']+sdma_events[0]['dur']) def test_graph_copy_bandwidth(self): sz = 256*1024*1024 dur = 10_000 prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)), ProfileDeviceEvent(device='NV:1:SDMA:0', tdiff=decimal.Decimal(-50)), ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV:1:SDMA:0', name=TracingKey("NV -> NV:1", ret=sz), st_id=0, en_id=1)], deps=[[]], sigs=[decimal.Decimal(1004), decimal.Decimal(1004+dur)])] j = load_profile(prof) sdma_events = j['layout']['NV:1:SDMA:0']['events'] self.assertEqual(sdma_events[0]["fmt"], {"B/s": sz/(dur*1e-6), "B": sz}) def test_block_ordering(self): prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)), ProfileDeviceEvent(device='NV:1', tdiff=decimal.Decimal(-500)), ProfileDeviceEvent(device='NV:SDMA:0', tdiff=decimal.Decimal(-100)), ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileRangeEvent(device='NV:1', name='E_3', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileRangeEvent(device='NV:SDMA:0', name='COPY', st=decimal.Decimal(1000), en=decimal.Decimal(1010)), ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV', name='E_2', st_id=0, en_id=1)], deps=[[]], sigs=[decimal.Decimal(1000), decimal.Decimal(1010)])] j = load_profile(prof) # graph grouped with its device, memory at the end self.assertListEqual(list(j['layout']), ['NV', 'NV Graph', 'NV:SDMA:0', 'NV:1']) @unittest.skipIf(sys.platform == 'win32', "TODO: ops_amd import fails on windows") def test_multi_sdma_ordering(self): props = {"gfx_target_version": 0} D, St, En = decimal.Decimal, decimal.Decimal(1000), decimal.Decimal(1010) prof = [# 2 AMD GPUs, 2 SDMA engines each ProfileDeviceEvent(device='AMD', tdiff=D(-1000), props=props), ProfileDeviceEvent(device='AMD:1', tdiff=D(-900), props=props), ProfileDeviceEvent(device='AMD:SDMA:0', tdiff=D(-100), props=props), ProfileDeviceEvent(device='AMD:SDMA:1', tdiff=D(-80), props=props), ProfileDeviceEvent(device='AMD:1:SDMA:0', tdiff=D(-60), props=props), ProfileDeviceEvent(device='AMD:1:SDMA:1', tdiff=D(-40), props=props), # compute + copy events ProfileRangeEvent(device='AMD', name='E_1', st=St, en=En), ProfileRangeEvent(device='AMD:1', name='E_2', st=St, en=En), ProfileRangeEvent(device='AMD:SDMA:0', name='COPY0', st=St, en=En), ProfileRangeEvent(device='AMD:SDMA:1', name='COPY1', st=St, en=En), ProfileRangeEvent(device='AMD:1:SDMA:0', name='COPY2', st=St, en=En), ProfileRangeEvent(device='AMD:1:SDMA:1', name='COPY3', st=St, en=En), # graph spanning compute + copy on GPU 0 ProfileGraphEvent(ents=[ProfileGraphEntry(device='AMD', name='E_1', st_id=0, en_id=1), ProfileGraphEntry(device='AMD:SDMA:0', name='COPY0', st_id=2, en_id=3)], deps=[[], [0]], sigs=[St, En, St, En]), # memory alloc on both GPUs ProfilePointEvent(device='AMD', name='alloc', key=0, arg={"sz":1024, "dtype":dtypes.float}, ts=St), ProfilePointEvent(device='AMD:1', name='alloc', key=1, arg={"sz":512, "dtype":dtypes.float}, ts=St)] j = load_profile(prof) # graph grouped with its device, memory at the end self.assertListEqual(list(j['layout']), ['AMD', 'AMD Graph', 'AMD:SDMA:0', 'AMD:SDMA:1', 'AMD:1', 'AMD:1:SDMA:0', 'AMD:1:SDMA:1', 'AMD Memory', 'AMD:1 Memory']) def test_bytes_per_kernel(self): step = 10 n_events = 1_000 prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)] sz = len(get_profile(VizData(), prof)) self.assertLessEqual(sz/n_events, 26) def test_calltrace(self): with save_viz() as viz: def fxn(): return Tensor.empty(10).mul(2).realize() with cpu_profile(TracingKey("test_fxn"), "CUSTOM"): fxn() codegen_trace = viz.list_items()[0]["steps"][0]["trace"] assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in codegen_trace), str(codegen_trace) profile_ret = load_profile(cpu_events) e = profile_ret["layout"]["CUSTOM"]["events"][0] self.assertEqual(e["name"], "test_fxn") runtime_trace = e["fmt"]["tb"] assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno+1 == l for f,l,*_ in runtime_trace), str(runtime_trace) # can pack up to 1hr 11 min of trace events def test_trace_duration(self): dur_mins = 72 n_events = 1_000 step = decimal.Decimal(dur_mins*60*1e6//n_events) prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)] with self.assertRaisesRegex(ValueError, "timestamp out of range"): get_profile(VizData(), prof) def test_python_marker(self): with save_viz(): a = Tensor.empty(1, device="NULL") b = Tensor.empty(1, device="NULL") (a+b).realize() profile_marker("test 1") (a*b).realize() profile_marker("test 2") profile_ret = load_profile(cpu_events) markers = profile_ret["markers"] kernels = profile_ret["layout"]["NULL"]["events"] self.assertEqual(len(markers), 2) assert kernels[0]["st"] <= markers[0]["ts"] <= kernels[1]["st"] assert markers[1]["ts"] >= kernels[1]["st"]+kernels[1]["dur"] def test_layout_order(self): with save_viz(): def fn(): return for dname in ["TINY", "USER", "TEST:1 N1", "TEST:2 N1", "TEST:1 N2", "TEST:1:ENGINE:0", "TEST:1:ENGINE:0 N1", "TEST:1"]: with cpu_profile("fn", dname): fn() layout = list(load_profile(cpu_events)["layout"]) self.assertListEqual(layout[:2], ["USER","TINY"]) self.assertListEqual(layout[2:], ["TEST:1", "TEST:1 N1", "TEST:1 N2", "TEST:1:ENGINE:0", "TEST:1:ENGINE:0 N1", "TEST:2 N1"]) def _alloc(b:int): a = Tensor.empty(b, device="NULL", dtype=dtypes.char) a.uop.buffer.allocate() return a class TestVizMemoryLayout(unittest.TestCase): def test_double_alloc(self): with save_viz(): a = _alloc(1) _b = _alloc(1) profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{a.device} Memory"] self.assertEqual(ret["peak"], 2) self.assertEqual(len(ret["events"]), 4) def test_del_once(self): with save_viz(): a = _alloc(1) del a b = _alloc(1) profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{b.device} Memory"] self.assertEqual(ret["peak"], 1) self.assertEqual(len(ret["events"]), 4) def test_alloc_free(self): with save_viz(): a = _alloc(1) _b = _alloc(1) del a c = _alloc(1) profile_ret = load_profile(Buffer.profile_events) ret = profile_ret["layout"][f"{c.device} Memory"] self.assertEqual(ret["peak"], 2) self.assertEqual(len(ret["events"]), 6) def test_free_last(self): with save_viz(): bufs = [] for _ in range(3): bufs.append(_alloc(1)) profile_marker("alloc") device = bufs[0].device while bufs: b = bufs.pop() del b profile_marker("free") profile = load_profile(cpu_events+Buffer.profile_events) ret = profile["layout"][f"{device} Memory"] self.assertEqual(ret["peak"], 3) self.assertEqual(len(ret["events"]), 6) self.assertEqual(len(profile["markers"]), 6) def test_producer_simple(self): with save_viz(): a = Tensor.ones(10, device="NULL") Tensor.realize(a.add(1).contiguous()) b = Tensor.ones(10, device="NULL") Tensor.realize(b.add(1).contiguous()) profile = load_profile(cpu_events+Buffer.profile_events) buffers = profile["layout"]["NULL Memory"]["events"] programs = profile["layout"]["NULL"]["events"] user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")] self.assertEqual(len(user_cnt), len(programs)) @unittest.skip("flaky") def test_inflight_buf(self): a = Tensor.empty(1, device="NULL") n = 4 for i in range(n): (a+i).realize() profile = load_profile(cpu_events+Buffer.profile_events) buffers = profile["layout"]["NULL Memory"]["events"] user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")] self.assertEqual(max(user_cnt), n) input_buf = buffers.pop() assert all(u[3] == 0 for u in input_buf["arg"]["users"]) def test_annotate_read_write(self): with save_viz(): a = Tensor.ones(4, device="NULL").contiguous().realize() b = a.assign(a+2) c = a+1 Tensor.realize(b, c) buf_events = load_profile(cpu_events+Buffer.profile_events)["layout"]["NULL Memory"]["events"] users = next((b["arg"]["users"] for b in buf_events if len(b["arg"].get("users",[])) == 3)) self.assertEqual(users[0][3], 1) # write Tensor.ones self.assertEqual(users[1][3], 2) # read+write Tensor.assign self.assertEqual(users[2][3], 0) # readonly def test_dedup_users(self): with save_viz(): a = Tensor.empty(1, device="NULL") for _ in range(n:=4): a.add(1).realize() profile = load_profile(cpu_events+Buffer.profile_events) programs = profile["layout"][a.device]["events"] users = profile["layout"][f"{a.device} Memory"]["events"].pop()["arg"]["users"] self.assertEqual(len(programs), len(set(users)), n) from tinygrad.uop.ops import KernelInfo from tinygrad.renderer.amd.dsl import s from tinygrad.runtime.autogen.amd.rdna3.ins import (s_add_u32, s_branch, s_cbranch_execz, s_cbranch_scc0, s_cbranch_scc1, s_cmp_eq_i32, s_cmp_eq_u64, s_code_end, s_endpgm, s_mov_b32, s_nop) from extra.gemm.amd_asm_matmul import Kernel @needs_tracked_pm class TestCfg(unittest.TestCase): def get_cfg(self, name:str, k:Kernel): insts = k.finalize() def fxn(out:UOp) -> UOp: lidx = UOp.special(1, "lidx0") gidx = UOp.special(1, "gidx0") sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name)) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts])))) with save_viz() as viz: with Context(DEV="NULL::gfx1100"): out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0] _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer) codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name) disasm = next(s for s in codegen_rewrites["steps"] if s["name"] == "View Disassembly") return get_render(viz.data, disasm["query"]) def test_simple(self): k = Kernel() k.label("entry") k.emit(s_branch(), target="bb1") k.label("bb1") k.emit(s_endpgm()) k.emit(s_code_end()) cfg = self.get_cfg("simple", k)["data"] self.assertEqual(len(cfg["blocks"]), 2) def test_diamond(self): k = Kernel() k.label("entry") k.emit(s_mov_b32(s[0], 0)) k.emit(s_mov_b32(s[1], 0)) k.emit(s_cmp_eq_u64(s[0:1], 0)) k.emit(s_cbranch_scc1(), target="if") k.emit(s_branch(), target="else") k.label("if") k.emit(s_nop(1)) k.emit(s_branch(), target="end") k.label("else") k.emit(s_nop(0)) k.label("end") k.emit(s_endpgm()) k.emit(s_code_end()) ret = self.get_cfg("diamond", k) cfg = ret["data"] self.assertEqual(len(cfg["blocks"]), 5) edge_count = sum(len(v) for v in cfg["paths"].values()) self.assertEqual(edge_count, 5) references:dict[str, list[str]] = {} for pc, tokens in cfg["pc_tokens"].items(): for t in tokens: for key in t["keys"]: references.setdefault(key, []).append(pc) self.assertEqual(len(references["r0"]), 2) insts = [cfg["pc_tokens"][pc][0]["st"] for pc in references["r0"]] self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64']) end_block = [" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1]] code_line = ret["src"].splitlines()[-1] self.assertEqual(len(end_block), 2) for st in [end_block[-1], code_line]: assert st.startswith("s_code_end") and st.endswith("x)"), st def test_loop(self): k = Kernel() k.label("entry") k.emit(s_mov_b32(s[1], 4)) k.label("loop") k.emit(s_add_u32(s[1], s[1], -1)) k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_cbranch_scc0(), target="loop") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("simple_loop", k) def test_loop_branch(self): k = Kernel() k.label("entry") k.emit(s_mov_b32(s[1], 4)) k.label("loop") k.emit(s_add_u32(s[1], s[1], -1)) k.emit(s_cmp_eq_i32(s[1], 2)) k.emit(s_cbranch_scc1(), target="cond") k.emit(s_branch(), target="cont") k.label("cond") k.emit(s_add_u32(s[1], s[1], -2)) k.label("cont") k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_cbranch_scc0(), target="loop") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("loop_if", k) def test_loop_break(self): k = Kernel() k.label("entry") k.emit(s_mov_b32(s[1], 8)) k.label("loop") k.emit(s_add_u32(s[1], s[1], -1)) k.emit(s_cmp_eq_i32(s[1], 5)) k.emit(s_cbranch_scc1(), target="break") k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_cbranch_scc0(), target="loop") k.label("break") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("loop_break", k) def test_switch(self): k = Kernel() k.label("entry") k.emit(s_cmp_eq_i32(s[0], 0)) k.emit(s_cbranch_scc1(), target="case0") k.emit(s_cmp_eq_i32(s[0], 1)) k.emit(s_cbranch_scc1(), target="case1") k.emit(s_branch(), target="case2") k.label("case0") k.emit(s_nop(0)) k.emit(s_branch(), target="join") k.label("case1") k.emit(s_nop(1)) k.emit(s_branch(), target="join") k.label("case2") k.emit(s_nop(2)) k.emit(s_branch(), target="join") k.label("join") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("switch_case", k) def test_ping_pong(self): k = Kernel() k.label("entry") k.emit(s_cmp_eq_i32(s[0], 0)) k.emit(s_cbranch_scc1(), target="ping") k.emit(s_branch(), target="pong") k.label("ping") k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_cbranch_scc1(), target="pong") k.emit(s_branch(), target="end") k.label("pong") k.emit(s_cmp_eq_i32(s[2], 0)) k.emit(s_cbranch_scc1(), target="ping") k.label("end") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("ping_pong", k) def test_colored_blocks(self): N = 10 k = Kernel() k.label("entry") k.emit(s_branch(), target="init0") for i in range(N): loop = f"loop{i}" k.label(f"init{i}") k.emit(s_mov_b32(s[1], i + 1)) k.emit(s_branch(), target=loop) k.label(loop) k.emit(s_nop(i & 7)) k.emit(s_add_u32(s[1], s[1], -1)) k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_cbranch_scc0(), target=loop) k.emit(s_branch(), target=f"init{i+1}" if i + 1 < N else "end") k.label("end") k.emit(s_endpgm()) k.emit(s_code_end()) self.get_cfg("test_colored_blocks", k) def test_jump_back_to_end(self): k = Kernel() k.label("entry") k.emit(s_mov_b32(s[1], 2)) k.emit(s_cbranch_execz(), target="loop") k.label("end") k.emit(s_endpgm()) k.label("loop") k.emit(s_add_u32(s[1], s[1], -1)) k.emit(s_cmp_eq_i32(s[1], 0)) k.emit(s_branch(), target="end") k.emit(s_code_end()) self.get_cfg("jump_back_to_end", k) # launch viz cli without subprocess def run_cli(*cli_args) -> list[dict]: from tinygrad.viz.cli import main, get_arg_parser args = get_arg_parser().parse_args(cli_args+("--json",)) with contextlib.redirect_stdout(buf:=io.StringIO()): main(args) return [json.loads(line) for line in buf.getvalue().strip().splitlines()] @contextlib.contextmanager def write_files(viz) -> list[str]: with tempfile.TemporaryDirectory() as tmpdir: (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace)) (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events)) yield ["--rewrites-path", str(r), "--profile-path", str(p)] class TestCLI(unittest.TestCase): @needs_tracked_pm def test_reconstruct_debug(self): with save_viz() as viz: Tensor.empty(1, device="NULL").add(2.0).realize() profile_marker("marker @ 1") Tensor.empty(1, device="NULL").add(3.0).realize() with write_files(viz) as files, Context(DEBUG=4): out = run_cli(*files, "-s", "NULL") assert any(s.get("value", "").startswith("void E") for s in out) assert any(s.get("name", "") == "marker @ 1" for s in out) def test_aggregate(self): N, CNT = 1024, 5 with save_viz() as viz: for _ in range(CNT): (Tensor.empty(N, N, device="NULL")@Tensor.empty(N, N, device="NULL")).realize() for _ in range(CNT): (Tensor.empty(N, N, device="NULL").assign(Tensor.empty(N, N, device="NULL"))).realize() with write_files(viz) as files, Context(NO_COLOR=1): kernels = run_cli(*files, "-s", "NULL", "-t") self.assertEqual(len(kernels), 2) gemm_summary = [s for s in kernels if s["name"].startswith("r_")][0] copy_summary = [s for s in kernels if s["name"].startswith("E_")][0] self.assertEqual(gemm_summary["count"], CNT) self.assertEqual(copy_summary["count"], CNT) def test_flops(self): test_n = [(8, 16), (16, 32), (32, 64)] with save_viz() as viz: @TinyJit def f(a, b): return (a@a.T), (b@b.T) a = Tensor.empty(64, 64, device="NULL") b = Tensor.empty(64, 64, device="NULL") for i_val, j_val in test_n: i = Variable("i", 1, 64).bind(i_val) j = Variable("j", 1, 64).bind(j_val) Tensor.realize(*f(a[:i], b[:j])) with write_files(viz) as files: out = run_cli(*files, "-s", "NULL") aggregate = run_cli(*files, "-s", "NULL", "-t") self.assertEqual(len(out), 3*2) # flops increases as N gets larger gflops = [row["fmt"]["FLOPS"] for row in out] self.assertGreater(gflops[4], gflops[2]) self.assertGreater(gflops[5], gflops[3]) # aggregate flops self.assertEqual(len(aggregate), 2) agg_gflops = [row["fmt"]["FLOPS"] for row in aggregate] assert all(min(gflops) < v < max(gflops) for v in agg_gflops), f"{agg_gflops}" def test_dedup(self): with save_viz() as viz: for _ in range(CNT:=4): # use kernel names unique to this test Tensor.custom_kernel(Tensor.empty(4, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k1_test_viz_dedup")))[0].realize() Tensor.custom_kernel(Tensor.empty(8, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo("k2_test_viz_dedup")))[0].realize() with write_files(viz) as files, Context(NO_COLOR=1): name = run_cli(*files, "-s", "NULL")[0]["name"] with Context(DEBUG=3): select = run_cli(*files, "-s", "NULL", name) self.assertEqual(len([s for s in select if s.get("value")]), 1, "debug output was not deduped") self.assertEqual(len([s for s in select if s.get("device") == "NULL"]), CNT, f"expected 4 runs for {name}") @needs_tracked_pm def test_call_graph(self): @function(precompile=True) def f(x): r = x.sum(axis=1).reshape(32, 1).expand(32, 32).contiguous() return x + r # turn off scache because this test requires a complete schedule rewrite with save_viz() as viz, Context(SCACHE=0): f(f(Tensor.empty(32, 32, device="NULL"))).realize() with write_files(viz) as files, Context(NO_COLOR=1): prgs = [s["name"] for s in run_cli(*files, "-s", "NULL")] with Context(DEBUG=5): out = run_cli(*files, "-s", "TINY") i = next(i for i,s in enumerate(out) if s.get("value", "").lstrip() == "View Kernel Graph") # next print is the CALL graph, CLI outputs exactly as web in TestVizIntegration.test_link_sched_codegen call_nodes = [n for n in out[i+1].values() if n["label"].startswith("CALL")] for i,n in enumerate(call_nodes): assert prgs[i] in n["label"], f"CALL must contain kernel name, got {n['label']}" def test_interval(self): def emit_kernel(name:str): Tensor.custom_kernel(Tensor.empty(1, device="NULL"), fxn=lambda _: UOp.sink(arg=KernelInfo(name=name)))[0].realize() with save_viz() as viz: emit_kernel("pre_1") emit_kernel("pre_2") profile_marker("interval_start") emit_kernel("target_1") emit_kernel("target_2") profile_marker("interval_end") emit_kernel("post_1") emit_kernel("post_2") with write_files(viz) as files, Context(NO_COLOR=1): flat = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end") aggregate = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end", "-t") final = run_cli(*files, "-s", "NULL", "--interval", "interval_end", "-t") self.assertEqual([s["name"] for s in flat], ["interval_start", "target_1", "target_2", "interval_end"]) self.assertEqual(sorted(s["name"] for s in aggregate), ["target_1", "target_2"]) assert all(s["name"].startswith("post_") for s in final), f"post_* kernels must be present in final, got {final}" if __name__ == "__main__": unittest.main()