[MLIR][Python] Implement pybind adapters for MlirBlock

[ftynse]

Is there any way to test this?

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

187

Nit: copy-pasta from pass manager.

[apaszke]

Addressed review comments

[apaszke]

I have tested this in my code, so I know it works (at least right now). Not sure if you have any tests for the other pybind adapters? All the methods in the Python bindings that take blocks use PyBlock, but that's not going to work for my own plugins that can only depend on the C types.

[jpienaar]

I think mlir/test/python/lib/PythonTestModule.cpp could fit here. If you could create a small test to avoid regressions that would be good

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

#define MAKE_MLIR_PYTHON_QUALNAME(local)                                       \
   MLIR_PYTHON_STRINGIZE_ARG(MLIR_PYTHON_PACKAGE_PREFIX) local
 
 #define MLIR_PYTHON_CAPSULE_AFFINE_EXPR                                        \
   MAKE_MLIR_PYTHON_QUALNAME("ir.AffineExpr._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_AFFINE_MAP                                         \
   MAKE_MLIR_PYTHON_QUALNAME("ir.AffineMap._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_ATTRIBUTE                                          \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Attribute._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_BLOCK MAKE_MLIR_PYTHON_QUALNAME("ir.Block._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_CONTEXT                                            \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Context._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY                                   \
   MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_EXECUTION_ENGINE                                   \
   MAKE_MLIR_PYTHON_QUALNAME("execution_engine.ExecutionEngine._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_INTEGER_SET                                        \
   MAKE_MLIR_PYTHON_QUALNAME("ir.IntegerSet._CAPIPtr")

  * a null attribute is returned (as checked via mlirAttributeIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirAttribute mlirPythonCapsuleToAttribute(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_ATTRIBUTE);
   MlirAttribute attr = {ptr};
   return attr;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirBlock.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the module in any way. */
+static inline PyObject *mlirPythonBlockToCapsule(MlirBlock block) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(block),
+                       MLIR_PYTHON_CAPSULE_BLOCK, NULL);
+}
+
+/** Extracts an MlirBlock from a capsule as produced from
+ * mlirPythonBlockToCapsule. If the capsule is not of the right type, then
+ * a null pass manager is returned (as checked via mlirBlockIsNull). */
+static inline MlirBlock mlirPythonCapsuleToBlock(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_BLOCK);
+  MlirBlock block = {ptr};
+  return block;
+}
+
 /** 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.
  */
 static inline PyObject *mlirPythonContextToCapsule(MlirContext context) {
   return PyCapsule_New(context.ptr, MLIR_PYTHON_CAPSULE_CONTEXT, NULL);
 }
 

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

static handle cast(MlirAttribute v, return_value_policy, handle) {
     return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
         .attr("Attribute")
         .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
         .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
         .release();
   }
 };
 
+/// Casts object -> MlirBlock.
+template <>
+struct type_caster<MlirBlock> {
+  PYBIND11_TYPE_CASTER(MlirBlock, _("MlirBlock"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToBlock(capsule.ptr());
+    return !mlirBlockIsNull(value);
+  }
+};
+
 /// Casts object -> MlirContext.
 template <>
 struct type_caster<MlirContext> {
   PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
   bool load(handle src, bool) {
     if (src.is_none()) {
       // Gets the current thread-bound context.
       // TODO: This raises an error of "No current context" currently.

mlir/lib/Bindings/Python/IRCore.cpp

static void bind(py::module &m) {
     py::class_<PyAttrBuilderMap>(m, "AttrBuilder", py::module_local())
         .def_static("contains", &PyAttrBuilderMap::dunderContains)
         .def_static("get", &PyAttrBuilderMap::dundeGetItemNamed)
         .def_static("insert", &PyAttrBuilderMap::dundeSetItemNamed);
   }
 };
 
 //------------------------------------------------------------------------------
+// PyBlock
+//------------------------------------------------------------------------------
+
+py::object PyBlock::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonBlockToCapsule(get()));
+}
+
+//------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:
   PyRegionIterator(PyOperationRef operation)
py::class_<PyRegion>(m, "Region", py::module_local())
              return self.get().ptr == other.get().ptr;
            })
       .def("__eq__", [](PyRegion &self, py::object &other) { return false; });
 
   //----------------------------------------------------------------------------
   // Mapping of PyBlock.
   //----------------------------------------------------------------------------
   py::class_<PyBlock>(m, "Block", py::module_local())
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyBlock::getCapsule)
       .def_property_readonly(
           "owner",
           [](PyBlock &self) {
             return self.getParentOperation()->createOpView();
           },
           "Returns the owning operation of this block.")
       .def_property_readonly(
           "region",

mlir/lib/Bindings/Python/IRModule.h

PyBlock(PyOperationRef parentOperation, MlirBlock block)
     assert(!mlirBlockIsNull(block) && "python block cannot be null");
   }
 
   MlirBlock get() { return block; }
   PyOperationRef &getParentOperation() { return parentOperation; }
 
   void checkValid() { return parentOperation->checkValid(); }
 
+  /// Gets a capsule wrapping the void* within the MlirBlock.
+  pybind11::object getCapsule();
+
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
 };
 
 /// An insertion point maintains a pointer to a Block and a reference operation.
 /// Calls to insert() will insert a new operation before the
 /// reference operation. If the reference operation is null, then appends to