Working loop interchange, initial checking

Author: christopherngutierrez

kerngen/const/options.py

@@ -0,0 +1,23 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Module for defining constants and enums used in the kernel generator"""
+from enum import Enum
+
+
+class LoopKey(Enum):
+    """Sort keys for PIsaOp instructions"""
+
+    RNS = "rns"
+    PART = "part"
+    UNIT = "unit"
+
+    @classmethod
+    def from_str(cls, value: str) -> "LoopKey":
+        """Convert a string to a LoopKey enum"""
+        if value is None:
+            raise ValueError("LoopKey cannot be None")
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(f"Invalid LoopKey: {value}") from None

kerngen/kernel_optimization/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0

kerngen/kernel_optimization/loops.py

@@ -0,0 +1,56 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Module for loop interchange optimization in P-ISA operations"""
+
+import re
+from const.options import LoopKey
+from high_parser.pisa_operations import PIsaOp, Comment
+
+
+def loop_interchange(
+    pisa_list: list[PIsaOp],
+    primary_key: LoopKey | None = LoopKey.PART,
+    secondary_key: LoopKey | None = LoopKey.RNS,
+) -> list[PIsaOp]:
+    """Batch pisa_list into groups and sort them by primary and optional secondary keys.
+
+    Args:
+        pisa_list: List of PIsaOp instructions
+        primary_key: Primary sort criterion from SortKey enum
+        secondary_key: Optional secondary sort criterion from SortKey enum
+
+    Returns:
+        List of processed PIsaOp instructions
+
+    Raises:
+        ValueError: If invalid sort key values provided
+    """
+    if primary_key is None and secondary_key is None:
+        return pisa_list
+
+    def get_sort_value(pisa: PIsaOp, key: LoopKey) -> int:
+        match key:
+            case LoopKey.RNS:
+                return pisa.q
+            case LoopKey.PART:
+                match = re.search(r"_(\d+)_", str(pisa))
+                return int(match[1]) if match else 0
+            case LoopKey.UNIT:
+                match = re.search(r"_(\d+),", str(pisa))
+                return int(match[1]) if match else 0
+            case _:
+                raise ValueError(f"Invalid sort key value: {key}")
+
+    def get_sort_key(pisa: PIsaOp) -> tuple:
+        primary_value = get_sort_value(pisa, primary_key)
+        if secondary_key:
+            secondary_value = get_sort_value(pisa, secondary_key)
+            return (primary_value, secondary_value)
+        return (primary_value,)
+
+    # Filter out comments
+    pisa_list_wo_comments = [p for p in pisa_list if not isinstance(p, Comment)]
+    # Sort based on primary and optional secondary keys
+    pisa_list_wo_comments.sort(key=get_sort_key)
+    return pisa_list_wo_comments

p-isa_tools/kerngen/kernel_parser/__init__.py

@@ -1,2 +1,2 @@
 # Copyright (C) 2024 Intel Corporation
-# SPDX-License-Identifier: Apache-2.0
+# SPDX-License-Identifier: Apache-2.0