Add optional arguments to passmanger ir printing

Author: Vrajang Parikh

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

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) {

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) {

mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi

@@ -4,6 +4,7 @@
 #   * Relative imports for cross-module references.
 #   * Add __all__
 
+from pickle import TRUE
 from typing import Any, Optional
 
 from . import ir as _ir
@@ -16,7 +17,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: ...