Adds MLIR C-API for marshaling Python capsules.

[stellaraccident]

Summary

• Providing stable, C-accessible definitions for bridging MLIR Python<->C APIs, we eliminate inter-extension dependencies (i.e. they can all share a diamond dependency on the MLIR C-API).
• Just provides accessors for context and module right now.
• Needed in NPComp in ~a week or so for high level Torch APIs.

[mehdi_amini]

mlir/include/mlir-c/PythonDefs.h

69

Isn't the isError redundant with returning a null context ?

94

Why isn't this file in the Python binding directory?

mlir/lib/Bindings/Python/IRModules.cpp

462

What about zero initializing explicitly? Reading such code is always suggesting me to read the implementation of the callee to make sure there isn't a path where it won't set the error flag.

mlir/lib/Bindings/Python/IRModules.h

210

Is this useful right now on the Module?

[stellaraccident]

Address comments.

[stellaraccident]

Thanks!

mlir/include/mlir-c/PythonDefs.h

69

Ah, we were missing an mlirContextIsNull and I just wasn't thinking (Alex still wants these structs to be mostly opaque from user's perspectives). Added it and removed the arg.

94

My mistake: I was considering the right place high order bit to put it in to be mlir-c but then neglected to create the sub-directory. Moved to mlir-c/Bindings/Python/Interop.h.

mlir/lib/Bindings/Python/IRModules.cpp

462

I think how this is now side-steps the concern, but let me know if not? (not quite following what you want to zero-initialize here)

mlir/lib/Bindings/Python/IRModules.h

210

I have a use for it in npcomp and am working around it not existing (basically, want to allocate a Python-visible Module and also interop with it at the MLIR-C API level).

[stellaraccident]

Rebase.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

462

Yes LG, it was the output parameter isError that is removed now.

mlir/include/mlir-c/PythonDefs.h

+/*===-- mlir-c/PythonDefs.h - Constants for Python interop --------*- C -*-===*\
+|*                                                                            *|
+|* Part of the LLVM Project, under the Apache License v2.0 with LLVM          *|
+|* Exceptions.                                                                *|
+|* See https://llvm.org/LICENSE.txt for license information.                  *|
+|* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception                    *|
+|*                                                                            *|
+|*===----------------------------------------------------------------------===*|
+|*                                                                            *|
+|* This header declares constants and helpers necessary for C-level           *|
+|* interop with the MLIR Python extension module. Since the Python bindings   *|
+|* are a thin wrapper around the MLIR C-API, a further C-API is not provided  *|
+|* specifically for the Python extension. Instead, simple facilities are      *|
+|* provided for translating between Python types and corresponding MLIR C-API *|
+|* types.
+|*                                                                            *|
+|* This header is standalone, requiring nothing beyond normal linking against *|
+|* the Python implementation.
+\*===----------------------------------------------------------------------===*/
+
+#ifndef MLIR_C_PYTHON_DEFS_H
+#define MLIR_C_PYTHON_DEFS_H
+
+#include <Python.h>
+
+#include "mlir-c/IR.h"
+
+#define MLIR_PYTHON_CAPSULE_CONTEXT "mlir.ir.Context._CAPIPtr"
+#define MLIR_PYTHON_CAPSULE_MODULE "mlir.ir.Module._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
+ * of the capsule are set. The caller is free to set a destructor and context
+ * as needed to manage anything further. */
+#define MLIR_PYTHON_CAPI_PTR_ATTR "_CAPIPtr"
+
+/** Attribute on MLIR Python objects that exposes a factory function for
+ * constructing the corresponding Python object from a type-specific
+ * capsule wrapping the C-API pointer. The signature of the function is:
+ *   def _CAPICreate(capsule) -> object
+ * Calling such a function implies a transfer of ownership of the object the
+ * capsule wraps: after such a call, the capsule should be considered invalid,
+ * and its wrapped pointer must not be destroyed.
+ *
+ * Only a very small number of Python objects can be created in such a fashion
+ * (i.e. top-level types such as Context where the lifetime can be cleanly
+ * delineated). */
+#define MLIR_PYTHON_CAPI_FACTORY_ATTR "_CAPICreate"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/** Creates a capsule object encapsulating the raw C-API MlirContext.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the context in any way.
+ */
+inline PyObject *mlirPythonContextToCapsule(MlirContext context) {
+  return PyCapsule_New(context.ptr, MLIR_PYTHON_CAPSULE_CONTEXT, NULL);
+}
+
+/** Extracts a MlirContext from a capsule as produced from
+ * mlirPythonContextToCapsule. If the capsule is not of the right type, then
+ * *isError is set to true and a null context is returned. In such a case, the
+ * Python APIs will have already set an error. */
+inline MlirContext mlirPythonCapsuleToContext(PyObject *capsule, int *isError) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_CONTEXT);
+  MlirContext context = {ptr};
+  *isError = !ptr;
+  return context;
+}
+
+/** Creates a capsule object encapsulating the raw C-API MlirModule.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the module in any way. */
+inline PyObject *mlirPythonModuleToCapsule(MlirModule module) {
+#ifdef __cplusplus
+  void *ptr = const_cast<void *>(module.ptr);
+#else
+  void *ptr = (void *)ptr;
+#endif
+  return PyCapsule_New(ptr, MLIR_PYTHON_CAPSULE_MODULE, NULL);
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

mlir/lib/Bindings/Python/IRModules.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
public:
   MlirContext get() { return context; }
 
   /// Gets a strong reference to this context, which will ensure it is kept
   /// alive for the life of the reference.
   PyMlirContextRef getRef() {
     return PyMlirContextRef(this, pybind11::cast(this));
   }
 
+  /// Gets a capsule wrapping the void* within the MlirContext.
+  pybind11::object getCapsule();
+
+  /// Creates a PyMlirContext from the MlirContext wrapped by a capsule.
+  /// Note that PyMlirContext instances are uniqued, so the returned object
+  /// may be a pre-existing object.
+  static pybind11::object createFromCapsule(pybind11::object capsule);
+
   /// Gets the count of live context objects. Used for testing.
   static size_t getLiveCount();
 
   /// Gets the count of live operations associated with this context.
   /// Used for testing.
   size_t getLiveOperationCount();
 
   /// Creates an operation. See corresponding python docstring.
public:
   MlirModule get() { return module; }
 
   /// Gets a strong reference to this module.
   PyModuleRef getRef() {
     return PyModuleRef(this,
                        pybind11::reinterpret_borrow<pybind11::object>(handle));
   }
 
+  /// Gets a capsule wrapping the void* within the MlirModule.
+  /// Note that the module does not (yet) provide a corresponding factory for
+  /// constructing from a capsule as that would require uniquing PyModule
+  /// instances, which is not currently done.
+  pybind11::object getCapsule();
+
 private:
   PyModule(PyMlirContextRef contextRef, MlirModule module)
       : BaseContextObject(std::move(contextRef)), module(module) {}
   MlirModule module;
   pybind11::handle handle;
 };
 
 /// Wrapper around PyOperation.

mlir/lib/Bindings/Python/IRModules.cpp

 //===- IRModules.cpp - IR Submodules of pybind module ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRModules.h"
 #include "PybindUtils.h"
 
+#include "mlir-c/PythonDefs.h"
 #include "mlir-c/Registration.h"
 #include "mlir-c/StandardAttributes.h"
 #include "mlir-c/StandardTypes.h"
 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
PyMlirContext::~PyMlirContext() {
   // Note that the only public way to construct an instance is via the
   // forContext method, which always puts the associated handle into
   // liveContexts.
   py::gil_scoped_acquire acquire;
   getLiveContexts().erase(context.ptr);
   mlirContextDestroy(context);
 }
 
+py::object PyMlirContext::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonContextToCapsule(get()));
+}
+
+py::object PyMlirContext::createFromCapsule(py::object capsule) {
+  int isError;
+  MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr(), &isError);
+  if (isError)
+    throw py::error_already_set();
+  return forContext(rawContext).releaseObject();
+}
+
 PyMlirContext *PyMlirContext::createNewContextForInit() {
   MlirContext context = mlirContextCreate();
   mlirRegisterAllDialects(context);
   return new PyMlirContext(context);
 }
 
 PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   py::gil_scoped_acquire acquire;
PyModuleRef PyModule::create(PyMlirContextRef contextRef, MlirModule module) {
   // which does not take ownership (delete will not be called).
   // Just be explicit.
   py::object pyRef =
       py::cast(unownedModule, py::return_value_policy::take_ownership);
   unownedModule->handle = pyRef;
   return PyModuleRef(unownedModule, std::move(pyRef));
 }
 
+py::object PyModule::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(get()));
+}
+
 //------------------------------------------------------------------------------
 // PyOperation
 //------------------------------------------------------------------------------
 
 PyOperation::PyOperation(PyMlirContextRef contextRef, MlirOperation operation)
     : BaseContextObject(std::move(contextRef)), operation(operation) {}
 
 PyOperation::~PyOperation() {
py::class_<PyMlirContext>(m, "Context")
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
              return ref.releaseObject();
            })
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
+                             &PyMlirContext::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyMlirContext::createFromCapsule)
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
py::class_<PyLocation>(m, "Location").def("__repr__", [](PyLocation &self) {
     PyPrintAccumulator printAccum;
     mlirLocationPrint(self.loc, printAccum.getCallback(),
                       printAccum.getUserData());
     return printAccum.join();
   });
 
   // Mapping of Module
   py::class_<PyModule>(m, "Module")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule)
       .def_property_readonly(
           "operation",
           [](PyModule &self) {
             return PyOperation::forOperation(self.getContext(),
                                              mlirModuleGetOperation(self.get()),
                                              self.getRef().releaseObject())
                 .releaseObject();
           },

mlir/test/Bindings/Python/context_lifecycle.py

 print("Free C1")
 c1 = None
 gc.collect()
 assert mlir.ir.Context._get_live_count() == 1
 print("Free C2")
 c2 = None
 gc.collect()
 assert mlir.ir.Context._get_live_count() == 0
+
+# Create a context, get its capsule and create from capsule.
+c4 = mlir.ir.Context()
+c4_capsule = c4._CAPIPtr
+assert '"mlir.ir.Context._CAPIPtr"' in repr(c4_capsule)
+c5 = mlir.ir.Context._CAPICreate(c4_capsule)
+assert c4 is c5

mlir/test/Bindings/Python/ir_module.py

def testModuleOperation():
   # Collect and verify lifetime.
   op1 = None
   op2 = None
   gc.collect()
   print("LIVE OPERATIONS:", ctx._get_live_operation_count())
   assert ctx._get_live_operation_count() == 0
 
 run(testModuleOperation)
+
+
+# CHECK-LABEL: TEST: testModuleCapsule
+def testModuleCapsule():
+  ctx = mlir.ir.Context()
+  module = ctx.parse_module(r"""module @successfulParse {}""")
+  # CHECK: "mlir.ir.Module._CAPIPtr"
+  print(module._CAPIPtr)
+
+run(testModuleCapsule)