[MLIR][Python] Add capsule methods for pybind11 to PyValue.

[mikeurbach]

Summary

Add the getCapsule() and createFromCapsule() methods to the
PyValue class, as well as the necessary interoperability.

[mikeurbach]

This is basically https://reviews.llvm.org/D99927 for Values instead of Operations. I've been testing it downstream in CIRCT. We have use-cases where we'd want a Python API that accepts a Value, or returns a Value, and this seemed to be missing for that.

Another thing we'll want is a nice casting helper, like this: https://github.com/llvm/circt/pull/940/commits/69f527952aa5acb2c73cc3b74e836483b5a5a5d8. It seems we have been adding this support downstream when needed. Does that kind of thing belong in the downstream project, or could I add it in MLIR? I don't see much pybind11 casting functionality in MLIR, so I'm curious if it's because it doesn't belong here or if we just haven't had the chance to move it upstream.

mlir/lib/Bindings/Python/IRModule.h

728

Frankly I'm in a little over my head here, is this the correct statement about ownership?

mlir/test/Bindings/Python/ir_value.py

25

Is this correct? It was confusing me last night... value2 is value is False, because value is an instance of OpResult not Value. But value2 == value is True, which I think is what we care about in this case.

[jdd]

Another thing we'll want is a nice casting helper, like this: https://github.com/llvm/circt/pull/940/commits/69f527952aa5acb2c73cc3b74e836483b5a5a5d8. It seems we have been adding this support downstream when needed. Does that kind of thing belong in the downstream project, or could I add it in MLIR? I don't see much pybind11 casting functionality in MLIR, so I'm curious if it's because it doesn't belong here or if we just haven't had the chance to move it upstream.

I originally stole this from npcomp. I think @stellaraccident said it was a candidate for upstreaming, which makes sense to me.

[mikeurbach]

Adding @ftynse for visibility.

[mikeurbach]

Rebase main

[mikeurbach]

@mehdi_amini since you approved https://reviews.llvm.org/D101398, do you mind taking a look at this one as well?

[stellaraccident]

mlir/lib/Bindings/Python/IRModule.h

728

Yes, the statement PyOperation::forOperation in the implementation is ensuring this.

mlir/test/Bindings/Python/ir_value.py

25

Yes, this is the best we can do here (values are just a wrapper and not uniqued between the portion/c side so eq vs is is all we can do).

[stellaraccident]

It would be nice to upstream the casting helpers once you all have them in a good state. Ideally, we have an out of tree sample that we use them in to avoid submitting dead code. One step at a time....

[mikeurbach]

Thanks Stella!

It would be nice to upstream the casting helpers once you all have them in a good state. Ideally, we have an out of tree sample that we use them in to avoid submitting dead code. One step at a time....

Sounds good to me.

mlir/include/mlir-c/Bindings/Python/Interop.h

 #define MLIR_PYTHON_CAPSULE_EXECUTION_ENGINE                                   \
   "mlir.execution_engine.ExecutionEngine._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_INTEGER_SET "mlir.ir.IntegerSet._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_LOCATION "mlir.ir.Location._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_MODULE "mlir.ir.Module._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_OPERATION "mlir.ir.Operation._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_TYPE "mlir.ir.Type._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_PASS_MANAGER "mlir.passmanager.PassManager._CAPIPtr"
+#define MLIR_PYTHON_CAPSULE_VALUE "mlir.ir.Value._CAPIPtr"
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
  * below.
  *
  * Ownership is not transferred by acquiring a capsule in this way: the
  * validity of the pointer wrapped by the capsule will be bounded by the
  * lifetime of the Python object that produced it. Only the name and pointer

 static inline MlirExecutionEngine
 mlirPythonCapsuleToExecutionEngine(PyObject *capsule) {
   void *ptr =
       PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_EXECUTION_ENGINE);
   MlirExecutionEngine jit = {ptr};
   return jit;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirValue.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the operation in any way.
+ */
+static inline PyObject *mlirPythonValueToCapsule(MlirValue value) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(value),
+                       MLIR_PYTHON_CAPSULE_VALUE, NULL);
+}
+
+/** Extracts an MlirValue from a capsule as produced from
+ * mlirPythonValueToCapsule. If the capsule is not of the right type, then a
+ * null type is returned (as checked via mlirValueIsNull). In such a case, the
+ * Python APIs will have already set an error. */
+static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_VALUE);
+  MlirValue value = {ptr};
+  return value;
+}
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_BINDINGS_PYTHON_INTEROP_H

mlir/lib/Bindings/Python/IRCore.cpp

 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
+#include "mlir-c/IR.h"
 #include "mlir-c/Registration.h"
 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
PyType PyType::createFromCapsule(py::object capsule) {
   return PyType(PyMlirContext::forContext(mlirTypeGetContext(rawType)),
                 rawType);
 }
 
 //------------------------------------------------------------------------------
 // PyValue and subclases.
 //------------------------------------------------------------------------------
 
+pybind11::object PyValue::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonValueToCapsule(get()));
+}
+
+PyValue PyValue::createFromCapsule(pybind11::object capsule) {
+  MlirValue value = mlirPythonCapsuleToValue(capsule.ptr());
+  if (mlirValueIsNull(value))
+    throw py::error_already_set();
+  MlirOperation owner;
+  if (mlirValueIsAOpResult(value))
+    owner = mlirOpResultGetOwner(value);
+  if (mlirValueIsABlockArgument(value))
+    owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(value));
+  if (mlirOperationIsNull(owner))
+    throw py::error_already_set();
+  MlirContext ctx = mlirOperationGetContext(owner);
+  PyOperationRef ownerRef =
+      PyOperation::forOperation(PyMlirContext::forContext(ctx), owner);
+  return PyValue(ownerRef, value);
+}
+
 namespace {
 /// CRTP base class for Python MLIR values that subclass Value and should be
 /// castable from it. The value hierarchy is one level deep and is not supposed
 /// to accommodate other levels unless core MLIR changes.
 template <typename DerivedTy>
 class PyConcreteValue : public PyValue {
 public:
   // Derived classes must define statics for:
py::class_<PyType>(m, "Type")
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Value.
   //----------------------------------------------------------------------------
   py::class_<PyValue>(m, "Value")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule)
       .def_property_readonly(
           "context",
           [](PyValue &self) { return self.getParentOperation()->getContext(); },
           "Context in which the value lives.")
       .def(
           "dump", [](PyValue &self) { mlirValueDump(self.get()); },
           kDumpDocstring)
       .def("__eq__",

mlir/lib/Bindings/Python/IRModule.h

public:
   PyValue(PyOperationRef parentOperation, MlirValue value)
       : parentOperation(parentOperation), value(value) {}
 
   MlirValue get() { return value; }
   PyOperationRef &getParentOperation() { return parentOperation; }
 
   void checkValid() { return parentOperation->checkValid(); }
 
+  /// Gets a capsule wrapping the void* within the MlirValue.
+  pybind11::object getCapsule();
+
+  /// Creates a PyValue from the MlirValue wrapped by a capsule. Ownership of
+  /// the underlying MlirValue is still tied to the owning operation.
+  static PyValue createFromCapsule(pybind11::object capsule);
+
 private:
   PyOperationRef parentOperation;
   MlirValue value;
 };
 
 /// Wrapper around MlirAffineExpr. Affine expressions are owned by the context.
 class PyAffineExpr : public BaseContextObject {
 public:

mlir/test/Bindings/Python/ir_value.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+from mlir.ir import *
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+
+
+# CHECK-LABEL: TEST: testCapsuleConversions
+def testCapsuleConversions():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    i32 = IntegerType.get_signless(32)
+    value = Operation.create("custom.op1", results=[i32]).result
+    value_capsule = value._CAPIPtr
+    assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule)
+    value2 = Value._CAPICreate(value_capsule)
+    assert value2 == value
+
+
+run(testCapsuleConversions)