sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity
Files
simple_rangeify
tinygrad/test/null/test_viz.py
T
George HotzandGitHub 581bfdd94f merge track_rewrites and profile_matches into rewrite_group [PR] (#17420) 
* merge track_rewrites and profile_matches into rewrite_group

* bug

* flip ctx polarity
2026-08-05 14:41:38 -07:00

1122 lines
47 KiB
Python
Raw Permalink Blame History

import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, math
from pathlib import Path
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):
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.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
a2 = uop_to_json(VizData(), a)[id(a)]
self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\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.CUSTOM, src=tuple(src), arg=arg)
exec_rewrite(a, [PatternMatcher([])])
a2 = next(viz.get_details(0, 0))["graph"][id(a)]
self.assertEqual(ansistrip(a2["label"]), "CUSTOM\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.FUNCTION, 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, a.dtype)
y = UOp.const(math.pi, dtypes.float)
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.src[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(z), id(y), id(ret)])
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)
b = UOp.variable("b",0,10)
c = UOp.variable("c",0,10)
d = UOp.variable("d",0,10)
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.CUSTOM, 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")
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)
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 setUp(self): self.arch = "gfx1100"
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) for x in insts]))))
with save_viz() as viz:
with Context(DEV=f"NULL::{self.arch}"):
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:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
(p:=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()
View Git Blame Copy Permalink