[MLIR][Python] make sure stubs get installed with LLVM_DISTRIBUTION_COMPONENTS (#168407)

Fixes https://github.com/llvm/llvm-project/issues/168393. Also adds
top-level `MLIR_PYTHON_STUBGEN_ENABLED` CMake option.

3 files changed, 28 insertions, 27 deletions

diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index 1a211f549576..570fb6f89dd1 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -210,6 +210,19 @@ set(MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES 0 CACHE BOOL
        'Development.Module' and ensure that find_package(pybind11) is \
        satisfied (and keep up to date as requirements evolve).")
 
+set(_mlir_python_stubgen_enabled ON)
+# Stubgen doesn't work when cross-compiling (stubgen will run in the host interpreter and then fail
+# to find the extension module for the host arch).
+# Note: Stubgen requires some extra handling to work properly when sanitizers are enabled,
+# so we skip running it in that case now.
+if(CMAKE_CROSSCOMPILING OR (NOT LLVM_USE_SANITIZER STREQUAL ""))
+  set(_mlir_python_stubgen_enabled OFF)
+endif()
+
+option(MLIR_PYTHON_STUBGEN_ENABLED
+       "Generate Python type stubs for the MLIR Python bindings."
+       ${_mlir_python_stubgen_enabled})
+
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   include(MLIRDetectPythonEnv)
   # Note that both upstream and downstreams often call this macro. It gates
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index df19fa86067e..8469bff0eda2 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -74,12 +74,7 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
 
 set(StandalonePythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
 
-set(_mlir_python_stubgen_enabled ON)
-if(CMAKE_CROSSCOMPILING OR (NOT LLVM_USE_SANITIZER STREQUAL ""))
-  set(_mlir_python_stubgen_enabled OFF)
-endif()
-
-if(_mlir_python_stubgen_enabled)
+if(MLIR_PYTHON_STUBGEN_ENABLED)
   # Everything here is very tightly coupled. See the ample descriptions at the bottom of
   # mlir/python/CMakeLists.txt.
 
@@ -146,7 +141,7 @@ set(_declared_sources
 )
 # For an external projects build, the MLIRPythonExtension.Core.type_stub_gen
 # target already exists and can just be added to DECLARED_SOURCES.
-if(EXTERNAL_PROJECT_BUILD AND _mlir_python_stubgen_enabled)
+if(EXTERNAL_PROJECT_BUILD AND MLIR_PYTHON_STUBGEN_ENABLED)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
 
@@ -158,7 +153,7 @@ add_mlir_python_modules(StandalonePythonModules
     StandalonePythonCAPI
 )
 
-if(_mlir_python_stubgen_enabled)
+if(MLIR_PYTHON_STUBGEN_ENABLED)
   if(NOT EXTERNAL_PROJECT_BUILD)
     add_dependencies(StandalonePythonModules "${_mlir_typestub_gen_target}")
   endif()
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 51c75764faf3..2acb6ee6cfda 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -895,24 +895,8 @@ add_mlir_python_common_capi_library(MLIRPythonCAPI
 ################################################################################
 
 _flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
-add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps})
-if(NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(install-mlir-python-sources
-    DEPENDS mlir-python-sources
-    COMPONENT mlir-python-sources
-  )
-endif()
-
-set(_mlir_python_stubgen_enabled ON)
-# Stubgen doesn't work when cross-compiling (stubgen will run in the host interpreter and then fail
-# to find the extension module for the host arch).
-# Note: Stubgen requires some extra handling to work properly when sanitizers are enabled,
-# so we skip running it in that case now.
-if(CMAKE_CROSSCOMPILING OR (NOT LLVM_USE_SANITIZER STREQUAL ""))
-  set(_mlir_python_stubgen_enabled OFF)
-endif()
 
-if(_mlir_python_stubgen_enabled)
+if(MLIR_PYTHON_STUBGEN_ENABLED)
   # _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.
@@ -965,6 +949,7 @@ if(_mlir_python_stubgen_enabled)
     ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
     SOURCES "${_core_type_stub_sources}"
   )
+  list(APPEND mlir_python_sources_deps MLIRPythonExtension.Core.type_stub_gen)
 
   # _mlirPythonTestNanobind stubgen
 
@@ -995,13 +980,21 @@ if(_mlir_python_stubgen_enabled)
   endif()
 endif()
 
+add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps})
+if(NOT LLVM_ENABLE_IDE)
+  add_llvm_install_targets(install-mlir-python-sources
+    DEPENDS mlir-python-sources
+    COMPONENT mlir-python-sources
+  )
+endif()
+
 ################################################################################
 # The fully assembled package of modules.
 # This must come last.
 ################################################################################
 
 set(_declared_sources MLIRPythonSources MLIRPythonExtension.RegisterEverything)
-if(_mlir_python_stubgen_enabled)
+if(MLIR_PYTHON_STUBGEN_ENABLED)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
 
@@ -1014,7 +1007,7 @@ add_mlir_python_modules(MLIRPythonModules
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )
-if(_mlir_python_stubgen_enabled)
+if(MLIR_PYTHON_STUBGEN_ENABLED)
   add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}")
   if(MLIR_INCLUDE_TESTS)
     add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")