New Feature: Kerngraph support for Relin (#109)

Refactored kerngraph to support grouping of p-isa lists to enable flexible reordering. Added regression testing for Relin Kernel.

Co-authored-by: Jose Rojas Chaves <[email protected]>

Author: christopherngutierrez

p-isa_tools/kerngen/kernel_optimization/loops.py

@@ -10,6 +10,36 @@
 from high_parser.pisa_operations import PIsaOp, Comment
 
 
+class PIsaOpGroup:
+    """A group of PIsaOp instructions with reorderable flag.
+
+    Attributes:
+        pisa_list: List of PIsaOp instructions
+        is_reorderable: Boolean indicating if the instructions can be reordered
+    """
+
+    pisa_list: list[PIsaOp]
+    is_reorderable: bool
+
+    def __init__(self, pisa_list: list[PIsaOp], is_reorderable: bool = False):
+        """Initialize PIsaOpGroup.
+
+        Args:
+            pisa_list: List of PIsaOp instructions
+            is_reorderable: Boolean indicating if instructions can be reordered
+        """
+        self.pisa_list = pisa_list
+        self.is_reorderable = is_reorderable
+
+    def __len__(self) -> int:
+        """Return the number of instructions in the group."""
+        return len(self.pisa_list)
+
+    def __iter__(self):
+        """Allow iteration over the PIsaOp instructions."""
+        return iter(self.pisa_list)
+
+
 def remove_comments(pisa_list: list[PIsaOp]) -> list[PIsaOp]:
     """Remove comments from a list of PIsaOp instructions.
 
@@ -22,40 +52,49 @@ def remove_comments(pisa_list: list[PIsaOp]) -> list[PIsaOp]:
     return [pisa for pisa in pisa_list if not isinstance(pisa, Comment)]
 
 
-def split_by_reorderable(pisa_list: list[PIsaOp]) -> tuple[list[PIsaOp], list[PIsaOp]]:
+def split_by_reorderable(pisa_list: list[PIsaOp]) -> list[PIsaOpGroup]:
     """Split a list of PIsaOp instructions into reorderable and non-reorderable groups.
 
     Args:
         pisa_list: List of PIsaOp instructions
 
     Returns:
-        Tuple containing two lists:
-            - reorderable: Instructions that can be reordered
-            - non_reorderable: Instructions that cannot be reordered
+        List of PIsaOpGroup objects containing grouped instructions with their reorderable status
     """
-
-    reorderable = []
-    non_reorderable = []
-    is_reorderable = False
+    groups = []
+    current_group = PIsaOpGroup([], is_reorderable=False)
+    no_reoderable_group = True
 
     for pisa in pisa_list:
         # if the pisa is a comment and it contains <reorderable> tag, treat the following pisa as reorderable until a </reorderable> tag is found.
         if isinstance(pisa, Comment):
             if "<reorderable>" in pisa.line:
-                is_reorderable = True
+                # If current group has instructions, append it to groups first
+                if current_group.pisa_list:
+                    groups.append(current_group)
+                # Create a new reorderable group
+                current_group = PIsaOpGroup([], is_reorderable=True)
+                no_reoderable_group = False
             elif "</reorderable>" in pisa.line:
-                is_reorderable = False
-
-        if is_reorderable:
-            reorderable.append(pisa)
+                # End reorderable section, append current group to groups
+                if current_group.pisa_list:
+                    groups.append(current_group)
+                # Create a new non-reorderable group
+                current_group = PIsaOpGroup([], is_reorderable=False)
         else:
-            non_reorderable.append(pisa)
+            # Add non-comment instruction to current group
+            current_group.pisa_list.append(pisa)
+
+    # Add any remaining instructions in current_group
+    if current_group.pisa_list:
+        groups.append(current_group)
+
+    # If there are no reorderable groups, set reorderable to True for the entire groups
+    if no_reoderable_group:
+        for group in groups:
+            group.is_reorderable = True
 
-    # if reoderable is empty, return non_reorderable as reorderable
-    if not reorderable:
-        reorderable = non_reorderable
-        non_reorderable = []
-    return remove_comments(reorderable), remove_comments(non_reorderable)
+    return groups
 
 
 def loop_interchange(

p-isa_tools/kerngen/kerngraph.py

@@ -52,7 +52,7 @@ def parse_args():
         nargs="*",
         default=[],
         # Composition high ops such are ntt, mod, and relin are not currently supported
-        choices=["add", "sub", "mul", "muli", "copy", "ntt", "intt", "mod"],
+        choices=["add", "sub", "mul", "muli", "copy", "ntt", "intt", "mod", "relin"],
         help="List of high_op names",
     )
     parser.add_argument(
@@ -78,47 +78,67 @@ def parse_args():
     return parser.parse_args()
 
 
-def main(args):
-    """Main function to read input and parse each line with KernelParser."""
-    input_lines = sys.stdin.read().strip().splitlines()
+def parse_kernels(input_lines, debug=False):
+    """Parse kernel strings from input lines."""
     valid_kernels = []
-    Config.legacy_mode = args.legacy
-
     for line in input_lines:
         try:
             kernel = KernelParser.parse_kernel(line)
             valid_kernels.append(kernel)
         except ValueError as e:
-            if args.debug:
+            if debug:
                 print(f"Error parsing line: {line}\nReason: {e}")
             continue  # Skip invalid lines
+    return valid_kernels
+
+
+def process_kernel_with_reordering(kernel, args):
+    """Process a kernel with reordering optimization."""
+    groups = split_by_reorderable(kernel.to_pisa())
+    processed_kernel = []
+    for group in groups:
+        if group.is_reorderable:
+            processed_kernel.append(
+                loop_interchange(
+                    group.pisa_list,
+                    primary_key=args.primary,
+                    secondary_key=args.secondary,
+                )
+            )
+        else:
+            processed_kernel.append(group.pisa_list)
+
+    for pisa in mixed_to_pisa_ops(processed_kernel):
+        print(pisa)
+
+
+def should_apply_reordering(kernel, targets):
+    """Check if reordering should be applied to this kernel."""
+    return targets and any(target.lower() in str(kernel).lower() for target in targets)
+
+
+def main(args):
+    """Main function to read input and parse each line with KernelParser."""
+    input_lines = sys.stdin.read().strip().splitlines()
+    Config.legacy_mode = args.legacy
+
+    valid_kernels = parse_kernels(input_lines, args.debug)
 
     if not valid_kernels:
         print("No valid kernel strings were parsed.")
-    else:
-        if args.debug:
-            print(
-                f"# Reordered targets {args.target} with primary key {args.primary} and secondary key {args.secondary}"
-            )
-        for kernel in valid_kernels:
-            if args.target and any(
-                target.lower() in str(kernel).lower() for target in args.target
-            ):
-                reorderable, non_reorderable = split_by_reorderable(kernel.to_pisa())
-                kernel = non_reorderable
-                kernel.append(
-                    loop_interchange(
-                        reorderable,
-                        primary_key=args.primary,
-                        secondary_key=args.secondary,
-                    )
-                )
-
-                for pisa in mixed_to_pisa_ops(kernel):
-                    print(pisa)
-            else:
-                for pisa in kernel.to_pisa():
-                    print(pisa)
+        return
+
+    if args.debug:
+        print(
+            f"# Reordered targets {args.target} with primary key {args.primary} and secondary key {args.secondary}"
+        )
+
+    for kernel in valid_kernels:
+        if should_apply_reordering(kernel, args.target):
+            process_kernel_with_reordering(kernel, args)
+        else:
+            for pisa in kernel.to_pisa():
+                print(pisa)
 
 
 if __name__ == "__main__":