[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());
----------------
PragmaTwice wrote:

Yeah `clone` is quite tricky since in the python side I'm not sure how to call it : )

>From the code, it seems that we first call `callback.clone`, and then call `callback.construct` in constructing an `ExternalPass`, so maybe here the pointer `dst.ptr()` will be dangling since the `nb::callable` is destructed (`callable::~callable` is called) and the ref_count become 0?

I think it is related to whether we `inc_ref` in the `callback.construct`. Maybe we can have such a design:
- in `callback.construct` we do not call `inc_ref`, it leaves to who pass the object into the `ExternalPass`;
- in `callback.clone`, we first call `inc_ref` and then pass the object (call `.release()` to prevent the `nb::object::~object` from decreasing the ref count), so that the ref count of this object can be 1, and in `callback.destruct` it is decreased to 0 and destructed;
- before the call to `mlirCreateExternalPass`, we first call `inc_ref` for the object to increase the ref count so that it should be valid at least before `callback.destruct`.

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