descript: use a tree structure to represent a loop nest

Author: qaco

src/xtc/schedules/descript.py

@@ -100,9 +100,6 @@ class ScheduleParser:
 
     _SPLIT_PATTERN = re.compile(r"^(.*)\[(-\d+|\d*)?:(-\d+|\d*)?\]$")
 
-    def __init__(self, abstract_axis: list[str]):
-        self.abstract_axis = abstract_axis
-
     def parse(self, spec: dict[str, Any]) -> ScheduleSpec:
         """Parse a schedule specification dict into an AST."""
         items: list[ScheduleItem] = []
@@ -224,31 +221,33 @@ def __init__(self, abstract_axis: list[str]):
 
     def interpret(self, spec: ScheduleSpec, root: str) -> LoopNest:
         """Interpret a schedule specification into a LoopNest."""
-        return self._interpret_spec(spec, root, head=[])
-
-    def _interpret_spec(
-        self, spec: ScheduleSpec, root: str, head: list[str]
-    ) -> LoopNest:
-        """Interpret a schedule spec recursively."""
         loop_nest = LoopNest(abstract_dims=self.abstract_axis)
-        slice = loop_nest.build_slice(root)
+        root_node = loop_nest.build_root_node(root)
+        self._interpret_spec_into_node(spec, root_node, root, head=[])
+        return loop_nest
 
+    def _interpret_spec_into_node(
+        self,
+        spec: ScheduleSpec,
+        node: LoopNestNode,
+        root: str,
+        head: list[str],
+    ) -> None:
+        """Interpret a schedule spec into an existing node (mutates node)."""
         # Track state during interpretation
         sizes: dict[str, int] = {}
         previous_cut: dict[str, int | None] = {a: 0 for a in self.abstract_axis}
         interchange: list[str] = list(head)
 
         for item in spec.items:
             if isinstance(item, SplitDecl):
-                self._interpret_split(
-                    item, slice, loop_nest, root, interchange, previous_cut
-                )
+                self._interpret_split(item, node, root, interchange, previous_cut)
             elif isinstance(item, TileDecl):
-                loop_name = self._interpret_tile(item, slice, interchange, sizes)
-                self._apply_annotations(item.annotations, loop_name, sizes, slice)
+                loop_name = self._interpret_tile(item, node, interchange, sizes)
+                self._apply_annotations(item.annotations, loop_name, sizes, node)
             elif isinstance(item, AxisDecl):
                 loop_name = self._interpret_axis(item, interchange)
-                self._apply_annotations(item.annotations, loop_name, sizes, slice)
+                self._apply_annotations(item.annotations, loop_name, sizes, node)
 
         # Check that all splits are complete
         for axis, cut in previous_cut.items():
@@ -257,14 +256,12 @@ def _interpret_spec(
                     f"Splitting of {axis} unachieved (stops at {cut})."
                 )
 
-        slice.interchange = interchange
-        return loop_nest
+        node.interchange = interchange
 
     def _interpret_split(
         self,
         item: SplitDecl,
-        slice: LoopNestSlice,
-        loop_nest: LoopNest,
+        node: LoopNestNode,
         root: str,
         interchange: list[str],
         previous_cut: dict[str, int | None],
@@ -279,7 +276,7 @@ def _interpret_split(
         # last one, so it cannot be the previous one.
         cut = previous_cut[axis_name]
 
-        # When x (the starting point of the slice) is not specified,
+        # When x (the starting point of the split) is not specified,
         # it is the previous cut
         if x is None:
             x = cut
@@ -291,34 +288,44 @@ def _interpret_split(
         previous_cut[axis_name] = y
 
         # Save the cutting points of the new dimensions
-        if axis_name not in slice.splits:
-            slice.splits[axis_name] = {}
-        new_dim_index = len(slice.splits[axis_name])
+        if axis_name not in node.splits:
+            node.splits[axis_name] = {}
+        new_dim_index = len(node.splits[axis_name])
         new_dim_name = f"{axis_name}[{new_dim_index}]"
         new_root_name = f"{root}/{new_dim_name}"
-        slice.splits[axis_name][new_dim_name] = x
+        node.splits[axis_name][new_dim_name] = x
         interchange.append(new_dim_name)
 
-        # Recursively interpret the nested schedule
-        inner_nest = self._interpret_spec(item.body, new_root_name, head=[axis_name])
-        loop_nest.slices += inner_nest.slices
+        # Create a child node for the nested schedule
+        child_node = LoopNestNode(
+            root=new_root_name, tiles={a: {} for a in self.abstract_axis}
+        )
+
+        # Create and attach the split child
+        split_child = SplitChild(axis=axis_name, start=x, end=y, node=child_node)
+        node.children.append(split_child)
+
+        # Recursively interpret the nested schedule into the child node
+        self._interpret_spec_into_node(
+            item.body, child_node, new_root_name, head=[axis_name]
+        )
 
     def _interpret_tile(
         self,
         item: TileDecl,
-        slice: LoopNestSlice,
+        node: LoopNestNode,
         interchange: list[str],
         sizes: dict[str, int],
     ) -> str:
         """Interpret a tile declaration. Returns the loop name."""
         self._check_axis_existence(item.axis)
-        tile_num = len(slice.tiles[item.axis])
+        tile_num = len(node.tiles[item.axis])
         loop_name = f"{item.axis}{tile_num}"
         if item.size <= 0:
             raise ScheduleInterpretError(
                 f"`{item}`: tile sizes should be strictly positive."
             )
-        slice.tiles[item.axis][loop_name] = item.size
+        node.tiles[item.axis][loop_name] = item.size
         sizes[loop_name] = item.size
         interchange.append(loop_name)
 
@@ -355,9 +362,9 @@ def _apply_annotations(
         annotations: Annotations,
         loop_name: str,
         sizes: dict[str, int],
-        slice: LoopNestSlice,
+        node: LoopNestNode,
     ) -> None:
-        """Apply annotations to a loop in the slice."""
+        """Apply annotations to a loop in the node."""
         if annotations.unroll_specified:
             unroll_factor = annotations.unroll_factor
             if unroll_factor is None:
@@ -371,13 +378,13 @@ def _apply_annotations(
                 raise ScheduleInterpretError(
                     f'`{{"unroll" = {unroll_factor}}}`: unroll parameter should be strictly positive.'
                 )
-            slice.unroll[loop_name] = unroll_factor
+            node.unroll[loop_name] = unroll_factor
 
         if annotations.vectorize:
-            slice.vectorize.append(loop_name)
+            node.vectorize.append(loop_name)
 
         if annotations.parallelize:
-            slice.parallelize.append(loop_name)
+            node.parallelize.append(loop_name)
 
     def _check_splitting_intervals(
         self,
@@ -431,17 +438,17 @@ def loops_to_axis(self) -> dict[str, str]:
         return loops_to_axis
 
     @staticmethod
-    def build_from_slices(slices: list["LoopNestSlice"]) -> "LoopsDimsMapper":
+    def build_from_nodes(nodes: list["LoopNestNode"]) -> "LoopsDimsMapper":
         tiles_to_axis = {}
         splits_to_axis = {}
         dims = set()
-        for slice in slices:
-            tiles_to_axis.update(LoopsDimsMapper._get_subloops_to_axis(slice.tiles))
-            splits_to_axis.update(LoopsDimsMapper._get_subloops_to_axis(slice.splits))
+        for node in nodes:
+            tiles_to_axis.update(LoopsDimsMapper._get_subloops_to_axis(node.tiles))
+            splits_to_axis.update(LoopsDimsMapper._get_subloops_to_axis(node.splits))
         refined_loops = list(tiles_to_axis) + list(splits_to_axis)
-        for slice in slices:
+        for node in nodes:
             dims.update(
-                [loop for loop in slice.interchange if loop not in refined_loops]
+                [loop for loop in node.interchange if loop not in refined_loops]
             )
             dims.update(tiles_to_axis.values())
             dims.update(splits_to_axis.values())
@@ -457,14 +464,32 @@ def _get_subloops_to_axis(subloops: dict[str, dict[str, Any]]) -> dict[str, str]
 
 
 @dataclass
-class LoopNestSlice:
-    """Represents a single slice of a loop nest with its transformations.
+class SplitChild:
+    """Represents a child node in the loop nest tree created by a split.
 
-    A slice describes the loops attached to a single root and
+    Attributes:
+        axis: The axis that was split to create this child.
+        start: The starting position of the split (inclusive), or None if unbounded.
+        end: The ending position of the split (exclusive), or None if unbounded.
+        node: The child LoopNestNode containing the nested schedule.
+    """
+
+    axis: str
+    start: int | None
+    end: int | None
+    node: "LoopNestNode"
+
+
+@dataclass
+class LoopNestNode:
+    """Represents a node in the loop nest tree with its transformations.
+
+    A node describes the loops attached to a single root and
     contains all the scheduling transformations applied to these loops.
+    Splits create child nodes, forming an explicit tree structure.
 
     Attributes:
-        root: Identifier of the the slice (either the base operation or
+        root: Identifier of the node (either the base operation or
             the content of a split).
         tiles: Tiling configuration per axis. Maps axis names to dicts of
             tile loop names and their sizes.
@@ -474,6 +499,7 @@ class LoopNestSlice:
         vectorize: List of loops to vectorize.
         parallelize: List of loops to parallelize.
         unroll: Maps loop names to their unroll factors.
+        children: List of SplitChild objects representing child nodes from splits.
     """
 
     root: str
@@ -483,6 +509,7 @@ class LoopNestSlice:
     vectorize: list[str] = field(default_factory=list)
     parallelize: list[str] = field(default_factory=list)
     unroll: dict[str, int] = field(default_factory=dict)
+    children: list[SplitChild] = field(default_factory=list)
 
     @property
     def splits_to_sizes(self) -> dict[str, int]:
@@ -509,25 +536,44 @@ def tiles_to_sizes(self) -> dict[str, int]:
 class LoopNest:
     """Represents a complete loop nest structure for scheduling.
 
-    A loop nest contains abstract dimensions and a collection of slices/
-    It provides validation to ensure consistency across all slices.
+    A loop nest contains abstract dimensions and a tree of nodes representing
+    the schedule. Splits create child nodes, forming an explicit tree structure.
 
     Attributes:
         abstract_dims: List of abstract dimension names for the loop nest.
-        slices: List of LoopNestSlice objects, one per scheduled operation.
+        root_node: The root node of the loop nest tree, or None if empty.
     """
 
     abstract_dims: list[str]
-    slices: list[LoopNestSlice] = field(default_factory=list)
+    root_node: LoopNestNode | None = None
 
     @property
-    def empty(self):
-        return not self.slices
+    def empty(self) -> bool:
+        return self.root_node is None
 
-    def build_slice(self, root: str) -> LoopNestSlice:
-        slice = LoopNestSlice(root=root, tiles={a: {} for a in self.abstract_dims})
-        self.slices = [slice] + self.slices
-        return slice
+    @property
+    def nodes(self) -> list[LoopNestNode]:
+        """Flatten the tree into a list of nodes.
+
+        Returns nodes in depth-first order, with the root node first,
+        followed by children in the order they were created.
+        """
+        if self.root_node is None:
+            return []
+        return self._collect_nodes_dfs(self.root_node)
+
+    def _collect_nodes_dfs(self, node: LoopNestNode) -> list[LoopNestNode]:
+        """Collect all nodes in the tree in depth-first order."""
+        result = [node]
+        for child in node.children:
+            result.extend(self._collect_nodes_dfs(child.node))
+        return result
+
+    def build_root_node(self, root: str) -> LoopNestNode:
+        """Build and set the root node of the loop nest tree."""
+        node = LoopNestNode(root=root, tiles={a: {} for a in self.abstract_dims})
+        self.root_node = node
+        return node
 
     def check(self):
         self._check_use_defined_dims()
@@ -536,13 +582,13 @@ def check(self):
         self._check_sizes()
 
     def _check_use_defined_dims(self):
-        mapper = LoopsDimsMapper.build_from_slices(self.slices)
+        mapper = LoopsDimsMapper.build_from_nodes(self.nodes)
         for dim in self.abstract_dims:
             if dim not in mapper.dims:
                 raise ScheduleValidationError(f"{dim} defined but never used")
 
     def _check_vectorization_consistency(self):
-        for sched in self.slices:
+        for sched in self.nodes:
             vect_above = False
             for loop_name in sched.interchange:
                 if loop_name in sched.vectorize:
@@ -553,9 +599,9 @@ def _check_vectorization_consistency(self):
                     )
 
     def _check_tiling_consistency(self) -> None:
-        mapper = LoopsDimsMapper.build_from_slices(self.slices)
+        mapper = LoopsDimsMapper.build_from_nodes(self.nodes)
         seen_axes: dict[str, int | None] = {}
-        for sched in self.slices:
+        for sched in self.nodes:
             for loop_name in sched.interchange:
                 if loop_name in mapper.dims:
                     seen_axes[loop_name] = None
@@ -571,9 +617,9 @@ def _check_tiling_consistency(self) -> None:
                     seen_axes[axis] = sched.tiles[axis][loop_name]
 
     def _check_sizes(self):
-        mapper = LoopsDimsMapper.build_from_slices(self.slices)
+        mapper = LoopsDimsMapper.build_from_nodes(self.nodes)
         current_size_of_split: dict[str, int | None] = {}
-        for sched in self.slices:
+        for sched in self.nodes:
             current_size_of_tile: dict[str, int] = {}
 
             for loop_name in sched.interchange:
@@ -677,7 +723,7 @@ def apply(self, node_name: str, spec: dict[str, dict[str, Any]]) -> None:
             ScheduleValidationError: If the resulting schedule is invalid.
         """
         # Parse the specification into an AST
-        parser = ScheduleParser(self.abstract_axis)
+        parser = ScheduleParser()
         ast = parser.parse(spec)
 
         # Interpret the AST into a LoopNest
@@ -694,16 +740,24 @@ def _apply_loop_nest(self, loop_nest: LoopNest) -> None:
         """Apply a LoopNest to the scheduler."""
         self.scheduler.set_dims(self.abstract_axis)
 
-        for slice in loop_nest.slices:
-            root = slice.root
+        if loop_nest.root_node is not None:
+            self._apply_node(loop_nest.root_node)
+
+    def _apply_node(self, node: LoopNestNode) -> None:
+        """Recursively apply a LoopNestNode and its children to the scheduler."""
+        root = node.root
+
+        for d, s in node.splits.items():
+            self.scheduler.split(d, s, root=root)
 
-            for d, s in slice.splits.items():
-                self.scheduler.split(d, s, root=root)
+        for d, s in node.tiles.items():
+            self.scheduler.tile(d, s, root=root)
 
-            for d, s in slice.tiles.items():
-                self.scheduler.tile(d, s, root=root)
+        self.scheduler.interchange(node.interchange, root=root)
+        self.scheduler.vectorize(node.vectorize, root=root)
+        self.scheduler.parallelize(node.parallelize, root=root)
+        self.scheduler.unroll(node.unroll, root=root)
 
-            self.scheduler.interchange(slice.interchange, root=root)
-            self.scheduler.vectorize(slice.vectorize, root=root)
-            self.scheduler.parallelize(slice.parallelize, root=root)
-            self.scheduler.unroll(slice.unroll, root=root)
+        # Recursively apply children
+        for child in node.children:
+            self._apply_node(child.node)