[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"
/>diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 7d97e34..85c8027 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -111,6 +111,11 @@
# 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 @@
--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 @@
"${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 @@
# 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 @@
${_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}")
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index a058cf9..8e79494 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -506,6 +506,10 @@
# 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 @@
MLIRCAPITransforms
MLIRCAPIRegisterEverything
GENERATE_TYPE_STUBS
+ "_mlirRegisterEverything.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -543,6 +548,7 @@
MLIRCAPIIR
MLIRCAPILinalg
GENERATE_TYPE_STUBS
+ "_mlirDialectsLinalg.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -558,6 +564,7 @@
MLIRCAPIIR
MLIRCAPIGPU
GENERATE_TYPE_STUBS
+ "_mlirDialectsGPU.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -573,6 +580,7 @@
MLIRCAPIIR
MLIRCAPILLVM
GENERATE_TYPE_STUBS
+ "_mlirDialectsLLVM.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -588,6 +596,7 @@
MLIRCAPIIR
MLIRCAPIQuant
GENERATE_TYPE_STUBS
+ "_mlirDialectsQuant.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -603,6 +612,7 @@
MLIRCAPIIR
MLIRCAPINVGPU
GENERATE_TYPE_STUBS
+ "_mlirDialectsNVGPU.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -618,6 +628,7 @@
MLIRCAPIIR
MLIRCAPIPDL
GENERATE_TYPE_STUBS
+ "_mlirDialectsPDL.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -633,6 +644,7 @@
MLIRCAPIIR
MLIRCAPISparseTensor
GENERATE_TYPE_STUBS
+ "_mlirDialectsSparseTensor.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -648,6 +660,7 @@
MLIRCAPIIR
MLIRCAPITransformDialect
GENERATE_TYPE_STUBS
+ "_mlirDialectsTransform.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -662,6 +675,7 @@
EMBED_CAPI_LINK_LIBS
MLIRCAPIAsync
GENERATE_TYPE_STUBS
+ "_mlirAsyncPasses.pyi"
)
if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -677,6 +691,7 @@
EMBED_CAPI_LINK_LIBS
MLIRCAPIExecutionEngine
GENERATE_TYPE_STUBS
+ "_mlirExecutionEngine.pyi"
)
endif()
@@ -692,6 +707,7 @@
EMBED_CAPI_LINK_LIBS
MLIRCAPIGPU
GENERATE_TYPE_STUBS
+ "_mlirGPUPasses.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -706,6 +722,7 @@
EMBED_CAPI_LINK_LIBS
MLIRCAPILinalg
GENERATE_TYPE_STUBS
+ "_mlirLinalgPasses.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -724,6 +741,7 @@
MLIRCAPISMT
MLIRCAPIExportSMTLIB
GENERATE_TYPE_STUBS
+ "_mlirDialectsSMT.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -738,6 +756,7 @@
EMBED_CAPI_LINK_LIBS
MLIRCAPISparseTensor
GENERATE_TYPE_STUBS
+ "_mlirSparseTensorPasses.pyi"
)
declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -752,6 +771,7 @@
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 @@
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 1d67658..8f0c82a 100644
--- a/mlir/python/mlir/_mlir_libs/.gitignore
+++ b/mlir/python/mlir/_mlir_libs/.gitignore
@@ -1 +1,2 @@
_mlir/**/*.pyi
+*.pyi
