descript: pretty-printer

Author: qaco

src/xtc/schedules/loop_nest.py

@@ -126,6 +126,83 @@ def tiles_to_sizes(self) -> dict[str, int]:
                 tiles_to_sizes[loop] = size
         return tiles_to_sizes
 
+    def pretty_print(self, indent: int = 0) -> str:
+        """Return a human-readable representation of the loop nest.
+
+        The output format resembles pseudocode:
+            for i in ...:
+              for j in ...:
+                for j0 in tile(j, 16):  // vectorized
+
+        Args:
+            indent: The initial indentation level (number of spaces).
+
+        Returns:
+            A multi-line string representing the loop nest structure.
+        """
+        lines: list[str] = []
+
+        # Build mapping from tile loop name to (axis, size)
+        tiles_info: dict[str, tuple[str, int]] = {}
+        for axis, tile_loops in self.tiles.items():
+            for loop_name, size in tile_loops.items():
+                tiles_info[loop_name] = (axis, size)
+
+        # Map split loop names to their child nodes (which contain split_origin)
+        split_to_child: dict[str, LoopNestNode] = {}
+        for child in self.children:
+            if child.split_origin is not None:
+                axis = child.split_origin.axis
+                if axis in self.splits:
+                    for loop_name, start in self.splits[axis].items():
+                        if start == child.split_origin.start:
+                            split_to_child[loop_name] = child
+                            break
+
+        current_indent = indent
+        for loop_name in self.interchange:
+            # Build the loop line based on loop type
+            if loop_name in tiles_info:
+                axis, size = tiles_info[loop_name]
+                line = f"for {loop_name} in tile({axis}, {size}):"
+            elif loop_name in split_to_child:
+                child = split_to_child[loop_name]
+                origin = child.split_origin
+                assert origin is not None
+                start = origin.start if origin.start is not None else "..."
+                end = origin.end if origin.end is not None else "..."
+                line = f"for {loop_name} in split({origin.axis}, {start}, {end}):"
+            else:
+                line = f"for {loop_name} in ...:"
+
+            # Add annotations as comments
+            annotations: list[str] = []
+            if loop_name in self.parallelize:
+                annotations.append("parallelized")
+            if loop_name in self.vectorize:
+                annotations.append("vectorized")
+            if loop_name in self.unroll:
+                factor = self.unroll[loop_name]
+                annotations.append(f"unroll({factor})")
+
+            if annotations:
+                line += "  // " + ", ".join(annotations)
+
+            lines.append(" " * current_indent + line)
+
+            # If this is a split loop with a child, recurse into the child
+            # Split loops don't increase indent for subsequent siblings
+            if loop_name in split_to_child:
+                child_output = split_to_child[loop_name].pretty_print(
+                    current_indent + 2
+                )
+                lines.append(child_output)
+            else:
+                # Regular loops increase nesting depth
+                current_indent += 2
+
+        return "\n".join(lines)
+
 
 @dataclass
 class LoopsDimsMapper:

tests/filecheck/schedules/test_descript_pretty_print.py

@@ -0,0 +1,83 @@
+# RUN: python %s --simple 2>&1 | filecheck %s --check-prefix=CHECK-SIMPLE
+# RUN: python %s --tiled 2>&1 | filecheck %s --check-prefix=CHECK-TILED
+# RUN: python %s --vectorized 2>&1 | filecheck %s --check-prefix=CHECK-VECTORIZED
+# RUN: python %s --full 2>&1 | filecheck %s --check-prefix=CHECK-FULL
+# RUN: python %s --split 2>&1 | filecheck %s --check-prefix=CHECK-SPLIT
+
+import sys
+from xtc.schedules.parsing import ScheduleParser
+from xtc.schedules.descript import ScheduleInterpreter
+
+parser = ScheduleParser()
+abstract_axis = ["i", "j", "k"]
+interpreter = ScheduleInterpreter(abstract_axis)
+
+if "--simple" in sys.argv:
+    spec = {"i": {}, "k": {}, "j": {}}
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
+elif "--tiled" in sys.argv:
+    spec = {"i": {}, "k": {}, "j": {}, "j#16": {}}
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
+elif "--vectorized" in sys.argv:
+    spec = {"i": {}, "k": {}, "j": {}, "j#16": {"vectorize": True}}
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
+elif "--full" in sys.argv:
+    spec = {
+        "i": {"parallelize": True},
+        "k": {},
+        "j": {},
+        "j#32": {},
+        "j#16": {"vectorize": True, "unroll": 4},
+    }
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
+elif "--split" in sys.argv:
+    spec = {
+        "i": {},
+        "j[:128]": {"k": {}, "k#32": {}},
+        "j[128:]": {"k": {}, "k#16": {"vectorize": True}},
+    }
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
+# CHECK-SIMPLE:      for i in ...:
+# CHECK-SIMPLE-NEXT:   for k in ...:
+# CHECK-SIMPLE-NEXT:     for j in ...:
+
+# CHECK-TILED:      for i in ...:
+# CHECK-TILED-NEXT:   for k in ...:
+# CHECK-TILED-NEXT:     for j in ...:
+# CHECK-TILED-NEXT:       for j0 in tile(j, 16):
+
+# CHECK-VECTORIZED:      for i in ...:
+# CHECK-VECTORIZED-NEXT:   for k in ...:
+# CHECK-VECTORIZED-NEXT:     for j in ...:
+# CHECK-VECTORIZED-NEXT:       for j0 in tile(j, 16):  // vectorized
+
+# CHECK-FULL:      for i in ...:  // parallelized
+# CHECK-FULL-NEXT:   for k in ...:
+# CHECK-FULL-NEXT:     for j in ...:
+# CHECK-FULL-NEXT:       for j0 in tile(j, 32):
+# CHECK-FULL-NEXT:         for j1 in tile(j, 16):  // vectorized, unroll(4)
+
+# CHECK-SPLIT:      for i in ...:
+# CHECK-SPLIT-NEXT:   for j[0] in split(j, 0, 128):
+# CHECK-SPLIT-NEXT:     for j in ...:
+# CHECK-SPLIT-NEXT:       for k in ...:
+# CHECK-SPLIT-NEXT:         for k0 in tile(k, 32):
+# CHECK-SPLIT-NEXT:   for j[1] in split(j, 128, ...):
+# CHECK-SPLIT-NEXT:     for j in ...:
+# CHECK-SPLIT-NEXT:       for k in ...:
+# CHECK-SPLIT-NEXT:         for k0 in tile(k, 16):  // vectorized