1//===- NanobindAdaptors.h - Interop with MLIR APIs via nanobind -----------===//
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, using nanobind. The facilities here do not
10// depend on implementation details of the MLIR Python API and do not introduce
11// C++-level dependencies with it (requiring only Python and CAPI-level
12// dependencies).
13//
14// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
15// cases, it should be used for dialect implementations (versus relying on
16// Pybind-based internals of the core libraries).
17//===----------------------------------------------------------------------===//
18
19#ifndef MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
20#define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
21
22#include <cstdint>
23
24#include "mlir-c/Diagnostics.h"
25#include "mlir-c/IR.h"
26// clang-format off
27#include "mlir/Bindings/Python/Nanobind.h"
28#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
29// clang-format on
30#include "llvm/ADT/Twine.h"
31
32// Raw CAPI type casters need to be declared before use, so always include them
33// first.
34namespace nanobind {
35namespace detail {
36
37/// Helper to convert a presumed MLIR API object to a capsule, accepting either
38/// an explicit Capsule (which can happen when two C APIs are communicating
39/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
40/// attribute (through which supported MLIR Python API objects export their
41/// contained API pointer as a capsule). Throws a type error if the object is
42/// neither. This is intended to be used from type casters, which are invoked
43/// with a raw handle (unowned). The returned object's lifetime may not extend
44/// beyond the apiObject handle without explicitly having its refcount increased
45/// (i.e. on return).
46static nanobind::object mlirApiObjectToCapsule(nanobind::handle apiObject) {
47 if (PyCapsule_CheckExact(apiObject.ptr()))
48 return nanobind::borrow<nanobind::object>(apiObject);
49 nanobind::object api =
50 nanobind::getattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR, nanobind::none());
51 if (api.is_none()) {
52 std::string repr = nanobind::cast<std::string>(nanobind::repr(apiObject));
53 throw nanobind::type_error(
54 (llvm::Twine("Expected an MLIR object (got ") + repr + ").")
55 .str()
56 .c_str());
57 }
58 return api;
59}
60
61// Note: Currently all of the following support cast from nanobind::object to
62// the 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.
67template <>
68struct type_caster<MlirAffineMap> {
69 NB_TYPE_CASTER(MlirAffineMap, const_name("MlirAffineMap"))
70 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
71 nanobind::object capsule = mlirApiObjectToCapsule(src);
72 value = mlirPythonCapsuleToAffineMap(capsule.ptr());
73 if (mlirAffineMapIsNull(value)) {
74 return false;
75 }
76 return !mlirAffineMapIsNull(value);
77 }
78 static handle from_cpp(MlirAffineMap v, rv_policy,
79 cleanup_list *cleanup) noexcept {
80 nanobind::object capsule =
81 nanobind::steal<nanobind::object>(mlirPythonAffineMapToCapsule(v));
82 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
83 .attr("AffineMap")
84 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
85 .release();
86 }
87};
88
89/// Casts object <-> MlirAttribute.
90template <>
91struct type_caster<MlirAttribute> {
92 NB_TYPE_CASTER(MlirAttribute, const_name("MlirAttribute"))
93 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
94 nanobind::object capsule = mlirApiObjectToCapsule(src);
95 value = mlirPythonCapsuleToAttribute(capsule.ptr());
96 return !mlirAttributeIsNull(value);
97 }
98 static handle from_cpp(MlirAttribute v, rv_policy,
99 cleanup_list *cleanup) noexcept {
100 nanobind::object capsule =
101 nanobind::steal<nanobind::object>(mlirPythonAttributeToCapsule(v));
102 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
103 .attr("Attribute")
104 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
105 .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
106 .release();
107 }
108};
109
110/// Casts object -> MlirBlock.
111template <>
112struct type_caster<MlirBlock> {
113 NB_TYPE_CASTER(MlirBlock, const_name("MlirBlock"))
114 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
115 nanobind::object capsule = mlirApiObjectToCapsule(src);
116 value = mlirPythonCapsuleToBlock(capsule.ptr());
117 return !mlirBlockIsNull(value);
118 }
119};
120
121/// Casts object -> MlirContext.
122template <>
123struct type_caster<MlirContext> {
124 NB_TYPE_CASTER(MlirContext, const_name("MlirContext"))
125 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
126 if (src.is_none()) {
127 // Gets the current thread-bound context.
128 // TODO: This raises an error of "No current context" currently.
129 // Update the implementation to pretty-print the helpful error that the
130 // core implementations print in this case.
131 src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
132 .attr("Context")
133 .attr("current");
134 }
135 nanobind::object capsule = mlirApiObjectToCapsule(src);
136 value = mlirPythonCapsuleToContext(capsule.ptr());
137 return !mlirContextIsNull(value);
138 }
139};
140
141/// Casts object <-> MlirDialectRegistry.
142template <>
143struct type_caster<MlirDialectRegistry> {
144 NB_TYPE_CASTER(MlirDialectRegistry, const_name("MlirDialectRegistry"))
145 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
146 nanobind::object capsule = mlirApiObjectToCapsule(src);
147 value = mlirPythonCapsuleToDialectRegistry(capsule.ptr());
148 return !mlirDialectRegistryIsNull(value);
149 }
150 static handle from_cpp(MlirDialectRegistry v, rv_policy,
151 cleanup_list *cleanup) noexcept {
152 nanobind::object capsule = nanobind::steal<nanobind::object>(
153 mlirPythonDialectRegistryToCapsule(v));
154 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
155 .attr("DialectRegistry")
156 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
157 .release();
158 }
159};
160
161/// Casts object <-> MlirLocation.
162template <>
163struct type_caster<MlirLocation> {
164 NB_TYPE_CASTER(MlirLocation, const_name("MlirLocation"))
165 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
166 if (src.is_none()) {
167 // Gets the current thread-bound context.
168 src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
169 .attr("Location")
170 .attr("current");
171 }
172 nanobind::object capsule = mlirApiObjectToCapsule(src);
173 value = mlirPythonCapsuleToLocation(capsule.ptr());
174 return !mlirLocationIsNull(value);
175 }
176 static handle from_cpp(MlirLocation v, rv_policy,
177 cleanup_list *cleanup) noexcept {
178 nanobind::object capsule =
179 nanobind::steal<nanobind::object>(mlirPythonLocationToCapsule(v));
180 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
181 .attr("Location")
182 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
183 .release();
184 }
185};
186
187/// Casts object <-> MlirModule.
188template <>
189struct type_caster<MlirModule> {
190 NB_TYPE_CASTER(MlirModule, const_name("MlirModule"))
191 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
192 nanobind::object capsule = mlirApiObjectToCapsule(src);
193 value = mlirPythonCapsuleToModule(capsule.ptr());
194 return !mlirModuleIsNull(value);
195 }
196 static handle from_cpp(MlirModule v, rv_policy,
197 cleanup_list *cleanup) noexcept {
198 nanobind::object capsule =
199 nanobind::steal<nanobind::object>(mlirPythonModuleToCapsule(v));
200 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
201 .attr("Module")
202 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
203 .release();
204 };
205};
206
207/// Casts object <-> MlirFrozenRewritePatternSet.
208template <>
209struct type_caster<MlirFrozenRewritePatternSet> {
210 NB_TYPE_CASTER(MlirFrozenRewritePatternSet,
211 const_name("MlirFrozenRewritePatternSet"))
212 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
213 nanobind::object capsule = mlirApiObjectToCapsule(src);
214 value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr());
215 return value.ptr != nullptr;
216 }
217 static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy, handle) {
218 nanobind::object capsule = nanobind::steal<nanobind::object>(
219 mlirPythonFrozenRewritePatternSetToCapsule(v));
220 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("rewrite"))
221 .attr("FrozenRewritePatternSet")
222 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
223 .release();
224 };
225};
226
227/// Casts object <-> MlirOperation.
228template <>
229struct type_caster<MlirOperation> {
230 NB_TYPE_CASTER(MlirOperation, const_name("MlirOperation"))
231 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
232 nanobind::object capsule = mlirApiObjectToCapsule(src);
233 value = mlirPythonCapsuleToOperation(capsule.ptr());
234 return !mlirOperationIsNull(value);
235 }
236 static handle from_cpp(MlirOperation v, rv_policy,
237 cleanup_list *cleanup) noexcept {
238 if (v.ptr == nullptr)
239 return nanobind::none();
240 nanobind::object capsule =
241 nanobind::steal<nanobind::object>(mlirPythonOperationToCapsule(v));
242 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
243 .attr("Operation")
244 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
245 .release();
246 };
247};
248
249/// Casts object <-> MlirValue.
250template <>
251struct type_caster<MlirValue> {
252 NB_TYPE_CASTER(MlirValue, const_name("MlirValue"))
253 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
254 nanobind::object capsule = mlirApiObjectToCapsule(src);
255 value = mlirPythonCapsuleToValue(capsule.ptr());
256 return !mlirValueIsNull(value);
257 }
258 static handle from_cpp(MlirValue v, rv_policy,
259 cleanup_list *cleanup) noexcept {
260 if (v.ptr == nullptr)
261 return nanobind::none();
262 nanobind::object capsule =
263 nanobind::steal<nanobind::object>(mlirPythonValueToCapsule(v));
264 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
265 .attr("Value")
266 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
267 .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
268 .release();
269 };
270};
271
272/// Casts object -> MlirPassManager.
273template <>
274struct type_caster<MlirPassManager> {
275 NB_TYPE_CASTER(MlirPassManager, const_name("MlirPassManager"))
276 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
277 nanobind::object capsule = mlirApiObjectToCapsule(src);
278 value = mlirPythonCapsuleToPassManager(capsule.ptr());
279 return !mlirPassManagerIsNull(value);
280 }
281};
282
283/// Casts object <-> MlirTypeID.
284template <>
285struct type_caster<MlirTypeID> {
286 NB_TYPE_CASTER(MlirTypeID, const_name("MlirTypeID"))
287 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
288 nanobind::object capsule = mlirApiObjectToCapsule(src);
289 value = mlirPythonCapsuleToTypeID(capsule.ptr());
290 return !mlirTypeIDIsNull(value);
291 }
292 static handle from_cpp(MlirTypeID v, rv_policy,
293 cleanup_list *cleanup) noexcept {
294 if (v.ptr == nullptr)
295 return nanobind::none();
296 nanobind::object capsule =
297 nanobind::steal<nanobind::object>(mlirPythonTypeIDToCapsule(v));
298 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
299 .attr("TypeID")
300 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
301 .release();
302 };
303};
304
305/// Casts object <-> MlirType.
306template <>
307struct type_caster<MlirType> {
308 NB_TYPE_CASTER(MlirType, const_name("MlirType"))
309 bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
310 nanobind::object capsule = mlirApiObjectToCapsule(src);
311 value = mlirPythonCapsuleToType(capsule.ptr());
312 return !mlirTypeIsNull(value);
313 }
314 static handle from_cpp(MlirType t, rv_policy,
315 cleanup_list *cleanup) noexcept {
316 nanobind::object capsule =
317 nanobind::steal<nanobind::object>(mlirPythonTypeToCapsule(t));
318 return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
319 .attr("Type")
320 .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
321 .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
322 .release();
323 }
324};
325
326/// Casts MlirStringRef -> object.
327template <>
328struct type_caster<MlirStringRef> {
329 NB_TYPE_CASTER(MlirStringRef, const_name("MlirStringRef"))
330 static handle from_cpp(MlirStringRef s, rv_policy,
331 cleanup_list *cleanup) noexcept {
332 return nanobind::str(s.data, s.length).release();
333 }
334};
335
336} // namespace detail
337} // namespace nanobind
338
339namespace mlir {
340namespace python {
341namespace nanobind_adaptors {
342
343/// Provides a facility like nanobind::class_ for defining a new class in a
344/// scope, but this allows extension of an arbitrary Python class, defining
345/// methods on it is a similar way. Classes defined in this way are very similar
346/// to if defined in Python in the usual way but use nanobind machinery to
347/// do it. These are not "real" nanobind classes but pure Python classes
348/// with no relation to a concrete C++ class.
349///
350/// Derived from a discussion upstream:
351/// https://github.com/pybind/pybind11/issues/1193
352/// (plus a fair amount of extra curricular poking)
353/// TODO: If this proves useful, see about including it in nanobind.
354class pure_subclass {
355public:
356 pure_subclass(nanobind::handle scope, const char *derivedClassName,
357 const nanobind::object &superClass) {
358 nanobind::object pyType =
359 nanobind::borrow<nanobind::object>((PyObject *)&PyType_Type);
360 nanobind::object metaclass = pyType(superClass);
361 nanobind::dict attributes;
362
363 thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass),
364 attributes);
365 scope.attr(derivedClassName) = thisClass;
366 thisClass.attr("__module__") = scope.attr("__name__");
367 }
368
369 template <typename Func, typename... Extra>
370 pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
371 nanobind::object cf = nanobind::cpp_function(
372 std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
373 nanobind::scope(thisClass), extra...);
374 thisClass.attr(name) = cf;
375 return *this;
376 }
377
378 template <typename Func, typename... Extra>
379 pure_subclass &def_property_readonly(const char *name, Func &&f,
380 const Extra &...extra) {
381 nanobind::object cf = nanobind::cpp_function(
382 std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
383 nanobind::scope(thisClass), extra...);
384 auto builtinProperty =
385 nanobind::borrow<nanobind::object>((PyObject *)&PyProperty_Type);
386 thisClass.attr(name) = builtinProperty(cf);
387 return *this;
388 }
389
390 template <typename Func, typename... Extra>
391 pure_subclass &def_staticmethod(const char *name, Func &&f,
392 const Extra &...extra) {
393 static_assert(!std::is_member_function_pointer<Func>::value,
394 "def_staticmethod(...) called with a non-static member "
395 "function pointer");
396 nanobind::object cf = nanobind::cpp_function(
397 std::forward<Func>(f),
398 nanobind::name(name), // nanobind::scope(thisClass),
399 extra...);
400 thisClass.attr(name) = cf;
401 return *this;
402 }
403
404 template <typename Func, typename... Extra>
405 pure_subclass &def_classmethod(const char *name, Func &&f,
406 const Extra &...extra) {
407 static_assert(!std::is_member_function_pointer<Func>::value,
408 "def_classmethod(...) called with a non-static member "
409 "function pointer");
410 nanobind::object cf = nanobind::cpp_function(
411 std::forward<Func>(f),
412 nanobind::name(name), // nanobind::scope(thisClass),
413 extra...);
414 thisClass.attr(name) =
415 nanobind::borrow<nanobind::object>(PyClassMethod_New(cf.ptr()));
416 return *this;
417 }
418
419 nanobind::object get_class() const { return thisClass; }
420
421protected:
422 nanobind::object superClass;
423 nanobind::object thisClass;
424};
425
426/// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
427/// constructor and type checking methods.
428class mlir_attribute_subclass : public pure_subclass {
429public:
430 using IsAFunctionTy = bool (*)(MlirAttribute);
431 using GetTypeIDFunctionTy = MlirTypeID (*)();
432
433 /// Subclasses by looking up the super-class dynamically.
434 mlir_attribute_subclass(nanobind::handle scope, const char *attrClassName,
435 IsAFunctionTy isaFunction,
436 GetTypeIDFunctionTy getTypeIDFunction = nullptr)
437 : mlir_attribute_subclass(
438 scope, attrClassName, isaFunction,
439 nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
440 .attr("Attribute"),
441 getTypeIDFunction) {}
442
443 /// Subclasses with a provided mlir.ir.Attribute super-class. This must
444 /// be used if the subclass is being defined in the same extension module
445 /// as the mlir.ir class (otherwise, it will trigger a recursive
446 /// initialization).
447 mlir_attribute_subclass(nanobind::handle scope, const char *typeClassName,
448 IsAFunctionTy isaFunction,
449 const nanobind::object &superCls,
450 GetTypeIDFunctionTy getTypeIDFunction = nullptr)
451 : pure_subclass(scope, typeClassName, superCls) {
452 // Casting constructor. Note that it is hard, if not impossible, to properly
453 // call chain to parent `__init__` in nanobind due to its special handling
454 // for init functions that don't have a fully constructed self-reference,
455 // which makes it impossible to forward it to `__init__` of a superclass.
456 // Instead, provide a custom `__new__` and call that of a superclass, which
457 // eventually calls `__init__` of the superclass. Since attribute subclasses
458 // have no additional members, we can just return the instance thus created
459 // without amending it.
460 std::string captureTypeName(
461 typeClassName); // As string in case if typeClassName is not static.
462 nanobind::object newCf = nanobind::cpp_function(
463 [superCls, isaFunction, captureTypeName](
464 nanobind::object cls, nanobind::object otherAttribute) {
465 MlirAttribute rawAttribute =
466 nanobind::cast<MlirAttribute>(otherAttribute);
467 if (!isaFunction(rawAttribute)) {
468 auto origRepr =
469 nanobind::cast<std::string>(nanobind::repr(otherAttribute));
470 throw std::invalid_argument(
471 (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
472 " (from " + origRepr + ")")
473 .str());
474 }
475 nanobind::object self = superCls.attr("__new__")(cls, otherAttribute);
476 return self;
477 },
478 nanobind::name("__new__"), nanobind::arg("cls"),
479 nanobind::arg("cast_from_attr"));
480 thisClass.attr("__new__") = newCf;
481
482 // 'isinstance' method.
483 static const char kIsinstanceSig[] =
484 "def isinstance(other_attribute: " MAKE_MLIR_PYTHON_QUALNAME(
485 "ir") ".Attribute) -> bool";
486 def_staticmethod(
487 "isinstance",
488 [isaFunction](MlirAttribute other) { return isaFunction(other); },
489 nanobind::arg("other_attribute"), nanobind::sig(kIsinstanceSig));
490 def("__repr__", [superCls, captureTypeName](nanobind::object self) {
491 return nanobind::repr(superCls(self))
492 .attr("replace")(superCls.attr("__name__"), captureTypeName);
493 });
494 if (getTypeIDFunction) {
495 def_staticmethod("get_static_typeid",
496 [getTypeIDFunction]() { return getTypeIDFunction(); });
497 nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
498 .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
499 getTypeIDFunction())(nanobind::cpp_function(
500 [thisClass = thisClass](const nanobind::object &mlirAttribute) {
501 return thisClass(mlirAttribute);
502 }));
503 }
504 }
505};
506
507/// Creates a custom subclass of mlir.ir.Type, implementing a casting
508/// constructor and type checking methods.
509class mlir_type_subclass : public pure_subclass {
510public:
511 using IsAFunctionTy = bool (*)(MlirType);
512 using GetTypeIDFunctionTy = MlirTypeID (*)();
513
514 /// Subclasses by looking up the super-class dynamically.
515 mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
516 IsAFunctionTy isaFunction,
517 GetTypeIDFunctionTy getTypeIDFunction = nullptr)
518 : mlir_type_subclass(
519 scope, typeClassName, isaFunction,
520 nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
521 .attr("Type"),
522 getTypeIDFunction) {}
523
524 /// Subclasses with a provided mlir.ir.Type super-class. This must
525 /// be used if the subclass is being defined in the same extension module
526 /// as the mlir.ir class (otherwise, it will trigger a recursive
527 /// initialization).
528 mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
529 IsAFunctionTy isaFunction,
530 const nanobind::object &superCls,
531 GetTypeIDFunctionTy getTypeIDFunction = nullptr)
532 : pure_subclass(scope, typeClassName, superCls) {
533 // Casting constructor. Note that it is hard, if not impossible, to properly
534 // call chain to parent `__init__` in nanobind due to its special handling
535 // for init functions that don't have a fully constructed self-reference,
536 // which makes it impossible to forward it to `__init__` of a superclass.
537 // Instead, provide a custom `__new__` and call that of a superclass, which
538 // eventually calls `__init__` of the superclass. Since attribute subclasses
539 // have no additional members, we can just return the instance thus created
540 // without amending it.
541 std::string captureTypeName(
542 typeClassName); // As string in case if typeClassName is not static.
543 nanobind::object newCf = nanobind::cpp_function(
544 [superCls, isaFunction, captureTypeName](nanobind::object cls,
545 nanobind::object otherType) {
546 MlirType rawType = nanobind::cast<MlirType>(otherType);
547 if (!isaFunction(rawType)) {
548 auto origRepr =
549 nanobind::cast<std::string>(nanobind::repr(otherType));
550 throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
551 captureTypeName + " (from " +
552 origRepr + ")")
553 .str());
554 }
555 nanobind::object self = superCls.attr("__new__")(cls, otherType);
556 return self;
557 },
558 nanobind::name("__new__"), nanobind::arg("cls"),
559 nanobind::arg("cast_from_type"));
560 thisClass.attr("__new__") = newCf;
561
562 // 'isinstance' method.
563 static const char kIsinstanceSig[] =
564 "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME(
565 "ir") ".Type) -> bool";
566 def_staticmethod(
567 "isinstance",
568 [isaFunction](MlirType other) { return isaFunction(other); },
569 nanobind::arg("other_type"), nanobind::sig(kIsinstanceSig));
570 def("__repr__", [superCls, captureTypeName](nanobind::object self) {
571 return nanobind::cast<std::string>(
572 nanobind::repr(superCls(self))
573 .attr("replace")(superCls.attr("__name__"), captureTypeName));
574 });
575 if (getTypeIDFunction) {
576 // 'get_static_typeid' method.
577 // This is modeled as a static method instead of a static property because
578 // `def_property_readonly_static` is not available in `pure_subclass` and
579 // we do not want to introduce the complexity that pybind uses to
580 // implement it.
581 def_staticmethod("get_static_typeid",
582 [getTypeIDFunction]() { return getTypeIDFunction(); });
583 nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
584 .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
585 getTypeIDFunction())(nanobind::cpp_function(
586 [thisClass = thisClass](const nanobind::object &mlirType) {
587 return thisClass(mlirType);
588 }));
589 }
590 }
591};
592
593/// Creates a custom subclass of mlir.ir.Value, implementing a casting
594/// constructor and type checking methods.
595class mlir_value_subclass : public pure_subclass {
596public:
597 using IsAFunctionTy = bool (*)(MlirValue);
598
599 /// Subclasses by looking up the super-class dynamically.
600 mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
601 IsAFunctionTy isaFunction)
602 : mlir_value_subclass(
603 scope, valueClassName, isaFunction,
604 nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
605 .attr("Value")) {}
606
607 /// Subclasses with a provided mlir.ir.Value super-class. This must
608 /// be used if the subclass is being defined in the same extension module
609 /// as the mlir.ir class (otherwise, it will trigger a recursive
610 /// initialization).
611 mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
612 IsAFunctionTy isaFunction,
613 const nanobind::object &superCls)
614 : pure_subclass(scope, valueClassName, superCls) {
615 // Casting constructor. Note that it is hard, if not impossible, to properly
616 // call chain to parent `__init__` in nanobind due to its special handling
617 // for init functions that don't have a fully constructed self-reference,
618 // which makes it impossible to forward it to `__init__` of a superclass.
619 // Instead, provide a custom `__new__` and call that of a superclass, which
620 // eventually calls `__init__` of the superclass. Since attribute subclasses
621 // have no additional members, we can just return the instance thus created
622 // without amending it.
623 std::string captureValueName(
624 valueClassName); // As string in case if valueClassName is not static.
625 nanobind::object newCf = nanobind::cpp_function(
626 [superCls, isaFunction, captureValueName](nanobind::object cls,
627 nanobind::object otherValue) {
628 MlirValue rawValue = nanobind::cast<MlirValue>(otherValue);
629 if (!isaFunction(rawValue)) {
630 auto origRepr =
631 nanobind::cast<std::string>(nanobind::repr(otherValue));
632 throw std::invalid_argument((llvm::Twine("Cannot cast value to ") +
633 captureValueName + " (from " +
634 origRepr + ")")
635 .str());
636 }
637 nanobind::object self = superCls.attr("__new__")(cls, otherValue);
638 return self;
639 },
640 nanobind::name("__new__"), nanobind::arg("cls"),
641 nanobind::arg("cast_from_value"));
642 thisClass.attr("__new__") = newCf;
643
644 // 'isinstance' method.
645 static const char kIsinstanceSig[] =
646 "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME(
647 "ir") ".Value) -> bool";
648 def_staticmethod(
649 "isinstance",
650 [isaFunction](MlirValue other) { return isaFunction(other); },
651 nanobind::arg("other_value"), nanobind::sig(kIsinstanceSig));
652 }
653};
654
655} // namespace nanobind_adaptors
656
657} // namespace python
658} // namespace mlir
659
660#endif // MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
661

source code of mlir/include/mlir/Bindings/Python/NanobindAdaptors.h