Skip to content

Conversation

makslevental
Copy link
Contributor

@makslevental makslevental commented Sep 3, 2025

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:

  1. Currently upstream "core" bindings and "dialect" bindings (all of the 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) and isinstance(..., 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!
  2. The dialect bindings incur a sizeable performance penalty relative to the core bindings in that every single trip across the wire (either python->cpp or cpp->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 the Mlir* opaque struct/ptr. This side isn't so onerous;
    b. When going from cpp->python we call long-hand call Python import APIs and construct the Python object using _CAPICreate. Note, there at least 2 attr 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.
  3. There is no technical reason for this division. The only reason I've ever seen/gotten for this division is that downstream users aren't sophisticated enough to use these types. That seems inaccurate since there are now many downstream users that were originally purely upstream;
    a. To wit: [mlir][Python] Make PyShapedType public #106105 already made one such Py* type public but only partially - that code actually can't work outside of upstream because python::PyConcreteType<...> is not visible anywhere outside of mlir/lib/Bindings/Python.

Thus, this is the first PR of three (or four):

  1. Move most of the mlir/lib/Bindings/Python headers into mlir/include/Bindings/Python (i.e., those headers which have Py* types);
  2. Move the Py* types in IRTypes.cpp and IRAttributes.cpp and etc. into their corresponding headers (this enables downstream to overload based on these types);
  3. Rewrite the upstream Dialect bindings in terms of PyConcreteType<...> and PyConcreteAttribute<...> and deprecate mlir_type_subclass and mlir_attribute_subclass;
  4. Eventually (say two releases), remove mlir_*_subclass and NanobindAdaptors.h

Note, I've already tested this E2E here, I'm just splitting it for easier review.

Footnotes

  1. Specifically when the modules are imported using importlib, which occurs with nanobind's stubgen;

  2. Python values being typed correctly when exiting from cpp;

  3. The workaround we implemented was a class method for the dialect bindings called Class.isinstance(...);

@makslevental makslevental force-pushed the users/makslevental/move-py-types branch from 566f457 to d7c734b Compare September 3, 2025 03:41
@makslevental makslevental changed the title [MLIR][Python] move Py* types [MLIR][Python] move Py* types in headers Sep 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant