print-after-all

Author: rengolin

lighthouse/execution/runner.py

@@ -212,6 +212,7 @@ def get_bench_wrapper_schedule(payload_func: str) -> ir.Module:
         """
         Get a schedule that wraps the payload function in a benchmarking function.
         The function name is defined in Runner and will be used by the runner benchmark method.
+        This schedule must apply to the module before any other in an optimizing pipeline.
         """
         with ir.Location.unknown():
             with schedule_boilerplate(result_types=[transform.any_op_t()]) as (

lighthouse/pipeline/driver.py

@@ -46,11 +46,13 @@ def add_stage(self, stage: lhs.Stage) -> None:
         # Users can derive their own classes from Stage and add them to the pipeline with this method.
         self.stages.append(stage)
 
-    def apply(self, module: ir.Module) -> ir.Module:
+    def apply(self, module: ir.Module, print_after_all: bool = False) -> ir.Module:
         if module.context != self.context:
             raise ValueError("Module context does not match driver context.")
         for stage in self.stages:
             module = stage.apply(module)
+            if print_after_all:
+                print(f"After stage {stage}:\n{module}")
         return module
 
     def __len__(self):
@@ -173,14 +175,14 @@ def reset(self) -> None:
         self.module = None
         self.pipeline_fixed = False
 
-    def run(self) -> ir.Module:
+    def run(self, print_after_all: bool = False) -> ir.Module:
         if self.module is None:
             raise ValueError("Module must not be empty.")
         if len(self.pipeline) == 0:
             raise ValueError("Pipeline must have at least one stage.")
 
         # Apply the whole pipeline.
-        self.pipeline.apply(self.module)
+        self.pipeline.apply(self.module, print_after_all=print_after_all)
 
         # The pipeline is now fixed and cannot be modified until reset is called.
         # This is to prevent accidental modifications to the pipeline after it has been run,

lighthouse/pipeline/stage.py

@@ -147,6 +147,7 @@ class PassStage(Stage):
     def __init__(self, passes: list[Pass], context: ir.Context):
         self.context = context
         self.pm = PassManager("builtin.module", self.context)
+        self.passes = passes
         add_bundle(self.pm, passes)
 
     def apply(self, module: ir.Module) -> ir.Module:
@@ -157,6 +158,8 @@ def apply(self, module: ir.Module) -> ir.Module:
         self.pm.run(module.operation)
         return module
 
+    def __str__(self) -> str:
+        return f"PassStage({[str(p) for p in self.passes]})"
 
 class TransformStage(Stage):
     """
@@ -218,3 +221,6 @@ def apply(self, module: ir.Module) -> ir.Module:
             raise ValueError("Missing module to apply transformations to.")
         self.schedule.apply(module.operation)
         return module
+
+    def __str__(self) -> str:
+        return f"TransformStage({self.module})"

tools/lh-run

@@ -111,6 +111,11 @@ if __name__ == "__main__":
         default=0,
         help="Print the Nth tensor. Default is 0 (no print).",
     )
+    Parser.add_argument(
+        "--print-mlir-after-all",
+        action=argparse.BooleanOptionalAction,
+        help="Whether to print the MLIR module after all stages. Default is False.",
+    )
     args = Parser.parse_args()
 
     # Initialize the random seed, for stable tests
@@ -125,6 +130,7 @@ if __name__ == "__main__":
 
     if args.benchmark:
         # Calling the benchmark wrapper, not the entry point.
+        # FIXME: Eliminate this cross-dependency between the Runner and the Driver.
         with driver.context:
             lh_dialects.register_and_load()
             bench_wrapper = Runner.get_bench_wrapper_schedule(args.entry_point)
@@ -140,7 +146,7 @@ if __name__ == "__main__":
     driver.add_stages(add_lower_to_llvm_stages())
 
     # Run the pipeline to get the optimized module.
-    optimized_module = driver.run()
+    optimized_module = driver.run(print_after_all=args.print_mlir_after_all)
     if args.print_optimized_module:
         print(optimized_module)