adding benchmark function through pattern rewrite schedule

Author: fschlimb

examples/feed-forward-mpi/feed-forward-mpi.py

@@ -13,8 +13,8 @@
 import ctypes
 from contextlib import contextmanager
 from typing import Optional
-
 import numpy as np
+
 from mlir import ir
 from mlir.dialects import transform
 from mlir.dialects.transform.bufferization import OneShotBufferizeOp
@@ -25,12 +25,13 @@
     make_nd_memref_descriptor,
     as_ctype,
 )
+
 from lighthouse.utils.memref import (
     to_ctype as memref_to_ctype,
     deallocate_memrefs_on_exit,
 )
 from lighthouse.pipeline.helper import apply_registered_pass, match
-from lighthouse.workload import Workload, benchmark
+from lighthouse.workload import Workload, benchmark, get_bench_wrapper_schedule
 from lighthouse.schedule.x86 import tile_and_vector_matmul
 from ff_weight_stationary import generate_ff_payload
 
@@ -103,8 +104,6 @@ class DistFF(Workload):
     where A, B, C are (M,K), (K,N), (N,K) matrices respectively.
     """
 
-    payload_function_name = "payload"
-
     def __init__(self, args, P: int, R: int):
         self.M = args.sizes[0]
         self.N = args.sizes[1]
@@ -308,8 +307,9 @@ def schedule_modules(
             transform.YieldOp()
             func = None
 
+        bench_schedule = get_bench_wrapper_schedule(self)
+
         tile_schedule = tile_and_vector_matmul.create()
-        tile_schedule.body.operations[0].verify()
 
         main_schedule = ir.Module.create()
         main_schedule.operation.attributes["transform.with_named_sequence"] = (
@@ -377,7 +377,7 @@ def schedule_modules(
                 transform.PrintOp(target=mod)
             transform.YieldOp()
 
-        return [pre_schedule, tile_schedule, main_schedule]
+        return [pre_schedule, bench_schedule, tile_schedule, main_schedule]
 
 
 if __name__ == "__main__":
@@ -394,7 +394,7 @@ def schedule_modules(
         # execute(wload, verbose=args.verbose)
         rprint(" Benchmark".center(60, "-"))
         times = benchmark(
-            wload, nruns=10, nwarmup=1, check_correctness=True, verbose=args.verbose
+            wload, nruns=100, nwarmup=10, check_correctness=True, verbose=args.verbose
         )
         # compute statistics
         times *= 1e6

examples/workload/example.py

@@ -19,7 +19,7 @@
 from mlir.execution_engine import ExecutionEngine
 
 from lighthouse.pipeline.helper import apply_registered_pass, canonicalize, match
-from lighthouse.workload import Workload, execute, benchmark
+from lighthouse.workload import Workload, execute, benchmark, get_bench_wrapper_schedule
 
 
 class ElementwiseSum(Workload):
@@ -158,7 +158,7 @@ def schedule_modules(
                 mod = apply_registered_pass(mod, "reconcile-unrealized-casts")
                 transform.YieldOp()
 
-        return [schedule_module]
+        return [get_bench_wrapper_schedule(self), schedule_module]
 
 
 if __name__ == "__main__":

examples/xegpu/matmul.py

@@ -17,7 +17,7 @@
 from mlir import ir
 from mlir.execution_engine import ExecutionEngine
 
-from lighthouse.workload import benchmark
+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
 from lighthouse.schedule.xegpu.mlp_schedule import get_schedule_module
@@ -195,14 +195,15 @@ def schedule_modules(
         self, stop_at_stage: Optional[str] = None, parameters: Optional[dict] = None
     ) -> list[ir.Module]:
         return [
+            get_bench_wrapper_schedule(self),
             get_schedule_module(
                 has_bias=self.has_bias,
                 has_relu=self.has_relu,
                 has_convert_c=False,
                 stop_at_stage=stop_at_stage,
                 nlayers=1,
                 params={"layer_0": parameters},
-            )
+            ),
         ]
 
     def shared_libs(self) -> list[str]:

examples/xegpu/mlp.py

@@ -21,7 +21,7 @@
 from mlir import ir
 from mlir.execution_engine import ExecutionEngine
 
-from lighthouse.workload import benchmark
+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
 from lighthouse.schedule.xegpu.mlp_schedule import get_schedule_module
@@ -256,14 +256,15 @@ def schedule_modules(
         self, stop_at_stage: Optional[str] = None, parameters: Optional[dict] = None
     ) -> list[ir.Module]:
         return [
+            get_bench_wrapper_schedule(self),
             get_schedule_module(
                 has_bias=self.has_bias,
                 has_relu=self.has_relu,
                 skip_final_layer_relu=True,
                 stop_at_stage=stop_at_stage,
                 nlayers=self.nlayers,
                 params=parameters,
-            )
+            ),
         ]
 
     def shared_libs(self) -> list[str]:

lighthouse/schedule/pattern_schedule.py

@@ -0,0 +1,70 @@
+from contextlib import contextmanager
+from mlir import rewrite, ir
+from mlir.dialects import ext, transform
+from mlir.dialects.transform import AnyOpType
+
+
+@ext.register_dialect
+class PatternDialect(ext.Dialect, name="lighthouse"):
+    pass
+
+
+def rewrite_pattern(patterns: dict, pname: str):
+    """Return a rewrite pattern class that can be registered with MLIR.
+    The patterns dict should map op names to their corresponding match and rewrite functions."""
+
+    @ext.register_operation(PatternDialect, replace=True)
+    class RewritePattern(PatternDialect.Operation, name=pname):
+        @classmethod
+        def attach_interface_impls(cls, ctx=None):
+            cls.PatternDescriptorOpInterfaceFallbackModel.attach(
+                cls.OPERATION_NAME, context=ctx
+            )
+
+        class PatternDescriptorOpInterfaceFallbackModel(
+            transform.PatternDescriptorOpInterface
+        ):
+            @staticmethod
+            def populate_patterns(
+                op: "RewritePattern",
+                patternset: rewrite.RewritePatternSet,
+            ) -> None:
+                for op_name, match_and_rewrite in patterns.items():
+                    patternset.add(op_name, match_and_rewrite, benefit=1)
+
+    return RewritePattern
+
+
+@contextmanager
+def schedule_boilerplate():
+    schedule = ir.Module.create()
+    schedule.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get()
+    with ir.InsertionPoint(schedule.body):
+        named_sequence = transform.NamedSequenceOp(
+            "__transform_main",
+            [AnyOpType.get()],
+            [AnyOpType.get()],
+            arg_attrs=[{"transform.consumed": ir.UnitAttr.get()}],
+        )
+        with ir.InsertionPoint(named_sequence.body):
+            yield schedule, named_sequence
+
+
+def pattern_rewrite_schedule(patterns: dict, pname: str = "rewrite_pattern"):
+    """Return a transform module that applies the given rewrite patterns.
+    patterns: dict mapping op names to match-and-rewrite functions.
+    pname: name for the generated rewrite pattern operation."""
+
+    rw_pattern = rewrite_pattern(patterns, pname)
+    PatternDialect.load(register=False, reload=False)
+    rw_pattern.attach_interface_impls()
+
+    with schedule_boilerplate() as (schedule, named_seq):
+        apply_patterns_op = transform.ApplyPatternsOp(named_seq.bodyTarget)
+        with ir.InsertionPoint(apply_patterns_op.patterns):
+            rw_pattern()
+        transform.yield_([named_seq.bodyTarget])
+        named_seq.verify()
+
+    schedule.body.operations[0].verify()
+    return schedule

lighthouse/schedule/x86/tile_and_vector_matmul.py

@@ -145,4 +145,6 @@ def create(tile_size=64) -> ir.Module:
         cleanup_func(named_seq.bodyTarget)
 
         transform.yield_()
+
+    schedule.body.operations[0].verify()
     return schedule

lighthouse/utils/memref.py

@@ -50,12 +50,17 @@ def get_packed_arg(
     return packed_args
 
 
-def to_packed_args(memref_descs) -> ctypes.Array[ctypes.c_void_p]:
+def to_packed_args(args) -> ctypes.Array[ctypes.c_void_p]:
     """
-    Convert a list of memref descriptors into packed ctype arguments.
+    Convert a list of memref descriptors and/or integers into packed ctype arguments.
 
     Args:
-        memref_descs: A list of memref descriptors.
+        args: A list of memref descriptors or integers.
     """
-    ctype_args = [to_ctype(memref) for memref in memref_descs]
+    ctype_args = []
+    for arg in args:
+        if isinstance(arg, int):
+            ctype_args.append(ctypes.pointer(ctypes.c_int64(arg)))
+        else:
+            ctype_args.append(to_ctype(arg))
     return get_packed_arg(ctype_args)

lighthouse/workload/__init__.py

@@ -1,4 +1,15 @@
 from .workload import Workload
-from .runner import execute, benchmark
+from .runner import (
+    execute,
+    benchmark,
+    bench_wrapper_pattern,
+    get_bench_wrapper_schedule,
+)
 
-__all__ = ["Workload", "benchmark", "execute"]
+__all__ = [
+    "Workload",
+    "bench_wrapper_pattern",
+    "benchmark",
+    "execute",
+    "get_bench_wrapper_schedule",
+]