Extend pipeline helper to schedules (#70)

So we can share the same bundles across pipelines and schedules.

Author: rengolin

examples/ingress/torch/compile_model.py

@@ -35,6 +35,9 @@ def lower_to_llvm(module: ir.Module) -> ir.Module:
     add_bundle(pm, PassBundles.BufferizationBundle)
     add_bundle(pm, PassBundles.CleanupBundle)
 
+    # Middle-end lowering.
+    add_bundle(pm, PassBundles.MLIRLoweringBundle)
+
     # Lower to LLVM.
     add_bundle(pm, PassBundles.LLVMLoweringBundle)
     add_bundle(pm, PassBundles.CleanupBundle)

examples/workload/example.py

@@ -18,7 +18,11 @@
 from mlir.dialects import transform
 from mlir.execution_engine import ExecutionEngine
 
-from lighthouse.pipeline.helper import apply_registered_pass, canonicalize, match
+from lighthouse.pipeline.helper import (
+    PassBundles,
+    apply_bundle,
+    match,
+)
 from lighthouse.workload import Workload, execute, benchmark
 
 
@@ -141,21 +145,15 @@ def schedule_module(
                     op_name="builtin.module",
                     deduplicate=True,
                 )
-                mod = apply_registered_pass(mod, "one-shot-bufferize")
-                mod = apply_registered_pass(mod, "convert-linalg-to-loops")
-                transform.apply_cse(mod)
-                canonicalize(mod)
+                mod = apply_bundle(mod, PassBundles.BufferizationBundle)
+                mod = apply_bundle(mod, PassBundles.MLIRLoweringBundle)
+                mod = apply_bundle(mod, PassBundles.CleanupBundle)
 
                 if stop_at_stage == "bufferized":
                     transform.YieldOp()
                     return schedule_module
 
-                mod = apply_registered_pass(mod, "convert-scf-to-cf")
-                mod = apply_registered_pass(mod, "finalize-memref-to-llvm")
-                mod = apply_registered_pass(mod, "convert-cf-to-llvm")
-                mod = apply_registered_pass(mod, "convert-arith-to-llvm")
-                mod = apply_registered_pass(mod, "convert-func-to-llvm")
-                mod = apply_registered_pass(mod, "reconcile-unrealized-casts")
+                mod = apply_bundle(mod, PassBundles.LLVMLoweringBundle)
                 transform.YieldOp()
 
         return schedule_module

lighthouse/pipeline/helper.py

@@ -4,6 +4,25 @@
 from mlir.dialects.transform import structured
 
 
+class Pass:
+    """
+    A simple wrapper class for MLIR passes.
+    The options can be serialized into a string for consumption by the PassManager.
+    Or used directly with the Transform Schedule by passing the options as a dictionary.
+    """
+
+    def __init__(self, name: str, options: dict = {}):
+        self.name = name
+        self.options = options
+
+    def __str__(self) -> str:
+        """serialize name + options dictionary for pass manager consumption"""
+        if not self.options:
+            return self.name
+        options_str = " ".join(f"{key}={value}" for key, value in self.options.items())
+        return f"{self.name}{{{options_str}}}"
+
+
 class PassBundles:
     """
     Predefined pass bundles for common transformations. These are not exhaustive and can be extended as needed.
@@ -12,30 +31,48 @@ class PassBundles:
 
     # All in one bufferization bundle. This is self consistent and should be used together.
     BufferizationBundle = [
-        "one-shot-bufferize{function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries}",
-        "drop-equivalent-buffer-results",
-        "buffer-deallocation-pipeline",
+        Pass(
+            "one-shot-bufferize",
+            {
+                "function-boundary-type-conversion": "identity-layout-map",
+                "bufferize-function-boundaries": True,
+            },
+        ),
+        Pass("drop-equivalent-buffer-results"),
+        # This last pass only works with the pass manager, not schedules.
+        # Pass("buffer-deallocation-pipeline"),
+    ]
+
+    # Lowers most dialects to basic control flow.
+    MLIRLoweringBundle = [
+        Pass("convert-linalg-to-loops"),
     ]
 
     # Lowers most dialects to LLVM Dialect
     LLVMLoweringBundle = [
-        "convert-linalg-to-loops",
-        "convert-scf-to-cf",
-        "convert-to-llvm",
-        "reconcile-unrealized-casts",
+        Pass("convert-scf-to-cf"),
+        Pass("convert-to-llvm"),
+        Pass("reconcile-unrealized-casts"),
     ]
 
     # Canonicalization bundle. This is a set of passes that can be used to clean up the IR after transformations.
     CleanupBundle = [
-        "cse",
-        "canonicalize",
+        Pass("cse"),
+        Pass("canonicalize"),
     ]
 
 
 # Utility function to add a bundle of passes to a PassManager. This can be used to easily add a predefined set of passes to a pipeline.
-def add_bundle(pm: PassManager, bundle: list[str]) -> None:
+def add_bundle(pm: PassManager, bundle: list[Pass]) -> None:
+    for p in bundle:
+        pm.add(str(p))
+
+
+# Utility function to add a bundle of passes to a Schedule. This can be used to easily add a predefined set of passes to a pipeline.
+def apply_bundle(op, bundle: list[Pass], *args, **kwargs) -> None:
     for p in bundle:
-        pm.add(p)
+        op = apply_registered_pass(op, p.name, options=p.options, *args, **kwargs)
+    return op
 
 
 # Utility function to add a bundle of passes to a Transform Schedule. This can be used to easily add a predefined set of passes to a pipeline.