MLIR  15.0.0git
PybindAdaptors.h
Go to the documentation of this file.
1 //===- PybindAdaptors.h - Adaptors for interop with MLIR APIs -------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 // This file contains adaptors for clients of the core MLIR Python APIs to
9 // interop via MLIR CAPI types. The facilities here do not depend on
10 // implementation details of the MLIR Python API and do not introduce C++-level
11 // dependencies with it (requiring only Python and CAPI-level dependencies).
12 //
13 // It is encouraged to be used both in-tree and out-of-tree. For in-tree use
14 // cases, it should be used for dialect implementations (versus relying on
15 // Pybind-based internals of the core libraries).
16 //===----------------------------------------------------------------------===//
17 
18 #ifndef MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
19 #define MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
20 
21 #include <pybind11/pybind11.h>
22 #include <pybind11/pytypes.h>
23 #include <pybind11/stl.h>
24 
26 #include "mlir-c/IR.h"
27 
28 #include "llvm/ADT/Optional.h"
29 #include "llvm/ADT/Twine.h"
30 
31 namespace py = pybind11;
32 
33 // Raw CAPI type casters need to be declared before use, so always include them
34 // first.
35 namespace pybind11 {
36 namespace detail {
37 
38 template <typename T>
39 struct type_caster<llvm::Optional<T>> : optional_caster<llvm::Optional<T>> {};
40 
41 /// Helper to convert a presumed MLIR API object to a capsule, accepting either
42 /// an explicit Capsule (which can happen when two C APIs are communicating
43 /// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
44 /// attribute (through which supported MLIR Python API objects export their
45 /// contained API pointer as a capsule). Throws a type error if the object is
46 /// neither. This is intended to be used from type casters, which are invoked
47 /// with a raw handle (unowned). The returned object's lifetime may not extend
48 /// beyond the apiObject handle without explicitly having its refcount increased
49 /// (i.e. on return).
50 static py::object mlirApiObjectToCapsule(py::handle apiObject) {
51  if (PyCapsule_CheckExact(apiObject.ptr()))
52  return py::reinterpret_borrow<py::object>(apiObject);
53  if (!py::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) {
54  auto repr = py::repr(apiObject).cast<std::string>();
55  throw py::type_error(
56  (llvm::Twine("Expected an MLIR object (got ") + repr + ").").str());
57  }
58  return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
59 }
60 
61 // Note: Currently all of the following support cast from py::object to the
62 // Mlir* C-API type, but only a few light-weight, context-bound ones
63 // implicitly cast the other way because the use case has not yet emerged and
64 // ownership is unclear.
65 
66 /// Casts object <-> MlirAffineMap.
67 template <>
68 struct type_caster<MlirAffineMap> {
69  PYBIND11_TYPE_CASTER(MlirAffineMap, _("MlirAffineMap"));
70  bool load(handle src, bool) {
71  py::object capsule = mlirApiObjectToCapsule(src);
72  value = mlirPythonCapsuleToAffineMap(capsule.ptr());
74  return false;
75  }
76  return !mlirAffineMapIsNull(value);
77  }
78  static handle cast(MlirAffineMap v, return_value_policy, handle) {
79  py::object capsule =
80  py::reinterpret_steal<py::object>(mlirPythonAffineMapToCapsule(v));
81  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
82  .attr("AffineMap")
83  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
84  .release();
85  }
86 };
87 
88 /// Casts object <-> MlirAttribute.
89 template <>
90 struct type_caster<MlirAttribute> {
91  PYBIND11_TYPE_CASTER(MlirAttribute, _("MlirAttribute"));
92  bool load(handle src, bool) {
93  py::object capsule = mlirApiObjectToCapsule(src);
94  value = mlirPythonCapsuleToAttribute(capsule.ptr());
95  return !mlirAttributeIsNull(value);
96  }
97  static handle cast(MlirAttribute v, return_value_policy, handle) {
98  py::object capsule =
99  py::reinterpret_steal<py::object>(mlirPythonAttributeToCapsule(v));
100  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
101  .attr("Attribute")
102  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
103  .release();
104  }
105 };
106 
107 /// Casts object -> MlirContext.
108 template <>
109 struct type_caster<MlirContext> {
110  PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
111  bool load(handle src, bool) {
112  if (src.is_none()) {
113  // Gets the current thread-bound context.
114  // TODO: This raises an error of "No current context" currently.
115  // Update the implementation to pretty-print the helpful error that the
116  // core implementations print in this case.
117  src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
118  .attr("Context")
119  .attr("current");
120  }
121  py::object capsule = mlirApiObjectToCapsule(src);
122  value = mlirPythonCapsuleToContext(capsule.ptr());
123  return !mlirContextIsNull(value);
124  }
125 };
126 
127 /// Casts object <-> MlirLocation.
128 template <>
129 struct type_caster<MlirLocation> {
130  PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
131  bool load(handle src, bool) {
132  if (src.is_none()) {
133  // Gets the current thread-bound context.
134  src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
135  .attr("Location")
136  .attr("current");
137  }
138  py::object capsule = mlirApiObjectToCapsule(src);
139  value = mlirPythonCapsuleToLocation(capsule.ptr());
140  return !mlirLocationIsNull(value);
141  }
142  static handle cast(MlirLocation v, return_value_policy, handle) {
143  py::object capsule =
144  py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(v));
145  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
146  .attr("Location")
147  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
148  .release();
149  }
150 };
151 
152 /// Casts object <-> MlirModule.
153 template <>
154 struct type_caster<MlirModule> {
155  PYBIND11_TYPE_CASTER(MlirModule, _("MlirModule"));
156  bool load(handle src, bool) {
157  py::object capsule = mlirApiObjectToCapsule(src);
158  value = mlirPythonCapsuleToModule(capsule.ptr());
159  return !mlirModuleIsNull(value);
160  }
161  static handle cast(MlirModule v, return_value_policy, handle) {
162  py::object capsule =
163  py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(v));
164  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
165  .attr("Module")
166  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
167  .release();
168  };
169 };
170 
171 /// Casts object <-> MlirOperation.
172 template <>
173 struct type_caster<MlirOperation> {
174  PYBIND11_TYPE_CASTER(MlirOperation, _("MlirOperation"));
175  bool load(handle src, bool) {
176  py::object capsule = mlirApiObjectToCapsule(src);
177  value = mlirPythonCapsuleToOperation(capsule.ptr());
178  return !mlirOperationIsNull(value);
179  }
180  static handle cast(MlirOperation v, return_value_policy, handle) {
181  if (v.ptr == nullptr)
182  return py::none();
183  py::object capsule =
184  py::reinterpret_steal<py::object>(mlirPythonOperationToCapsule(v));
185  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
186  .attr("Operation")
187  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
188  .release();
189  };
190 };
191 
192 /// Casts object -> MlirPassManager.
193 template <>
194 struct type_caster<MlirPassManager> {
195  PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
196  bool load(handle src, bool) {
197  py::object capsule = mlirApiObjectToCapsule(src);
198  value = mlirPythonCapsuleToPassManager(capsule.ptr());
199  return !mlirPassManagerIsNull(value);
200  }
201 };
202 
203 /// Casts object <-> MlirType.
204 template <>
205 struct type_caster<MlirType> {
206  PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
207  bool load(handle src, bool) {
208  py::object capsule = mlirApiObjectToCapsule(src);
209  value = mlirPythonCapsuleToType(capsule.ptr());
210  return !mlirTypeIsNull(value);
211  }
212  static handle cast(MlirType t, return_value_policy, handle) {
213  py::object capsule =
214  py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
215  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
216  .attr("Type")
217  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
218  .release();
219  }
220 };
221 
222 } // namespace detail
223 } // namespace pybind11
224 
225 namespace mlir {
226 namespace python {
227 namespace adaptors {
228 
229 /// Provides a facility like py::class_ for defining a new class in a scope,
230 /// but this allows extension of an arbitrary Python class, defining methods
231 /// on it is a similar way. Classes defined in this way are very similar to
232 /// if defined in Python in the usual way but use Pybind11 machinery to do
233 /// it. These are not "real" Pybind11 classes but pure Python classes with no
234 /// relation to a concrete C++ class.
235 ///
236 /// Derived from a discussion upstream:
237 /// https://github.com/pybind/pybind11/issues/1193
238 /// (plus a fair amount of extra curricular poking)
239 /// TODO: If this proves useful, see about including it in pybind11.
241 public:
242  pure_subclass(py::handle scope, const char *derivedClassName,
243  const py::object &superClass) {
244  py::object pyType =
245  py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
246  py::object metaclass = pyType(superClass);
247  py::dict attributes;
248 
249  thisClass =
250  metaclass(derivedClassName, py::make_tuple(superClass), attributes);
251  scope.attr(derivedClassName) = thisClass;
252  }
253 
254  template <typename Func, typename... Extra>
255  pure_subclass &def(const char *name, Func &&f, const Extra &... extra) {
256  py::cpp_function cf(
257  std::forward<Func>(f), py::name(name), py::is_method(thisClass),
258  py::sibling(py::getattr(thisClass, name, py::none())), extra...);
259  thisClass.attr(cf.name()) = cf;
260  return *this;
261  }
262 
263  template <typename Func, typename... Extra>
264  pure_subclass &def_property_readonly(const char *name, Func &&f,
265  const Extra &... extra) {
266  py::cpp_function cf(
267  std::forward<Func>(f), py::name(name), py::is_method(thisClass),
268  py::sibling(py::getattr(thisClass, name, py::none())), extra...);
269  auto builtinProperty =
270  py::reinterpret_borrow<py::object>((PyObject *)&PyProperty_Type);
271  thisClass.attr(name) = builtinProperty(cf);
272  return *this;
273  }
274 
275  template <typename Func, typename... Extra>
276  pure_subclass &def_staticmethod(const char *name, Func &&f,
277  const Extra &... extra) {
279  "def_staticmethod(...) called with a non-static member "
280  "function pointer");
281  py::cpp_function cf(
282  std::forward<Func>(f), py::name(name), py::scope(thisClass),
283  py::sibling(py::getattr(thisClass, name, py::none())), extra...);
284  thisClass.attr(cf.name()) = py::staticmethod(cf);
285  return *this;
286  }
287 
288  template <typename Func, typename... Extra>
289  pure_subclass &def_classmethod(const char *name, Func &&f,
290  const Extra &... extra) {
292  "def_classmethod(...) called with a non-static member "
293  "function pointer");
294  py::cpp_function cf(
295  std::forward<Func>(f), py::name(name), py::scope(thisClass),
296  py::sibling(py::getattr(thisClass, name, py::none())), extra...);
297  thisClass.attr(cf.name()) =
298  py::reinterpret_borrow<py::object>(PyClassMethod_New(cf.ptr()));
299  return *this;
300  }
301 
302  py::object get_class() const { return thisClass; }
303 
304 protected:
305  py::object superClass;
306  py::object thisClass;
307 };
308 
309 /// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
310 /// constructor and type checking methods.
312 public:
313  using IsAFunctionTy = bool (*)(MlirAttribute);
314 
315  /// Subclasses by looking up the super-class dynamically.
316  mlir_attribute_subclass(py::handle scope, const char *attrClassName,
317  IsAFunctionTy isaFunction)
319  scope, attrClassName, isaFunction,
320  py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
321  .attr("Attribute")) {}
322 
323  /// Subclasses with a provided mlir.ir.Attribute super-class. This must
324  /// be used if the subclass is being defined in the same extension module
325  /// as the mlir.ir class (otherwise, it will trigger a recursive
326  /// initialization).
327  mlir_attribute_subclass(py::handle scope, const char *typeClassName,
328  IsAFunctionTy isaFunction, const py::object &superCls)
329  : pure_subclass(scope, typeClassName, superCls) {
330  // Casting constructor. Note that it hard, if not impossible, to properly
331  // call chain to parent `__init__` in pybind11 due to its special handling
332  // for init functions that don't have a fully constructed self-reference,
333  // which makes it impossible to forward it to `__init__` of a superclass.
334  // Instead, provide a custom `__new__` and call that of a superclass, which
335  // eventually calls `__init__` of the superclass. Since attribute subclasses
336  // have no additional members, we can just return the instance thus created
337  // without amending it.
338  std::string captureTypeName(
339  typeClassName); // As string in case if typeClassName is not static.
340  py::cpp_function newCf(
341  [superCls, isaFunction, captureTypeName](py::object cls,
342  py::object otherAttribute) {
343  MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherAttribute);
344  if (!isaFunction(rawAttribute)) {
345  auto origRepr = py::repr(otherAttribute).cast<std::string>();
346  throw std::invalid_argument(
347  (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
348  " (from " + origRepr + ")")
349  .str());
350  }
351  py::object self = superCls.attr("__new__")(cls, otherAttribute);
352  return self;
353  },
354  py::name("__new__"), py::arg("cls"), py::arg("cast_from_attr"));
355  thisClass.attr("__new__") = newCf;
356 
357  // 'isinstance' method.
358  def_staticmethod(
359  "isinstance",
360  [isaFunction](MlirAttribute other) { return isaFunction(other); },
361  py::arg("other_attribute"));
362  }
363 };
364 
365 /// Creates a custom subclass of mlir.ir.Type, implementing a casting
366 /// constructor and type checking methods.
368 public:
369  using IsAFunctionTy = bool (*)(MlirType);
370 
371  /// Subclasses by looking up the super-class dynamically.
372  mlir_type_subclass(py::handle scope, const char *typeClassName,
373  IsAFunctionTy isaFunction)
375  scope, typeClassName, isaFunction,
376  py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Type")) {}
377 
378  /// Subclasses with a provided mlir.ir.Type super-class. This must
379  /// be used if the subclass is being defined in the same extension module
380  /// as the mlir.ir class (otherwise, it will trigger a recursive
381  /// initialization).
382  mlir_type_subclass(py::handle scope, const char *typeClassName,
383  IsAFunctionTy isaFunction, const py::object &superCls)
384  : pure_subclass(scope, typeClassName, superCls) {
385  // Casting constructor. Note that it hard, if not impossible, to properly
386  // call chain to parent `__init__` in pybind11 due to its special handling
387  // for init functions that don't have a fully constructed self-reference,
388  // which makes it impossible to forward it to `__init__` of a superclass.
389  // Instead, provide a custom `__new__` and call that of a superclass, which
390  // eventually calls `__init__` of the superclass. Since attribute subclasses
391  // have no additional members, we can just return the instance thus created
392  // without amending it.
393  std::string captureTypeName(
394  typeClassName); // As string in case if typeClassName is not static.
395  py::cpp_function newCf(
396  [superCls, isaFunction, captureTypeName](py::object cls,
397  py::object otherType) {
398  MlirType rawType = py::cast<MlirType>(otherType);
399  if (!isaFunction(rawType)) {
400  auto origRepr = py::repr(otherType).cast<std::string>();
401  throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
402  captureTypeName + " (from " +
403  origRepr + ")")
404  .str());
405  }
406  py::object self = superCls.attr("__new__")(cls, otherType);
407  return self;
408  },
409  py::name("__new__"), py::arg("cls"), py::arg("cast_from_type"));
410  thisClass.attr("__new__") = newCf;
411 
412  // 'isinstance' method.
413  def_staticmethod(
414  "isinstance",
415  [isaFunction](MlirType other) { return isaFunction(other); },
416  py::arg("other_type"));
417  }
418 };
419 
420 } // namespace adaptors
421 } // namespace python
422 } // namespace mlir
423 
424 #endif // MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
TODO: Remove this file when SCCP and integer range analysis have been ported to the new framework...
mlir_attribute_subclass(py::handle scope, const char *attrClassName, IsAFunctionTy isaFunction)
Subclasses by looking up the super-class dynamically.
Creates a custom subclass of mlir.ir.Type, implementing a casting constructor and type checking metho...
static PyObject * mlirPythonLocationToCapsule(MlirLocation loc)
Creates a capsule object encapsulating the raw C-API MlirLocation.
Definition: Interop.h:178
static bool mlirAffineMapIsNull(MlirAffineMap affineMap)
Checks whether an affine map is null.
Definition: AffineMap.h:47
static MlirModule mlirPythonCapsuleToModule(PyObject *capsule)
Extracts an MlirModule from a capsule as produced from mlirPythonModuleToCapsule. ...
Definition: Interop.h:205
The OpAsmOpInterface, see OpAsmInterface.td for more details.
Definition: CallGraph.h:221
static MlirContext mlirPythonCapsuleToContext(PyObject *capsule)
Extracts a MlirContext from a capsule as produced from mlirPythonContextToCapsule.
Definition: Interop.h:169
Creates a custom subclass of mlir.ir.Attribute, implementing a casting constructor and type checking ...
mlir_type_subclass(py::handle scope, const char *typeClassName, IsAFunctionTy isaFunction, const py::object &superCls)
Subclasses with a provided mlir.ir.Type super-class.
static MlirOperation mlirPythonCapsuleToOperation(PyObject *capsule)
Extracts an MlirOperations from a capsule as produced from mlirPythonOperationToCapsule.
Definition: Interop.h:241
pure_subclass(py::handle scope, const char *derivedClassName, const py::object &superClass)
static bool mlirAttributeIsNull(MlirAttribute attr)
Checks whether an attribute is null.
Definition: IR.h:725
static bool mlirPassManagerIsNull(MlirPassManager passManager)
Checks if a PassManager is null.
Definition: Pass.h:62
static MlirType mlirPythonCapsuleToType(PyObject *capsule)
Extracts an MlirType from a capsule as produced from mlirPythonTypeToCapsule.
Definition: Interop.h:260
static handle cast(MlirModule v, return_value_policy, handle)
static handle cast(MlirAttribute v, return_value_policy, handle)
static handle cast(MlirOperation v, return_value_policy, handle)
static constexpr const bool value
static py::object mlirApiObjectToCapsule(py::handle apiObject)
Helper to convert a presumed MLIR API object to a capsule, accepting either an explicit Capsule (whic...
static PyObject * mlirPythonModuleToCapsule(MlirModule module)
Creates a capsule object encapsulating the raw C-API MlirModule.
Definition: Interop.h:196
pure_subclass & def_classmethod(const char *name, Func &&f, const Extra &... extra)
static PyObject * mlirPythonTypeToCapsule(MlirType type)
Creates a capsule object encapsulating the raw C-API MlirType.
Definition: Interop.h:251
static PyObject * mlirPythonAffineMapToCapsule(MlirAffineMap affineMap)
Creates a capsule object encapsulating the raw C-API MlirAffineMap.
Definition: Interop.h:270
Provides a facility like py::class_ for defining a new class in a scope, but this allows extension of...
pure_subclass & def_staticmethod(const char *name, Func &&f, const Extra &... extra)
#define MAKE_MLIR_PYTHON_QUALNAME(local)
Definition: Interop.h:56
mlir_attribute_subclass(py::handle scope, const char *typeClassName, IsAFunctionTy isaFunction, const py::object &superCls)
Subclasses with a provided mlir.ir.Attribute super-class.
#define MLIR_PYTHON_CAPI_PTR_ATTR
Attribute on MLIR Python objects that expose their C-API pointer.
Definition: Interop.h:91
pure_subclass & def(const char *name, Func &&f, const Extra &... extra)
static MlirLocation mlirPythonCapsuleToLocation(PyObject *capsule)
Extracts an MlirLocation from a capsule as produced from mlirPythonLocationToCapsule.
Definition: Interop.h:187
static handle cast(MlirLocation v, return_value_policy, handle)
static handle cast(MlirAffineMap v, return_value_policy, handle)
static handle cast(MlirType t, return_value_policy, handle)
pure_subclass & def_property_readonly(const char *name, Func &&f, const Extra &... extra)
static PyObject * mlirPythonOperationToCapsule(MlirOperation operation)
Creates a capsule object encapsulating the raw C-API MlirOperation.
Definition: Interop.h:233
mlir_type_subclass(py::handle scope, const char *typeClassName, IsAFunctionTy isaFunction)
Subclasses by looking up the super-class dynamically.
static bool mlirContextIsNull(MlirContext context)
Checks whether a context is null.
Definition: IR.h:91
static bool mlirLocationIsNull(MlirLocation location)
Checks if the location is null.
Definition: IR.h:207
static PyObject * mlirPythonAttributeToCapsule(MlirAttribute attribute)
Creates a capsule object encapsulating the raw C-API MlirAttribute.
Definition: Interop.h:142
static MlirPassManager mlirPythonCapsuleToPassManager(PyObject *capsule)
Extracts an MlirPassManager from a capsule as produced from mlirPythonPassManagerToCapsule.
Definition: Interop.h:223
static MlirAffineMap mlirPythonCapsuleToAffineMap(PyObject *capsule)
Extracts an MlirAffineMap from a capsule as produced from mlirPythonAffineMapToCapsule.
Definition: Interop.h:279
#define MLIR_PYTHON_CAPI_FACTORY_ATTR
Attribute on MLIR Python objects that exposes a factory function for constructing the corresponding P...
Definition: Interop.h:104
static bool mlirOperationIsNull(MlirOperation op)
Checks whether the underlying operation is null.
Definition: IR.h:376
static bool mlirTypeIsNull(MlirType type)
Checks whether a type is null.
Definition: IR.h:693
static bool mlirModuleIsNull(MlirModule module)
Checks whether a module is null.
Definition: IR.h:239
static MlirAttribute mlirPythonCapsuleToAttribute(PyObject *capsule)
Extracts an MlirAttribute from a capsule as produced from mlirPythonAttributeToCapsule.
Definition: Interop.h:151