[MLIR][Python] Add optional arguments to PassManager IR printing #89301

vraspar · 2024-04-18T20:26:38Z

This MR adds support for more granular IR printing for python pass manager by exposing existing parameters Pass Manager available through the C++ API.

Specifically, it adds following arguments:

print_before_pass
print_after_pass
print_after_only_on_change
print_after_only_on_failure

github-actions · 2024-04-18T20:26:54Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-04-18T20:27:26Z

@llvm/pr-subscribers-mlir

Author: None (vraspar)

Changes

This MR adds support for more granular ir printing for python pass manager by exposing existing parameters Pass Manager CPP API.

Specifically, it adds followng arguments:

print_before_pass
print_after_pass
print_after_only_on_change
print_after_only_on_failure

Full diff: https://github.com/llvm/llvm-project/pull/89301.diff

4 Files Affected:

(modified) mlir/include/mlir-c/Pass.h (+9-1)
(modified) mlir/lib/Bindings/Python/Pass.cpp (+19-2)
(modified) mlir/lib/CAPI/IR/Pass.cpp (+17-2)
(modified) mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi (+9-1)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 35db138305d1e2..47edaecd5ed8d8 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -76,7 +76,15 @@ mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 
 /// Enable mlir-print-ir-after-all.
 MLIR_CAPI_EXPORTED void
-mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
+mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
+                                bool shouldPrintBeforePass,
+                                bool shouldPrintAfterPass,
+                                bool printAfterOnlyOnChange,
+                                bool printAfterOnlyOnFailure);
+
+// Enable timing of passes
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableTiming(MlirPassManager passManager);
 
 /// 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..60694fa00782a3 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -74,10 +74,27 @@ 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 print_before_pass,
+            bool print_after_pass,
+            bool print_after_only_on_change,
+            bool print_after_only_on_failure) {
+            mlirPassManagerEnableIRPrinting(passManager.get(),
+                                            print_before_pass,
+                                            print_after_pass,
+                                            print_after_only_on_change,
+                                            print_after_only_on_failure);
           },
+          "print_before_pass"_a = true, "print_after_pass"_a = true, 
+          "print_after_only_on_change"_a = true, "print_after_only_on_failure"_a = false,
           "Enable mlir-print-ir-after-all.")
+      .def(
+          "enable_timing",
+          [](PyPassManager &PyPassManager) {
+            mlirPassManagerEnableTiming(PyPassManager.get());
+          },
+          "Enable timing of passes"
+      )
       .def(
           "enable_verifier",
           [](PyPassManager &passManager, bool enable) {
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index d242baae99c086..a499b1cdbf501f 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -44,8 +44,23 @@ MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
   return wrap(unwrap(passManager)->run(unwrap(op)));
 }
 
-void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
-  return unwrap(passManager)->enableIRPrinting();
+void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
+                                bool shouldPrintBeforePass,
+                                bool shouldPrintAfterPass,
+                                bool printAfterOnlyOnChange,
+                                bool printAfterOnlyOnFailure) {
+  auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass *, Operation *) {return shouldPrintBeforePass;};
+  auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass *, Operation *) {return shouldPrintAfterPass;};
+ 
+  return unwrap(passManager)->enableIRPrinting(shouldPrintBeforeFn,
+                                                shouldPrintAfterFn,
+                                                true,
+                                                printAfterOnlyOnChange,
+                                                printAfterOnlyOnFailure); 
+}
+
+void mlirPassManagerEnableTiming(MlirPassManager passManager) {
+  return unwrap(passManager)->enableTiming();
 }
 
 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..1c0e7d4d0f6151 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -4,6 +4,7 @@
 #   * Relative imports for cross-module references.
 #   * Add __all__
 
+from pickle import TRUE
 from typing import Any, Optional
 
 from . import ir as _ir
@@ -16,7 +17,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_pass=True,
+        print_after_pass=True,
+        print_after_only_on_change=True,
+        print_after_only_on_failure=False,
+    ) -> None: ...
+    def enable_timing(self) -> None: ...
     def enable_verifier(self, enable: bool) -> None: ...
     @staticmethod
     def parse(pipeline: str, context: Optional[_ir.Context] = None) -> PassManager: ...

mlir/lib/CAPI/IR/Pass.cpp

mlir/include/mlir-c/Pass.h

mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi

github-actions · 2024-04-18T21:04:59Z

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff 4759890f859277cd798648a9a333573cd088d98a dd062a077de06018cd765526696f6e7200b66c10 -- mlir/include/mlir-c/Pass.h mlir/lib/Bindings/Python/Pass.cpp mlir/lib/CAPI/IR/Pass.cpp

View the diff from clang-format here.

diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index 6d63b6519bd..fdee1077526 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -50,8 +50,12 @@ void mlirPassManagerEnableIRPrinting(MlirPassManager passManager,
                                      bool printModuleScope,
                                      bool printAfterOnlyOnChange,
                                      bool printAfterOnlyOnFailure) {
-  auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass *, Operation *) {return shouldPrintBeforePass;};
-  auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass *, Operation *) {return shouldPrintAfterPass;};
+  auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass *, Operation *) {
+    return shouldPrintBeforePass;
+  };
+  auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass *, Operation *) {
+    return shouldPrintAfterPass;
+  };
 
   return unwrap(passManager)
       ->enableIRPrinting(shouldPrintBeforeFn, shouldPrintAfterFn,

makslevental · 2024-04-18T21:07:41Z

Overall this is cool - just those small things and I'll approve (and then we wait for 2 days for Windows pre-commit CI workers lol). But I'll keep hitting "approve and run" for you.

ftynse

I wonder if this would be the right time to introduce a PassManager::PrintConfig struct, or even a more general PassManagerConfig, that would take care of all these API flags. This way, we won't break C API unnecessarily and sparse C and C++ clients the pain on passing a list of mostly useless options.

mlir/lib/Bindings/Python/Pass.cpp

ftynse · 2024-04-22T07:59:20Z

mlir/lib/CAPI/IR/Pass.cpp

+  auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass *, Operation *) {return shouldPrintBeforePass;};
+  auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass *, Operation *) {return shouldPrintAfterPass;};


Can't we actually take a C function pointer instead? On the Python side, we can then do a dispatch based on the actual type and accept bool as well as a py::function.

Hi! I am so sorry for the delay.

I am new to the MLIR world, so my understanding might be incorrect. But, I believe useful implementation from py::function to std::function<bool(Pass *, Operation *)> will need additional wrappers around the Pass and Operation class. Should we add these changes to this pull request?

Thank you so much!

Wrappers are already available. This can take a bool (*shouldPrintAfterPass)(MlirPass, MlirOperation, void *userData). Python binding side can define two overloads: one that takes a bool and another that takes a std::function<bool (MlirPass, MlirOperation)>, both using lambdas to delegate further. I believe we have type casters in place for this to work pretty much out of the box.

mlir/lib/Bindings/Python/Pass.cpp

vraspar · 2024-04-22T16:28:04Z

The implementation of PassManager::PrintConfig sounds like a great idea. I am planning of adding more options for ir printing from IRPrinterConfing in future.

Andrew-Cain · 2024-05-22T04:09:02Z

Hi @vraspar, just checking in to see if you're still working on this PR. It's been a month since the last update. Enabling the additional IR printing functionality through the Python extension is something we care about. Let me know if you need any help.

vraspar · 2024-05-26T23:51:55Z

Hi @vraspar, just checking in to see if you're still working on this PR. It's been a month since the last update. Enabling the additional IR printing functionality through the Python extension is something we care about. Let me know if you need any help.

Hi @Andrew-Cain. I am so sorry, I just saw this message. I was on vacation since beginnning of May and did not bring my laptop. I will implement suggested changes and update the PR. Thank you!

makslevental · 2024-05-27T01:54:32Z

vacation since beginnning of May and did not bring my laptop

man what a great vacation!

vraspar requested review from ftynse, makslevental and stellaraccident as code owners April 18, 2024 20:26

llvmbot added mlir:python mlir labels Apr 18, 2024

makslevental reviewed Apr 18, 2024

View reviewed changes

mlir/lib/CAPI/IR/Pass.cpp Outdated Show resolved Hide resolved

makslevental reviewed Apr 18, 2024

View reviewed changes

mlir/include/mlir-c/Pass.h Outdated Show resolved Hide resolved

makslevental reviewed Apr 18, 2024

View reviewed changes

mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi Outdated Show resolved Hide resolved

vraspar force-pushed the vrajangp/pass-manager branch from caff7d4 to 43a4862 Compare April 18, 2024 21:02

joker-eph changed the title ~~Add optional arguments to passmanger ir printing~~ [MLIR][Python] Add optional arguments to PassManager IR printing Apr 18, 2024

Add optional arguments to passmanger ir printing

Unverified

The committer email address is not verified.

Learn about vigilant mode

dd062a0

vraspar force-pushed the vrajangp/pass-manager branch from 43a4862 to dd062a0 Compare April 21, 2024 23:47

ftynse requested changes Apr 22, 2024

View reviewed changes

vraspar added 2 commits May 27, 2024 19:44

Merge branch 'llvm:main' into vrajangp/pass-manager

0613984

Refactor enable_ir_printing function in Python Pass.cpp

f915c47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][Python] Add optional arguments to PassManager IR printing #89301

[MLIR][Python] Add optional arguments to PassManager IR printing #89301

vraspar commented Apr 18, 2024 •

edited by joker-eph

Loading

github-actions bot commented Apr 18, 2024

llvmbot commented Apr 18, 2024

github-actions bot commented Apr 18, 2024 •

edited

Loading

makslevental commented Apr 18, 2024 •

edited

Loading

ftynse left a comment

ftynse Apr 22, 2024

vraspar May 31, 2024

ftynse Jun 7, 2024

vraspar commented Apr 22, 2024

Andrew-Cain commented May 22, 2024

vraspar commented May 26, 2024

makslevental commented May 27, 2024

		auto shouldPrintBeforeFn = [shouldPrintBeforePass](Pass , Operation ) {return shouldPrintBeforePass;};
		auto shouldPrintAfterFn = [shouldPrintAfterPass](Pass , Operation ) {return shouldPrintAfterPass;};

[MLIR][Python] Add optional arguments to PassManager IR printing #89301

Are you sure you want to change the base?

[MLIR][Python] Add optional arguments to PassManager IR printing #89301

Conversation

vraspar commented Apr 18, 2024 • edited by joker-eph Loading

github-actions bot commented Apr 18, 2024

llvmbot commented Apr 18, 2024

github-actions bot commented Apr 18, 2024 • edited Loading

makslevental commented Apr 18, 2024 • edited Loading

ftynse left a comment

Choose a reason for hiding this comment

ftynse Apr 22, 2024

Choose a reason for hiding this comment

vraspar May 31, 2024

Choose a reason for hiding this comment

ftynse Jun 7, 2024

Choose a reason for hiding this comment

vraspar commented Apr 22, 2024

Andrew-Cain commented May 22, 2024

vraspar commented May 26, 2024

makslevental commented May 27, 2024

vraspar commented Apr 18, 2024 •

edited by joker-eph

Loading

github-actions bot commented Apr 18, 2024 •

edited

Loading

makslevental commented Apr 18, 2024 •

edited

Loading