[Mlir-commits] [mlir] [mlir][Python] Generate type stubs for dialect extensions (PR #236657)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 02:11:11 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/236657
>From 9f5044b0657ec0762619125bd5f126e401b78b2a Mon Sep 17 00:00:00 2001
From: Twice <twice at apache.org>
Date: Thu, 29 Oct 2026 04:41:51 +0000
Subject: [PATCH 1/2] [mlir][Python] Generate type stubs for dialect extensions
Continue #175403 by generating extension stubs for each concrete Python
package. Keep imports and build dependencies local to that package, handle
optional extensions, and normalize native IR type references.
Co-authored-by: makslevental <maksim.levental at gmail.com>
---
mlir/cmake/modules/AddMLIRPython.cmake | 102 ++++++++++++-
mlir/cmake/modules/CMakeLists.txt | 1 +
.../modules/MLIRPythonTypeStubFixup.cmake | 21 +++
mlir/python/CMakeLists.txt | 64 +++-----
.../mlir/_mlir_libs/_mlir/dialects/pdl.pyi | 63 --------
.../mlir/_mlir_libs/_mlir/dialects/quant.pyi | 142 ------------------
.../_mlir/dialects/transform/__init__.pyi | 25 ---
.../mlir/_mlir_libs/_mlirExecutionEngine.pyi | 24 ---
8 files changed, 139 insertions(+), 303 deletions(-)
create mode 100644 mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
delete mode 100644 mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi
delete mode 100644 mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi
delete mode 100644 mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/__init__.pyi
delete mode 100644 mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 0511116789fd1a..4f7f50ce3621c3 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -114,12 +114,13 @@ endfunction()
# IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen.
# PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
# VERBOSE: Emit logging/status messages during stub generation (default: OFF).
+# TARGET_NAME: Optional custom target name, instead of concatenating all deps.
# Outputs:
# NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
function(mlir_generate_type_stubs)
cmake_parse_arguments(ARG
"VERBOSE"
- "MODULE_NAME;OUTPUT_DIR;PATTERN_FILE"
+ "MODULE_NAME;OUTPUT_DIR;PATTERN_FILE;TARGET_NAME"
"IMPORT_PATHS;DEPENDS_TARGETS;OUTPUTS;DEPENDS_TARGET_SRC_DEPS"
${ARGN})
@@ -232,8 +233,12 @@ function(mlir_generate_type_stubs)
COMMENT "Generating type stubs using: ${_nb_stubgen_cmd}"
)
- list(JOIN ARG_DEPENDS_TARGETS "." _name)
- set(_name "${_name}.${ARG_MODULE_NAME}.type_stubs")
+ if(ARG_TARGET_NAME)
+ set(_name "${ARG_TARGET_NAME}")
+ else()
+ list(JOIN ARG_DEPENDS_TARGETS "." _name)
+ set(_name "${_name}.${ARG_MODULE_NAME}.type_stubs")
+ endif()
add_custom_target("${_name}" DEPENDS ${_generated_type_stubs})
set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
endfunction()
@@ -253,9 +258,11 @@ endfunction()
# EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
# on. These will be collected for all extensions and put into an
# aggregate dylib that is linked against.
+# GENERATE_TYPE_STUBS: Generate a stub for this leaf extension in each package
+# that builds it, using that package's module name and native extension.
function(declare_mlir_python_extension name)
cmake_parse_arguments(ARG
- "_PRIVATE_SUPPORT_LIB"
+ "_PRIVATE_SUPPORT_LIB;GENERATE_TYPE_STUBS"
"ROOT_DIR;MODULE_NAME;ADD_TO_PARENT"
"SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
${ARGN})
@@ -274,11 +281,12 @@ function(declare_mlir_python_extension name)
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
- EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS"
+ EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS;mlir_python_GENERATE_TYPE_STUBS"
mlir_python_SOURCES_TYPE "${SOURCES_TYPE}"
mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
mlir_python_DEPENDS ""
+ mlir_python_GENERATE_TYPE_STUBS "${ARG_GENERATE_TYPE_STUBS}"
)
# Set the interface source and link_libs properties of the target
@@ -440,6 +448,51 @@ function(build_nanobind_lib)
)
endfunction()
+# Generate a leaf extension stub after its concrete package has been assembled.
+# Source declarations can be embedded by several packages with different names;
+# generated files and dependencies must belong to each concrete package.
+function(_mlir_python_add_extension_type_stub modules_target sources_target)
+ cmake_parse_arguments(ARG "" "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX"
+ "DEPENDS_TARGETS" ${ARGN})
+ get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
+ get_target_property(_extension_srcs ${sources_target} INTERFACE_SOURCES)
+ set(_import_root "${ARG_ROOT_PREFIX}")
+ string(REPLACE "." ";" _package_components "${ARG_PACKAGE_PREFIX}")
+ foreach(_component IN LISTS _package_components)
+ cmake_path(GET _import_root PARENT_PATH _import_root)
+ endforeach()
+ set(_raw_dir "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/${modules_target}/_mlir_libs/nonfixup")
+ mlir_generate_type_stubs(
+ TARGET_NAME ${modules_target}.generate_type_stubs.${_module_name}
+ MODULE_NAME ${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}
+ DEPENDS_TARGETS ${ARG_DEPENDS_TARGETS}
+ OUTPUT_DIR "${_raw_dir}"
+ OUTPUTS ${_module_name}.pyi
+ DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}"
+ IMPORT_PATHS "${_import_root}"
+ )
+ set(_raw_stub "${_raw_dir}/${_module_name}.pyi")
+ set(_stub "${ARG_ROOT_PREFIX}/_mlir_libs/${_module_name}.pyi")
+ set(_fixup "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/MLIRPythonTypeStubFixup.cmake")
+ add_custom_command(
+ OUTPUT "${_stub}"
+ COMMAND ${CMAKE_COMMAND}
+ "-DINPUT_FILE=${_raw_stub}"
+ "-DOUTPUT_FILE=${_stub}"
+ "-DMLIR_PYTHON_PACKAGE_PREFIX=${ARG_PACKAGE_PREFIX}"
+ -P "${_fixup}"
+ DEPENDS "${_raw_stub}" "${_fixup}"
+ COMMENT "Generating ${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name} type stub"
+ VERBATIM
+ )
+ set(_stub_target "${modules_target}.type_stubs.${_module_name}")
+ add_custom_target(${_stub_target} DEPENDS "${_stub}")
+ add_dependencies(${modules_target} ${_stub_target})
+ install(FILES "${_stub}"
+ DESTINATION "${ARG_INSTALL_PREFIX}/_mlir_libs"
+ COMPONENT ${modules_target})
+endfunction()
+
# Function: add_mlir_python_modules
# Adds python modules to a project, building them from a list of declared
# source groupings (see declare_mlir_python_sources and
@@ -456,6 +509,11 @@ endfunction()
# will share nanobind types with other bindings packages. Expected to be unique
# per project (and per specific set of bindings, for projects with multiple
# bindings packages).
+# MLIR_PYTHON_PACKAGE_PREFIX: Import name of the concrete package. Defaults to
+# MLIR_PYTHON_PACKAGE_PREFIX; embedders must pass their own package name for
+# stubgen, matching their C++ MLIR_PYTHON_PACKAGE_PREFIX definition.
+# STUBGEN_DEPENDS_TARGETS: Extra package preparation targets needed before
+# importing extensions (for example, colocating DLLs in the package).
# DECLARED_SOURCES: List of declared source groups to include. The entire
# DAG of source modules is included.
# COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
@@ -471,11 +529,13 @@ endfunction()
function(add_mlir_python_modules name)
cmake_parse_arguments(ARG
""
- "ROOT_PREFIX;INSTALL_PREFIX;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
- "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_LINK_LIBS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS"
+ "ROOT_PREFIX;INSTALL_PREFIX;MLIR_BINDINGS_PYTHON_NB_DOMAIN;MLIR_PYTHON_PACKAGE_PREFIX"
+ "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES;STUBGEN_DEPENDS_TARGETS;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_LINK_LIBS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS"
${ARGN})
- # TODO(max): do the same for MLIR_PYTHON_PACKAGE_PREFIX?
+ if(NOT ARG_MLIR_PYTHON_PACKAGE_PREFIX)
+ set(ARG_MLIR_PYTHON_PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}")
+ endif()
if((NOT ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN) AND MLIR_BINDINGS_PYTHON_NB_DOMAIN)
set(ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN ${MLIR_BINDINGS_PYTHON_NB_DOMAIN})
endif()
@@ -598,6 +658,32 @@ function(add_mlir_python_modules name)
_process_target(${name} ${sources_target} "${_mlir_python_support_libs}")
endforeach()
+ if(MLIR_PYTHON_STUBGEN_ENABLED)
+ # Import package initializers only after its native extensions and Python
+ # sources are ready. Depend on concrete targets, not the package aggregate:
+ # the latter also depends on the generated stubs and would form a cycle.
+ set(_stubgen_inputs ${ARG_STUBGEN_DEPENDS_TARGETS})
+ foreach(sources_target ${_flat_targets})
+ get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)
+ if(_source_type STREQUAL "extension")
+ get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
+ list(APPEND _stubgen_inputs "${name}.extension.${_module_name}.dso")
+ elseif(_source_type STREQUAL "pure")
+ list(APPEND _stubgen_inputs "${name}.sources.${sources_target}")
+ endif()
+ endforeach()
+ foreach(sources_target ${_flat_targets})
+ get_target_property(_generate_stub ${sources_target} mlir_python_GENERATE_TYPE_STUBS)
+ if(_generate_stub)
+ _mlir_python_add_extension_type_stub(${name} ${sources_target}
+ ROOT_PREFIX "${ARG_ROOT_PREFIX}"
+ INSTALL_PREFIX "${ARG_INSTALL_PREFIX}"
+ PACKAGE_PREFIX "${ARG_MLIR_PYTHON_PACKAGE_PREFIX}"
+ DEPENDS_TARGETS ${_stubgen_inputs})
+ endif()
+ endforeach()
+ endif()
+
# Create an install target.
if(NOT LLVM_ENABLE_IDE)
add_llvm_install_targets(
diff --git a/mlir/cmake/modules/CMakeLists.txt b/mlir/cmake/modules/CMakeLists.txt
index 4fed889eb2ee1d..8270a12e23d9ea 100644
--- a/mlir/cmake/modules/CMakeLists.txt
+++ b/mlir/cmake/modules/CMakeLists.txt
@@ -168,6 +168,7 @@ if (NOT LLVM_INSTALL_TOOLCHAIN_ONLY AND
${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/MLIRConfigVersion.cmake
${CMAKE_CURRENT_SOURCE_DIR}/AddMLIR.cmake
${CMAKE_CURRENT_SOURCE_DIR}/AddMLIRPython.cmake
+ ${CMAKE_CURRENT_SOURCE_DIR}/MLIRPythonTypeStubFixup.cmake
${CMAKE_CURRENT_SOURCE_DIR}/IRDLToCpp.cmake
${CMAKE_CURRENT_SOURCE_DIR}/MLIRDetectPythonEnv.cmake
DESTINATION ${MLIR_INSTALL_PACKAGE_DIR}
diff --git a/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake b/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
new file mode 100644
index 00000000000000..880a45d4530ee2
--- /dev/null
+++ b/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
@@ -0,0 +1,21 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+file(READ "${INPUT_FILE}" _content)
+string(REPLACE "\r\n" "\n" _content "${_content}")
+# DefaultingPyMlirContext and DefaultingPyLocation describe their types using
+# bare _mlir.ir names even when the extension is imported under a package.
+# Normalize both those names and fully-qualified native IR references.
+foreach(_module
+ "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs._mlir.ir"
+ "_mlir.ir"
+ "${MLIR_PYTHON_PACKAGE_PREFIX}.ir")
+ string(REPLACE "import ${_module}\n" "" _content "${_content}")
+ string(REPLACE "${_module}" "_ir" _content "${_content}")
+endforeach()
+if(NOT _content MATCHES "(^|[^A-Za-z0-9_.])_mlir\\.")
+ string(REPLACE "import _mlir\n" "" _content "${_content}")
+endif()
+file(WRITE "${OUTPUT_FILE}"
+ "from ${MLIR_PYTHON_PACKAGE_PREFIX} import ir as _ir\n${_content}")
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index d7a57a15f82420..71f4bfa1020d8e 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -48,7 +48,6 @@ set(_execution_engine_python_sources)
if(MLIR_ENABLE_EXECUTION_ENGINE)
set(_execution_engine_python_sources
execution_engine.py
- _mlir_libs/_mlirExecutionEngine.pyi
)
endif()
@@ -228,7 +227,6 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/TransformOps.td
SOURCES
dialects/transform/__init__.py
- _mlir_libs/_mlir/dialects/transform/__init__.pyi
DIALECT_NAME transform
GEN_ENUM_BINDINGS_TD_FILE
"../../include/mlir/Dialect/Transform/IR/TransformAttrs.td"
@@ -436,8 +434,7 @@ declare_mlir_python_sources(
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
GEN_ENUM_BINDINGS
SOURCES
- dialects/quant.py
- _mlir_libs/_mlir/dialects/quant.pyi)
+ dialects/quant.py)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -453,7 +450,6 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/PDLOps.td
SOURCES
dialects/pdl.py
- _mlir_libs/_mlir/dialects/pdl.pyi
DIALECT_NAME pdl)
declare_mlir_dialect_python_bindings(
@@ -606,6 +602,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
MODULE_NAME _mlirDialectsLinalg
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -617,6 +614,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
MODULE_NAME _mlirDialectsGPU
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -628,6 +626,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
MODULE_NAME _mlirDialectsLLVM
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -641,6 +640,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
MODULE_NAME _mlirDialectsQuant
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.quant
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -652,6 +652,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
MODULE_NAME _mlirDialectsNVGPU
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -663,6 +664,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
MODULE_NAME _mlirDialectsPDL
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -674,6 +676,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
MODULE_NAME _mlirDialectsSparseTensor
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -685,6 +688,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
MODULE_NAME _mlirDialectsTransform
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.transform
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -699,6 +703,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Nanobind
MODULE_NAME _mlirDialectsIRDL
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.irdl
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -721,6 +726,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
if(MLIR_ENABLE_EXECUTION_ENGINE)
declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
MODULE_NAME _mlirExecutionEngine
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -752,6 +758,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Nanobind
MODULE_NAME _mlirDialectsSMT
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.smt
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -774,6 +781,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
MODULE_NAME _mlirTransformInterpreter
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.transform
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -784,6 +792,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
declare_mlir_python_extension(MLIRPythonExtension.Dialects.AMDGPU.Nanobind
MODULE_NAME _mlirDialectsAMDGPU
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonSources.Dialects.amdgpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
SOURCES
@@ -829,6 +838,7 @@ if(MLIR_INCLUDE_TESTS)
declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
MODULE_NAME _mlirPythonTestNanobind
+ GENERATE_TYPE_STUBS
ADD_TO_PARENT MLIRPythonTestSources.Dialects
ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
SOURCES
@@ -884,6 +894,8 @@ add_mlir_python_common_capi_library(MLIRPythonCAPI
_flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
if(MLIR_PYTHON_STUBGEN_ENABLED)
+ set(_mlir_typestub_gen_targets)
+
# _mlir stubgen
# Note: All this needs to come before add_mlir_python_modules(MLIRPythonModules so that the install targets for the
# generated type stubs get created.
@@ -904,14 +916,15 @@ if(MLIR_PYTHON_STUBGEN_ENABLED)
# **NOT** MLIRPythonModules.extension._mlir.dso. So be sure to use the correct target!
get_target_property(_core_extension_srcs MLIRPythonExtension.Core INTERFACE_SOURCES)
- # Why is MODULE_NAME _mlir here but mlir._mlir_libs._mlirPythonTestNanobind below???
+ # The core stub generator imports _mlir directly. Dialect stub generators in
+ # add_mlir_python_modules instead use the full package-qualified module name.
# The _mlir extension can be imported independently of any other python code and/or extension modules.
# I.e., you could do `cd $MLIRPythonModules_ROOT_PREFIX/_mlir_libs && python -c "import _mlir"` (try it!).
# _mlir is also (currently) the only extension for which this is possible because dialect extensions modules,
# which generally make use of `mlir_value_subclass/mlir_type_subclass/mlir_attribute_subclass`, perform an
# `import mlir` right when they're loaded (see the mlir_*_subclass ctors in NanobindAdaptors.h).
- # Note, this also why IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" here while below
- # "${MLIRPythonModules_ROOT_PREFIX}/.." (because MLIR_BINDINGS_PYTHON_INSTALL_PREFIX, by default, ends at mlir).
+ # This also explains the _mlir_libs import root here, versus the parent of the
+ # concrete package for dialect extensions.
#
# Further note: this function creates file targets like
# "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs/_mlir/__init__.pyi". These must match the file targets
@@ -925,7 +938,7 @@ if(MLIR_PYTHON_STUBGEN_ENABLED)
DEPENDS_TARGET_SRC_DEPS "${_core_extension_srcs}"
IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs"
)
- set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}")
+ list(APPEND _mlir_typestub_gen_targets ${NB_STUBGEN_CUSTOM_TARGET})
list(TRANSFORM _core_type_stub_sources PREPEND "_mlir_libs/")
# Note, we do not do ADD_TO_PARENT here so that the type stubs are not associated (as mlir_DEPENDS) with
@@ -937,34 +950,6 @@ if(MLIR_PYTHON_STUBGEN_ENABLED)
SOURCES "${_core_type_stub_sources}"
)
list(APPEND mlir_python_sources_deps MLIRPythonExtension.Core.type_stub_gen)
-
- # _mlirPythonTestNanobind stubgen
-
- if(MLIR_INCLUDE_TESTS)
- get_target_property(_test_extension_srcs MLIRPythonTestSources.PythonTestExtensionNanobind INTERFACE_SOURCES)
- mlir_generate_type_stubs(
- # This is the FQN path because dialect modules import _mlir when loaded. See above.
- MODULE_NAME mlir._mlir_libs._mlirPythonTestNanobind
- DEPENDS_TARGETS
- # You need both _mlir and _mlirPythonTestNanobind because dialect modules import _mlir when loaded
- # (so _mlir needs to be built before calling stubgen).
- MLIRPythonModules.extension._mlir.dso
- MLIRPythonModules.extension._mlirPythonTestNanobind.dso
- # You need this one so that ir.py "built" because mlir._mlir_libs.__init__.py import mlir.ir in _site_initialize.
- MLIRPythonModules.sources.MLIRPythonSources.Core.Python
- OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
- OUTPUTS _mlirPythonTestNanobind.pyi
- DEPENDS_TARGET_SRC_DEPS "${_test_extension_srcs}"
- IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/.."
- )
- set(_mlirPythonTestNanobind_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}")
- declare_mlir_python_sources(
- MLIRPythonTestSources.PythonTestExtensionNanobind.type_stub_gen
- ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
- ADD_TO_PARENT MLIRPythonTestSources.Dialects
- SOURCES _mlir_libs/_mlirPythonTestNanobind.pyi
- )
- endif()
endif()
add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps})
@@ -995,8 +980,5 @@ add_mlir_python_modules(MLIRPythonModules
MLIRPythonCAPI
)
if(MLIR_PYTHON_STUBGEN_ENABLED)
- add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}")
- if(MLIR_INCLUDE_TESTS)
- add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")
- endif()
+ add_dependencies(MLIRPythonModules ${_mlir_typestub_gen_targets})
endif()
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi b/mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi
deleted file mode 100644
index d12c6839deabaf..00000000000000
--- a/mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi
+++ /dev/null
@@ -1,63 +0,0 @@
-# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-# See https://llvm.org/LICENSE.txt for license information.
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-
-from mlir.ir import Type, Context
-
-__all__ = [
- 'PDLType',
- 'AttributeType',
- 'OperationType',
- 'RangeType',
- 'TypeType',
- 'ValueType',
-]
-
-
-class PDLType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
-
-class AttributeType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(context: Context | None = None) -> AttributeType: ...
-
-
-class OperationType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(context: Context | None = None) -> OperationType: ...
-
-
-class RangeType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(element_type: Type) -> RangeType: ...
-
- @property
- def element_type(self) -> Type: ...
-
-
-class TypeType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(context: Context | None = None) -> TypeType: ...
-
-
-class ValueType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(context: Context | None = None) -> ValueType: ...
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi b/mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi
deleted file mode 100644
index 3f5304584edeff..00000000000000
--- a/mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi
+++ /dev/null
@@ -1,142 +0,0 @@
-# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-# See https://llvm.org/LICENSE.txt for license information.
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-
-from mlir.ir import DenseElementsAttr, Type
-
-__all__ = [
- "QuantizedType",
- "AnyQuantizedType",
- "UniformQuantizedType",
- "UniformQuantizedPerAxisType",
- "CalibratedQuantizedType",
-]
-
-class QuantizedType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def default_minimum_for_integer(is_signed: bool, integral_width: int) -> int:
- ...
-
- @staticmethod
- def default_maximum_for_integer(is_signed: bool, integral_width: int) -> int:
- ...
-
- @property
- def expressed_type(self) -> Type: ...
-
- @property
- def flags(self) -> int: ...
-
- @property
- def is_signed(self) -> bool: ...
-
- @property
- def storage_type(self) -> Type: ...
-
- @property
- def storage_type_min(self) -> int: ...
-
- @property
- def storage_type_max(self) -> int: ...
-
- @property
- def storage_type_integral_width(self) -> int: ...
-
- def is_compatible_expressed_type(self, candidate: Type) -> bool: ...
-
- @property
- def quantized_element_type(self) -> Type: ...
-
- def cast_from_storage_type(self, candidate: Type) -> Type: ...
-
- @staticmethod
- def cast_to_storage_type(type: Type) -> Type: ...
-
- def cast_from_expressed_type(self, candidate: Type) -> Type: ...
-
- @staticmethod
- def cast_to_expressed_type(type: Type) -> Type: ...
-
- def cast_expressed_to_storage_type(self, candidate: Type) -> Type: ...
-
-
-class AnyQuantizedType(QuantizedType):
-
- @classmethod
- def get(cls, flags: int, storage_type: Type, expressed_type: Type,
- storage_type_min: int, storage_type_max: int) -> Type:
- ...
-
-
-class UniformQuantizedType(QuantizedType):
-
- @classmethod
- def get(cls, flags: int, storage_type: Type, expressed_type: Type,
- scale: float, zero_point: int, storage_type_min: int,
- storage_type_max: int) -> Type: ...
-
- @property
- def scale(self) -> float: ...
-
- @property
- def zero_point(self) -> int: ...
-
- @property
- def is_fixed_point(self) -> bool: ...
-
-
-class UniformQuantizedPerAxisType(QuantizedType):
-
- @classmethod
- def get(cls, flags: int, storage_type: Type, expressed_type: Type,
- scales: list[float], zero_points: list[int], quantized_dimension: int,
- storage_type_min: int, storage_type_max: int):
- ...
-
- @property
- def scales(self) -> list[float]: ...
-
- @property
- def zero_points(self) -> list[int]: ...
-
- @property
- def quantized_dimension(self) -> int: ...
-
- @property
- def is_fixed_point(self) -> bool: ...
-
-class UniformQuantizedSubChannelType(QuantizedType):
-
- @classmethod
- def get(cls, flags: int, storage_type: Type, expressed_type: Type,
- scales: DenseElementsAttr, zero_points: DenseElementsAttr,
- quantized_dimensions: list[int], block_sizes: list[int],
- storage_type_min: int, storage_type_max: int):
- ...
-
- @property
- def quantized_dimensions(self) -> list[int]: ...
-
- @property
- def block_sizes(self) -> list[int]: ...
-
- @property
- def scales(self) -> DenseElementsAttr: ...
-
- @property
- def zero_points(self) -> DenseElementsAttr: ...
-
-def CalibratedQuantizedType(QuantizedType):
-
- @classmethod
- def get(cls, expressed_type: Type, min: float, max: float): ...
-
- @property
- def min(self) -> float: ...
-
- @property
- def max(self) -> float: ...
\ No newline at end of file
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/__init__.pyi b/mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/__init__.pyi
deleted file mode 100644
index a3f1b09102379f..00000000000000
--- a/mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/__init__.pyi
+++ /dev/null
@@ -1,25 +0,0 @@
-# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-# See https://llvm.org/LICENSE.txt for license information.
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-
-from mlir.ir import Type, Context
-
-
-class AnyOpType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(context: Context | None = None) -> AnyOpType: ...
-
-
-class OperationType(Type):
- @staticmethod
- def isinstance(type: Type) -> bool: ...
-
- @staticmethod
- def get(operation_name: str, context: Context | None = None) -> OperationType: ...
-
- @property
- def operation_name(self) -> str: ...
diff --git a/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi b/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
deleted file mode 100644
index 4b82c784892958..00000000000000
--- a/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
+++ /dev/null
@@ -1,24 +0,0 @@
-# Originally imported via:
-# stubgen {...} -m mlir._mlir_libs._mlirExecutionEngine
-# Local modifications:
-# * Relative imports for cross-module references.
-# * Add __all__
-
-from collections.abc import Sequence
-
-from ._mlir import ir as _ir
-
-__all__ = [
- "ExecutionEngine",
-]
-
-class ExecutionEngine:
- def __init__(self, module: _ir.Module, opt_level: int = 2, shared_libs: Sequence[str] = ...) -> None: ...
- def _CAPICreate(self) -> object: ...
- def _testing_release(self) -> None: ...
- def dump_to_object_file(self, file_name: str) -> None: ...
- def raw_lookup(self, func_name: str) -> int: ...
- def raw_register_runtime(self, name: str, callback: object) -> None: ...
- def init() -> None: ...
- @property
- def _CAPIPtr(self) -> object: ...
>From 3986f9d138f4f665274ff763e77582c5cd861eb1 Mon Sep 17 00:00:00 2001
From: Twice <twice at apache.org>
Date: Thu, 29 Oct 2026 04:49:13 -0400
Subject: [PATCH 2/2] [mlir][Python] Fix defaulting type names without stub
postprocessing
Use the package prefix in the Context and Location caster descriptions.
Generate extension stubs directly in their package and remove the text
rewrite step.
Signed-off-by: Twice <twice at apache.org>
---
mlir/cmake/modules/AddMLIRPython.cmake | 24 ++++---------------
mlir/cmake/modules/CMakeLists.txt | 1 -
.../modules/MLIRPythonTypeStubFixup.cmake | 21 ----------------
mlir/include/mlir/Bindings/Python/IRCore.h | 7 ++++--
4 files changed, 9 insertions(+), 44 deletions(-)
delete mode 100644 mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 4f7f50ce3621c3..b73bcadfa24729 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -461,34 +461,18 @@ function(_mlir_python_add_extension_type_stub modules_target sources_target)
foreach(_component IN LISTS _package_components)
cmake_path(GET _import_root PARENT_PATH _import_root)
endforeach()
- set(_raw_dir "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/${modules_target}/_mlir_libs/nonfixup")
+ set(_stub_dir "${ARG_ROOT_PREFIX}/_mlir_libs")
mlir_generate_type_stubs(
TARGET_NAME ${modules_target}.generate_type_stubs.${_module_name}
MODULE_NAME ${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}
DEPENDS_TARGETS ${ARG_DEPENDS_TARGETS}
- OUTPUT_DIR "${_raw_dir}"
+ OUTPUT_DIR "${_stub_dir}"
OUTPUTS ${_module_name}.pyi
DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}"
IMPORT_PATHS "${_import_root}"
)
- set(_raw_stub "${_raw_dir}/${_module_name}.pyi")
- set(_stub "${ARG_ROOT_PREFIX}/_mlir_libs/${_module_name}.pyi")
- set(_fixup "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/MLIRPythonTypeStubFixup.cmake")
- add_custom_command(
- OUTPUT "${_stub}"
- COMMAND ${CMAKE_COMMAND}
- "-DINPUT_FILE=${_raw_stub}"
- "-DOUTPUT_FILE=${_stub}"
- "-DMLIR_PYTHON_PACKAGE_PREFIX=${ARG_PACKAGE_PREFIX}"
- -P "${_fixup}"
- DEPENDS "${_raw_stub}" "${_fixup}"
- COMMENT "Generating ${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name} type stub"
- VERBATIM
- )
- set(_stub_target "${modules_target}.type_stubs.${_module_name}")
- add_custom_target(${_stub_target} DEPENDS "${_stub}")
- add_dependencies(${modules_target} ${_stub_target})
- install(FILES "${_stub}"
+ add_dependencies(${modules_target} ${NB_STUBGEN_CUSTOM_TARGET})
+ install(FILES "${_stub_dir}/${_module_name}.pyi"
DESTINATION "${ARG_INSTALL_PREFIX}/_mlir_libs"
COMPONENT ${modules_target})
endfunction()
diff --git a/mlir/cmake/modules/CMakeLists.txt b/mlir/cmake/modules/CMakeLists.txt
index 8270a12e23d9ea..4fed889eb2ee1d 100644
--- a/mlir/cmake/modules/CMakeLists.txt
+++ b/mlir/cmake/modules/CMakeLists.txt
@@ -168,7 +168,6 @@ if (NOT LLVM_INSTALL_TOOLCHAIN_ONLY AND
${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/MLIRConfigVersion.cmake
${CMAKE_CURRENT_SOURCE_DIR}/AddMLIR.cmake
${CMAKE_CURRENT_SOURCE_DIR}/AddMLIRPython.cmake
- ${CMAKE_CURRENT_SOURCE_DIR}/MLIRPythonTypeStubFixup.cmake
${CMAKE_CURRENT_SOURCE_DIR}/IRDLToCpp.cmake
${CMAKE_CURRENT_SOURCE_DIR}/MLIRDetectPythonEnv.cmake
DESTINATION ${MLIR_INSTALL_PACKAGE_DIR}
diff --git a/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake b/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
deleted file mode 100644
index 880a45d4530ee2..00000000000000
--- a/mlir/cmake/modules/MLIRPythonTypeStubFixup.cmake
+++ /dev/null
@@ -1,21 +0,0 @@
-# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-# See https://llvm.org/LICENSE.txt for license information.
-# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-file(READ "${INPUT_FILE}" _content)
-string(REPLACE "\r\n" "\n" _content "${_content}")
-# DefaultingPyMlirContext and DefaultingPyLocation describe their types using
-# bare _mlir.ir names even when the extension is imported under a package.
-# Normalize both those names and fully-qualified native IR references.
-foreach(_module
- "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs._mlir.ir"
- "_mlir.ir"
- "${MLIR_PYTHON_PACKAGE_PREFIX}.ir")
- string(REPLACE "import ${_module}\n" "" _content "${_content}")
- string(REPLACE "${_module}" "_ir" _content "${_content}")
-endforeach()
-if(NOT _content MATCHES "(^|[^A-Za-z0-9_.])_mlir\\.")
- string(REPLACE "import _mlir\n" "" _content "${_content}")
-endif()
-file(WRITE "${OUTPUT_FILE}"
- "from ${MLIR_PYTHON_PACKAGE_PREFIX} import ir as _ir\n${_content}")
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 3edf12fbafee1e..956413004b98e9 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -20,6 +20,7 @@
#include "NanobindUtils.h"
#include "mlir-c/AffineExpr.h"
#include "mlir-c/AffineMap.h"
+#include "mlir-c/Bindings/Python/Interop.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/Debug.h"
#include "mlir-c/Diagnostics.h"
@@ -332,7 +333,8 @@ class MLIR_PYTHON_API_EXPORTED DefaultingPyMlirContext
: public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
public:
using Defaulting::Defaulting;
- static constexpr const char kTypeDescription[] = "_mlir.ir.Context";
+ static constexpr const char kTypeDescription[] =
+ MAKE_MLIR_PYTHON_QUALNAME("ir.Context");
static PyMlirContext &resolve();
};
@@ -618,7 +620,8 @@ class MLIR_PYTHON_API_EXPORTED DefaultingPyLocation
: public Defaulting<DefaultingPyLocation, PyLocation> {
public:
using Defaulting::Defaulting;
- static constexpr const char kTypeDescription[] = "_mlir.ir.Location";
+ static constexpr const char kTypeDescription[] =
+ MAKE_MLIR_PYTHON_QUALNAME("ir.Location");
static PyLocation &resolve();
operator MlirLocation() const { return *get(); }
More information about the Mlir-commits
mailing list