[Mlir-commits] [mlir] [MLIR][Python][NO MERGE] Support Python-defined passes in MLIR (PR #157369)

[Maksim Levental]

================
@@ -157,6 +159,44 @@ 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_python_pass",
+          [](PyPassManager &passManager, const std::string &name,
+             const std::string &argument, const std::string &description,
+             const std::string &opName, const nb::callable &run) {
+            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.clone = [](void *obj) {
+              auto src = nb::handle(static_cast<PyObject *>(obj));
+              nb::callable dst;
+              nb::inst_copy(dst, src);
+              return static_cast<void *>(dst.ptr());
----------------
makslevental wrote:

this is the only part i'm not 100% on: i'm not sure if `inst_copy` copies the ref count or what (and whether `construct` is called before or after). need to find a way to test this path.

https://github.com/llvm/llvm-project/pull/157369