descript: expose pack_at

Author: qaco

src/xtc/backends/tvm/TVMScheduler.py

@@ -538,6 +538,11 @@ def get_loop_nest(self) -> LoopNest:
         # Build buffer_at mapping
         root_node.buffer_at = {axis: None for axis in self.write_caches}
 
+        # Build pack_at mapping
+        root_node.pack_at = {
+            axis: (input_idx, None, pad) for axis, input_idx, pad in self.read_buffers
+        }
+
         return loop_nest
 
 

src/xtc/schedules/descript.py

@@ -213,6 +213,9 @@ def _apply_annotations(
         if annotations.buffer_specified:
             node.buffer_at[loop_name] = annotations.buffer
 
+        if annotations.pack_specified and annotations.pack is not None:
+            node.pack_at[loop_name] = annotations.pack
+
     def _check_splitting_intervals(
         self,
         item: SplitDecl,
@@ -305,6 +308,9 @@ def _apply_node(self, node: LoopNestNode) -> None:
         for axis, mtype in node.buffer_at.items():
             self.scheduler.buffer_at(axis, mtype=mtype, root=root)
 
+        for axis, (input_idx, mtype, pad) in node.pack_at.items():
+            self.scheduler.pack_at(axis, input_idx, mtype=mtype, pad=pad, root=root)
+
         # Recursively apply children
         for child in node.children:
             self._apply_node(child)

src/xtc/schedules/loop_nest.py

@@ -94,6 +94,9 @@ class LoopNestNode(Node["LoopNestNode"]):
         unroll: Maps loop names to their unroll factors.
         buffer_at: Buffer configuration per axis. Maps axis names to optional
             memory types (mtype). None means default memory type.
+        pack_at: Pack configuration per axis. Maps axis names to tuples of
+            (input_idx, mtype, pad). input_idx is the input buffer index,
+            mtype is the memory type (None for default), pad enables padding.
     """
 
     root: str
@@ -104,6 +107,7 @@ class LoopNestNode(Node["LoopNestNode"]):
     parallelize: list[str] = field(default_factory=list)
     unroll: dict[str, int] = field(default_factory=dict)
     buffer_at: dict[str, str | None] = field(default_factory=dict)
+    pack_at: dict[str, tuple[int, str | None, bool]] = field(default_factory=dict)
 
     def pretty_print(self, indent: int = 0) -> str:
         """Return a human-readable representation of the loop nest.
@@ -204,7 +208,7 @@ def pretty_print(self, indent: int = 0) -> str:
         return "\n".join(lines)
 
     def _add_annotations(self, line: str, loop_name: str) -> str:
-        """Add annotations (parallelized, vectorized, unroll, buffer) to a loop line."""
+        """Add annotations (parallelized, vectorized, unroll, buffer, pack) to a loop line."""
         annotations: list[str] = []
         if loop_name in self.parallelize:
             annotations.append("parallelized")
@@ -218,6 +222,14 @@ def _add_annotations(self, line: str, loop_name: str) -> str:
                 annotations.append(f"buffer({mtype})")
             else:
                 annotations.append("buffer")
+        if loop_name in self.pack_at:
+            input_idx, mtype, pad = self.pack_at[loop_name]
+            parts = [str(input_idx)]
+            if mtype is not None:
+                parts.append(mtype)
+            if pad:
+                parts.append("pad")
+            annotations.append(f"pack({', '.join(parts)})")
         if annotations:
             line += "  // " + ", ".join(annotations)
         return line

src/xtc/schedules/parsing.py

@@ -25,6 +25,9 @@ class Annotations:
         buffer: The memory type for the buffer. None means default memory type.
             Only meaningful when buffer_specified is True.
         buffer_specified: True if buffer was explicitly requested.
+        pack: Pack configuration as (input_idx, mtype, pad). mtype is None for default.
+            Only meaningful when pack_specified is True.
+        pack_specified: True if pack was explicitly requested.
     """
 
     unroll_factor: int | None = None
@@ -33,6 +36,8 @@ class Annotations:
     parallelize: bool = False
     buffer: str | None = None
     buffer_specified: bool = False
+    pack: tuple[int, str | None, bool] | None = None
+    pack_specified: bool = False
 
 
 @dataclass(frozen=True)
@@ -152,6 +157,8 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations
         parallelize = False
         buffer: str | None = None
         buffer_specified = False
+        pack: tuple[int, str | None, bool] | None = None
+        pack_specified = False
 
         for key, param in value.items():
             if key == "unroll":
@@ -186,6 +193,9 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations
                     )
                 buffer = None if param == "default" else param
                 buffer_specified = True
+            elif key == "pack":
+                pack = self._parse_pack_param(param, context)
+                pack_specified = True
             else:
                 raise ScheduleParseError(f"Unknown annotation on {context}: {key}")
 
@@ -196,8 +206,42 @@ def _parse_annotations(self, value: dict[str, Any], context: str) -> Annotations
             parallelize=parallelize,
             buffer=buffer,
             buffer_specified=buffer_specified,
+            pack=pack,
+            pack_specified=pack_specified,
         )
 
+    def _parse_pack_param(
+        self, param: Any, context: str
+    ) -> tuple[int, str | None, bool]:
+        """Parse pack parameter into (input_idx, mtype, pad) tuple."""
+        if not isinstance(param, (list, tuple)) or len(param) != 3:
+            raise ScheduleParseError(
+                f'`{{"pack" = {param}}}` on {context}: pack parameter should be a tuple (input_idx, mtype, pad).'
+            )
+
+        input_idx, mtype, pad = param
+
+        if not isinstance(input_idx, int):
+            raise ScheduleParseError(
+                f'`{{"pack" = {param}}}` on {context}: input_idx should be an integer.'
+            )
+
+        if mtype is not None and not isinstance(mtype, str):
+            raise ScheduleParseError(
+                f'`{{"pack" = {param}}}` on {context}: mtype should be a string or None.'
+            )
+
+        if not isinstance(pad, bool):
+            raise ScheduleParseError(
+                f'`{{"pack" = {param}}}` on {context}: pad should be a boolean.'
+            )
+
+        # Convert "default" to None for mtype
+        if mtype == "default":
+            mtype = None
+
+        return (input_idx, mtype, pad)
+
     def _parse_split_syntax(
         self, declaration: str
     ) -> tuple[str, int | None, int | None]:

tests/filecheck/schedules/test_descript_pretty_print.py

@@ -4,6 +4,7 @@
 # RUN: python %s --full 2>&1 | filecheck %s --check-prefix=CHECK-FULL
 # RUN: python %s --split 2>&1 | filecheck %s --check-prefix=CHECK-SPLIT
 # RUN: python %s --buffer 2>&1 | filecheck %s --check-prefix=CHECK-BUFFER
+# RUN: python %s --pack 2>&1 | filecheck %s --check-prefix=CHECK-PACK
 
 import sys
 from xtc.schedules.parsing import ScheduleParser
@@ -64,6 +65,17 @@
     loop_nest = interpreter.interpret(ast, root="C")
     print(loop_nest.root_node.pretty_print())
 
+elif "--pack" in sys.argv:
+    spec = {
+        "i": {"parallelize": True},
+        "k": {"pack": (0, "default", False)},
+        "j": {"pack": (1, "shared", True)},
+        "j#16": {"vectorize": True},
+    }
+    ast = parser.parse(spec)
+    loop_nest = interpreter.interpret(ast, root="C")
+    print(loop_nest.root_node.pretty_print())
+
 # CHECK-SIMPLE:      loop i
 # CHECK-SIMPLE-NEXT:   loop k
 # CHECK-SIMPLE-NEXT:     loop j
@@ -105,3 +117,9 @@
 # CHECK-BUFFER-NEXT:     loop j  // buffer(shared)
 # CHECK-BUFFER-NEXT:       tile(j, 16)  // vectorized
 # CHECK-BUFFER-NEXT:         ...
+
+# CHECK-PACK:      loop i  // parallelized
+# CHECK-PACK-NEXT:   loop k  // pack(0)
+# CHECK-PACK-NEXT:     loop j  // pack(1, shared, pad)
+# CHECK-PACK-NEXT:       tile(j, 16)  // vectorized
+# CHECK-PACK-NEXT:         ...

tests/filecheck/schedules/test_get_descript.py

@@ -33,6 +33,7 @@
 sch.vectorize(["J0"])
 if "--tvm" in sys.argv:
     sch.buffer_at("J")
+    sch.pack_at("I", 0, pad=True)
 
 loop_nest = sch.get_loop_nest()
 print(loop_nest.root_node.pretty_print())
@@ -45,7 +46,7 @@
 # CHECK-MLIR-NEXT:           ...
 
 # CHECK-TVM:      loop K
-# CHECK-TVM-NEXT:   loop I
+# CHECK-TVM-NEXT:   loop I  // pack(0, pad)
 # CHECK-TVM-NEXT:     loop J  // buffer
 # CHECK-TVM-NEXT:       tile(I, 2)  // unroll(2)
 # CHECK-TVM-NEXT:         tile(J, 16)  // vectorized