[mlir][Python] create MLIRPythonSupport

[makslevental]

What

This PR adds a shared library MLIRPythonSupport which contains all of the CRTP classes ike PyConcreteValue, PyConcreteType, PyConcreteAttribute, as well as other useful code like Defaulting* and etc enabling their reuse in downstream projects. Downstream projects can now do

struct PyTestType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyTestType> {
  ...
};

class PyTestAttr : public mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteAttribute<PyTestAttr> {
  ...
}

NB_MODULE(_mlirPythonTestNanobind, m) {
  PyTestType::bind(m);
  PyTestAttr::bind(m);
}

instead of using the discordant alternative mlir_type_subclass/mlir_attr_subclass (same goes for PyConcreteValue/mlir_value_subclass).

Why

This PR is mostly code motion (along with CMake) but before I describe the changes I want to state the goals/benefits:

1. Currently upstream "core" extensions and "dialect" extensions (all of the Dialect* extensions here) are a two-tier system;
   a. core extensions enjoy first class support as far as type inference¹, type stub generation, and ease of implementation, while dialect extensions 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 extensions²;
   c. Since these types are not exposed in public headers, downstream users (dialect extensions or not) cannot write functions that overload on e.g. PyFloat8*Type - that's quite a useful feature!
2. The dialect extensions incur a sizeable performance penalty relative to the core extensions 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; this is already much more efficiently handled by nanobind itself!
3. This division blocks various features: in some configurations³ 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). This blocks type stub generation for dialect extensions (i.e., the reason we currently only generate type stubs for _mlir).

How

Prior this was not done/possible because of "ODR" issues but I have resolved those issues; the basic idea for how we solve this is "move things we want to share into shared libraries":

1. Move IRCore (stuff like PyConcreteValue, PyConcreteType, PyConcreteAttribute) into MLIRPythonSupport;
   • Note, we move the rest of the things in IRModule.h (renamed to IRCore.h) because PyConcreteValue, PyConcreteType, PyConcreteAttribute depend on them. This makes for a bigger PR than one would hope for but ultimately I think we should give people access to these classes to use as they see fit (specifically inherit from, but also liberally use in bindings signatures instead of the opaque Mlir* struct wrappers).
   • Note, we move the classes into MLIRPythonSupport but we leave the bind calls in the _mlir extension. Users should not need to bind/re-bind these bindings (and they couldn't anyway - nanobind will warn if you try to bind the same class twice).
2. Put all of this code into a nested namespace MLIR_BINDINGS_PYTHON_DOMAIN which is determined by a compile time define (and tied to MLIR_BINDINGS_PYTHON_NB_DOMAIN). This is necessary in order to prevent conflicts on both symbol name and typeid (necessary for nanobind to not double register binded types) between multiple bindings libraries (e.g., torch-mlir, and jax). Note nanobind doesn't support module_local like pybind11. It does support NB_DOMAIN but that is not sufficient for disambiguating typeids across projects (to wit: we currently define NB_DOMAIN and it was still necessary to move everything to a nested namespace);
3. Build the nanobind library itself as a shared object (and link it to both the extensions and MLIRPythonSupport).
4. CMake to make this work, in-tree, out-of-tree, downstream, upstream, etc.

Testing

Three tests are added here

1. PythonTestModuleNanobind is ported to use PyConcreteType<PyTestType> instead of mlir_type_subclass and PyConcreteAttribute<PyTestAttr> instead of mlir_atrr_subclass, verifying this works for non-core extensions in-tree;
2. StandaloneExtensionNanobind is ported to use struct PyCustomType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyCustomType> instead of mlir_type_subclass verifying this works for non-core extensions out-of-tree;
3. StandaloneExtensionNanobind's smoketest is extended to also load another bindings package (namely mlir) verifying MLIR_BINDINGS_PYTHON_DOMAIN successfully disambiguates symbols and typeids.

I have also tested this downstream: llvm/eudsl#287 as well run the following builder bots:

mlir-gcc: https://lab.llvm.org/buildbot/#/builders/116/builds/22692

I have also tested against IREE: iree-org/iree#21916

Integration

It is highly recommended to set the CMake var MLIR_BINDINGS_PYTHON_NB_DOMAIN (which will also determine MLIR_BINDINGS_PYTHON_DOMAIN) to something unique for each downstream. This can also be passed explicitly to add_mlir_python_modules if your project builds multiple bindings packages. I added a WARNING to this effect in AddMLIRPython.cmake.

Notes

This is the first PR of a few - the follow-ups will move more "support" types into respective headers (e.g., RankedTensorType into IRTypes.h so that TestIntegerRankedTensorType can be ported).

Currently this PR has all the impls (members, methods, etc) in IRCore.h for easier review (so that the diff registers as much as possible as code motion). After review I will move all the impls into IRCore.cpp. Done.

Footnotes

1. Python values being typed correctly when exiting from cpp; ↩

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

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

[jpienaar]

The cmake change seems fine, it is non-trivial and does depend on specifics nanobind side a bit. But not sure that's avoidable/can be approved.

[jpienaar]

How would formatv approach look here?

[makslevental]

done

[makslevental]

ah this doesn't work because you start getting complaints like

ld64.lld: error: undefined symbol: typeinfo for llvm::support::detail::format_adapter

[makslevental]

The cmake change seems fine, it is non-trivial and does depend on specifics nanobind side a bit. But not sure that's avoidable/can be approved.

It would all be much nicer if we could just configure nanobind_add_module to either build a module or shared object (and maybe they would take the PR) but I'm not allowed to send them PRs currently 🤷‍♂️