[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());
+            };
+            callbacks.run = [](MlirOperation op, MlirExternalPass,
+                               void *userData) {
+              nb::steal<nb::callable>(static_cast<PyObject *>(userData))(op);
----------------
makslevental wrote:

Really? Where does it say the ref count will be decreased? Either way sure nb::handle works too I think. It's funny - I wish there were a good way to print ref counts for things so you could figure out if you are leaking references.

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