[MLIR][Python] reland stubgen #157853

makslevental · 2025-09-10T13:12:55Z

This PR relands stubgen. It works via an updated API for declare_mlir_python_extension:

declare_mlir_python_extension(
#   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.
)

Note: this patch includes a bump of nanobind so you will need to do pip install -r mlir/python/requirements.txt --force.

EDIT:

see #157930 for the current version of this feature.

makslevental · 2025-09-10T13:21:14Z

This was thoroughly tested here #157589 with all the problem scenarios people mentioned but it would be good if it was tested again. Please if I've tagged you, can you try the patch locally. Note you can do so with something like

cd $LLVM_ROOT
git fetch origin main
git checkout 3c810b76b97456e4e1c115dcf3238a799067c466
wget https://github.com/llvm/llvm-project/pull/157853.patch
git apply 157853.patch

Note: this patch includes a bump of nanobind so you will need to do pip install -r mlir/python/requirements.txt --force.

llvmbot · 2025-09-10T13:25:30Z

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

This PR relands stubgen. It works via an updated API for declare_mlir_python_extension:

declare_mlir_python_extension(
#   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.
)

and add_mlir_python_modules:

add_mlir_python_modules(
#   PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`.
#     This is used to determine type stub generation python module names.
)

Patch is 121.91 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/157853.diff

15 Files Affected:

(modified) .ci/all_requirements.txt (+7-3)
(modified) mlir/cmake/modules/AddMLIRPython.cmake (+154-21)
(modified) mlir/examples/standalone/CMakeLists.txt (+1-1)
(modified) mlir/examples/standalone/python/CMakeLists.txt (+5-1)
(modified) mlir/python/CMakeLists.txt (+27-13)
(added) mlir/python/mlir/_mlir_libs/.gitignore (+2)
(removed) mlir/python/mlir/_mlir_libs/_mlir/init.pyi (-12)
(removed) mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi (-63)
(removed) mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi (-142)
(removed) mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/init.pyi (-25)
(removed) mlir/python/mlir/_mlir_libs/_mlir/ir.pyi (-2846)
(removed) mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi (-36)
(removed) mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi (-24)
(modified) mlir/python/requirements.txt (+3-2)
(modified) mlir/test/python/CMakeLists.txt (+1-1)

diff --git a/.ci/all_requirements.txt b/.ci/all_requirements.txt
index dea9646f648ad..ac9682a09bec1 100644
--- a/.ci/all_requirements.txt
+++ b/.ci/all_requirements.txt
@@ -194,9 +194,9 @@ ml-dtypes==0.5.1 ; python_version < "3.13" \
     --hash=sha256:d13755f8e8445b3870114e5b6240facaa7cb0c3361e54beba3e07fa912a6e12b \
     --hash=sha256:fd918d4e6a4e0c110e2e05be7a7814d10dc1b95872accbf6512b80a109b71ae1
     # via -r mlir/python/requirements.txt
-nanobind==2.7.0 \
-    --hash=sha256:73b12d0e751d140d6c1bf4b215e18818a8debfdb374f08dc3776ad208d808e74 \
-    --hash=sha256:f9f1b160580c50dcf37b6495a0fd5ec61dc0d95dae5f8004f87dd9ad7eb46b34
+nanobind==2.9.2 \
+    --hash=sha256:c37957ffd5eac7eda349cff3622ecd32e5ee1244ecc912c99b5bc8188bafd16e \
+    --hash=sha256:e7608472de99d375759814cab3e2c94aba3f9ec80e62cfef8ced495ca5c27d6e
     # via -r mlir/python/requirements.txt
 numpy==2.0.2 \
     --hash=sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a \
@@ -383,6 +383,10 @@ swig==4.3.1 \
     --hash=sha256:efec16327029f682f649a26da726bb0305be8800bd0f1fa3e81bf0769cf5b476 \
     --hash=sha256:fc496c0d600cf1bb2d91e28d3d6eae9c4301e5ea7a0dec5a4281b5efed4245a8
     # via -r lldb/test/requirements.txt
+typing-extensions==4.15.0 \
+    --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \
+    --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548
+    # via -r mlir/python/requirements.txt
 urllib3==2.5.0 \
     --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \
     --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 2b883558d33c6..b74828fa8cdba 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -99,6 +99,78 @@ function(declare_mlir_python_sources name)
   endif()
 endfunction()
 
+# Function: generate_type_stubs
+# Turns on automatic type stub generation (via nanobind's stubgen) for extension modules.
+# Arguments:
+#   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)
+#   CORE_MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module
+#     (e.g., something like StandalonePythonModules.extension._mlir.dso)
+#   OUTPUT_DIR: The root output directory to emit the type stubs into.
+#   OUTPUTS: List of expected outputs.
+#   DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
+# Outputs:
+#   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
+function(generate_type_stubs )
+  cmake_parse_arguments(ARG
+    ""
+    "FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR"
+    "OUTPUTS;DEPENDS_TARGET_SRC_DEPS"
+    ${ARGN})
+  # for people doing find_package(nanobind)
+  if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
+    set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
+  elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
+    set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py")
+  # for people using FetchContent_Declare and FetchContent_MakeAvailable
+  elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py)
+    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py")
+  elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py)
+    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/stubgen.py")
+  else()
+    message(FATAL_ERROR "generate_type_stubs(): could not locate 'stubgen.py'!")
+  endif()
+  file(REAL_PATH "${NB_STUBGEN}" NB_STUBGEN)
+
+  file(REAL_PATH "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." _import_path)
+
+  set(_nb_stubgen_cmd
+      "${Python_EXECUTABLE}"
+      "${NB_STUBGEN}"
+      --module
+      "${ARG_FQ_MODULE_NAME}"
+      -i
+      "${_import_path}"
+      --recursive
+      --include-private
+      --output-dir
+      "${ARG_OUTPUT_DIR}"
+      --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}
+    WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
+    DEPENDS
+      "${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})
+  set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
+endfunction()
+
 # Function: declare_mlir_python_extension
 # Declares a buildable python extension from C++ source files. The built
 # module is considered a python source file and included as everything else.
@@ -115,11 +187,15 @@ 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: 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
     ""
     "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
-    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
+    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -132,15 +208,33 @@ function(declare_mlir_python_extension name)
   endif()
 
   add_library(${name} INTERFACE)
+
+  if(NOT ARG_GENERATE_TYPE_STUBS)
+    set(ARG_GENERATE_TYPE_STUBS OFF)
+  endif()
+  if("${ARG_GENERATE_TYPE_STUBS}" STREQUAL "ON")
+    set(ARG_GENERATE_TYPE_STUBS "${ARG_MODULE_NAME}.pyi")
+  endif()
+  if(ARG_GENERATE_TYPE_STUBS)
+    list(TRANSFORM ARG_GENERATE_TYPE_STUBS PREPEND "_mlir_libs/")
+    declare_mlir_python_sources(
+      "${name}.type_stub_gen"
+      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+      ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
+      SOURCES "${ARG_GENERATE_TYPE_STUBS}"
+    )
+  endif()
+
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
-    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS;mlir_python_BINDINGS_LIBRARY"
+    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS;mlir_python_BINDINGS_LIBRARY;mlir_python_GENERATE_TYPE_STUBS"
     mlir_python_SOURCES_TYPE extension
     mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
     mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     mlir_python_DEPENDS ""
     mlir_python_BINDINGS_LIBRARY "${ARG_PYTHON_BINDINGS_LIBRARY}"
+    mlir_python_GENERATE_TYPE_STUBS "${ARG_GENERATE_TYPE_STUBS}"
   )
 
   # Set the interface source and link_libs properties of the target
@@ -206,10 +300,12 @@ endfunction()
 #     DAG of source modules is included.
 #   COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
 #     extension depend on (see mlir_python_add_common_capi_library).
+#   PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`.
+#     This is used to determine type stub generation python module names.
 function(add_mlir_python_modules name)
   cmake_parse_arguments(ARG
     ""
-    "ROOT_PREFIX;INSTALL_PREFIX"
+    "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX"
     "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES"
     ${ARGN})
   # Helper to process an individual target.
@@ -243,6 +339,41 @@ function(add_mlir_python_modules name)
       )
       add_dependencies(${modules_target} ${_extension_target})
       mlir_python_setup_extension_rpath(${_extension_target})
+      # NOTE: `sources_target` (naturally) lists all the sources (it's the INTERFACE
+      # target defined above in declare_mlir_python_extension). It's also the name of the
+      # target that gets exported (i.e., is populated as an INTERFACE IMPORTED library in MLIRTargets.cmake).
+      # This is why all metadata is queried from `sources_target`. On the other hand
+      # `_extension_target` is the actual dylib target that's built just above with `add_mlir_python_extension`.
+      # That's why dependencies are in terms of `_extension_target`.
+      get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS)
+      if(_generate_type_stubs)
+        if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL ""))
+          message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}")
+        endif()
+        # TL;DR: all paths here are load bearing and annoyingly coupled. Changing paths here
+        # (or related code in declare_mlir_python_extension) will break either in-tree or out-of-tree generation.
+        #
+        # We remove _mlir_libs here because OUTPUT_DIR already includes it.
+        # Specifically OUTPUT_DIR already includes it because that's the actual directory
+        # where we want stubgen to dump the emitted sources. This is load bearing because up above
+        # (in declare_mlir_python_extension) we prefixed all the paths with _mlir_libs and
+        # we specified declare_mlir_python_sources with ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs".
+        #
+        # NOTE: INTERFACE_SOURCES is a genex in the build dir ($<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}"
+          DEPENDS_TARGET ${_extension_target}
+          CORE_MLIR_DEPENDS_TARGET ${name}
+          OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
+          OUTPUTS "${_generate_type_stubs}"
+          DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}"
+        )
+        add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
+      endif()
     else()
       message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
       return()
@@ -678,26 +809,28 @@ function(add_mlir_python_extension libname extname)
       # the super project handle compile options as it wishes.
       get_property(NB_LIBRARY_TARGET_NAME TARGET ${libname} PROPERTY LINK_LIBRARIES)
       target_compile_options(${NB_LIBRARY_TARGET_NAME}
-	PRIVATE
-	  -Wall -Wextra -Wpedantic
-	  -Wno-c++98-compat-extra-semi
-	  -Wno-cast-qual
-	  -Wno-covered-switch-default
-	  -Wno-nested-anon-types
-	  -Wno-unused-parameter
-	  -Wno-zero-length-array
-	  ${eh_rtti_enable})
+        PRIVATE
+          -Wall -Wextra -Wpedantic
+          -Wno-c++98-compat-extra-semi
+          -Wno-cast-qual
+          -Wno-covered-switch-default
+          -Wno-deprecated-literal-operator
+          -Wno-nested-anon-types
+          -Wno-unused-parameter
+          -Wno-zero-length-array
+          ${eh_rtti_enable})
 
       target_compile_options(${libname}
-	PRIVATE
-	  -Wall -Wextra -Wpedantic
-	  -Wno-c++98-compat-extra-semi
-	  -Wno-cast-qual
-	  -Wno-covered-switch-default
-	  -Wno-nested-anon-types
-	  -Wno-unused-parameter
-	  -Wno-zero-length-array
-	  ${eh_rtti_enable})
+        PRIVATE
+          -Wall -Wextra -Wpedantic
+          -Wno-c++98-compat-extra-semi
+          -Wno-cast-qual
+          -Wno-covered-switch-default
+          -Wno-deprecated-literal-operator
+          -Wno-nested-anon-types
+          -Wno-unused-parameter
+          -Wno-zero-length-array
+          ${eh_rtti_enable})
     endif()
 
     if(APPLE)
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 a0eca9c095775..8a5672dcf82cf 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -39,6 +39,9 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
   EMBED_CAPI_LINK_LIBS
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
+  GENERATE_TYPE_STUBS
+    "_standaloneDialectsNanobind/__init__.pyi"
+    "_standaloneDialectsNanobind/standalone.pyi"
 )
 
 
@@ -76,4 +79,5 @@ add_mlir_python_modules(StandalonePythonModules
     MLIRPythonSources.Dialects.builtin
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
-  )
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
+)
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index c983914722ce1..80cd3a2b76e7e 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -1,5 +1,9 @@
 include(AddMLIRPython)
 
+# Specifies that all MLIR packages are co-located under the `mlir_standalone`
+# top level package (the API has been embedded in a relocatable way).
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=${MLIR_PYTHON_PACKAGE_PREFIX}.")
+
 ################################################################################
 # Structural groupings.
 ################################################################################
@@ -23,11 +27,6 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
     passmanager.py
     rewrite.py
     dialects/_ods_common.py
-
-    # The main _mlir module has submodules: include stubs from each.
-    _mlir_libs/_mlir/__init__.pyi
-    _mlir_libs/_mlir/ir.pyi
-    _mlir_libs/_mlir/passmanager.pyi
 )
 
 declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras
@@ -43,7 +42,6 @@ declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
   ADD_TO_PARENT MLIRPythonSources
   SOURCES
     execution_engine.py
-    _mlir_libs/_mlirExecutionEngine.pyi
   SOURCES_GLOB
     runtime/*.py
 )
@@ -195,7 +193,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"
@@ -367,8 +364,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
@@ -384,7 +380,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(
@@ -510,6 +505,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
 
     # Dialects
     MLIRCAPIFunc
+  GENERATE_TYPE_STUBS
+    "_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,6 +531,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPIConversion
     MLIRCAPITransforms
     MLIRCAPIRegisterEverything
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -545,6 +546,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPILinalg
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -559,6 +561,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIGPU
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -573,6 +576,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPILLVM
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -587,6 +591,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIQuant
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -601,6 +606,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPINVGPU
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -615,6 +621,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIPDL
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -629,6 +636,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPISparseTensor
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -643,6 +651,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPITransformDialect
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -663,13 +672,14 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
     MODULE_NAME _mlirExecutionEngine
     ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
     ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  PYTHON_BINDINGS_LIBRARY nanobind
+    PYTHON_BINDINGS_LIBRARY nanobind
     SOURCES
       ExecutionEngineModule.cpp
     PRIVATE_LINK_LIBS
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIExecutionEngine
+    GENERATE_TYPE_STUBS ON
   )
 endif()
 
@@ -714,6 +724,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPIIR
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
+  GENERATE_TYPE_STUBS ON
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -740,6 +751,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransformDialectTransforms
+  GENERATE_TYPE_STUBS ON
 )
 
 # TODO: Figure out how to put this in the test tree.
@@ -798,6 +810,7 @@ if(MLIR_INCLUDE_TESTS)
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
+    GENERATE_TYPE_STUBS ON
   )
 endif()
 
@@ -817,7 +830,7 @@ endif()
 add_mlir_python_common_capi_library(MLIRPythonCAPI
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
-  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_HEADERS
     MLIRPythonCAPI.HeaderSou...
[truncated]

kuhar

The cmake output looks clean on my machine

joker-eph · 2025-09-10T18:51:41Z

Note you can do so with something like [...]

FYI, in general I test PR that way:

$ git fetch origin pull/157853/head
$ git checkout FETCH_HEAD

I checked and locally I see still the standalone example being broken:

# | ninja: error: 'build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/__init__.pyi', needed by 'python_packages/standalone/mlir_standalone/_mlir_libs/_mlir/__init__.pyi', missing and no known rule to make it

makslevental · 2025-09-10T18:54:18Z

# | ninja: error: 'build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/__init__.pyi', needed by 'python_packages/standalone/mlir_standalone/_mlir_libs/_mlir/__init__.pyi', missing and no known rule to make it

Can you check your nanobind version: pip freeze | grep nanobind; you need >=2.9.0. I (and others) observed this failure mode due to a non-updated version of nanobind (which dumps the stubs differently).

joker-eph · 2025-09-10T19:18:55Z

Actually it seems like I see this issue in the current main branch, so something is already broken somehow.

joker-eph · 2025-09-10T19:20:04Z

$ pip install nanobind --upgrade

# | ninja: error: 'build/tools/mlir/python/type_stubs/_mlir_libs/_mlir/__init__.pyi', needed by 'python_packages/standalone/mlir_standalone/_mlir_libs/_mlir/__init__.pyi', missing and no known rule to make it
Can you check your nanobind version: pip freeze | grep nanobind; you need >=2.9.0. I (and others) observed this failure mode due to a non-updated version of nanobind (which dumps the stubs differently).

Thanks! $ pip install nanobind --upgrade did the trick!
Is this something we could detect and error out in CMake?

joker-eph · 2025-09-10T19:36:38Z

Also: does it mean that the standalone example didn't catch any issue? Can we update it to provide more coverage representative of the issues that the other folks hit?

makslevental · 2025-09-10T20:07:55Z

Thanks! $ pip install nanobind --upgrade did the trick! Is this something we could detect and error out in CMake?

The error, which probably surfaced by nanobind itself, was hidden because in #157583 I --quieted output to prevent spamming during build 😄 . The "correct" solution is to not --quiet that output and instead fix the spamming using DEPFILE (which is done in this patch).

Also: does it mean that the standalone example didn't catch any issue? Can we update it to provide more coverage representative of the issues that the other folks hit?

Yes this has all been very unfortunate and partially due to standalone not being a robust enough test of build paths. I fixed some stuff here #156197 but I would like to fix more (see #157944, which isn't ready yet).

FYI: this patch can't land due to complication with IREEs so I'm reverting to draft. I've narrowed the feature and will shortly release #157930 for review.

makslevental force-pushed the users/makslevental/reland-stubgen branch from cb26723 to a1ea24f Compare September 10, 2025 13:15

makslevental marked this pull request as ready for review September 10, 2025 13:24

makslevental requested review from ftynse, stellaraccident and rolfmorel as code owners September 10, 2025 13:24

llvmbot added mlir:python MLIR Python bindings mlir labels Sep 10, 2025

makslevental requested review from jpienaar, rkayaith and joker-eph September 10, 2025 13:25

makslevental force-pushed the users/makslevental/reland-stubgen branch from a1ea24f to ba5a9c9 Compare September 10, 2025 13:27

makslevental requested a review from ingomueller-net September 10, 2025 13:27

makslevental force-pushed the users/makslevental/reland-stubgen branch 2 times, most recently from 60b8cbb to 4ba5645 Compare September 10, 2025 13:31

makslevental mentioned this pull request Sep 10, 2025

[MLIR][Python] fix stubgen for downstream users #157589

Closed

makslevental requested review from christopherbate, boomanaiden154, grypp, ScottTodd, kuhar, chsigg and Groverkss September 10, 2025 13:41

kuhar reviewed Sep 10, 2025

View reviewed changes

makslevental force-pushed the users/makslevental/reland-stubgen branch 4 times, most recently from 9f4c643 to 56e85b6 Compare September 10, 2025 14:43

ScottTodd removed their request for review September 10, 2025 15:22

[MLIR][Python] reland stubgen

2e59800

makslevental force-pushed the users/makslevental/reland-stubgen branch from 56e85b6 to 2e59800 Compare September 10, 2025 15:43

Hardcode84 mentioned this pull request Sep 10, 2025

[python] Test upstream bindings change iree-org/iree#21916

Draft

makslevental marked this pull request as draft September 10, 2025 20:08

makslevental closed this Sep 12, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][Python] reland stubgen #157853

[MLIR][Python] reland stubgen #157853

makslevental commented Sep 10, 2025 •

edited

Loading

Uh oh!

makslevental commented Sep 10, 2025 •

edited

Loading

Uh oh!

llvmbot commented Sep 10, 2025

Uh oh!

kuhar left a comment

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

makslevental commented Sep 10, 2025 •

edited

Loading

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

makslevental commented Sep 10, 2025 •

edited

Loading

Uh oh!

Uh oh!

[MLIR][Python] reland stubgen #157853

[MLIR][Python] reland stubgen #157853

Conversation

makslevental commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

makslevental commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 10, 2025

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

makslevental commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

joker-eph commented Sep 10, 2025

Uh oh!

makslevental commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

makslevental commented Sep 10, 2025 •

edited

Loading

makslevental commented Sep 10, 2025 •

edited

Loading

makslevental commented Sep 10, 2025 •

edited

Loading

makslevental commented Sep 10, 2025 •

edited

Loading