[transform_ext] Move populate_pattern op to dialects dir

Introduces op
`transform_ext.populate_pattern TARGET_OP_KIND PAT_NAME PRIORITY`
where patterns can be registered on `PopulatePatternOp` via its
`name_to_pattern_rewrite` class member.

Fixes #80.

Author: Rolf Morel

examples/workload/example.py

@@ -18,6 +18,7 @@
 from mlir.dialects import transform
 from mlir.execution_engine import ExecutionEngine
 
+from lighthouse import dialects as lh_dialects
 from lighthouse.pipeline.helper import match
 from lighthouse.pipeline.opt import PassBundles, apply_bundle
 
@@ -159,6 +160,8 @@ def schedule_modules(
 
 if __name__ == "__main__":
     with ir.Context(), ir.Location.unknown():
+        lh_dialects.register_and_load()
+
         wload = ElementwiseSum(400, 400)
 
         print(" Dump kernel ".center(60, "-"))

examples/workload/example_mlir.py

@@ -21,6 +21,7 @@
 from mlir.dialects import func, linalg, arith, memref
 from mlir.execution_engine import ExecutionEngine
 
+from lighthouse import dialects as lh_dialects
 from lighthouse.workload import execute, benchmark
 import lighthouse.utils as lh_utils
 
@@ -195,6 +196,8 @@ def payload_module(self):
 
 if __name__ == "__main__":
     with ir.Context(), ir.Location.unknown():
+        lh_dialects.register_and_load()
+
         wload = ElementwiseSumMLIRAlloc(400, 400)
 
         print(" Dump kernel ".center(60, "-"))

examples/xegpu/matmul.py

@@ -19,6 +19,7 @@
 from mlir import ir
 from mlir.execution_engine import ExecutionEngine
 
+from lighthouse import dialects as lh_dialects
 from lighthouse.workload import benchmark, get_bench_wrapper_schedule
 from lighthouse.utils.memref import to_ctype as memref_to_ctype
 from lighthouse.utils.numpy import numpy_to_ctype
@@ -360,6 +361,8 @@ def parse_cli():
     c_type = "f32"
 
     with ir.Context(), ir.Location.unknown():
+        lh_dialects.register_and_load()
+
         wload = XeGPUMatMul(
             M=M,
             N=N,

examples/xegpu/mlp.py

@@ -24,6 +24,7 @@
 from mlir import ir
 from mlir.execution_engine import ExecutionEngine
 
+from lighthouse import dialects as lh_dialects
 from lighthouse.workload import benchmark, get_bench_wrapper_schedule
 from lighthouse.utils.memref import to_ctype as memref_to_ctype
 from lighthouse.utils.numpy import numpy_to_ctype
@@ -375,6 +376,8 @@ def parse_cli():
     identity_weights = args.check_result
 
     with ir.Context(), ir.Location.unknown():
+        lh_dialects.register_and_load()
+
         wload = XeGPUMLP(
             batch_size=args.batch_size,
             input_size=args.input_size,

lighthouse/dialects/__init__.py

@@ -0,0 +1,4 @@
+def register_and_load():
+    from . import transform_ext
+
+    transform_ext.register_and_load()

lighthouse/dialects/transform_ext.py

@@ -0,0 +1,60 @@
+from mlir import rewrite, ir
+from mlir.dialects import ext, transform
+
+
+def register_and_load(context=None):
+    TransformExtensionDialect.load()
+
+
+class TransformExtensionDialect(ext.Dialect, name="transform_ext"):
+    @classmethod
+    def load(cls, *args, **kwargs):
+        super().load(*args, **kwargs)
+        for op_cls in cls.operations:
+            if hasattr(op_cls, "attach_interface_impls"):
+                op_cls.attach_interface_impls()
+
+
+class PopulatePatternOp(TransformExtensionDialect.Operation, name="populate_pattern"):
+    """An operation to populate a pattern set with a specific pattern.
+
+    To be used in the region of `transform.apply_patterns`."""
+
+    op_kind: ir.StringAttr
+    pattern_name: ir.StringAttr
+    priority: ir.IntegerAttr
+
+    # A mapping from pattern names to their corresponding rewrite functions.
+    # This should be populated by the users of this operation. In effect serves
+    # as a registry for rewrite patterns that can be applied by this operation.
+    name_to_rewrite_pattern = {}
+
+    @classmethod
+    def attach_interface_impls(cls, context=None):
+        cls.PatternDescriptorOpInterfaceModel.attach(
+            cls.OPERATION_NAME, context=context
+        )
+
+    class PatternDescriptorOpInterfaceModel(transform.PatternDescriptorOpInterface):
+        @staticmethod
+        def populate_patterns(
+            op: "PopulatePatternOp",
+            patternset: rewrite.RewritePatternSet,
+        ) -> None:
+            patternset.add(
+                op.op_kind.value,
+                op.name_to_rewrite_pattern[op.pattern_name.value],
+                benefit=op.priority.value,
+            )
+
+
+def populate_pattern(
+    op_kind: str, pattern_name: str, priority: int
+) -> PopulatePatternOp:
+    """Camelcase constructor for PopulatePatternOp."""
+    priority_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), priority)
+    return PopulatePatternOp(
+        op_kind=ir.StringAttr.get(op_kind),
+        pattern_name=ir.StringAttr.get(pattern_name),
+        priority=priority_attr,
+    )

lighthouse/schedule/pattern_schedule.py

@@ -5,10 +5,11 @@
 import numpy as np
 import os
 from mlir import ir
-from mlir.dialects import func, arith, scf, memref
+from mlir.dialects import func, arith, scf, memref, transform
 from mlir.execution_engine import ExecutionEngine
 from mlir.runtime.np_to_memref import get_ranked_memref_descriptor
-from lighthouse.schedule.pattern_schedule import pattern_rewrite_schedule
+from lighthouse.dialects.transform_ext import PopulatePatternOp, populate_pattern
+from lighthouse.schedule.utils import schedule_boilerplate
 from lighthouse.utils.mlir import func_cif, get_mlir_library_path
 from lighthouse.utils.memref import to_packed_args
 from lighthouse.workload import Workload
@@ -121,16 +122,24 @@ def bench(*args):
 
 
 def get_bench_wrapper_schedule(workload: Workload):
-    return pattern_rewrite_schedule(
-        {
-            "func.func": bench_wrapper_pattern(
-                workload.payload_function_name,
-                workload.benchmark_function_name,
-            )
-        },
-        "add_bench_pattern",
+    PopulatePatternOp.name_to_rewrite_pattern["bench_wrapper"] = bench_wrapper_pattern(
+        workload.payload_function_name,
+        workload.benchmark_function_name,
     )
 
+    with schedule_boilerplate() as (schedule, named_seq):
+        apply_patterns_op = transform.apply_patterns(named_seq.bodyTarget)
+        with ir.InsertionPoint(apply_patterns_op.patterns):
+            populate_pattern(
+                op_kind="func.func",
+                pattern_name="bench_wrapper",
+                priority=1,
+            )
+        transform.yield_([named_seq.bodyTarget])
+
+    schedule.body.operations[0].verify()
+    return schedule
+
 
 def benchmark(
     workload: Workload,

lighthouse/workload/runner.py

@@ -3,7 +3,7 @@ name = "lighthouse"
 dynamic = ["version"]
 requires-python = ">=3.10,<3.13"  # Bounds are due to torch-mlir's packaging
 dependencies = [
-    "mlir-python-bindings==20260315+69780be1d"
+    "mlir-python-bindings==20260316+f46a51538"
 ]
 
 [dependency-groups]