[MLIR][Python] fix stubgen #157583

makslevental · 2025-09-09T00:28:38Z

In #155741 I broke the cardinal rule of CMake: nothing happens when you think it happens 🤷. Meaning: declare_mlir_python_sources(SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi") wasn't picking up any sources because they aren't generated yet. This of course makes sense in retrospect (the stubs are generated as part of the build process, post extension compile, rather than the configure process).

Thus, the API needs to be:

GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen, relative to _mlir_libs.

Partially as a result of this omission, the stubs weren't being installed into either the build dir nor the install dir. That is also fixed now:

Source dir (for easy reference):

Build dir (for forthcoming typechecker tests):

Install dir:

llvmbot · 2025-09-09T00:55:00Z

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

In #155741 I broke the cardinal rule of CMake: nothing happens when you think it happens 🤷. Meaning: declare_mlir_python_sources(SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi") wasn't picking up any sources because they aren't generated yet. This of course makes sense in retrospect (the stubs are generated as part of the build process, post extension compile, rather than the configure process).

Thus, the API needs to be:

GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen, relative to _mlir_libs.

Partially as a result of this omission, the stubs weren't being installed into either the build dir nor the install dir. That is also fixed now:

Source dir (for easy reference):

Build dir (for forthcoming typechecker tests):

Install dir:

Full diff: https://github.com/llvm/llvm-project/pull/157583.diff

3 Files Affected:

(modified) mlir/cmake/modules/AddMLIRPython.cmake (+26-11)
(modified) mlir/python/CMakeLists.txt (+21)
(modified) mlir/python/mlir/_mlir_libs/.gitignore (+1)

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 7d97e34147f06..85c80276c1bcf 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -111,6 +111,11 @@ endfunction()
 # Outputs:
 #   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
 function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR)
+  cmake_parse_arguments(ARG
+    ""
+    ""
+    "OUTPUTS"
+    ${ARGN})
   if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
     set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
   elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
@@ -135,9 +140,9 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
       --output-dir
       "${OUTPUT_DIR}")
 
-  set(NB_STUBGEN_OUTPUT "${OUTPUT_DIR}/${MODULE_NAME}.pyi")
+  list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
   add_custom_command(
-    OUTPUT ${NB_STUBGEN_OUTPUT}
+    OUTPUT ${_generated_type_stubs}
     COMMAND ${NB_STUBGEN_CMD}
     WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
     DEPENDS
@@ -146,7 +151,7 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
       "${DEPENDS_TARGET}"
   )
   set(_name "MLIRPythonModuleStubs_${_module}")
-  add_custom_target("${_name}" ALL DEPENDS ${NB_STUBGEN_OUTPUT})
+  add_custom_target("${_name}" ALL DEPENDS ${_generated_type_stubs})
   set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
 endfunction()
 
@@ -166,12 +171,12 @@ endfunction()
 #     on. These will be collected for all extensions and put into an
 #     aggregate dylib that is linked against.
 #   PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
-#   GENERATE_TYPE_STUBS: Enable type stub generation.
+#   GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen relative to _mlir_libs.
 function(declare_mlir_python_extension name)
   cmake_parse_arguments(ARG
-    "GENERATE_TYPE_STUBS"
+    ""
     "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
-    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
+    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -302,15 +307,25 @@ function(add_mlir_python_modules name)
           ${_module_name}
           ${_extension_target}
           ${name}
-          "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs/_mlir"
+          "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
+          OUTPUTS "${_generate_type_stubs}"
         )
+        add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
+        set(_stubgen_target "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen")
         declare_mlir_python_sources(
-          "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen"
-          ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+          ${_stubgen_target}
+          ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
           ADD_TO_PARENT "${sources_target}"
-          SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi"
+          SOURCES "${_generate_type_stubs}"
         )
-        add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
+        set(_pure_sources_target "${modules_target}.sources.${sources_target}_type_stub_gen")
+        add_mlir_python_sources_target(${_pure_sources_target}
+          INSTALL_COMPONENT ${modules_target}
+          INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
+          OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+          SOURCES_TARGETS ${_stubgen_target}
+        )
+        add_dependencies(${modules_target} ${_pure_sources_target})
       endif()
     else()
       message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index a058cf926abdc..8e7949480f21e 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -506,6 +506,10 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     # Dialects
     MLIRCAPIFunc
   GENERATE_TYPE_STUBS
+    "_mlir/__init__.pyi"
+    "_mlir/ir.pyi"
+    "_mlir/passmanager.pyi"
+    "_mlir/rewrite.pyi"
 )
 
 # This extension exposes an API to register all dialects, extensions, and passes
@@ -528,6 +532,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPITransforms
     MLIRCAPIRegisterEverything
   GENERATE_TYPE_STUBS
+    "_mlirRegisterEverything.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -543,6 +548,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
     MLIRCAPIIR
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
+    "_mlirDialectsLinalg.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -558,6 +564,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
     MLIRCAPIIR
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
+    "_mlirDialectsGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -573,6 +580,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
     MLIRCAPIIR
     MLIRCAPILLVM
   GENERATE_TYPE_STUBS
+    "_mlirDialectsLLVM.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -588,6 +596,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
     MLIRCAPIIR
     MLIRCAPIQuant
   GENERATE_TYPE_STUBS
+    "_mlirDialectsQuant.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -603,6 +612,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
     MLIRCAPIIR
     MLIRCAPINVGPU
   GENERATE_TYPE_STUBS
+    "_mlirDialectsNVGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -618,6 +628,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
     MLIRCAPIIR
     MLIRCAPIPDL
   GENERATE_TYPE_STUBS
+    "_mlirDialectsPDL.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -633,6 +644,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
     MLIRCAPIIR
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
+    "_mlirDialectsSparseTensor.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -648,6 +660,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
     MLIRCAPIIR
     MLIRCAPITransformDialect
   GENERATE_TYPE_STUBS
+    "_mlirDialectsTransform.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -662,6 +675,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIAsync
   GENERATE_TYPE_STUBS
+    "_mlirAsyncPasses.pyi"
 )
 
 if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -677,6 +691,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIExecutionEngine
     GENERATE_TYPE_STUBS
+      "_mlirExecutionEngine.pyi"
   )
 endif()
 
@@ -692,6 +707,7 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
+    "_mlirGPUPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -706,6 +722,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
+    "_mlirLinalgPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -724,6 +741,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
   GENERATE_TYPE_STUBS
+    "_mlirDialectsSMT.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -738,6 +756,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
+    "_mlirSparseTensorPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -752,6 +771,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransformDialectTransforms
   GENERATE_TYPE_STUBS
+    "_mlirTransformInterpreter.pyi"
 )
 
 # TODO: Figure out how to put this in the test tree.
@@ -811,6 +831,7 @@ if(MLIR_INCLUDE_TESTS)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
     GENERATE_TYPE_STUBS
+      "_mlirPythonTestNanobind.pyi"
   )
 endif()
 
diff --git a/mlir/python/mlir/_mlir_libs/.gitignore b/mlir/python/mlir/_mlir_libs/.gitignore
index 1d6765818d346..8f0c82ab0aac9 100644
--- a/mlir/python/mlir/_mlir_libs/.gitignore
+++ b/mlir/python/mlir/_mlir_libs/.gitignore
@@ -1 +1,2 @@
 _mlir/**/*.pyi
+*.pyi

rkayaith · 2025-09-09T01:03:48Z

mlir/cmake/modules/AddMLIRPython.cmake

 #     aggregate dylib that is linked against.
 #   PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
-#   GENERATE_TYPE_STUBS: Enable type stub generation.
+#   GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen relative to _mlir_libs.


out of curiosity, what happens if someone gets this list wrong?

CMake will try to create dependencies on files that are never generated - you'll get (admittedly confusing) errors like

CMake Error at /Users/maksimlevental/dev_projects/llvm-project/mlir/cmake/modules/AddMLIRPython.cmake:79 (target_sources): Cannot find source file: /Users/maksimlevental/dev_projects/llvm-project/mlir/python/mlir/_mlirAsyncPasses.pyi

even though that file doesn't exist until after the first generation even if you get everything right.

It's bad UX but there's no way around it; CMake sets up a task graph during the configure step right, where "edges" are built files. Since it can't know what stubgen will generate before the extension itself is built you have to supply these "edges" yourself.

an error sounds better than subtle silent issues 👍

krzysz00

Seems fine, I'm trusting this is a reasonable emergency fix

This reverts commit d02c7ae.

ftynse · 2025-09-10T09:42:18Z

I seem to be still getting these at d4f7995

CMake Error at cmake_install.cmake:2070 (file):
  file INSTALL cannot find
  "~/git/llvm-project/build/tools/mlir/python/_mlir_libs/_mlir/rewrite.pyi":
  No such file or directory.
Call Stack (most recent call first):
  ~/git/llvm-project/build/tools/mlir/cmake_install.cmake:102 (include)
  ~/git/llvm-project/build/tools/cmake_install.cmake:82 (include)
  ~/git/llvm-project/build/cmake_install.cmake:132 (include)

Only for top-level _mlir_libs/_mlir though.

I'm running the install-distribution target with

-DLLVM_DISTRIBUTION_COMPONENTS="llvm-headers;llvm-libraries;cmake-exports;FileCheck;count;not;mlir-headers;mlir-libraries;mlir-cmake-exports;mlir-tblgen;mlir-python-sources"

in cmake initially.

ftynse · 2025-09-10T09:52:12Z

When I work around the missing files, building a downstream from the installation hits further problems:

FAILED: python/_mlir_libs/_mlirDialectsGPU.pyi ~/git/downstream/build/python/_mlir_libs/_mlirDialectsGPU.pyi 
cd ~/usr-llvm/lib/cmake/mlir && ~/git/downstream/.venv/bin/python3.11 ~/git/downstream/.venv/lib/python3.11/site-packages/nanobind/stubgen.py --module ._mlir_libs._mlirDialectsGPU -i ~/git/downstream --recursive --include-private --output-dir ~/git/downstream/build/python/_mlir_libs --quiet
Traceback (most recent call last):
  File "~/git/downstream/.venv/lib/python3.11/site-packages/nanobind/stubgen.py", line 1444, in <module>
    main()
  File "~/git/downstream/.venv/lib/python3.11/site-packages/nanobind/stubgen.py", line 1371, in main
    mod_imported = importlib.import_module(mod)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/importlib/__init__.py", line 121, in import_module
    raise TypeError(msg.format(name))
TypeError: the 'package' argument is required to perform a relative import for '._mlir_libs._mlirDialectsGPU'

This reverts commit d02c7ae.

Despite several hotfixes, things remain broken, in particular: - installation/distribution (`ninja install / install-distribution`); - downstream projects with bindings exposed. See #157583 (comment) for more details. Reverts #155741, #157583, #157697. Let's make sure things are fixed and re-land as a unit.

Despite several hotfixes, things remain broken, in particular: - installation/distribution (`ninja install / install-distribution`); - downstream projects with bindings exposed. See llvm/llvm-project#157583 (comment) for more details. Reverts #155741, #157583, #157697. Let's make sure things are fixed and re-land as a unit.

makslevental force-pushed the users/makslevental/fix-stubgen branch 2 times, most recently from 94db74b to aaa1368 Compare September 9, 2025 00:42

[MLIR][Python] fix stubgen

bf5663e

makslevental force-pushed the users/makslevental/fix-stubgen branch from aaa1368 to bf5663e Compare September 9, 2025 00:45

makslevental marked this pull request as ready for review September 9, 2025 00:54

makslevental requested review from ftynse, stellaraccident and rolfmorel as code owners September 9, 2025 00:54

llvmbot added mlir:python MLIR Python bindings mlir labels Sep 9, 2025

rkayaith approved these changes Sep 9, 2025

View reviewed changes

krzysz00 approved these changes Sep 9, 2025

View reviewed changes

makslevental merged commit d02c7ae into main Sep 9, 2025
13 checks passed

makslevental deleted the users/makslevental/fix-stubgen branch September 9, 2025 01:12

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

041ad97

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

2147c8a

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

cfadbe7

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

30f4847

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

0b9a8aa

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

b36a41c

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

f94f0af

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

861b5fd

This reverts commit d02c7ae.

krzysz00 pushed a commit to iree-org/llvm-project that referenced this pull request Sep 9, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

3106ffe

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 10, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

32e0f3f

This reverts commit d02c7ae.

rkayaith added a commit to iree-org/llvm-project that referenced this pull request Sep 10, 2025

Revert "[MLIR][Python] fix stubgen (llvm#157583)"

2f4aa33

This reverts commit d02c7ae.

ftynse added a commit that referenced this pull request Sep 10, 2025

Revert "[MLIR][Python] fix stubgen (#157583)"

6b6fab0

This reverts commit d02c7ae.

ftynse mentioned this pull request Sep 10, 2025

Revert stubgen-related patches #157831

Merged

makslevental mentioned this pull request Sep 10, 2025

[MLIR][Python] reland stubgen #157853

Closed

makslevental mentioned this pull request Sep 11, 2025

[MLIR][Standalone] test Standalone against installed distribution #157944

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][Python] fix stubgen #157583

[MLIR][Python] fix stubgen #157583

makslevental commented Sep 9, 2025 •

edited

Loading

Uh oh!

llvmbot commented Sep 9, 2025

Uh oh!

rkayaith Sep 9, 2025

Uh oh!

makslevental Sep 9, 2025

Uh oh!

rkayaith Sep 9, 2025

Uh oh!

krzysz00 left a comment

Uh oh!

Uh oh!

ftynse commented Sep 10, 2025 •

edited

Loading

Uh oh!

ftynse commented Sep 10, 2025

Uh oh!

Uh oh!

[MLIR][Python] fix stubgen #157583

[MLIR][Python] fix stubgen #157583

Conversation

makslevental commented Sep 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 9, 2025

Uh oh!

rkayaith Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

makslevental Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

rkayaith Sep 9, 2025

Choose a reason for hiding this comment

Uh oh!

krzysz00 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ftynse commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftynse commented Sep 10, 2025

Uh oh!

Uh oh!

makslevental commented Sep 9, 2025 •

edited

Loading

ftynse commented Sep 10, 2025 •

edited

Loading