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

In https://github.com/llvm/llvm-project/pull/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):** <img width="300" height="674" alt="image" src="https://github.com/user-attachments/assets/a569f066-c2bd-4361-91f3-1c75181e51da" /> **Build dir (for forthcoming typechecker tests):** <img width="398" height="551" alt="image" src="https://github.com/user-attachments/assets/017859f9-fddb-49ee-85e5-915f5b5f7377" /> **Install dir:** <img width="456" height="884" alt="image" src="https://github.com/user-attachments/assets/8051be7e-898c-4ec8-a11e-e2408b241a56" />
2026-01-14 03:50:17 +08:00 · 2025-09-08 18:12:24 -07:00
parent 7d04e37904
commit d02c7aeb35
3 changed files with 50 additions and 13 deletions
--- 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"
-        )
-        declare_mlir_python_sources(
-          "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen"
-          ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-          ADD_TO_PARENT "${sources_target}"
-          SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi"
+          "${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(
+          ${_stubgen_target}
+          ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
+          ADD_TO_PARENT "${sources_target}"
+          SOURCES "${_generate_type_stubs}"
+        )
+        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}")
--- 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()

--- a/mlir/python/mlir/_mlir_libs/.gitignore
+++ b/mlir/python/mlir/_mlir_libs/.gitignore
@@ -1 +1,2 @@
 _mlir/**/*.pyi
+*.pyi