Improves graph rendering (#110)

* Adds input/output to the text rendering

* mypy

* improve graph rendering

Author: xadupre

_unittests/ut_helpers/test_graph_helper.py

@@ -1,7 +1,9 @@
 import textwrap
 import unittest
+import numpy as np
 import onnx
 import onnx.helper as oh
+import onnx.numpy_helper as onh
 from onnx_diagnostic.ext_test_case import ExtTestCase
 from onnx_diagnostic.helpers.graph_helper import GraphRendering
 
@@ -126,36 +128,35 @@ def test_text_rendering(self):
         )
         graph = GraphRendering(proto)
         text = textwrap.dedent(graph.text_rendering(prefix="|")).strip("\n")
-        print()
-        print(text)
-        expected = textwrap.dedent(
-            """
-            |
-            |
-            |
-            |
-            |   X            Y
-            |   |            |
-            |   +------+-----+-----------+
-            |          |                 |
-            |         Add               Neg
-            |          |                 |
-            |          +-----+-----------+
-            |                |
-            |               Mul
-            |                |
-            |                +-----------+
-            |                            |
-            |                           Mul
-            |                            |
-            |                            +----------------------+
-            |                                                   |
-            |                                                   Z
-            |
-            |
-            """
-        ).strip("\n")
-        self.assertEqual(expected, text)
+        self.assertIn("└-------┬---┼---┘", text)
+
+    def test_text_rendering_more(self):
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Add", ["X", "Y"], ["xy"]),
+                    oh.make_node("Neg", ["Y"], ["ny"]),
+                    oh.make_node("Mul", ["xy", "ny"], ["a"]),
+                    oh.make_node("Div", ["xy", "two"], ["b"]),
+                    oh.make_node("Add", ["b", "Y"], ["by"]),
+                    oh.make_node("Mod", ["a", "ny"], ["ay"]),
+                    oh.make_node("Sub", ["ay", "by"], ["Z"]),
+                ],
+                "-nd-",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+                [onh.from_array(np.array([2], dtype=np.float32), name="two")],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+        onnx.checker.check_model(proto)
+        graph = GraphRendering(proto)
+        text = textwrap.dedent(graph.text_rendering(prefix="|")).strip("\n")
+        self.assertIn(" └-------┬---┼---┴-------┐", text)
 
 
 if __name__ == "__main__":

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -14,11 +14,13 @@ def dummy_path(self):
         )
 
     def test_parser_print(self):
-        st = StringIO()
-        with redirect_stdout(st):
-            main(["print", "raw", self.dummy_path])
-        text = st.getvalue()
-        self.assertIn("Add", text)
+        for fmt in ["raw", "text", "pretty", "printer"]:
+            with self.subTest(format=fmt):
+                st = StringIO()
+                with redirect_stdout(st):
+                    main(["print", fmt, self.dummy_path])
+                text = st.getvalue()
+                self.assertIn("Add", text)
 
     def test_parser_stats(self):
         output = self.get_dump_file("test_parser_stats.xlsx")

onnx_diagnostic/_command_lines_parser.py

@@ -126,9 +126,23 @@ def get_parser_print() -> ArgumentParser:
         """
         ),
         epilog="To show a model.",
+        formatter_class=RawTextHelpFormatter,
     )
     parser.add_argument(
-        "fmt", choices=["pretty", "raw"], help="Format to use.", default="pretty"
+        "fmt",
+        choices=["pretty", "raw", "text", "printer"],
+        default="pretty",
+        help=textwrap.dedent(
+            """
+            Prints out a model on the standard output.
+            raw     - just prints the model with print(...)
+            printer - onnx.printer.to_text(...)
+            pretty  - an improved rendering
+            text    - uses GraphRendering
+            """.strip(
+                "\n"
+            )
+        ),
     )
     parser.add_argument("input", type=str, help="onnx model to load")
     return parser
@@ -144,6 +158,12 @@ def _cmd_print(argv: List[Any]):
         from .helpers.onnx_helper import pretty_onnx
 
         print(pretty_onnx(onx))
+    elif args.fmt == "printer":
+        print(onnx.printer.to_text(onx))
+    elif args.fmt == "text":
+        from .helpers.graph_helper import GraphRendering
+
+        print(GraphRendering(onx).text_rendering())
     else:
         raise ValueError(f"Unexpected value fmt={args.fmt!r}")
 

onnx_diagnostic/helpers/graph_helper.py

@@ -1,3 +1,4 @@
+import pprint
 from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
 import onnx
 import onnx.helper as oh
@@ -45,8 +46,12 @@ def computation_order(
                 f"Missing input in node {i_node} type={node.op_type}: "
                 f"{[i for i in node.input if i not in number]}"
             )
-            mx = max(number[i] for i in node.input) + 1
-            results[i_node] = mx
+            if node.input:
+                mx = max(number[i] for i in node.input) + 1
+                results[i_node] = mx
+            else:
+                # A constant
+                mx = max(number.values()) if number else 0
             for i in node.output:
                 number[i] = mx
         return results
@@ -80,7 +85,7 @@ def graph_positions(
             indices = [i for i, o in enumerate(order) if o == row]
             assert indices, f"indices cannot be empty for row={row}, order={order}"
             ns = [nodes[i] for i in indices]
-            mx = [max(names.get(i, 0) for i in n.input) for n in ns]
+            mx = [(max(names.get(i, 0) for i in n.input) if n.input else 0) for n in ns]
             mix = [(m, i) for i, m in enumerate(mx)]
             mix.sort()
             for c, (_m, i) in enumerate(mix):
@@ -108,14 +113,18 @@ def text_positions(
         for i, (_row, col) in enumerate(new_positions):
             size = len(nodes[i].op_type) + 5
             column_size[col] = max(column_size[col], size)
+            assert column_size[col] < 200, (
+                f"column_size[{col}]={column_size[col]}, this is quite big, i={i}, "
+                f"nodes[i].op_type={nodes[i].op_type}"
+            )
 
         # cumulated
         sort = sorted(column_size.items())
         cumul = dict(sort[:1])
         results = {sort[0][0]: sort[0][1] // 2}
         for col, size in sort[1:]:
-            c = sum(cumul.values())
-            cumul[col] = c
+            c = max(cumul.values())
+            cumul[col] = c + size
             results[col] = c + size // 2
         return [(row, results[col]) for row, col in new_positions]
 
@@ -190,6 +199,49 @@ def build_node_edges(cls, nodes: Sequence[onnx.NodeProto]) -> Set[Tuple[int, int
                     edges.add(edge)
         return edges
 
+    ADD_RULES = {
+        ("┴", "┘"): "┴",
+        ("┴", "└"): "┴",
+        ("┬", "┐"): "┬",
+        ("┬", "┌"): "┬",
+        ("-", "└"): "┴",
+        ("-", "|"): "┼",
+        ("-", "┐"): "┬",
+        ("┐", "-"): "┬",
+        ("┘", "-"): "┴",
+        ("┴", "-"): "┴",
+        ("-", "┘"): "┴",
+        ("┌", "-"): "┬",
+        ("┬", "-"): "┬",
+        ("-", "┌"): "┬",
+        ("|", "-"): "┼",
+        ("└", "-"): "┴",
+        ("|", "└"): "├",
+        ("|", "┘"): "┤",
+        ("┐", "|"): "┤",
+        ("┬", "|"): "┼",
+        ("|", "┐"): "┤",
+        ("|", "┌"): "├",
+        ("├", "-"): "┼",
+        ("└", "|"): "├",
+        ("┤", "┐"): "┤",
+        ("┤", "|"): "┤",
+        ("├", "|"): "├",
+        ("┴", "┌"): "┼",
+        ("┐", "┌"): "┬",
+        ("┌", "┐"): "┬",
+        ("┌", "|"): "┼",
+        ("┴", "┐"): "┼",
+        ("┐", "└"): "┼",
+        ("┬", "┘"): "┼",
+        ("├", "└"): "├",
+        ("┤", "┌"): "┼",
+        ("┘", "|"): "┤",
+        ("┴", "|"): "┼",
+        ("┤", "-"): "┼",
+        ("┘", "└"): "┴",
+    }
+
     @classmethod
     def text_grid(cls, grid: List[List[str]], position: Tuple[int, int], text: str):
         """
@@ -230,24 +282,44 @@ def text_edge(
         assert (
             0 <= p2[1] < min(len(g) for g in grid)
         ), f"p2={p2}, sizes={[len(g) for g in grid]}"
-        grid[p1[0] + 1][p1[1]] = "|"
-        grid[p1[0] + 2][p1[1]] = "+"
+
+        def add(s1, s2):
+            assert s2 != " ", f"s1={s1!r}, s2={s2!r}"
+            if s1 == " " or s1 == s2:
+                return s2
+            if s1 == "┼" or s2 == "┼":
+                return "┼"
+            if (s1, s2) in cls.ADD_RULES:
+                return cls.ADD_RULES[s1, s2]
+            raise NotImplementedError(f"Unable to add: ({s1!r},{s2!r}): '',")
+
+        def place(grid, x, y, symbol):
+            grid[x][y] = add(grid[x][y], symbol)
+
+        place(grid, p1[0] + 1, p1[1], "|")
+        place(grid, p1[0] + 2, p1[1], "└" if p1[1] < p2[1] else "┘")
 
         if p1[0] + 2 == p2[0] - 2:
             a, b = (p1[1] + 1, p2[1] - 1) if p1[1] < p2[1] else (p2[1] + 1, p1[1] - 1)
-            grid[p1[0] + 2][a : b + 1] = ["-"] * (b - a + 1)
+            for i in range(a, b + 1):
+                place(grid, p1[0] + 2, i, "-")
         else:
             middle = (p1[1] + p2[1]) // 2
             a, b = (p1[1] + 1, middle - 1) if p1[1] < middle else (middle + 1, p1[1] - 1)
-            grid[p1[0] + 2][a : b + 1] = ["-"] * (b - a + 1)
+            for i in range(a, b + 1):
+                place(grid, p1[0] + 2, i, "-")
             a, b = (p1[1] + 1, middle - 1) if p1[1] < middle else (middle + 1, p1[1] - 1)
-            grid[p1[0] + 2][a : b + 1] = ["-"] * (b - a + 1)
+            for i in range(a, b + 1):
+                place(grid, p1[0] + 2, i, "-")
+
+            place(grid, p1[0] + 2, middle, "┐" if p1[1] < p2[1] else "┌")
+            place(grid, p2[0] - 2, middle, "└" if p1[1] < p2[1] else "┘")
 
-            grid[p1[0] + 2][middle] = "+"
-            grid[p2[0] - 2][middle] = "+"
+            for i in range(p1[0] + 2 + 1, p2[0] - 2):
+                place(grid, i, middle, "|")
 
-        grid[p2[0] - 2][p2[1]] = "+"
-        grid[p2[0] - 1][p2[1]] = "|"
+        place(grid, p2[0] - 2, p2[1], "┐" if p1[1] < p2[1] else "┌")
+        place(grid, p2[0] - 1, p2[1], "|")
 
     def text_rendering(self, prefix="") -> str:
         """
@@ -298,6 +370,7 @@ def text_rendering(self, prefix="") -> str:
         text_pos = self.text_positions(nodes, positions)
         edges = self.build_node_edges(nodes)
         max_len = max(col for _, col in text_pos) + max(len(n.op_type) for n in nodes)
+        assert max_len < 1e6, f"max_len={max_len}, text_pos=\n{pprint.pformat(text_pos)}"
         max_row = max(row for row, _ in text_pos) + 2
         grid = [[" " for i in range(max_len + 1)] for _ in range(max_row + 1)]