diff --git a/mlir/include/mlir-c/Bindings/Python/Interop.h b/mlir/include/mlir-c/Bindings/Python/Interop.h --- a/mlir/include/mlir-c/Bindings/Python/Interop.h +++ b/mlir/include/mlir-c/Bindings/Python/Interop.h @@ -62,6 +62,7 @@ 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 \ @@ -175,6 +176,23 @@ 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. diff --git a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h --- a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h +++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h @@ -102,6 +102,17 @@ } }; +/// Casts object -> MlirBlock. +template <> +struct type_caster { + 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 { diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -257,6 +257,14 @@ } }; +//------------------------------------------------------------------------------ +// PyBlock +//------------------------------------------------------------------------------ + +py::object PyBlock::getCapsule() { + return py::reinterpret_steal(mlirPythonBlockToCapsule(get())); +} + //------------------------------------------------------------------------------ // Collections. //------------------------------------------------------------------------------ @@ -2968,6 +2976,7 @@ // Mapping of PyBlock. //---------------------------------------------------------------------------- py::class_(m, "Block", py::module_local()) + .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyBlock::getCapsule) .def_property_readonly( "owner", [](PyBlock &self) { diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -763,6 +763,9 @@ void checkValid() { return parentOperation->checkValid(); } + /// Gets a capsule wrapping the void* within the MlirBlock. + pybind11::object getCapsule(); + private: PyOperationRef parentOperation; MlirBlock block;