From 907c7b626c0b9cbb8298fca33470e04b2eee3dc2 Mon Sep 17 00:00:00 2001 From: makslevental Date: Mon, 8 Sep 2025 18:21:49 -0700 Subject: [PATCH 1/9] [MLIR][Python] fix stubgen for FetchContent users --- mlir/cmake/modules/AddMLIRPython.cmake | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index ee6c87a8b635e..8d2a73f385b9d 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -116,10 +116,16 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP "" "OUTPUTS" ${ARGN}) + # for people doing find_package(nanobind) if(EXISTS ${nanobind_DIR}/../src/stubgen.py) set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py") elseif(EXISTS ${nanobind_DIR}/../stubgen.py) set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py") + # for people using FetchContent_Declare and FetchContent_MakeAvailable + elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py) + set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py") + elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py) + set(NB_STUBGEN "${nanobind_SOURCE_DIR}/stubgen.py") else() message(FATAL_ERROR "generate_type_stubs(): could not locate 'stubgen.py'!") endif() From ab5f2425ffab1bab8e7e3969fb60bcf174868dfb Mon Sep 17 00:00:00 2001 From: makslevental Date: Tue, 9 Sep 2025 12:23:25 -0700 Subject: [PATCH 2/9] fix a bunch of stuff --- mlir/cmake/modules/AddMLIRPython.cmake | 31 ++++++++++++------- mlir/examples/standalone/CMakeLists.txt | 2 +- .../examples/standalone/python/CMakeLists.txt | 3 ++ mlir/python/CMakeLists.txt | 2 ++ 4 files changed, 25 insertions(+), 13 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 8d2a73f385b9d..8564be133a66b 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -102,7 +102,7 @@ endfunction() # Function: generate_type_stubs # Turns on automatic type stub generation (via nanobind's stubgen) for extension modules. # Arguments: -# MODULE_NAME: The name of the extension module as specified in declare_mlir_python_extension. +# FQ_MODULE_NAME: The fully-qualified name of the extension module (used for importing in python). # DEPENDS_TARGET: The dso target corresponding to the extension module # (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso) # MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module @@ -110,7 +110,7 @@ endfunction() # OUTPUT_DIR: The root output directory to emit the type stubs into. # 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) +function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR) cmake_parse_arguments(ARG "" "" @@ -131,14 +131,13 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP endif() file(REAL_PATH "${NB_STUBGEN}" NB_STUBGEN) - set(_module "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs.${MODULE_NAME}") file(REAL_PATH "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." _import_path) set(NB_STUBGEN_CMD "${Python_EXECUTABLE}" "${NB_STUBGEN}" --module - "${_module}" + "${FQ_MODULE_NAME}" -i "${_import_path}" --recursive @@ -157,8 +156,8 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP "${MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python" "${DEPENDS_TARGET}" ) - set(_name "MLIRPythonModuleStubs_${_module}") - add_custom_target("${_name}" ALL DEPENDS ${_generated_type_stubs}) + set(_name "${FQ_MODULE_NAME}.type_stubs") + add_custom_target("${_name}" DEPENDS ${_generated_type_stubs}) set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE) endfunction() @@ -271,10 +270,14 @@ endfunction() # DAG of source modules is included. # COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every # extension depend on (see mlir_python_add_common_capi_library). +# GENERATE_TYPE_STUBS: Enable type stub generation for all modules +# which have enabled type stub generation. +# PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX. This is used +# to determine type stub generation python module names. function(add_mlir_python_modules name) cmake_parse_arguments(ARG - "" - "ROOT_PREFIX;INSTALL_PREFIX" + "GENERATE_TYPE_STUBS" + "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX" "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES" ${ARGN}) # Helper to process an individual target. @@ -309,23 +312,27 @@ function(add_mlir_python_modules name) add_dependencies(${modules_target} ${_extension_target}) mlir_python_setup_extension_rpath(${_extension_target}) get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS) - if(_generate_type_stubs) + if(ARG_GENERATE_TYPE_STUBS AND _generate_type_stubs) + if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) + message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX") + endif() + set(_fully_qualified_module_name "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}") generate_type_stubs( - ${_module_name} + ${_fully_qualified_module_name} ${_extension_target} ${name} "${CMAKE_CURRENT_BINARY_DIR}/_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") + set(_stubgen_target "${_fully_qualified_module_name}.type_stub_gen") declare_mlir_python_sources( ${_stubgen_target} ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs" ADD_TO_PARENT "${sources_target}" SOURCES "${_generate_type_stubs}" ) - set(_pure_sources_target "${modules_target}.sources.${sources_target}_type_stub_gen") + 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" diff --git a/mlir/examples/standalone/CMakeLists.txt b/mlir/examples/standalone/CMakeLists.txt index 88dfa3e5d57a3..d403e6dc3709e 100644 --- a/mlir/examples/standalone/CMakeLists.txt +++ b/mlir/examples/standalone/CMakeLists.txt @@ -55,7 +55,7 @@ if(MLIR_ENABLE_BINDINGS_PYTHON) include(MLIRDetectPythonEnv) mlir_configure_python_dev_packages() set(MLIR_PYTHON_PACKAGE_PREFIX "mlir_standalone" CACHE STRING "" FORCE) - set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/mlir_standalone" CACHE STRING "" FORCE) + set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/${MLIR_PYTHON_PACKAGE_PREFIX}" CACHE STRING "" FORCE) add_subdirectory(python) endif() add_subdirectory(test) diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index cb10518e94e33..cd1b070adc53c 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -40,6 +40,7 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension StandaloneCAPI PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS + "_standaloneDialectsNanobind.pyi" ) @@ -77,4 +78,6 @@ add_mlir_python_modules(StandalonePythonModules MLIRPythonSources.Dialects.builtin COMMON_CAPI_LINK_LIBS StandalonePythonCAPI + GENERATE_TYPE_STUBS + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 8e7949480f21e..65189a888ae65 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -888,4 +888,6 @@ add_mlir_python_modules(MLIRPythonModules ${_ADDL_TEST_SOURCES} COMMON_CAPI_LINK_LIBS MLIRPythonCAPI + GENERATE_TYPE_STUBS + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) From 5f4f855aa71dc564756c12ae7c3651cafe2b7991 Mon Sep 17 00:00:00 2001 From: makslevental Date: Tue, 9 Sep 2025 17:32:34 -0700 Subject: [PATCH 3/9] god what a pain --- mlir/cmake/modules/AddMLIRPython.cmake | 34 +++++------ .../examples/standalone/python/CMakeLists.txt | 5 +- mlir/python/CMakeLists.txt | 60 ++++++++++++------- 3 files changed, 59 insertions(+), 40 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 8564be133a66b..45ebfa818cb45 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -108,6 +108,7 @@ endfunction() # MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module # (e.g., something like StandalonePythonModules.extension._mlir.dso) # OUTPUT_DIR: The root output directory to emit the type stubs into. +# OUTPUTS: List of expected outputs. # Outputs: # NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on. function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR) @@ -181,7 +182,7 @@ endfunction() function(declare_mlir_python_extension name) cmake_parse_arguments(ARG "" - "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY" + "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY;PACKAGE_PREFIX" "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS" ${ARGN}) @@ -207,6 +208,19 @@ function(declare_mlir_python_extension name) mlir_python_GENERATE_TYPE_STUBS "${ARG_GENERATE_TYPE_STUBS}" ) + if(ARG_GENERATE_TYPE_STUBS) + if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) + message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}") + endif() + set(_stubgen_target "${ARG_PACKAGE_PREFIX}._mlir_libs.${name}.type_stub_gen") + declare_mlir_python_sources( + ${_stubgen_target} + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + ADD_TO_PARENT "${ARG_ADD_TO_PARENT}" + SOURCES "${ARG_GENERATE_TYPE_STUBS}" + ) + endif() + # Set the interface source and link_libs properties of the target # These properties support generator expressions and are automatically exported list(TRANSFORM ARG_SOURCES PREPEND "${ARG_ROOT_DIR}/" OUTPUT_VARIABLE _build_sources) @@ -317,29 +331,15 @@ function(add_mlir_python_modules name) message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX") endif() set(_fully_qualified_module_name "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}") + list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "") generate_type_stubs( ${_fully_qualified_module_name} ${_extension_target} ${name} - "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs" + "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" OUTPUTS "${_generate_type_stubs}" ) add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}") - set(_stubgen_target "${_fully_qualified_module_name}.type_stub_gen") - declare_mlir_python_sources( - ${_stubgen_target} - ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/_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/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index cd1b070adc53c..82fe9d1c0b95d 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -40,7 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension StandaloneCAPI PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS - "_standaloneDialectsNanobind.pyi" + "_mlir_libs/_standaloneDialectsNanobind.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) @@ -80,4 +81,4 @@ add_mlir_python_modules(StandalonePythonModules StandalonePythonCAPI GENERATE_TYPE_STUBS PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" - ) +) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 65189a888ae65..d48e238991d7f 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -506,10 +506,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Core # Dialects MLIRCAPIFunc GENERATE_TYPE_STUBS - "_mlir/__init__.pyi" - "_mlir/ir.pyi" - "_mlir/passmanager.pyi" - "_mlir/rewrite.pyi" + "_mlir_libs/_mlir/__init__.pyi" + "_mlir_libs/_mlir/ir.pyi" + "_mlir_libs/_mlir/passmanager.pyi" + "_mlir_libs/_mlir/rewrite.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) # This extension exposes an API to register all dialects, extensions, and passes @@ -532,7 +533,8 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MLIRCAPITransforms MLIRCAPIRegisterEverything GENERATE_TYPE_STUBS - "_mlirRegisterEverything.pyi" + "_mlir_libs/_mlirRegisterEverything.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind @@ -548,7 +550,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind MLIRCAPIIR MLIRCAPILinalg GENERATE_TYPE_STUBS - "_mlirDialectsLinalg.pyi" + "_mlir_libs/_mlirDialectsLinalg.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind @@ -564,7 +567,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind MLIRCAPIIR MLIRCAPIGPU GENERATE_TYPE_STUBS - "_mlirDialectsGPU.pyi" + "_mlir_libs/_mlirDialectsGPU.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind @@ -580,7 +584,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind MLIRCAPIIR MLIRCAPILLVM GENERATE_TYPE_STUBS - "_mlirDialectsLLVM.pyi" + "_mlir_libs/_mlirDialectsLLVM.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind @@ -596,7 +601,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind MLIRCAPIIR MLIRCAPIQuant GENERATE_TYPE_STUBS - "_mlirDialectsQuant.pyi" + "_mlir_libs/_mlirDialectsQuant.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind @@ -612,7 +618,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind MLIRCAPIIR MLIRCAPINVGPU GENERATE_TYPE_STUBS - "_mlirDialectsNVGPU.pyi" + "_mlir_libs/_mlirDialectsNVGPU.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind @@ -628,7 +635,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind MLIRCAPIIR MLIRCAPIPDL GENERATE_TYPE_STUBS - "_mlirDialectsPDL.pyi" + "_mlir_libs/_mlirDialectsPDL.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind @@ -644,7 +652,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind MLIRCAPIIR MLIRCAPISparseTensor GENERATE_TYPE_STUBS - "_mlirDialectsSparseTensor.pyi" + "_mlir_libs/_mlirDialectsSparseTensor.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind @@ -660,7 +669,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind MLIRCAPIIR MLIRCAPITransformDialect GENERATE_TYPE_STUBS - "_mlirDialectsTransform.pyi" + "_mlir_libs/_mlirDialectsTransform.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses @@ -675,7 +685,8 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses EMBED_CAPI_LINK_LIBS MLIRCAPIAsync GENERATE_TYPE_STUBS - "_mlirAsyncPasses.pyi" + "_mlir_libs/_mlirAsyncPasses.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) if(MLIR_ENABLE_EXECUTION_ENGINE) @@ -691,7 +702,8 @@ if(MLIR_ENABLE_EXECUTION_ENGINE) EMBED_CAPI_LINK_LIBS MLIRCAPIExecutionEngine GENERATE_TYPE_STUBS - "_mlirExecutionEngine.pyi" + "_mlir_libs/_mlirExecutionEngine.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) endif() @@ -707,7 +719,8 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses EMBED_CAPI_LINK_LIBS MLIRCAPIGPU GENERATE_TYPE_STUBS - "_mlirGPUPasses.pyi" + "_mlir_libs/_mlirGPUPasses.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses @@ -722,7 +735,8 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses EMBED_CAPI_LINK_LIBS MLIRCAPILinalg GENERATE_TYPE_STUBS - "_mlirLinalgPasses.pyi" + "_mlir_libs/_mlirLinalgPasses.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind @@ -741,7 +755,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind MLIRCAPISMT MLIRCAPIExportSMTLIB GENERATE_TYPE_STUBS - "_mlirDialectsSMT.pyi" + "_mlir_libs/_mlirDialectsSMT.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses @@ -756,7 +771,8 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses EMBED_CAPI_LINK_LIBS MLIRCAPISparseTensor GENERATE_TYPE_STUBS - "_mlirSparseTensorPasses.pyi" + "_mlir_libs/_mlirSparseTensorPasses.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter @@ -771,7 +787,8 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter EMBED_CAPI_LINK_LIBS MLIRCAPITransformDialectTransforms GENERATE_TYPE_STUBS - "_mlirTransformInterpreter.pyi" + "_mlir_libs/_mlirTransformInterpreter.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) # TODO: Figure out how to put this in the test tree. @@ -831,7 +848,8 @@ if(MLIR_INCLUDE_TESTS) EMBED_CAPI_LINK_LIBS MLIRCAPIPythonTestDialect GENERATE_TYPE_STUBS - "_mlirPythonTestNanobind.pyi" + "_mlir_libs/_mlirPythonTestNanobind.pyi" + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) endif() From 2ba888ebc3266cb400113533b9f621f8a978989f Mon Sep 17 00:00:00 2001 From: makslevental Date: Tue, 9 Sep 2025 18:06:14 -0700 Subject: [PATCH 4/9] remove PACKAGE_PREFIX --- mlir/cmake/modules/AddMLIRPython.cmake | 25 ++++++++----------- .../examples/standalone/python/CMakeLists.txt | 1 - mlir/python/CMakeLists.txt | 18 ------------- 3 files changed, 11 insertions(+), 33 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 45ebfa818cb45..7051213fdde6d 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -182,7 +182,7 @@ endfunction() function(declare_mlir_python_extension name) cmake_parse_arguments(ARG "" - "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY;PACKAGE_PREFIX" + "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY" "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS" ${ARGN}) @@ -196,6 +196,16 @@ function(declare_mlir_python_extension name) endif() add_library(${name} INTERFACE) + + if(ARG_GENERATE_TYPE_STUBS) + declare_mlir_python_sources( + "${name}.type_stub_gen" + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + ADD_TO_PARENT "${ARG_ADD_TO_PARENT}" + SOURCES "${ARG_GENERATE_TYPE_STUBS}" + ) + endif() + set_target_properties(${name} PROPERTIES # Yes: Leading-lowercase property names are load bearing and the recommended # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261 @@ -208,19 +218,6 @@ function(declare_mlir_python_extension name) mlir_python_GENERATE_TYPE_STUBS "${ARG_GENERATE_TYPE_STUBS}" ) - if(ARG_GENERATE_TYPE_STUBS) - if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) - message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}") - endif() - set(_stubgen_target "${ARG_PACKAGE_PREFIX}._mlir_libs.${name}.type_stub_gen") - declare_mlir_python_sources( - ${_stubgen_target} - ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" - ADD_TO_PARENT "${ARG_ADD_TO_PARENT}" - SOURCES "${ARG_GENERATE_TYPE_STUBS}" - ) - endif() - # Set the interface source and link_libs properties of the target # These properties support generator expressions and are automatically exported list(TRANSFORM ARG_SOURCES PREPEND "${ARG_ROOT_DIR}/" OUTPUT_VARIABLE _build_sources) diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index 82fe9d1c0b95d..24a0dee9299c3 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -41,7 +41,6 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS "_mlir_libs/_standaloneDialectsNanobind.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index d48e238991d7f..d1e1b08cd8607 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -510,7 +510,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core "_mlir_libs/_mlir/ir.pyi" "_mlir_libs/_mlir/passmanager.pyi" "_mlir_libs/_mlir/rewrite.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) # This extension exposes an API to register all dialects, extensions, and passes @@ -534,7 +533,6 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MLIRCAPIRegisterEverything GENERATE_TYPE_STUBS "_mlir_libs/_mlirRegisterEverything.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind @@ -551,7 +549,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind MLIRCAPILinalg GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsLinalg.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind @@ -568,7 +565,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind MLIRCAPIGPU GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsGPU.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind @@ -585,7 +581,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind MLIRCAPILLVM GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsLLVM.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind @@ -602,7 +597,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind MLIRCAPIQuant GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsQuant.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind @@ -619,7 +613,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind MLIRCAPINVGPU GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsNVGPU.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind @@ -636,7 +629,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind MLIRCAPIPDL GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsPDL.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind @@ -653,7 +645,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind MLIRCAPISparseTensor GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsSparseTensor.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind @@ -670,7 +661,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind MLIRCAPITransformDialect GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsTransform.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses @@ -686,7 +676,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses MLIRCAPIAsync GENERATE_TYPE_STUBS "_mlir_libs/_mlirAsyncPasses.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) if(MLIR_ENABLE_EXECUTION_ENGINE) @@ -703,7 +692,6 @@ if(MLIR_ENABLE_EXECUTION_ENGINE) MLIRCAPIExecutionEngine GENERATE_TYPE_STUBS "_mlir_libs/_mlirExecutionEngine.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) endif() @@ -720,7 +708,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses MLIRCAPIGPU GENERATE_TYPE_STUBS "_mlir_libs/_mlirGPUPasses.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses @@ -736,7 +723,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses MLIRCAPILinalg GENERATE_TYPE_STUBS "_mlir_libs/_mlirLinalgPasses.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind @@ -756,7 +742,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind MLIRCAPIExportSMTLIB GENERATE_TYPE_STUBS "_mlir_libs/_mlirDialectsSMT.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses @@ -772,7 +757,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses MLIRCAPISparseTensor GENERATE_TYPE_STUBS "_mlir_libs/_mlirSparseTensorPasses.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter @@ -788,7 +772,6 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter MLIRCAPITransformDialectTransforms GENERATE_TYPE_STUBS "_mlir_libs/_mlirTransformInterpreter.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) # TODO: Figure out how to put this in the test tree. @@ -849,7 +832,6 @@ if(MLIR_INCLUDE_TESTS) MLIRCAPIPythonTestDialect GENERATE_TYPE_STUBS "_mlir_libs/_mlirPythonTestNanobind.pyi" - PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) endif() From 2483348b7233e8b071cfcc0479de05c0388ed2e7 Mon Sep 17 00:00:00 2001 From: makslevental Date: Tue, 9 Sep 2025 18:13:20 -0700 Subject: [PATCH 5/9] remove mlir_libs prefix --- mlir/cmake/modules/AddMLIRPython.cmake | 61 ++++++++++++------- .../examples/standalone/python/CMakeLists.txt | 3 +- mlir/python/CMakeLists.txt | 55 ++++++----------- mlir/test/python/CMakeLists.txt | 2 +- 4 files changed, 59 insertions(+), 62 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 7051213fdde6d..338dd1e7ae98f 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -105,16 +105,16 @@ endfunction() # FQ_MODULE_NAME: The fully-qualified name of the extension module (used for importing in python). # DEPENDS_TARGET: The dso target corresponding to the extension module # (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso) -# MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module +# CORE_MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module # (e.g., something like StandalonePythonModules.extension._mlir.dso) # OUTPUT_DIR: The root output directory to emit the type stubs into. # OUTPUTS: List of expected outputs. # Outputs: # NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on. -function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR) +function(generate_type_stubs ) cmake_parse_arguments(ARG "" - "" + "FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR" "OUTPUTS" ${ARGN}) # for people doing find_package(nanobind) @@ -134,30 +134,30 @@ function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET O file(REAL_PATH "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." _import_path) - set(NB_STUBGEN_CMD + set(_nb_stubgen_cmd "${Python_EXECUTABLE}" "${NB_STUBGEN}" --module - "${FQ_MODULE_NAME}" + "${ARG_FQ_MODULE_NAME}" -i "${_import_path}" --recursive --include-private --output-dir - "${OUTPUT_DIR}" + "${ARG_OUTPUT_DIR}" --quiet) - list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs) + list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs) add_custom_command( OUTPUT ${_generated_type_stubs} - COMMAND ${NB_STUBGEN_CMD} + COMMAND ${_nb_stubgen_cmd} WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}" DEPENDS - "${MLIR_DEPENDS_TARGET}.extension._mlir.dso" - "${MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python" - "${DEPENDS_TARGET}" + "${ARG_CORE_MLIR_DEPENDS_TARGET}.extension._mlir.dso" + "${ARG_CORE_MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python" + "${ARG_DEPENDS_TARGET}" ) - set(_name "${FQ_MODULE_NAME}.type_stubs") + set(_name "${ARG_FQ_MODULE_NAME}.type_stubs") add_custom_target("${_name}" DEPENDS ${_generated_type_stubs}) set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE) endfunction() @@ -178,7 +178,10 @@ 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: List of generated type stubs expected from stubgen relative to _mlir_libs. +# GENERATE_TYPE_STUBS: Either +# 1. OFF (default) +# 2. ON if ${MODULE_NAME}.pyi is the only stub +# 3. A list of generated type stubs expected from stubgen relative to _mlir_libs. function(declare_mlir_python_extension name) cmake_parse_arguments(ARG "" @@ -197,7 +200,14 @@ function(declare_mlir_python_extension name) add_library(${name} INTERFACE) + if(NOT ARG_GENERATE_TYPE_STUBS) + set(ARG_GENERATE_TYPE_STUBS OFF) + endif() + if("${ARG_GENERATE_TYPE_STUBS}" STREQUAL "ON") + set(ARG_GENERATE_TYPE_STUBS "${ARG_MODULE_NAME}.pyi") + endif() if(ARG_GENERATE_TYPE_STUBS) + list(TRANSFORM ARG_GENERATE_TYPE_STUBS PREPEND "_mlir_libs/") declare_mlir_python_sources( "${name}.type_stub_gen" ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" @@ -282,9 +292,9 @@ endfunction() # COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every # extension depend on (see mlir_python_add_common_capi_library). # GENERATE_TYPE_STUBS: Enable type stub generation for all modules -# which have enabled type stub generation. -# PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX. This is used -# to determine type stub generation python module names. +# which have declared type stub generation with `declare_mlir_python_extension`. +# PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`. +# This is used to determine type stub generation python module names. function(add_mlir_python_modules name) cmake_parse_arguments(ARG "GENERATE_TYPE_STUBS" @@ -325,15 +335,22 @@ function(add_mlir_python_modules name) get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS) if(ARG_GENERATE_TYPE_STUBS AND _generate_type_stubs) if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) - message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX") + message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}") endif() - set(_fully_qualified_module_name "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}") + # TL;DR: everything here is load bearing and annoyingly coupled. Changing anything here + # (or in declare_mlir_python_extension) will break either in-tree or out-of-tree generation. + # + # We remove _mlir_libs here because OUTPUT_DIR already includes it. + # Specifically OUTPUT_DIR already includes it because that's the actual directory + # where we want stubgen to dump the emitted sources. This is load bearing because up above + # (in declare_mlir_python_extension) we prefixed all the paths with _mlir_libs and + # we specified declare_mlir_python_sources with ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs". list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "") generate_type_stubs( - ${_fully_qualified_module_name} - ${_extension_target} - ${name} - "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" + FQ_MODULE_NAME "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}" + DEPENDS_TARGET ${_extension_target} + CORE_MLIR_DEPENDS_TARGET ${name} + OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" OUTPUTS "${_generate_type_stubs}" ) add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}") diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index 24a0dee9299c3..d9afd6f70edd5 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -40,7 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension StandaloneCAPI PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS - "_mlir_libs/_standaloneDialectsNanobind.pyi" + "_standaloneDialectsNanobind/__init__.pyi" + "_standaloneDialectsNanobind/standalone.pyi" ) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index d1e1b08cd8607..cf6679db7b684 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -506,10 +506,10 @@ declare_mlir_python_extension(MLIRPythonExtension.Core # Dialects MLIRCAPIFunc GENERATE_TYPE_STUBS - "_mlir_libs/_mlir/__init__.pyi" - "_mlir_libs/_mlir/ir.pyi" - "_mlir_libs/_mlir/passmanager.pyi" - "_mlir_libs/_mlir/rewrite.pyi" + "_mlir/__init__.pyi" + "_mlir/ir.pyi" + "_mlir/passmanager.pyi" + "_mlir/rewrite.pyi" ) # This extension exposes an API to register all dialects, extensions, and passes @@ -531,8 +531,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MLIRCAPIConversion MLIRCAPITransforms MLIRCAPIRegisterEverything - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirRegisterEverything.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind @@ -547,8 +546,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPILinalg - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsLinalg.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind @@ -563,8 +561,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIGPU - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsGPU.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind @@ -579,8 +576,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPILLVM - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsLLVM.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind @@ -595,8 +591,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIQuant - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsQuant.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind @@ -611,8 +606,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPINVGPU - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsNVGPU.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind @@ -627,8 +621,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIPDL - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsPDL.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind @@ -643,8 +636,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPISparseTensor - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsSparseTensor.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind @@ -659,8 +651,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPITransformDialect - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsTransform.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses @@ -674,8 +665,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIAsync - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirAsyncPasses.pyi" ) if(MLIR_ENABLE_EXECUTION_ENGINE) @@ -690,8 +679,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE) LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIExecutionEngine - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirExecutionEngine.pyi" + GENERATE_TYPE_STUBS ON ) endif() @@ -706,8 +694,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIGPU - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirGPUPasses.pyi" ) declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses @@ -721,8 +707,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPILinalg - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirLinalgPasses.pyi" ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind @@ -740,8 +724,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind MLIRCAPIIR MLIRCAPISMT MLIRCAPIExportSMTLIB - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirDialectsSMT.pyi" + GENERATE_TYPE_STUBS ON ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses @@ -755,8 +738,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPISparseTensor - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirSparseTensorPasses.pyi" ) declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter @@ -770,8 +751,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPITransformDialectTransforms - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirTransformInterpreter.pyi" + GENERATE_TYPE_STUBS ON ) # TODO: Figure out how to put this in the test tree. @@ -830,8 +810,7 @@ if(MLIR_INCLUDE_TESTS) LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIPythonTestDialect - GENERATE_TYPE_STUBS - "_mlir_libs/_mlirPythonTestNanobind.pyi" + GENERATE_TYPE_STUBS ON ) endif() diff --git a/mlir/test/python/CMakeLists.txt b/mlir/test/python/CMakeLists.txt index d68f3ff82e883..3459a9aa4be7a 100644 --- a/mlir/test/python/CMakeLists.txt +++ b/mlir/test/python/CMakeLists.txt @@ -13,7 +13,7 @@ add_subdirectory(lib) set(MLIR_PYTHON_TEST_DEPENDS MLIRPythonModules) if(NOT MLIR_STANDALONE_BUILD) - list(APPEND MLIR_PYTHON_TEST_DEPENDS FileCheck count) + list(APPEND MLIR_PYTHON_TEST_DEPENDS FileCheck count not split-file yaml2obj) endif() add_lit_testsuite(check-mlir-python "Running the MLIR Python regression tests" ${CMAKE_CURRENT_BINARY_DIR} From 8da646294473addac14682684a59e3b1625d8c89 Mon Sep 17 00:00:00 2001 From: makslevental Date: Tue, 9 Sep 2025 22:30:20 -0700 Subject: [PATCH 6/9] add DEPFILEs --- mlir/cmake/modules/AddMLIRPython.cmake | 27 +++++++++++++++++++++++--- mlir/python/CMakeLists.txt | 2 +- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 338dd1e7ae98f..26ae884d9cfcd 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -109,13 +109,14 @@ endfunction() # (e.g., something like StandalonePythonModules.extension._mlir.dso) # OUTPUT_DIR: The root output directory to emit the type stubs into. # OUTPUTS: List of expected outputs. +# DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE). # Outputs: # NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on. function(generate_type_stubs ) cmake_parse_arguments(ARG "" "FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR" - "OUTPUTS" + "OUTPUTS;DEPENDS_TARGET_SRC_DEPS" ${ARGN}) # for people doing find_package(nanobind) if(EXISTS ${nanobind_DIR}/../src/stubgen.py) @@ -148,6 +149,13 @@ function(generate_type_stubs ) --quiet) list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs) + set(_depfile "${ARG_OUTPUT_DIR}/${ARG_FQ_MODULE_NAME}.d") + if ((NOT EXISTS ${_depfile}) AND ARG_DEPENDS_TARGET_SRC_DEPS) + list(JOIN ARG_DEPENDS_TARGET_SRC_DEPS " " _depfiles) + list(TRANSFORM _generated_type_stubs APPEND ": ${_depfiles}" OUTPUT_VARIABLE _depfiles) + list(JOIN _depfiles "\n" _depfiles) + file(GENERATE OUTPUT "${_depfile}" CONTENT "${_depfiles}") + endif() add_custom_command( OUTPUT ${_generated_type_stubs} COMMAND ${_nb_stubgen_cmd} @@ -156,6 +164,7 @@ function(generate_type_stubs ) "${ARG_CORE_MLIR_DEPENDS_TARGET}.extension._mlir.dso" "${ARG_CORE_MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python" "${ARG_DEPENDS_TARGET}" + DEPFILE "${_depfile}" ) set(_name "${ARG_FQ_MODULE_NAME}.type_stubs") add_custom_target("${_name}" DEPENDS ${_generated_type_stubs}) @@ -332,19 +341,30 @@ function(add_mlir_python_modules name) ) add_dependencies(${modules_target} ${_extension_target}) mlir_python_setup_extension_rpath(${_extension_target}) + # NOTE: `sources_target` (naturally) lists all the sources (it's the INTERFACE + # target defined above in declare_mlir_python_extension). It's also the name of the + # target that gets exported (i.e., is populated as an INTERFACE IMPORTED library in MLIRTargets.cmake). + # This is why all metadata is queried from `sources_target`. On the other hand + # `_extension_target` is the actual dylib target that's built just above with `add_mlir_python_extension`. + # That's why dependencies are in terms of `_extension_target`. get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS) if(ARG_GENERATE_TYPE_STUBS AND _generate_type_stubs) if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}") endif() - # TL;DR: everything here is load bearing and annoyingly coupled. Changing anything here - # (or in declare_mlir_python_extension) will break either in-tree or out-of-tree generation. + # TL;DR: all paths here are load bearing and annoyingly coupled. Changing paths here + # (or related code in declare_mlir_python_extension) will break either in-tree or out-of-tree generation. # # We remove _mlir_libs here because OUTPUT_DIR already includes it. # Specifically OUTPUT_DIR already includes it because that's the actual directory # where we want stubgen to dump the emitted sources. This is load bearing because up above # (in declare_mlir_python_extension) we prefixed all the paths with _mlir_libs and # we specified declare_mlir_python_sources with ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs". + # + # NOTE: INTERFACE_SOURCES is a genex in the build dir ($ and $) + # which will be evaluated by file(GENERATE ...). In the install dir it's a conventional path + # (see install/lib/cmake/mlir/MLIRTargets.cmake). + get_target_property(_extension_srcs ${sources_target} INTERFACE_SOURCES) list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "") generate_type_stubs( FQ_MODULE_NAME "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}" @@ -352,6 +372,7 @@ function(add_mlir_python_modules name) CORE_MLIR_DEPENDS_TARGET ${name} OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" OUTPUTS "${_generate_type_stubs}" + DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}" ) add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}") endif() diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index cf6679db7b684..8b8aa488abdc3 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -672,7 +672,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE) MODULE_NAME _mlirExecutionEngine ADD_TO_PARENT MLIRPythonSources.ExecutionEngine ROOT_DIR "${PYTHON_SOURCE_DIR}" - PYTHON_BINDINGS_LIBRARY nanobind + PYTHON_BINDINGS_LIBRARY nanobind SOURCES ExecutionEngineModule.cpp PRIVATE_LINK_LIBS From f539569983b6be0a5fba3e6581e13a45beb76303 Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Wed, 10 Sep 2025 00:17:59 -0700 Subject: [PATCH 7/9] Update AddMLIRPython.cmake Remove GENERATE from add_python_modules --- mlir/cmake/modules/AddMLIRPython.cmake | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 26ae884d9cfcd..b74828fa8cdba 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -300,13 +300,11 @@ endfunction() # DAG of source modules is included. # COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every # extension depend on (see mlir_python_add_common_capi_library). -# GENERATE_TYPE_STUBS: Enable type stub generation for all modules -# which have declared type stub generation with `declare_mlir_python_extension`. # PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`. # This is used to determine type stub generation python module names. function(add_mlir_python_modules name) cmake_parse_arguments(ARG - "GENERATE_TYPE_STUBS" + "" "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX" "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES" ${ARGN}) @@ -348,7 +346,7 @@ function(add_mlir_python_modules name) # `_extension_target` is the actual dylib target that's built just above with `add_mlir_python_extension`. # That's why dependencies are in terms of `_extension_target`. get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS) - if(ARG_GENERATE_TYPE_STUBS AND _generate_type_stubs) + if(_generate_type_stubs) if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL "")) message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}") endif() From 182c493ad7f25938982c47d6ea4d97c9be992b2c Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Wed, 10 Sep 2025 00:18:35 -0700 Subject: [PATCH 8/9] Update CMakeLists.txt Remove GENERATE from add_python_modules --- mlir/examples/standalone/python/CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index d9afd6f70edd5..8a5672dcf82cf 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -79,6 +79,5 @@ add_mlir_python_modules(StandalonePythonModules MLIRPythonSources.Dialects.builtin COMMON_CAPI_LINK_LIBS StandalonePythonCAPI - GENERATE_TYPE_STUBS PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) From c82b91db924beeaf3caba91114e43ab280eb294b Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Wed, 10 Sep 2025 00:19:02 -0700 Subject: [PATCH 9/9] Update CMakeLists.txt --- mlir/python/CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 8b8aa488abdc3..80cd3a2b76e7e 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -867,6 +867,5 @@ add_mlir_python_modules(MLIRPythonModules ${_ADDL_TEST_SOURCES} COMMON_CAPI_LINK_LIBS MLIRPythonCAPI - GENERATE_TYPE_STUBS PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" )