[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