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

[Yuanqiang Liu]

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

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

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

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