[Mlir-commits] [mlir] [MLIR][Python] fix stubgen for downstream users (PR #157589)
Maksim Levental
llvmlistbot at llvm.org
Tue Sep 9 14:07:01 PDT 2025
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/157589
>From 907c7b626c0b9cbb8298fca33470e04b2eee3dc2 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Mon, 8 Sep 2025 18:21:49 -0700
Subject: [PATCH 1/2] [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 ae0eaf379b193f2be798194de6fcf0e3dce6cba5 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 9 Sep 2025 12:23:25 -0700
Subject: [PATCH 2/2] fix a bunch of stuff
---
mlir/cmake/modules/AddMLIRPython.cmake | 29 ++++++++++++-------
mlir/examples/standalone/CMakeLists.txt | 2 +-
.../examples/standalone/python/CMakeLists.txt | 3 ++
mlir/python/CMakeLists.txt | 2 ++
4 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 8d2a73f385b9d..2f3ad060767c0 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,7 +156,7 @@ 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}")
+ set(_name "${FQ_MODULE_NAME}.type_stubs")
add_custom_target("${_name}" ALL 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}"
)
More information about the Mlir-commits
mailing list