From 3a04e518ecf3ac239587e84b91f1963ccde1576b Mon Sep 17 00:00:00 2001 From: kormann <49917710+DKormann@users.noreply.github.com> Date: Thu, 27 Jun 2024 00:02:19 +0200 Subject: [PATCH] print_tree UPat +fix (#5132) * fix and extend print_tree * typing * typing * fix upat * fix none * ws * rm prefix * mv luop dag * typo * test print_tree --- test/test_print_tree.py | 66 ++++++++++++++++++++++++++++++++++++++++ tinygrad/engine/graph.py | 17 ++++++----- 2 files changed, 75 insertions(+), 8 deletions(-) create mode 100644 test/test_print_tree.py diff --git a/test/test_print_tree.py b/test/test_print_tree.py new file mode 100644 index 0000000000..db58fa4366 --- /dev/null +++ b/test/test_print_tree.py @@ -0,0 +1,66 @@ +#%% +import unittest +from tinygrad.engine.graph import print_tree + +from tinygrad import Tensor, dtypes +from tinygrad.codegen.uops import UOps, UOp, UPat +from tinygrad.ops import BinaryOps + +import sys, io + +class TestPrintTree(unittest.TestCase): + + def _capture_print(self, fn): + capturedOutput = io.StringIO() + sys.stdout = capturedOutput + fn() + sys.stdout = sys.__stdout__ + return capturedOutput.getvalue() + + def test_print_uop(self): + x = Tensor.arange(10).schedule()[-1].ast[0] + output = self._capture_print(lambda: print_tree(x)) + assert output == '\ + 0 ━┳ BufferOps.STORE MemBuffer(idx=0, dtype=dtypes.int, \ +st=ShapeTracker(views=(View(shape=(10, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))\n\ + 1 ┗━┳ BinaryOps.ADD None\n\ + 2 ┣━┳ ReduceOps.SUM (1,)\n\ + 3 ┃ ┗━━ BufferOps.CONST ConstBuffer(val=1, dtype=dtypes.int, st=ShapeTrac\ +ker(views=(View(shape=(11, 19), strides=(0, 0), offset=0, mask=((0, 11), (9, 19))\ +, contiguous=False), View(shape=(10, 10), strides=(1, 20), offset=0, mask=None, contiguous=False))))\n\ + 4 ┗━━ BufferOps.CONST ConstBuffer(val=-1, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(10,\ + 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)))\n' + + x = UOp.var("x", dtypes.int) + x = (x + x) - UOp.const(dtypes.int, 2) + output = self._capture_print(lambda: print_tree(x)) + assert output == '\ + 0 ━┳ UOps.ALU BinaryOps.ADD\n\ + 1 ┣━┳ UOps.ALU BinaryOps.ADD\n\ + 2 ┃ ┣━━ UOps.VAR x\n\ + 3 ┃ ┗━━ UOps.VAR x\n\ + 4 ┗━┳ UOps.ALU UnaryOps.NEG\n\ + 5 ┗━━ UOps.CONST 2\n' + + x = UPat(UOp.alu(BinaryOps.ADD, UOp.var("x", dtypes.int), UOp.var("x", dtypes.int))) + assert self._capture_print(lambda: print_tree(x)) == '\ + 0 ━━ UOps.ALU : dtypes.int [, ] BinaryOps.ADD None\n' + + x = UPat.compile(UOp.store(UOp.var("buf"), UOp.var("idx"), + UOp(UOps.CAST, src=tuple(UOp(UOps.GEP, arg=i, src=(UOp.var("val"),)) for i in range(4)))), UOp.store) + assert self._capture_print(lambda: print_tree(x)) == '\ + 0 ━┳ UOps.STORE None\n\ + 1 ┣━━ None None\n\ + 2 ┣━━ None None\n\ + 3 ┗━┳ UOps.CAST None\n\ + 4 ┣━┳ UOps.GEP 0\n\ + 5 ┃ ┗━━ None None\n\ + 6 ┣━┳ UOps.GEP 1\n\ + 7 ┃ ┗━━ None None\n\ + 8 ┣━┳ UOps.GEP 2\n\ + 9 ┃ ┗━━ None None\n\ + 10 ┗━┳ UOps.GEP 3\n\ + 11 ┗━━ None None\n' + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tinygrad/engine/graph.py b/tinygrad/engine/graph.py index d382b68642..f4924d366b 100644 --- a/tinygrad/engine/graph.py +++ b/tinygrad/engine/graph.py @@ -4,7 +4,7 @@ from typing import List, Any, DefaultDict, Union from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, LoadOps, BufferOps, TernaryOps, LazyOp from tinygrad.device import Device from tinygrad.helpers import GRAPHPATH, DEBUG, GlobalCounters, getenv -from tinygrad.codegen.uops import UOps, UOp +from tinygrad.codegen.uops import UOps, UOp, UPat from tinygrad.shape.symbolic import NumNode from tinygrad.lazy import LazyBuffer @@ -75,18 +75,19 @@ def log_lazybuffer(lb:'LazyBuffer', scheduled=False): # realized but unseen? G.add_node(nm(lb), label=f'"{str(lb.base.realized)[5:-1].replace(" ", chr(10))}\nb:{nm(lb.realized)}"', style='filled', fillcolor="#f0c08080") -def _tree(luop:Union[LazyOp,UOp], cycles, cnt, prefix=""): +def _tree(dag:Union[LazyOp, UOp, UPat], cycles, cnt): cnt[0] += 1 - if len(luop.src) == 0: return [f"━━ {prefix}{luop.op.name} {luop.arg if luop.arg else ''}"] - if (lid := id(luop)) in cycles and cycles[lid][1] > (tcnt := getenv("TREE_CYCLE_CNT", 5)) and tcnt >= 0: - return [f"━⬆︎ goto {cycles[id(luop)][0]}: {luop.op.name}"] + src = dag.src if isinstance(dag.src, (list, tuple)) else [] if dag.src is None else [dag.src] + if len(src) == 0: return [f"━━ {dag.op} {dag.arg}"] + if (lid := id(dag)) in cycles and cycles[lid][1] > (tcnt := getenv("TREE_CYCLE_CNT", 5)) and tcnt >= 0: + return [f"━⬆︎ goto {cycles[id(dag)][0]}: {dag.op}"] cycles[lid] = (cnt[0], 1 if lid not in cycles else cycles[lid][1]+1) - lines = [f"━┳ {prefix}{luop.op.name} {luop.arg if luop.arg else ''}"] - childs = [_tree(c, cycles, cnt) for c in luop.src[:]] + lines = [f"━┳ {dag.op} {dag.arg}"] + childs = [_tree(c, cycles, cnt) for c in src] for c in childs[:-1]: lines += [f" ┣{c[0]}"] + [f" ┃{l}" for l in c[1:]] return lines + [" ┗"+childs[-1][0]] + [" "+l for l in childs[-1][1:]] -def print_tree(luop:Union[LazyOp,UOp]): print("\n".join([f"{str(i).rjust(3)} {s}" for i,s in enumerate(_tree(luop, {}, [-1]))])) +def print_tree(dag:Union[LazyOp, UOp, UPat]): print("\n".join([f"{str(i).rjust(3)} {s}" for i,s in enumerate(_tree(dag, {}, [-1]))])) def graph_uops(uops:List[UOp]): colors = {UOps.ALU: "#ffffc0", UOps.LOAD: "#ffc0c0", UOps.STORE: "#c0ffc0", UOps.SPECIAL: "#c0c0ff", UOps.CONST: "#e0e0e0",