Skip to content

[MLIR][Python] Expose PassManager::enableStatistics to CAPI and Python #162591

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Oct 9, 2025
+62 −0
Merged

[MLIR][Python] Expose PassManager::enableStatistics to CAPI and Python #162591

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
16b033d
[MLIR][Python] Expose PassManager::enableStatistics to CAPI and Python
PragmaTwice Oct 8, 2025
ce2b2ce
fix format
PragmaTwice Oct 9, 2025
aa6a57e
Merge branch 'main' into mlir-python-stat
PragmaTwice Oct 9, 2025
157bb40
fix clang warning
PragmaTwice Oct 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions mlir/include/mlir-c/Pass.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,18 @@ mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable);
MLIR_CAPI_EXPORTED void
mlirPassManagerEnableTiming(MlirPassManager passManager);

/// Enumerated type of pass display modes.
/// Mainly used in mlirPassManagerEnableStatistics.
typedef enum {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way to disable statistics?

Copy link
Member Author

@PragmaTwice PragmaTwice Oct 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the answer is no if we don't change the C++ API. (see https://mlir.llvm.org/doxygen/classmlir_1_1PassManager.html)

In C++, the only public API for statistics is like void enableStatistics(PassDisplayMode displayMode) and internally it is assigned to a std::optional<PassDisplayMode>, and the bool(optional) can indicate if it is enabled.

MLIR_PASS_DISPLAY_MODE_LIST,
MLIR_PASS_DISPLAY_MODE_PIPELINE,
} MlirPassDisplayMode;

/// Enable pass statistics.
MLIR_CAPI_EXPORTED void
mlirPassManagerEnableStatistics(MlirPassManager passManager,
MlirPassDisplayMode displayMode);

/// Nest an OpPassManager under the top-level PassManager, the nested
/// passmanager will only run on operations matching the provided name.
/// The returned OpPassManager will be destroyed when the parent is destroyed.
Expand Down
15 changes: 15 additions & 0 deletions mlir/lib/Bindings/Python/Pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,13 @@ class PyPassManager {

/// Create the `mlir.passmanager` here.
void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
//----------------------------------------------------------------------------
// Mapping of enumerated types
//----------------------------------------------------------------------------
nb::enum_<MlirPassDisplayMode>(m, "PassDisplayMode")
.value("LIST", MLIR_PASS_DISPLAY_MODE_LIST)
.value("PIPELINE", MLIR_PASS_DISPLAY_MODE_PIPELINE);

//----------------------------------------------------------------------------
// Mapping of MlirExternalPass
//----------------------------------------------------------------------------
Expand Down Expand Up @@ -138,6 +145,14 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
mlirPassManagerEnableTiming(passManager.get());
},
"Enable pass timing.")
.def(
"enable_statistics",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naming not: "enable statistics" calls for a boolean argument. I'd consider something like set pass display mode instead, or a pass display mode property if there is a way to fetch back the current value

Copy link
Member Author

@PragmaTwice PragmaTwice Oct 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It indeed looks good but I think currently we cannot do it before some modification to the C++ API. I'll try to look into the C++ APIs and see if we can change later : )

Currently these names are just C++ ones with underscores.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd also suggest to rename the C++ function. It may have been enable_statistics(bool) at some point...

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me take this task and see what we can do : )

[](PyPassManager &passManager, MlirPassDisplayMode displayMode) {
mlirPassManagerEnableStatistics(passManager.get(), displayMode);
},
"displayMode"_a =
MlirPassDisplayMode::MLIR_PASS_DISPLAY_MODE_PIPELINE,
"Enable pass statistics.")
.def_static(
"parse",
[](const std::string &pipeline, DefaultingPyMlirContext context) {
Expand Down
15 changes: 15 additions & 0 deletions mlir/lib/CAPI/IR/Pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include "mlir/CAPI/Support.h"
#include "mlir/CAPI/Utils.h"
#include "mlir/Pass/PassManager.h"
#include "llvm/Support/ErrorHandling.h"
#include <optional>

using namespace mlir;
Expand Down Expand Up @@ -79,6 +80,20 @@ void mlirPassManagerEnableTiming(MlirPassManager passManager) {
unwrap(passManager)->enableTiming();
}

void mlirPassManagerEnableStatistics(MlirPassManager passManager,
MlirPassDisplayMode displayMode) {
PassDisplayMode mode;
switch (displayMode) {
case MLIR_PASS_DISPLAY_MODE_LIST:
mode = PassDisplayMode::List;
break;
case MLIR_PASS_DISPLAY_MODE_PIPELINE:
mode = PassDisplayMode::Pipeline;
break;
}
unwrap(passManager)->enableStatistics(mode);
}

MlirOpPassManager mlirPassManagerGetNestedUnder(MlirPassManager passManager,
MlirStringRef operationName) {
return wrap(&unwrap(passManager)->nest(unwrap(operationName)));
Expand Down
20 changes: 20 additions & 0 deletions mlir/test/python/pass_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,3 +435,23 @@ def print_file_tree(directory, prefix=""):

print_file_tree(temp_dir)
log("// Tree printing end")


# CHECK-LABEL: TEST: testEnableStatistics
@run
def testEnableStatistics():
with Context() as ctx:
module = ModuleOp.parse(
"""
module {
func.func @main() {
%0 = arith.constant 10
return
}
}
"""
)
pm = PassManager.parse("builtin.module(canonicalize)")
pm.enable_statistics()
# CHECK: Pass statistics report
pm.run(module)