Add FunctionType to MLIR C and Python bindings.

[ftynse]

mlir/lib/CAPI/IR/StandardTypes.cpp

321–335 ↗

(On Diff #294694)

Nit: consider adding explicit casts between signed (intptr_t) and unsigned integers to avoid warnings.

[stellaraccident]

Address comments.

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.cpp

1316

I notice that some def/def_static take py::arg, while others don't take. When do we need to use py::arg? What is the purpose of using py::arg?

[stellaraccident]

Rebase.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

1316

It is generally a good idea to include py::arg wherever there is an API where people may want to invoke things with keywords or have other hints to that effect (i.e. most of them). Any that are missing should not be considered intentional unless if it is a purely positional/internal/dunder API.

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

+/*===-- mlir-c/Interop.h - Constants for Python/C-API 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_BINDINGS_PYTHON_INTEROP_H
+#define MLIR_C_BINDINGS_PYTHON_INTEROP_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
+ * a null context is returned (as checked via mlirContextIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+inline MlirContext mlirPythonCapsuleToContext(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_CONTEXT);
+  MlirContext context = {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_C_BINDINGS_PYTHON_INTEROP_H

mlir/include/mlir-c/IR.h

 /*============================================================================*/
 
 /** Creates an MLIR context and transfers its ownership to the caller. */
 MlirContext mlirContextCreate();
 
 /** Checks if two contexts are equal. */
 int mlirContextEqual(MlirContext ctx1, MlirContext ctx2);
 
+/** Checks whether a context is null. */
+inline int mlirContextIsNull(MlirContext context) { return !context.ptr; }
+
 /** Takes an MLIR context owned by the caller and destroys it. */
 void mlirContextDestroy(MlirContext context);
 
 /** Sets whether unregistered dialects are allowed in this context. */
 void mlirContextSetAllowUnregisteredDialects(MlirContext context, int allow);
 
 /** Returns whether the context allows unregistered dialects. */
 int mlirContextGetAllowUnregisteredDialects(MlirContext context);

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/Bindings/Python/Interop.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) {
+  MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr());
+  if (mlirContextIsNull(rawContext))
+    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() {
c.def_static(
            std::vector<PyType> results) {
           SmallVector<MlirType, 4> inputsRaw(inputs.begin(), inputs.end());
           SmallVector<MlirType, 4> resultsRaw(results.begin(), results.end());
           MlirType t = mlirFunctionTypeGet(context.get(), inputsRaw.size(),
                                            inputsRaw.data(), resultsRaw.size(),
                                            resultsRaw.data());
           return PyFunctionType(context.getRef(), t);
         },
         py::arg("context"), py::arg("inputs"), py::arg("results"),
         "Gets a FunctionType from a list of input and result types");
     c.def_property_readonly(
         "inputs",
         [](PyFunctionType &self) {
           MlirType t = self.type;
           auto contextRef = self.getContext();
           py::list types;
           for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self.type);
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)