Skip to content

[MLIR][Python] fix stubgen #157583

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 9, 2025
Merged

[MLIR][Python] fix stubgen #157583

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 26 additions & 11 deletions mlir/cmake/modules/AddMLIRPython.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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()

Expand All @@ -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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

an error sounds better than subtle silent issues 👍

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)
Expand Down Expand Up @@ -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}")
Expand Down
21 changes: 21 additions & 0 deletions mlir/python/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -677,6 +691,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
EMBED_CAPI_LINK_LIBS
MLIRCAPIExecutionEngine
GENERATE_TYPE_STUBS
"_mlirExecutionEngine.pyi"
)
endif()

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -811,6 +831,7 @@ if(MLIR_INCLUDE_TESTS)
EMBED_CAPI_LINK_LIBS
MLIRCAPIPythonTestDialect
GENERATE_TYPE_STUBS
"_mlirPythonTestNanobind.pyi"
)
endif()

Expand Down
1 change: 1 addition & 0 deletions mlir/python/mlir/_mlir_libs/.gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
_mlir/**/*.pyi
*.pyi