[AutoBump] Merge with fixes of f8eceb45 (Sep 17, needs torch bump) (6)

mlir/include/mlir-c/Pass.h

@@ -74,9 +74,11 @@ mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
 MLIR_CAPI_EXPORTED MlirLogicalResult
 mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 
-/// Enable mlir-print-ir-after-all.
-MLIR_CAPI_EXPORTED void
-mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
+/// Enable IR printing.
+MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
+    MlirPassManager passManager, bool printBeforeAll, bool printAfterAll,
+    bool printModuleScope, bool printAfterOnlyOnChange,
+    bool printAfterOnlyOnFailure);
 
 /// Enable lir-reproducer-before-all.
 MLIR_CAPI_EXPORTED void

mlir/lib/Bindings/Python/Pass.cpp

@@ -74,10 +74,17 @@ 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 printBeforeAll,
+             bool printAfterAll, bool printModuleScope, bool printAfterChange,
+             bool printAfterFailure) {
+            mlirPassManagerEnableIRPrinting(
+                passManager.get(), printBeforeAll, printAfterAll,
+                printModuleScope, printAfterChange, printAfterFailure);
           },
-          "Enable mlir-print-ir-after-all.")
+          "print_before_all"_a = false, "print_after_all"_a = true,
+          "print_module_scope"_a = false, "print_after_change"_a = false,
+          "print_after_failure"_a = false,
+          "Enable IR printing, default as mlir-print-ir-after-all.")
       .def(
           "enable_reproducer_before_all",
           [](PyPassManager &passManager, const std::string &outputDir) {

mlir/lib/CAPI/IR/Pass.cpp

@@ -44,8 +44,21 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
   return wrap(unwrap(passManager)->run(unwrap(op)));
 }
 
-void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
-  return unwrap(passManager)->enableIRPrinting();
+void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
+                                     bool printBeforeAll, bool printAfterAll,
+                                     bool printModuleScope,
+                                     bool printAfterOnlyOnChange,
+                                     bool printAfterOnlyOnFailure) {
+  auto shouldPrintBeforePass = [printBeforeAll](Pass *, Operation *) {
+    return printBeforeAll;
+  };
+  auto shouldPrintAfterPass = [printAfterAll](Pass *, Operation *) {
+    return printAfterAll;
+  };
+  return unwrap(passManager)
+      ->enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
+                         printModuleScope, printAfterOnlyOnChange,
+                         printAfterOnlyOnFailure);
 }
 
 void mlirPassManagerEnableReproducerBeforeAll(MlirPassManager passManager,

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_all: bool = False,
+        print_after_all: bool = True,
+        print_module_scope: bool = False,
+        print_after_change: bool = False,
+        print_after_failure: bool = False,
+    ) -> None: ...
     def enable_reproducer_before_all(self, output_dir: str) -> None: ...
     def enable_verifier(self, enable: bool) -> None: ...
     @staticmethod

mlir/test/python/pass_manager.py

@@ -300,14 +300,40 @@ def testPrintIrAfterAll():
         pm = PassManager.parse("builtin.module(canonicalize)")
         ctx.enable_multithreading(False)
         pm.enable_ir_printing()
-        # CHECK: // -----// IR Dump Before Canonicalizer (canonicalize) ('builtin.module' operation) //----- //
+        # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) //----- //
+        # CHECK: module {
+        # CHECK:   func.func @main() {
+        # CHECK:     return
+        # CHECK:   }
+        # CHECK: }
+        pm.run(module)
+
+
+# CHECK-LABEL: TEST: testPrintIrBeforeAndAfterAll
+@run
+def testPrintIrBeforeAndAfterAll():
+    with Context() as ctx:
+        module = ModuleOp.parse(
+            """
+          module {
+            func.func @main() {
+              %0 = arith.constant 10
+              return
+            }
+          }
+        """
+        )
+        pm = PassManager.parse("builtin.module(canonicalize)")
+        ctx.enable_multithreading(False)
+        pm.enable_ir_printing(print_before_all=True, print_after_all=True)
+        # CHECK: // -----// IR Dump Before Canonicalizer (canonicalize) //----- //
         # CHECK: module {
         # CHECK:   func.func @main() {
         # CHECK:     %[[C10:.*]] = arith.constant 10 : i64
         # CHECK:     return
         # CHECK:   }
         # CHECK: }
-        # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) ('builtin.module' operation) //----- //
+        # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) //----- //
         # CHECK: module {
         # CHECK:   func.func @main() {
         # CHECK:     return