[MLIR] [Python] Add capsule methods for pybind11 to PyOperation

[jdd]

Summary

Add the getCapsule() and createFromCapsule() methods to the PyOperation class.

[jdd]

@stellaraccident I really don't understand what I've done here, but it works... not that that means much when it comes to memory management. Am I doing it correctly or is there something I'm missing about Operation ownership semantics and interactions with PyObjectRef?

[stellaraccident]

Thanks! You may not have known why, but this is the right way to do it. The unit test has a leak test which verifies that everything gets freed. The factory functions you are using maintain 1:1 correspondence between context and operation objects and the backing capi instances. Let me know if you run in to further issues: ownership of operations is complicated and this stuff doesn't have a lot of mileage on it.

[jdd]

Thanks! I think my success is more a testament to the ownership management infra already there. Landing this now.

mlir/lib/Bindings/Python/IRCore.cpp

 
 PyBlock PyOperation::getBlock() {
   PyOperationRef parentOperation = getParentOperation();
   MlirBlock block = mlirOperationGetBlock(get());
   assert(!mlirBlockIsNull(block) && "Attached operation has null parent");
   return PyBlock{std::move(parentOperation), block};
 }
 
+py::object PyOperation::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonOperationToCapsule(get()));
+}
+
+py::object PyOperation::createFromCapsule(py::object capsule) {
+  MlirOperation rawOperation = mlirPythonCapsuleToOperation(capsule.ptr());
+  if (mlirOperationIsNull(rawOperation))
+    throw py::error_already_set();
+  MlirContext rawCtxt = mlirOperationGetContext(rawOperation);
+  return forOperation(PyMlirContext::forContext(rawCtxt), rawOperation)
+      .releaseObject();
+}
+
 py::object PyOperation::create(
     std::string name, llvm::Optional<std::vector<PyType *>> results,
     llvm::Optional<std::vector<PyValue *>> operands,
     llvm::Optional<py::dict> attributes,
     llvm::Optional<std::vector<PyBlock *>> successors, int regions,
     DefaultingPyLocation location, py::object maybeIp) {
   llvm::SmallVector<MlirValue, 4> mlirOperands;
   llvm::SmallVector<MlirType, 4> mlirResults;
void mlir::python::populateIRCore(py::module &m) {
   py::class_<PyOperation, PyOperationBase>(m, "Operation")
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
                   kOperationCreateDocstring)
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
+                             &PyOperation::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyOperation::createFromCapsule)
       .def_property_readonly("name",
                              [](PyOperation &self) {
                                MlirOperation operation = self.get();
                                MlirStringRef name = mlirIdentifierStr(
                                    mlirOperationGetName(operation));
                                return py::str(name.data, name.length);
                              })
       .def_property_readonly(

mlir/lib/Bindings/Python/IRModule.h

public:
   /// Gets the owning block or raises an exception if the operation has no
   /// owning block.
   PyBlock getBlock();
 
   /// Gets the parent operation or raises an exception if the operation has
   /// no parent.
   PyOperationRef getParentOperation();
 
+  /// Gets a capsule wrapping the void* within the MlirOperation.
+  pybind11::object getCapsule();
+
+  /// Creates a PyOperation from the MlirOperation wrapped by a capsule.
+  /// Ownership of the underlying MlirOperation is taken by calling this
+  /// function.
+  static pybind11::object createFromCapsule(pybind11::object capsule);
+
   /// Creates an operation. See corresponding python docstring.
   static pybind11::object
   create(std::string name, llvm::Optional<std::vector<PyType *>> results,
          llvm::Optional<std::vector<PyValue *>> operands,
          llvm::Optional<pybind11::dict> attributes,
          llvm::Optional<std::vector<PyBlock *>> successors, int regions,
          DefaultingPyLocation location, pybind11::object ip);
 

mlir/test/Bindings/Python/ir_operation.py

def testOperationName():
 
   # CHECK: custom.op1
   # CHECK: custom.op2
   # CHECK: custom.op1
   for op in module.body.operations:
     print(op.operation.name)
 
 run(testOperationName)
+
+# CHECK-LABEL: TEST: testCapsuleConversions
+def testCapsuleConversions():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    m = Operation.create("custom.op1").operation
+    m_capsule = m._CAPIPtr
+    assert '"mlir.ir.Operation._CAPIPtr"' in repr(m_capsule)
+    m2 = Operation._CAPICreate(m_capsule)
+    assert m2 is m
+
+run(testCapsuleConversions)