[Mlir-commits] [mlir] [MLIR][Python] fix stubgen (PR #157583)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 8 17:55:06 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Maksim Levental (makslevental)
<details>
<summary>Changes</summary>
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" />
---
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
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/157583
More information about the Mlir-commits
mailing list