From d02c7aeb35a2ac92cd61545329bf8814c3b88de5 Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Mon, 8 Sep 2025 18:12:24 -0700 Subject: [PATCH] [MLIR][Python] fix stubgen (#157583) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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):** image **Build dir (for forthcoming typechecker tests):** image **Install dir:** image --- mlir/cmake/modules/AddMLIRPython.cmake | 41 ++++++++++++++++++-------- mlir/python/CMakeLists.txt | 21 +++++++++++++ mlir/python/mlir/_mlir_libs/.gitignore | 1 + 3 files changed, 50 insertions(+), 13 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 7d97e34147f0..85c80276c1bc 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" - ) - 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 a058cf926abd..8e7949480f21 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 1d6765818d34..8f0c82ab0aac 100644 --- a/mlir/python/mlir/_mlir_libs/.gitignore +++ b/mlir/python/mlir/_mlir_libs/.gitignore @@ -1 +1,2 @@ _mlir/**/*.pyi +*.pyi