[Mlir-commits] [mlir] 2e51e15 - [MLIR][Python] enhance python ir printing with pringing flags (#117836)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Dec 5 01:31:09 PST 2024
Author: Yuanqiang Liu
Date: 2024-12-05T10:31:04+01:00
New Revision: 2e51e150e161bd5fb5b8adb8655744a672ced002
URL: https://github.com/llvm/llvm-project/commit/2e51e150e161bd5fb5b8adb8655744a672ced002
DIFF: https://github.com/llvm/llvm-project/commit/2e51e150e161bd5fb5b8adb8655744a672ced002.diff
LOG: [MLIR][Python] enhance python ir printing with pringing flags (#117836)
Close https://github.com/llvm/llvm-project/pull/65854
Added:
Modified:
mlir/include/mlir-c/Pass.h
mlir/lib/Bindings/Python/Pass.cpp
mlir/lib/CAPI/IR/Pass.cpp
mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
mlir/test/python/pass_manager.py
Removed:
################################################################################
diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 6019071cfdaa29..8fd8e9956a65a3 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -81,7 +81,8 @@ mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
MlirPassManager passManager, bool printBeforeAll, bool printAfterAll,
bool printModuleScope, bool printAfterOnlyOnChange,
- bool printAfterOnlyOnFailure, MlirStringRef treePrintingPath);
+ bool printAfterOnlyOnFailure, MlirOpPrintingFlags flags,
+ MlirStringRef treePrintingPath);
/// 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 e8d28abe6d583a..e991deaae2daa5 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -76,20 +76,33 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
"enable_ir_printing",
[](PyPassManager &passManager, bool printBeforeAll,
bool printAfterAll, bool printModuleScope, bool printAfterChange,
- bool printAfterFailure,
+ bool printAfterFailure, std::optional<int64_t> largeElementsLimit,
+ bool enableDebugInfo, bool printGenericOpForm,
std::optional<std::string> optionalTreePrintingPath) {
+ MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
+ if (largeElementsLimit)
+ mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
+ *largeElementsLimit);
+ if (enableDebugInfo)
+ mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
+ /*prettyForm=*/false);
+ if (printGenericOpForm)
+ mlirOpPrintingFlagsPrintGenericOpForm(flags);
std::string treePrintingPath = "";
if (optionalTreePrintingPath.has_value())
treePrintingPath = optionalTreePrintingPath.value();
mlirPassManagerEnableIRPrinting(
passManager.get(), printBeforeAll, printAfterAll,
- printModuleScope, printAfterChange, printAfterFailure,
+ printModuleScope, printAfterChange, printAfterFailure, flags,
mlirStringRefCreate(treePrintingPath.data(),
treePrintingPath.size()));
+ mlirOpPrintingFlagsDestroy(flags);
},
"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,
+ "large_elements_limit"_a = py::none(), "enable_debug_info"_a = false,
+ "print_generic_op_form"_a = false,
"tree_printing_dir_path"_a = py::none(),
"Enable IR printing, default as mlir-print-ir-after-all.")
.def(
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index 01151eafeb5268..883b7e8bb832d2 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -49,6 +49,7 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
bool printModuleScope,
bool printAfterOnlyOnChange,
bool printAfterOnlyOnFailure,
+ MlirOpPrintingFlags flags,
MlirStringRef treePrintingPath) {
auto shouldPrintBeforePass = [printBeforeAll](Pass *, Operation *) {
return printBeforeAll;
@@ -60,13 +61,14 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
return unwrap(passManager)
->enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
printModuleScope, printAfterOnlyOnChange,
- printAfterOnlyOnFailure);
+ printAfterOnlyOnFailure, /*out=*/llvm::errs(),
+ *unwrap(flags));
unwrap(passManager)
->enableIRPrintingToFileTree(shouldPrintBeforePass, shouldPrintAfterPass,
printModuleScope, printAfterOnlyOnChange,
printAfterOnlyOnFailure,
- unwrap(treePrintingPath));
+ unwrap(treePrintingPath), *unwrap(flags));
}
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 229979ae33608c..0d2eaffe16d3ec 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -22,6 +22,10 @@ class PassManager:
print_module_scope: bool = False,
print_after_change: bool = False,
print_after_failure: bool = False,
+ large_elements_limit: int | None = None,
+ enable_debug_info: bool = False,
+ print_generic_op_form: bool = False,
+ tree_printing_dir_path: str | None = None,
) -> None: ...
def enable_verifier(self, enable: bool) -> None: ...
@staticmethod
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index a794a3fc6fa006..ecac57e3302f01 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -342,6 +342,27 @@ def testPrintIrBeforeAndAfterAll():
pm.run(module)
+# CHECK-LABEL: TEST: testPrintIrLargeLimitElements
+ at run
+def testPrintIrLargeLimitElements():
+ with Context() as ctx:
+ module = ModuleOp.parse(
+ """
+ module {
+ func.func @main() -> tensor<3xi64> {
+ %0 = arith.constant dense<[1, 2, 3]> : tensor<3xi64>
+ return %0 : tensor<3xi64>
+ }
+ }
+ """
+ )
+ pm = PassManager.parse("builtin.module(canonicalize)")
+ ctx.enable_multithreading(False)
+ pm.enable_ir_printing(large_elements_limit=2)
+ # CHECK: %[[CST:.*]] = arith.constant dense_resource<__elided__> : tensor<3xi64>
+ pm.run(module)
+
+
# CHECK-LABEL: TEST: testPrintIrTree
@run
def testPrintIrTree():
More information about the Mlir-commits
mailing list