[Mlir-commits] [mlir] [mlir][Python] generate type stubs for dialect extensions (PR #175403)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/175403

>From 87e6dfcd47f52be4f60625195d8a5874841d47e2 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Sat, 10 Jan 2026 18:11:42 -0800
Subject: [PATCH] [mlir][Python] generate type stubs for dialect extensions

---
 mlir/python/CMakeLists.txt | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 003a06b16daac..a7cece87a11d6 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -924,7 +924,25 @@ 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}")
+  set(_mlir_typestub_gen_targets "${NB_STUBGEN_CUSTOM_TARGET}")
+
+  get_target_property(_linalg_extension_srcs MLIRPythonExtension.Dialects.Linalg.Nanobind INTERFACE_SOURCES)
+  mlir_generate_type_stubs(
+    MODULE_NAME ${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs._mlirDialectsLinalg
+    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._mlirDialectsLinalg.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 _mlirDialectsLinalg.pyi
+    DEPENDS_TARGET_SRC_DEPS "${_linalg_extension_srcs}"
+    IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/.."
+  )
+  list(APPEND _mlir_typestub_gen_targets "${NB_STUBGEN_CUSTOM_TARGET}")
+  list(APPEND _core_type_stub_sources "_mlirDialectsLinalg.pyi")
 
   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
@@ -943,7 +961,7 @@ if(MLIR_PYTHON_STUBGEN_ENABLED)
     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
+      MODULE_NAME ${MLIR_PYTHON_PACKAGE_PREFIX}._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).
@@ -994,7 +1012,7 @@ add_mlir_python_modules(MLIRPythonModules
     MLIRPythonCAPI
 )
 if(MLIR_PYTHON_STUBGEN_ENABLED)
-  add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}")
+  add_dependencies(MLIRPythonModules ${_mlir_typestub_gen_targets})
   if(MLIR_INCLUDE_TESTS)
     add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")
   endif()