[Mlir-commits] [mlir] [MLIR] [Python] align python ir printing with mlir-print-ir-after-all (PR #107522)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Sep 17 19:06:09 PDT 2024


https://github.com/xurui1995 updated https://github.com/llvm/llvm-project/pull/107522

>From b6ad0449e0ccb546be5b55b3b9abcdf427d39602 Mon Sep 17 00:00:00 2001
From: "Xu, Rui" <rui.xu at intel.com>
Date: Thu, 5 Sep 2024 22:15:36 -0700
Subject: [PATCH 1/5] align python ir printing with mlir-print-ir-after-all

---
 mlir/lib/CAPI/IR/Pass.cpp        | 7 ++++++-
 mlir/test/python/pass_manager.py | 7 -------
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index d242baae99c086..4d49e2ad949756 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -45,7 +45,12 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
 }
 
 void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
-  return unwrap(passManager)->enableIRPrinting();
+  auto shouldPrintBeforePass = [](Pass *, Operation *) { return false; };
+  auto shouldPrintAfterPass = [](Pass *, Operation *) { return true; };
+  bool printModuleScope = true, printAfterOnlyOnChange = false;
+  return unwrap(passManager)
+      ->enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
+                         printModuleScope, printAfterOnlyOnChange);
 }
 
 void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index 43af80b53166cc..fc7225f920f730 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -300,13 +300,6 @@ 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: 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: module {
         # CHECK:   func.func @main() {

>From ca245fe9f6bf8d1855c3ccd7cc802e8c2a27cb02 Mon Sep 17 00:00:00 2001
From: "Xu, Rui" <rui.xu at intel.com>
Date: Sun, 8 Sep 2024 21:16:37 -0700
Subject: [PATCH 2/5] add options for enable_ir_printing

---
 mlir/include/mlir-c/Pass.h                    |  8 +++--
 mlir/lib/Bindings/Python/Pass.cpp             | 13 ++++++--
 mlir/lib/CAPI/IR/Pass.cpp                     | 18 +++++++---
 .../mlir/_mlir_libs/_mlir/passmanager.pyi     |  9 ++++-
 mlir/test/python/pass_manager.py              | 33 +++++++++++++++++++
 5 files changed, 69 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 35db138305d1e2..fc45a47bcf7ffc 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/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, default as mlir-print-ir-after-all.
+MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
+    MlirPassManager passManager, bool printBeforeAll = false,
+    bool printAfterAll = true, bool printModuleScope = false,
+    bool printAfterOnlyOnChange = false, bool printAfterOnlyOnFailure = false);
 
 /// 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..1d0e5ce2115a0a 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/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_verifier",
           [](PyPassManager &passManager, bool enable) {
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index 4d49e2ad949756..a6c9fbd08d45a6 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -44,13 +44,21 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
   return wrap(unwrap(passManager)->run(unwrap(op)));
 }
 
-void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
-  auto shouldPrintBeforePass = [](Pass *, Operation *) { return false; };
-  auto shouldPrintAfterPass = [](Pass *, Operation *) { return true; };
-  bool printModuleScope = true, printAfterOnlyOnChange = false;
+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);
+                         printModuleScope, printAfterOnlyOnChange,
+                         printAfterOnlyOnFailure);
 }
 
 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..5d115e8222d730 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_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_verifier(self, enable: bool) -> None: ...
     @staticmethod
     def parse(pipeline: str, context: Optional[_ir.Context] = None) -> PassManager: ...
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index fc7225f920f730..119cc8a0692fd9 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -307,3 +307,36 @@ def testPrintIrAfterAll():
         # CHECK:   }
         # CHECK: }
         pm.run(module)
+
+
+# CHECK-LABEL: TEST: testPrintIrBeforeAndAfterAll
+ at 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) ('builtin.module' operation) //----- //
+        # 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: module {
+        # CHECK:   func.func @main() {
+        # CHECK:     return
+        # CHECK:   }
+        # CHECK: }
+        pm.run(module)
\ No newline at end of file

>From be503dc9518973a468521e1101f31efcfb9527da Mon Sep 17 00:00:00 2001
From: "Xu, Rui" <rui.xu at intel.com>
Date: Sun, 8 Sep 2024 21:19:25 -0700
Subject: [PATCH 3/5] rm default value in CAPI

---
 mlir/include/mlir-c/Pass.h | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index fc45a47bcf7ffc..2218ec0f47d199 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -74,11 +74,11 @@ mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
 MLIR_CAPI_EXPORTED MlirLogicalResult
 mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 
-/// Enable IR printing, default as mlir-print-ir-after-all.
+/// Enable IR printing.
 MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
-    MlirPassManager passManager, bool printBeforeAll = false,
-    bool printAfterAll = true, bool printModuleScope = false,
-    bool printAfterOnlyOnChange = false, bool printAfterOnlyOnFailure = false);
+    MlirPassManager passManager, bool printBeforeAll, bool printAfterAll,
+    bool printModuleScope, bool printAfterOnlyOnChange,
+    bool printAfterOnlyOnFailure);
 
 /// Enable / disable verify-each.
 MLIR_CAPI_EXPORTED void

>From 0594942e75af55777d23d84788b0694d2b8bfd97 Mon Sep 17 00:00:00 2001
From: "Xu, Rui" <rui.xu at intel.com>
Date: Sun, 8 Sep 2024 22:28:08 -0700
Subject: [PATCH 4/5] fix ut

---
 mlir/test/python/pass_manager.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index 119cc8a0692fd9..fec5cea31c04c0 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -300,7 +300,7 @@ def testPrintIrAfterAll():
         pm = PassManager.parse("builtin.module(canonicalize)")
         ctx.enable_multithreading(False)
         pm.enable_ir_printing()
-        # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) ('builtin.module' operation) //----- //
+        # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) //----- //
         # CHECK: module {
         # CHECK:   func.func @main() {
         # CHECK:     return
@@ -326,14 +326,14 @@ def testPrintIrBeforeAndAfterAll():
         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) ('builtin.module' operation) //----- //
+        # 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

>From c7fa4c28c98ca6da07e7ff984a1f10fc1a874367 Mon Sep 17 00:00:00 2001
From: "Xu, Rui" <rui.xu at intel.com>
Date: Wed, 18 Sep 2024 10:02:29 +0800
Subject: [PATCH 5/5] add newline

---
 mlir/test/python/pass_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index fec5cea31c04c0..74967032562351 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -339,4 +339,4 @@ def testPrintIrBeforeAndAfterAll():
         # CHECK:     return
         # CHECK:   }
         # CHECK: }
-        pm.run(module)
\ No newline at end of file
+        pm.run(module)



More information about the Mlir-commits mailing list