sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

viz/cli: select all calls in graph by default (#16214)

Browse Source
This commit is contained in:
qazal
2026-05-15 21:01:44 +09:00
committed by GitHub
parent 28b98e529d
commit d54fa86b71
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/viz/kernel_graph.py
+4 -4
View File
@@ -7,15 +7,15 @@ def get_node(graph:dict, key): return graph[str(key)]
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="print CALL graph from DEBUG=5 tinygrad.viz.cli --json output")
parser.add_argument("kernel", type=str, default=None, help="Kernel name to stop at (default: print all kernels)")
parser.add_argument("kernel", type=str, nargs="?", default="ALL", metavar="NAME", help="Kernel name to stop at (default: print all kernels)")
args = parser.parse_args()
ref:int|None = None
for line in sys.stdin:
if not line.strip(): continue
graph = json.loads(line)
if ref is not None and graph.get("ref") == ref:
if graph.get("ref") is not None and (args.kernel == "ALL" or graph["ref"] == ref):
print(graph)
if (v:=json.loads(next(sys.stdin)).get("value")): print(v)
if (v:=json.loads(next(sys.stdin, "{}")).get("value")): print(v)
if ref is not None or not isinstance(rec:=next(iter(graph.values()), {}), dict) or "label" not in rec: continue
for v in graph.values():
if not v["label"].startswith("CALL"): continue
@@ -39,6 +39,6 @@ if __name__ == "__main__":
src_str = ["SRC"]+get_node(graph, get_node(graph, s)["src"][0][1])["label"].splitlines()[1:]
print(" ".join(idx_str+src_str))
ss += [x[1] for x in get_node(graph, s)["src"]]
if args.kernel is not None and args.kernel in ansistrip(v["label"]):
if args.kernel != "ALL" and args.kernel in ansistrip(v["label"]):
ref = v["ref"]
break