[MLIR][Python][NO MERGE] Support Python-defined passes in MLIR

[makslevental]

based heavily on #156000

[makslevental]

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.

[PragmaTwice]

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.

[makslevental]

so maybe here the pointer dst.ptr() will be dangling since the nb::callable is destructed

Seems not possible since the tests pass? Are you sure that clone is called for our use?

[PragmaTwice]

The patch sounds good to me and definitely much cleaner and simpler than mine! I'll try to merge it into my PR. Just left a few minor comments to discuss some details.

[PragmaTwice]

If possible maybe we can make argument, description and opName keyword arguments? So that we can use it like pm.add_python_pass(name, run).

[makslevental]

Ya 100% that's a good idea (like you have in your PR). But I don't think you actually need to make them kw - you just have to make them have default args. Anyway either form is good with me.

[PragmaTwice]

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.