[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)

[joker-eph]

It seems that this brought back the verbose logging in the build output:

[4707/4761] Generating type stubs using: /usr/bin/python3.10;/home/mamini/.local/lib/python3.10/site-packages/nanobind/stubgen.py;--module;mlir._mlir_libs._mlirPythonTestNanobin...vm-project/build/tools/mlir/python_packages/mlir_core;--recursive;--include-private;--output-dir;/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs
Module "mlir._mlir_libs._mlirPythonTestNanobind" ..
  - importing ..
  - analyzing ..
  - writing stub "/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs/_mlirPythonTestNanobind.pyi" ..
[4709/4761] Generating type stubs using: /usr/bin/python3.10;/home/mamini/.local/lib/python3.10/site-packages/nanobind/stubgen.py;--module;_mlir;-i;/llvm-pr.../tools/mlir/python_packages/mlir_core/mlir/_mlir_libs;--recursive;--include-private;--output-dir;/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs
Module "_mlir" ..
  - importing ..
  - analyzing ..
  - writing stub "/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/ir.pyi" ..
  - writing stub "/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/rewrite.pyi" ..
  - writing stub "/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/passmanager.pyi" ..
  - writing stub "/llvm-project/build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/__init__.pyi" ..
[4760/4761] Running the MLIR regression tests

Can you look into this please?

[makslevental]

It seems that this brought back the verbose logging in the build output:

@joker-eph only the first time right? the dep files are being written successfully

so you should only see this on a clean build.

[joker-eph]

But why do we print this in the first place? This is neither an error nor a warning or anything that would require the user attention, or is it?

[makslevental]

But why do we print this in the first place? This is neither an error nor a warning or anything that would require the user attention, or is it?

Because you asked for a mechanism for reporting failures in the type generation process itself:

#157853 (comment)

Prior when it was quiet, your error was hidden.

[joker-eph]

I don't follow:

1. there is no error here, so nothing to diagnose, I expect it to stay quiet.
2. my comment that you're linking to was describing the need for a version check of the dependency in cmake: that means that running cmake I would expect a message(FATAL_ERROR "Nanobind is too old, minimum version should be xx.yy").

[nikic]

Does MLIR have any kind of policy on bumping versions for Python dependencies? The LLVM project is generally very conservative about dependency requirements -- you can easily build LLVM against LTS distro packages. However, the MLIR nanobind requirement in particular is being bumped extremely aggressively -- the version required here is so new that it's not even part of the upcoming Fedora 43 release (a leading edge distro).

[kuhar]

Do distros prefer/expect to use system python packages? AFAIK, the general recommendation for MLIR-based projects is to use a virtual environment.

[nikic]

Yes, distros always use system packages.

And as this is a requirement of the LLVM build itself (not just for downstream dependents), I'd have to build LLVM inside a venv -- is this what people working on MLIR do? That sounds like really bad ergonomics.

[makslevental]

Echoing @kuhar - nanobind is a source dependency (it's always compiled during build of nanobind bindings) so there's no reason to depend on a distro's distribution and thus we advise users to either pip install or to use ExternalProject (or whatever source dependency mechanism they prefer).

I'd have to build LLVM inside a venv -- is this what people working on MLIR do? That sounds like really bad ergonomics.

If you build the Python bindings, you point the build (using Pyrhon_EXECUTABLE) to a venv yes. Not sure what's unergonomic about that?

[makslevental]

Speaking of ExternalProject, I wouldn't be opposed to migrating to using that instead of pip install (there are actually a whole mess of hacks in our CMake around finding nanobind and pybind in the venv).

[nikic]

If you build the Python bindings, you point the build (using Pyrhon_EXECUTABLE) to a venv yes. Not sure what's unergonomic about that?

I was going to say that adding mlir to the project list should work without having to perform any other arcane invocations, but apparently MLIR_ENABLE_BINDINGS_PYTHON is already disabled by default. This being a non-default configuration does make the requirement less bad than I initially thought.

But I feel like the fact that you have to add docs like https://mlir.llvm.org/docs/Bindings/Python/#recommended-development-practices indicates that something is going wrong here. Building LLVM should not require this kind of additional setup. I should not have to remember to activate a venv any time I work on LLVM.

[joker-eph]

You don't have to use a venv, you can install the pip requirements for the current user I believe (this is what I have setup locally)

[makslevental]

I should not have to remember to activate a venv any time I work on LLVM.

I guess I don't disagree with the essence of this but this is just by virtue of Python and the nature of Python dev, not nanobind. To wit: nanobind is not the only package in our requirements.txt you would need to get from pip:

https://packages.fedoraproject.org/search?query=ml-dtypes&releases=Fedora+41&releases=Fedora+42&releases=Fedora+43&start=0

[rkayaith]

But I feel like the fact that you have to add docs like https://mlir.llvm.org/docs/Bindings/Python/#recommended-development-practices indicates that something is going wrong here. Building LLVM should not require this kind of additional setup. I should not have to remember to activate a venv any time I work on LLVM.

This is pretty standard fare for python development. Also I believe you should only need to activate the venv when configuring (or alternatively it should be sufficient to just specify Python3_EXECUTABLE, without having to activate, as mentioned above)

[kwk]

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.

In other words, what is the feature/change from 2.9 that you depend on?

[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.