[Mlir-commits] [mlir] [MLIR][Python] fix stubgen for downstream users (PR #157589)
Maksim Levental
llvmlistbot at llvm.org
Wed Sep 10 00:19:13 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/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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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 ($<BUILD_INTERFACE> and $<INSTALL_INTERFACE>)
+ # 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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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 <maksim.levental at gmail.com>
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}"
)
More information about the Mlir-commits
mailing list