[mlir] Add Python bindings to enable default passmanager timing

[maerhart]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Martin Erhart (maerhart)

Changes

---

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

3 Files Affected:

• (modified) mlir/include/mlir-c/Pass.h (+4)
• (modified) mlir/lib/Bindings/Python/Pass.cpp (+6)
• (modified) mlir/lib/CAPI/IR/Pass.cpp (+4)

diff --git a/mlir/include/mlir-c/Pass.h b/mlir/include/mlir-c/Pass.h
index 8fd8e9956a65a..0d2e19ee7fb0a 100644
--- a/mlir/include/mlir-c/Pass.h
+++ b/mlir/include/mlir-c/Pass.h
@@ -88,6 +88,10 @@ MLIR_CAPI_EXPORTED void mlirPassManagerEnableIRPrinting(
 MLIR_CAPI_EXPORTED void
 mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable);
 
+/// Enable pass timing.
+MLIR_CAPI_EXPORTED void
+mlirPassManagerEnableTiming(MlirPassManager passManager);
+
 /// 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.
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 858c3bd5745fe..8d84864b9db4d 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -112,6 +112,12 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
             mlirPassManagerEnableVerifier(passManager.get(), enable);
           },
           "enable"_a, "Enable / disable verify-each.")
+      .def(
+          "enable_timing",
+          [](PyPassManager &passManager) {
+            mlirPassManagerEnableTiming(passManager.get());
+          },
+          "Enable pass timing.")
       .def_static(
           "parse",
           [](const std::string &pipeline, DefaultingPyMlirContext context) {
diff --git a/mlir/lib/CAPI/IR/Pass.cpp b/mlir/lib/CAPI/IR/Pass.cpp
index 883b7e8bb832d..3c499c3e4974d 100644
--- a/mlir/lib/CAPI/IR/Pass.cpp
+++ b/mlir/lib/CAPI/IR/Pass.cpp
@@ -75,6 +75,10 @@ void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
   unwrap(passManager)->enableVerifier(enable);
 }
 
+void mlirPassManagerEnableTiming(MlirPassManager passManager) {
+  unwrap(passManager)->enableTiming();
+}
+
 MlirOpPassManager mlirPassManagerGetNestedUnder(MlirPassManager passManager,
                                                 MlirStringRef operationName) {
   return wrap(&unwrap(passManager)->nest(unwrap(operationName)));

[makslevental]

Should we have disable too? Or is that unneccessary for some reason that I can't think of right now?

[maerhart]

Should we have disable too? Or is that unneccessary for some reason that I can't think of right now?

Enabling timing adds an instrumentation and AFAIK we don't even have a way to remove instrumentations in the C++ implementation of the passmanager, so it would be more than just a CAPI and Python bindings change.
Also, I never had the need to disable timing again after enabling it. Usually it's enough to just not enable it in the first place. However, if the need arises, we can always add it later on.