[Mlir-commits] [mlir] c83318e - [MLIR][Python] Implement pybind adapters for MlirBlock

[Jacques Pienaar]

Author: Adam Paszke
Date: 2023-07-12T22:27:01-07:00
New Revision: c83318e3e07c40dae7cf28839be3fe62ebcde369

URL: https://github.com/llvm/llvm-project/commit/c83318e3e07c40dae7cf28839be3fe62ebcde369
DIFF: https://github.com/llvm/llvm-project/commit/c83318e3e07c40dae7cf28839be3fe62ebcde369.diff

LOG: [MLIR][Python] Implement pybind adapters for MlirBlock

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D155092

Added: 
    

Modified: 
    mlir/include/mlir-c/Bindings/Python/Interop.h
    mlir/include/mlir/Bindings/Python/PybindAdaptors.h
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/Bindings/Python/IRModule.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/Bindings/Python/Interop.h b/mlir/include/mlir-c/Bindings/Python/Interop.h
index 33332d6a34bb88..f79c10cb938382 100644
--- 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 @@ static inline MlirAttribute mlirPythonCapsuleToAttribute(PyObject *capsule) {
   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
index 44a10d619d029a..49680c8b79b135 100644
--- a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
@@ -102,6 +102,17 @@ struct type_caster<MlirAttribute> {
   }
 };
 
+/// 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> {

diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 6c0b4a0604e31d..39049f387eda70 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -257,6 +257,14 @@ struct PyAttrBuilderMap {
   }
 };
 
+//------------------------------------------------------------------------------
+// PyBlock
+//------------------------------------------------------------------------------
+
+py::object PyBlock::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonBlockToCapsule(get()));
+}
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
@@ -2968,6 +2976,7 @@ void mlir::python::populateIRCore(py::module &m) {
   // 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) {

diff  --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 76acfe5e7790be..5da1d7d25ebe9c 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -763,6 +763,9 @@ class PyBlock {
 
   void checkValid() { return parentOperation->checkValid(); }
 
+  /// Gets a capsule wrapping the void* within the MlirBlock.
+  pybind11::object getCapsule();
+
 private:
   PyOperationRef parentOperation;
   MlirBlock block;