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

[llvmlistbot at llvm.org]

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

When using the `enable_ir_printing` API from Python, it invokes IR printing with default args, printing the IR before each pass and printing IR after pass only if there have been changes. This PR attempts to align the `enable_ir_printing` API with the documentation

* default arg of `enableIRPrinting`
https://github.com/llvm/llvm-project/blob/093b8bfe6b64c916647ae64af8066df22bb6ea65/mlir/include/mlir/Pass/PassManager.h#L380-L387

* PyPassManager
https://github.com/llvm/llvm-project/blob/77f1b481b884621c12cde5f2ce6f080f11dbbbcc/mlir/lib/Bindings/Python/Pass.cpp#L75-L80

>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] 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() {