[MLIR][Python] move Py* types in headers #156575
Draft
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This refactors the Python bindings to expose the various
Py*
types as public API. There are micro and macro reasons for doing this.The micro reason (local problem we're facing downstream) is that in some configurations1 we trigger a circular import bug because "dialect" types and attributes perform an import of the root
_mlir
module when they are created (the types themselves, not even instances of those types).The macro reasons are several:
Dialect*
bindings here) are a two-tier system;a. core bindings enjoy first class support as far as type inference2, type stub generation, and ease of implementation, while dialect bindings have poorer support, incorrect type stub generation much more tedious (boilerplate) implementation;
b. Crucially, this two-tiered system is reflected in the fact that the two sets of types/attributes are not in the same Python object hierarchy. To wit:
isinstance(..., Type)
andisinstance(..., Attribute)
are not supported for the dialect bindings3;c. Since these types are not exposed in public headers, downstream users (dialect bindings or not) cannot write functions that overload on e.g.
PyFloat8*Type
- that's quite a useful feature!python->cpp
orcpp->python
) requires work in addition to nanobind's own casting/construction pipeline;a. When going from
python->cpp
, we extract the capsule object from the Python object and then extract from the capsule theMlir*
opaque struct/ptr. This side isn't so onerous;b. When going from
cpp->python
we call long-hand call Pythonimport
APIs and construct the Python object using_CAPICreate
. Note, there at least 2attr
calls incurred in addition to_CAPICreate
;c. All of this is much more efficiently handled by nanobind itself, as is done for the various
Py*
types.a. To wit: [mlir][Python] Make
PyShapedType
public #106105 already made one suchPy*
type public but only partially - that code actually can't work outside of upstream becausepython::PyConcreteType<...>
is not visible anywhere outside ofmlir/lib/Bindings/Python
.Thus, this is the first PR of three (or four):
mlir/lib/Bindings/Python
headers intomlir/include/Bindings/Python
(i.e., those headers which havePy*
types);Py*
types inIRTypes.cpp
andIRAttributes.cpp
and etc. into their corresponding headers (this enables downstream to overload based on these types);PyConcreteType<...>
andPyConcreteAttribute<...>
and deprecatemlir_type_subclass
andmlir_attribute_subclass
;mlir_*_subclass
andNanobindAdaptors.h
Note, I've already tested this E2E here, I'm just splitting it for easier review.
Footnotes
Specifically when the modules are imported using
importlib
, which occurs with nanobind's stubgen; ↩Python values being typed correctly when exiting from cpp; ↩
The workaround we implemented was a class method for the dialect bindings called
Class.isinstance(...)
; ↩