diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index ee6c87a8b635e..b74828fa8cdba 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -102,57 +102,72 @@ 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 +# 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. +# 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 MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR) +function(generate_type_stubs ) cmake_parse_arguments(ARG "" - "" - "OUTPUTS" + "FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR" + "OUTPUTS;DEPENDS_TARGET_SRC_DEPS" ${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() 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 + set(_nb_stubgen_cmd "${Python_EXECUTABLE}" "${NB_STUBGEN}" --module - "${_module}" + "${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) + 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} + 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}" + DEPFILE "${_depfile}" ) - set(_name "MLIRPythonModuleStubs_${_module}") - add_custom_target("${_name}" ALL DEPENDS ${_generated_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() @@ -172,7 +187,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 "" @@ -190,6 +208,23 @@ function(declare_mlir_python_extension name) endif() 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" + 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 @@ -265,10 +300,12 @@ 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). +# 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 "" - "ROOT_PREFIX;INSTALL_PREFIX" + "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX" "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES" ${ARGN}) # Helper to process an individual target. @@ -302,31 +339,40 @@ 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(_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: 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( - ${_module_name} - ${_extension_target} - ${name} - "${CMAKE_CURRENT_BINARY_DIR}/_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}" + DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}" ) 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_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/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..8a5672dcf82cf 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -40,6 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension StandaloneCAPI PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS + "_standaloneDialectsNanobind/__init__.pyi" + "_standaloneDialectsNanobind/standalone.pyi" ) @@ -77,4 +79,5 @@ add_mlir_python_modules(StandalonePythonModules MLIRPythonSources.Dialects.builtin COMMON_CAPI_LINK_LIBS StandalonePythonCAPI - ) + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" +) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 8e7949480f21e..80cd3a2b76e7e 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -531,8 +531,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MLIRCAPIConversion MLIRCAPITransforms MLIRCAPIRegisterEverything - GENERATE_TYPE_STUBS - "_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 - "_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 - "_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 - "_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 - "_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 - "_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 - "_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 - "_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 - "_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 - "_mlirAsyncPasses.pyi" ) if(MLIR_ENABLE_EXECUTION_ENGINE) @@ -683,15 +672,14 @@ 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 LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIExecutionEngine - GENERATE_TYPE_STUBS - "_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 - "_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 - "_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 - "_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 - "_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 - "_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 - "_mlirPythonTestNanobind.pyi" + GENERATE_TYPE_STUBS ON ) endif() @@ -888,4 +867,5 @@ add_mlir_python_modules(MLIRPythonModules ${_ADDL_TEST_SOURCES} COMMON_CAPI_LINK_LIBS MLIRPythonCAPI + PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}" ) 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}