[MLIR][Python] Support Python-defined passes in MLIR (#156000)

It closes #155996.

This PR added a method `add(callable, ..)` to
`mlir.passmanager.PassManager` to accept a callable object for defining
passes in the Python side.

This is a simple example of a Python-defined pass.
```python
from mlir.passmanager import PassManager

def demo_pass_1(op):
    # do something with op
    pass

class DemoPass:
    def __init__(self, ...):
        pass
    def __call__(op):
        # do something
        pass

demo_pass_2 = DemoPass(..)

pm = PassManager('any', ctx)
pm.add(demo_pass_1)
pm.add(demo_pass_2)
pm.add("registered-passes")
pm.run(..)
```

---------

Co-authored-by: cnb.bsD2OPwAgEA <[email protected]>
Co-authored-by: Maksim Levental <[email protected]>

Author: 3 people

mlir/lib/Bindings/Python/MainModule.cpp

@@ -136,7 +136,7 @@ NB_MODULE(_mlir, m) {
   populateRewriteSubmodule(rewriteModule);
 
   // Define and populate PassManager submodule.
-  auto passModule =
+  auto passManagerModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
-  populatePassManagerSubmodule(passModule);
+  populatePassManagerSubmodule(passManagerModule);
 }

mlir/lib/Bindings/Python/Pass.cpp

@@ -10,8 +10,10 @@
 
 #include "IRModule.h"
 #include "mlir-c/Pass.h"
+// clang-format off
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
+// clang-format on
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -157,6 +159,45 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
           "pipeline"_a,
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
+      .def(
+          "add",
+          [](PyPassManager &passManager, const nb::callable &run,
+             std::optional<std::string> &name, const std::string &argument,
+             const std::string &description, const std::string &opName) {
+            if (!name.has_value()) {
+              name = nb::cast<std::string>(
+                  nb::borrow<nb::str>(run.attr("__name__")));
+            }
+            MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
+            MlirTypeID passID =
+                mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+            MlirExternalPassCallbacks callbacks;
+            callbacks.construct = [](void *obj) {
+              (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();
+            };
+            callbacks.destruct = [](void *obj) {
+              (void)nb::handle(static_cast<PyObject *>(obj)).dec_ref();
+            };
+            callbacks.initialize = nullptr;
+            callbacks.clone = [](void *) -> void * {
+              throw std::runtime_error("Cloning Python passes not supported");
+            };
+            callbacks.run = [](MlirOperation op, MlirExternalPass,
+                               void *userData) {
+              nb::borrow<nb::callable>(static_cast<PyObject *>(userData))(op);
+            };
+            auto externalPass = mlirCreateExternalPass(
+                passID, mlirStringRefCreate(name->data(), name->length()),
+                mlirStringRefCreate(argument.data(), argument.length()),
+                mlirStringRefCreate(description.data(), description.length()),
+                mlirStringRefCreate(opName.data(), opName.size()),
+                /*nDependentDialects*/ 0, /*dependentDialects*/ nullptr,
+                callbacks, /*userData*/ run.ptr());
+            mlirPassManagerAddOwnedPass(passManager.get(), externalPass);
+          },
+          "run"_a, "name"_a.none() = nb::none(), "argument"_a.none() = "",
+          "description"_a.none() = "", "op_name"_a.none() = "",
+          "Add a python-defined pass to the pass manager.")
       .def(
           "run",
           [](PyPassManager &passManager, PyOperationBase &op) {

mlir/lib/CAPI/IR/Pass.cpp

@@ -145,10 +145,14 @@ class ExternalPass : public Pass {
       : Pass(passID, opName), id(passID), name(name), argument(argument),
         description(description), dependentDialects(dependentDialects),
         callbacks(callbacks), userData(userData) {
-    callbacks.construct(userData);
+    if (callbacks.construct)
+      callbacks.construct(userData);
   }
 
-  ~ExternalPass() override { callbacks.destruct(userData); }
+  ~ExternalPass() override {
+    if (callbacks.destruct)
+      callbacks.destruct(userData);
+  }
 
   StringRef getName() const override { return name; }
   StringRef getArgument() const override { return argument; }