[Mlir-commits] [mlir] Add optional arguments to passmanger ir printing (PR #89301)

[llvmlistbot at llvm.org]

https://github.com/vraspar updated https://github.com/llvm/llvm-project/pull/89301

>From 43a486292501d59bed599310090681fdc1d95300 Mon Sep 17 00:00:00 2001
From: Vrajang Parikh <vrajangp at d-matrix.ai>
Date: Mon, 8 Apr 2024 21:07:00 +0000
Subject: [PATCH] Add optional arguments to passmanger ir printing

---
 mlir/include/mlir-c/Pass.h                    | 10 ++++++++-
 mlir/lib/Bindings/Python/Pass.cpp             | 21 +++++++++++++++++--
 mlir/lib/CAPI/IR/Pass.cpp                     | 19 +++++++++++++++--
 .../mlir/_mlir_libs/_mlir/passmanager.pyi     |  9 +++++++-
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 35db138305d1e2..47edaecd5ed8d8 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -76,7 +76,15 @@ mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 
 /// Enable mlir-print-ir-after-all.
 MLIR_CAPI_EXPORTED void
-mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
+mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
+                                bool shouldPrintBeforePass,
+                                bool shouldPrintAfterPass,
+                                bool printAfterOnlyOnChange,
+                                bool printAfterOnlyOnFailure);
+
+// Enable timing of passes
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableTiming(MlirPassManager passManager);
 
 /// Enable / disable verify-each.
 MLIR_CAPI_EXPORTED void
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index a68421b61641f6..60694fa00782a3 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -74,10 +74,27 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
            "Releases (leaks) the backing pass manager (testing)")
       .def(
           "enable_ir_printing",
-          [](PyPassManager &passManager) {
-            mlirPassManagerEnableIRPrinting(passManager.get());
+          [](PyPassManager &passManager,
+            bool print_before_pass,
+            bool print_after_pass,
+            bool print_after_only_on_change,
+            bool print_after_only_on_failure) {
+            mlirPassManagerEnableIRPrinting(passManager.get(),
+                                            print_before_pass,
+                                            print_after_pass,
+                                            print_after_only_on_change,
+                                            print_after_only_on_failure);
           },
+          "print_before_pass"_a = true, "print_after_pass"_a = true, 
+          "print_after_only_on_change"_a = true, "print_after_only_on_failure"_a = false,
           "Enable mlir-print-ir-after-all.")
+      .def(
+          "enable_timing",
+          [](PyPassManager &PyPassManager) {
+            mlirPassManagerEnableTiming(PyPassManager.get());
+          },
+          "Enable timing of passes"
+      )
       .def(
           "enable_verifier",
           [](PyPassManager &passManager, bool enable) {
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index d242baae99c086..a499b1cdbf501f 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -44,8 +44,23 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
   return wrap(unwrap(passManager)->run(unwrap(op)));
 }
 
-void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
-  return unwrap(passManager)->enableIRPrinting();
+void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
+                                bool shouldPrintBeforePass,
+                                bool shouldPrintAfterPass,
+                                bool printAfterOnlyOnChange,
+                                bool printAfterOnlyOnFailure) {
+  auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass *, Operation *) {return shouldPrintBeforePass;};
+  auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass *, Operation *) {return shouldPrintAfterPass;};
+ 
+  return unwrap(passManager)->enableIRPrinting(shouldPrintBeforeFn,
+                                                shouldPrintAfterFn,
+                                                true,
+                                                printAfterOnlyOnChange,
+                                                printAfterOnlyOnFailure); 
+}
+
+void mlirPassManagerEnableTiming(MlirPassManager passManager) {
+  return unwrap(passManager)->enableTiming();
 }
 
 void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
index c072d5e0fb86f3..9f049c1135bbdf 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -16,7 +16,14 @@ class PassManager:
     def __init__(self, context: Optional[_ir.Context] = None) -> None: ...
     def _CAPICreate(self) -> object: ...
     def _testing_release(self) -> None: ...
-    def enable_ir_printing(self) -> None: ...
+    def enable_ir_printing(
+        self,
+        print_before_pass=True,
+        print_after_pass=True,
+        print_after_only_on_change=True,
+        print_after_only_on_failure=False,
+    ) -> None: ...
+    def enable_timing(self) -> None: ...
     def enable_verifier(self, enable: bool) -> None: ...
     @staticmethod
     def parse(pipeline: str, context: Optional[_ir.Context] = None) -> PassManager: ...