[MLIR][Python] reland (narrower) type stub generation

[makslevental]

This a reland of #155741 which was reverted at #157831. This version is narrower in scope - it only turns on automatic stub generation for MLIRPythonExtension.Core._mlir and does not do anything automatically. Specifically, the only CMake code added to AddMLIRPython.cmake is the mlir_generate_type_stubs function which is then used only in a manual way. The API for mlir_generate_type_stubs is:

Arguments:
  MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
  DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the
    specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
    and the core bindings extension module (e.g., something like StandalonePythonModules.extension._mlir.dso).
  OUTPUT_DIR: The root output directory to emit the type stubs into.
  OUTPUTS: List of expected outputs.
  DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
  IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen.
  PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
Outputs:
  NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.

Downstream users should use mlir_generate_type_stubs in coordination with declare_mlir_python_sources to turn on stub generation for their own downstream dialect extensions and upstream dialect extensions if they so choose. Standalone example shows an example.

Note, downstream will also need to set -DMLIR_PYTHON_PACKAGE_PREFIX=... correctly for their bindings.

[makslevental]

Here's a brief retrospective of the failures involved in the previous PR:

1. We don't have robust testing of all the various types of install/build configs people use;
   • This has been tested against IREE:
     • [python] Test upstream bindings change iree-org/iree#21916
   • In an effort to improve upstream I have:
     • [MLIR][Standalone] test Standalone against installed distribution #157944
     • [MLIR] enable Standalone example test for Windows #158183
2. The reason why the core bindings are safe to generate is because they are independent of all other modules/packages/etc and thus they can be loaded directly from their build/install dir (_mlir_libs). On the contrary, the reason the dialect bindings are not safe is they use mlir_attribute_subclass, mlir_type_subclass, etc. which all perform an import mlir when the types are created, i.e., at module load time. What that means is the location of the core _mlir module needs to be known and it must be importable when the dialect modules are loaded;
   • This is a design flaw of dialect bindings (amongst others), one that I've discussed before, and one that can be fixed but will take more work.
3. In an ideal world we wouldn't need this upstream at all (it's just a shell-out to nanobind's stubgen.py). The problem is our CMake is fairly spaghetti-fied (the Python bindings stuff specifically). That's why the first PR did things automatically and why this narrow, less automatic, PR has so many comment lines. The solution is to refactor the various CMake functions (like declare_*) to have less abstract interfaces, to perform less magic, to enable downstream to extend. This PR is styled as such (a reusable function not plugged into any of the other functions but used manually).

[makslevental]

Note, this has been tested against IREE with @Hardcode84's help iree-org/iree#21916 where it pass PkgCI / Build Packages / Linux Release (x86_64)

[makslevental]

I see that this PR introduces find_package(nanobind 2.9 CONFIG REQUIRED). Before the dependency was on version 2.4. Is 2.9 really required or does 2.8 suffice? I cannot tell if 2.9 is really needed or not. But since there was no update from 2.4 to 2.6 or 2.8 before moving to 2.9 my gut feeling is that you simply took the latest version you can find upstream.

Your "gut feeling" is incorrect - 2.9 was specifically the first release that had this fix wjakob/nanobind#1124 which is what was needed for this PR.

[ingomueller-net]

A bunch of pyright checks are failing after this commit. I believe at least the following are due to problems in the MLIR code base, though I don't understand exactly how they are related to this PR:

• error: "__getitem__" method not defined on type "RegionSequence" (reportIndexIssue)
  Indeed, RegionSequence inherits from Slicable, which says in its documentation that it binds __len__ and __getitem__ but not through nanobind. I guess that this the reason why no stubs are generated for these two dunder methods.
• error: Object of type "object" is not callable
  Indeed, register_operation returns an nb::object. Since it is meant to be a decorator, I think it should return a Callable. It returns an object created from nb::cpp_function, which (curiously, IMO) does not return an nb::callable but an nb::object. I tried fixing this with a manually crafted nb::sig(.) -- that makes the error go away and -- curiously -- all other errors as well. This could mean that pyright is broken by the new signature, so it's not clear to me if that's a proper fix.
• error: Import symbol "X" has type "type[Y]", which is not assignable to declared type "type[Z]"
  This is almost certainly a consequence of the previous problem but since my attempt to fix it removed all error reports, I don't fully trust it.

@makslevental: What's your take? What's the best way to deal with this?

[ingomueller-net]

Minimal progress: I found that the following manually-specified signature allows pyright to continue to report all other errors (most of which I consider correctly detected problems, exception below):

def register_operation(dialect_class: type[T], *, replace: bool = False) -> type[T]:

I don't know how to handle the closely related error: Import symbol "X" has type "type[Y]", which is not assignable to declared type "type[Z]", though, which is also still reported.

[makslevental]

error: Import symbol "PlanRelOp" has type "type[substrait_mlir.dialects._substrait_ops_gen.PlanRelOp]", which is not assignable to declared type "type[substrait_mlir.dialects.substrait.PlanRelOp]"

suggests to me the signature for register_operation should be

def register_operation(dialect_class: type[T], *, replace: bool = False) -> type[U]:

?

[ingomueller-net]

I think you are right in that T and U are different classes (namely those for the dialect and the op, respectively). But I think the correct annotation is actually even closer to the following (because register_operation returns a function):

def register_operation(dialect_class: type[T], *, replace: bool = False) -> Callable[[type[U]], type[U]]:

However, neither your proposal nor my new one make the error: Import symbol "PlanRelOp" has type... go away. The way I can make it go away is to refrain from using the wildcard import, import the base class with from ... import PlanRelOp as _PlanRelOpBase, and then class PlanRelOp(_PlanRelOpBase). Unfortunately, this means that I have to list all generated ops manually for import, which isn't great. My plan was to "solve" this with:

from ._substrait_ops_gen import *  # type: ignore

If the annotation with Callable sounds like a good idea, I can create a PR.

For the error: "__getitem__" method not defined ..., my next idea would be to use a pattern file. Again, if you think that's a good idea, I can have a stab.

[makslevental]

For the error: "__getitem__" method not defined ..., my next idea would be to use a pattern file. Again, if you think that's a good idea, I can have a stab.

I believe this is because pyright expects some type to be generic - is it the builder classes - in which case a pattern file (or an explicit signature won't help) because they're not generic?

[ingomueller-net]

For the error: "__getitem__" method not defined ..., my next idea would be to use a pattern file. Again, if you think that's a good idea, I can have a stab.

I believe this is because pyright expects some type to be generic - is it the builder classes - in which case a pattern file (or an explicit signature won't help) because they're not generic?

I am not sure I understand this explanation. My understanding is that PyRegionList inherits from Slicable, which implements the __length__ and __getitem__ functions manually rather than through nanobind for performance reasons and does simply not provide type annotations that way. If I add the following annotation manually to ir.pyi, the error goes away, so I suppose a pattern file could do that as well:

def __getitem__(self, index: int) -> Region: ...

On a second though, I guess it would be cleaner to implement a solution for any class inhereting from Slicable, so I should probably investigate that.

[makslevental]

My understanding is that PyRegionList

Err sorry I lost track of the context. Yes your explanation makes sense. So then we should just add a generic type hint to sliceable? That CRTP already captures the name of the instantiating class.

[ingomueller-net]

If a generic type hint is possible, that'd be great. Unfortunately, I haven't found a way, though I stopped investigating after Gemini told me it couldn't be done.

I have looked into using pattern files. Unfortunately, I ran into problems with the CMake rule that MLIR uses: part of the rule expects the pattern file to be a target name (i.e., without slashes), while other parts expect it to be a path (which contains slashes). I experimented with some work-around and run into a limitation of nanobind, which doesn't allow to add type annotations; pattern files today only replace existing ones. I have created wjakob/nanobind#1235 to overcome this, though.

TL;DR: If someone (not me 😞) manages to make the generic type hints for Sliceable work, that'd probably be best. Otherwise, with some pending work, the pattern files should work as well.

[Ahajha]

Does nb::sig work? It completely replaces the type signature, which might be clunky, but it's an option and doesn't involve a separate file.

[Ahajha]

Does nb::sig work? It completely replaces the type signature, which might be clunky, but it's an option and doesn't involve a separate file.

[makslevental]

generic type hints for Sliceable

by generic i just mean parameterized by the CRTP instantiator. I can do it later today.

[makslevental]

prototype here #170551

[ingomueller-net]

Does nb::sig work? It completely replaces the type signature, which might be clunky, but it's an option and doesn't involve a separate file.

Unfortunately, that doesn't work because it can only be used for functions created by nanobind and the functions in question are not only not created by nanobind but not even PyFunctionObjects but plain C function pointers, so it's not possible to attach a typing annotation there. But I think Maks' #170551 works nicely 😃

[ingomueller-net]

... I think the correct annotation is actually even closer to the following (because register_operation returns a function):

def register_operation(dialect_class: type[T], *, replace: bool = False) -> Callable[[type[U]], type[U]]:

...
If the annotation with Callable sounds like a good idea, I can create a PR.

--> #170627