geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

do not require clearing method_cache in viz tests (#16428)

Browse Source 
* update

* update test_dedup
This commit is contained in:
qazal
2026-05-29 18:12:34 +09:00
committed by GitHub
parent 54cfb794b8
commit 81bc00c006
1 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_viz.py
+7 -8
View File
@@ -12,8 +12,7 @@ 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 to_program_cache
from tinygrad.codegen import to_program
from tinygrad.codegen import do_to_program
@track_rewrites(name=True)
def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
@@ -41,7 +40,6 @@ class VizTrace:
@contextlib.contextmanager
def save_viz():
for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
to_program_cache.clear()
Buffer.profile_events.clear()
cpu_events.clear()
viz = VizTrace()
@@ -342,7 +340,7 @@ class TestVizIntegration(unittest.TestCase):
def test_codegen_tracing(self):
with save_viz() as viz:
ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0]
prg = to_program(ast, Device[Device.DEFAULT].renderer)
prg = do_to_program(ast, Device[Device.DEFAULT].renderer)
lst = viz.list_items()
self.assertEqual(len(lst), 3)
self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
@@ -362,7 +360,7 @@ class TestVizIntegration(unittest.TestCase):
from tinygrad.engine.realize import compile_linear
sched = compile_linear(sched)
with Context(NO_COLOR=0):
prgs = [to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
prgs = [do_to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
lst = viz.list_items()
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")
@@ -769,7 +767,7 @@ class TestCfg(unittest.TestCase):
with save_viz() as viz:
with Context(DEV=f"NULL::{self.arch}"):
out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
_ = to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
_ = 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"])
@@ -1006,8 +1004,9 @@ class TestCLI(unittest.TestCase):
def test_dedup(self):
with save_viz() as viz:
for _ in range(CNT:=4):
Tensor.empty(4, device="NULL").add(1).realize()
Tensor.empty(8, device="NULL").add(1).realize()
# 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):