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

[llvmlistbot at llvm.org]

================
@@ -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);
----------------
PragmaTwice wrote:

It seems that, `nb::steal` will construct a temporary `nb::object` without increasing the ref count, but when the object is destructed (`object::~object`), `dec_ref` will always be called, so here the ref count of `userData` will be decreased?

Maybe here we can just use `nb::handle` to avoid any `inc_ref` and `dec_ref`?

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