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

makslevental · web-flow · commit d02c7aeb35a2 · 2025-09-08T18:12:24.000-07:00
In llvm#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" />
diff --git 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
@@ -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
@@ -1 +1,2 @@
 _mlir/**/*.pyi
+*.pyi
