[Mlir-commits] [mlir] [MLIR][Python] enhance python ir printing with pringing flags (PR #117836)

[Yuanqiang Liu]

https://github.com/qingyunqu created https://github.com/llvm/llvm-project/pull/117836

Close https://github.com/llvm/llvm-project/pull/65854

>From 05aef3b09a0d8b149c7a093c8577372c31dc3c9e Mon Sep 17 00:00:00 2001
From: Yuanqiang Liu <liuyuanqiang.yqliu at bytedance.com>
Date: Wed, 27 Nov 2024 11:11:31 +0800
Subject: [PATCH] [MLIR][Python] enhance python ir printing with pringing flags

---
 mlir/include/mlir-c/Pass.h                      |  2 +-
 mlir/lib/Bindings/Python/Pass.cpp               | 17 +++++++++++++++--
 mlir/lib/CAPI/IR/Pass.cpp                       |  6 ++++--
 .../mlir/_mlir_libs/_mlir/passmanager.pyi       |  3 +++
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 2218ec0f47d199..f5d431b48f1c6a 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -78,7 +78,7 @@ mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
     MlirPassManager passManager, bool printBeforeAll, bool printAfterAll,
     bool printModuleScope, bool printAfterOnlyOnChange,
-    bool printAfterOnlyOnFailure);
+    bool printAfterOnlyOnFailure, MlirOpPrintingFlags flags);
 
 /// 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 1d0e5ce2115a0a..7f135031ab45ea 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -76,14 +76,27 @@ 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) {
+            MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
+            if (largeElementsLimit)
+              mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
+                                                         *largeElementsLimit);
+            if (enableDebugInfo)
+              mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
+                                                 /*prettyForm=*/false);
+            if (printGenericOpForm)
+              mlirOpPrintingFlagsPrintGenericOpForm(flags);
             mlirPassManagerEnableIRPrinting(
                 passManager.get(), printBeforeAll, printAfterAll,
-                printModuleScope, printAfterChange, printAfterFailure);
+                printModuleScope, printAfterChange, printAfterFailure, flags);
+            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,
           "Enable IR printing, default as mlir-print-ir-after-all.")
       .def(
           "enable_verifier",
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index a6c9fbd08d45a6..84da467215ddff 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -48,7 +48,8 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
                                      bool printBeforeAll, bool printAfterAll,
                                      bool printModuleScope,
                                      bool printAfterOnlyOnChange,
-                                     bool printAfterOnlyOnFailure) {
+                                     bool printAfterOnlyOnFailure,
+                                     MlirOpPrintingFlags flags) {
   auto shouldPrintBeforePass = [printBeforeAll](Pass *, Operation *) {
     return printBeforeAll;
   };
@@ -58,7 +59,8 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
   return unwrap(passManager)
       ->enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
                          printModuleScope, printAfterOnlyOnChange,
-                         printAfterOnlyOnFailure);
+                         printAfterOnlyOnFailure, /*out=*/llvm::errs(),
+                         *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..701de885a71544 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -22,6 +22,9 @@ 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,
     ) -> None: ...
     def enable_verifier(self, enable: bool) -> None: ...
     @staticmethod