[Mlir-commits] [mlir] [MLIR][Python] fix stubgen for downstream users (PR #157589)
Maksim Levental
llvmlistbot at llvm.org
Tue Sep 9 18:06:28 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/4] [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/4] 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/4] 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/4] 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()
More information about the Mlir-commits
mailing list