[MLIR][Python] Skip stubgen while any sanitizer is enabled (llvm#164661)

The intention of this PR is described
llvm#164197 (comment)
(and
llvm#164197 (comment)).

When sanitizers are enabled, some additional setup (such as preloading
certain libraries) seems to be required for the stubgen step to work
properly
(llvm#164197 (comment)).
In this case, I chose to simply skip the stubgen process, as supporting
it would likely require some extra CMake logic, and type stubs don’t
appear to be particularly necessary in this configuration.

Author: PragmaTwice

mlir/examples/standalone/python/CMakeLists.txt

@@ -74,7 +74,12 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
 
 set(StandalonePythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
 
-if(NOT CMAKE_CROSSCOMPILING)
+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)
   # Everything here is very tightly coupled. See the ample descriptions at the bottom of
   # mlir/python/CMakeLists.txt.
 
@@ -141,7 +146,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 (NOT CMAKE_CROSSCOMPILING))
+if(EXTERNAL_PROJECT_BUILD AND _mlir_python_stubgen_enabled)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
 
@@ -153,7 +158,7 @@ add_mlir_python_modules(StandalonePythonModules
     StandalonePythonCAPI
 )
 
-if(NOT CMAKE_CROSSCOMPILING)
+if(_mlir_python_stubgen_enabled)
   if(NOT EXTERNAL_PROJECT_BUILD)
     add_dependencies(StandalonePythonModules "${_mlir_typestub_gen_target}")
   endif()

mlir/python/CMakeLists.txt

@@ -894,9 +894,16 @@ if(NOT LLVM_ENABLE_IDE)
   )
 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).
-if(NOT CMAKE_CROSSCOMPILING)
+# 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)
   # _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.
@@ -985,7 +992,7 @@ endif()
 ################################################################################
 
 set(_declared_sources MLIRPythonSources MLIRPythonExtension.RegisterEverything)
-if(NOT CMAKE_CROSSCOMPILING)
+if(_mlir_python_stubgen_enabled)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
 
@@ -998,7 +1005,7 @@ add_mlir_python_modules(MLIRPythonModules
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )
-if(NOT CMAKE_CROSSCOMPILING)
+if(_mlir_python_stubgen_enabled)
   add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}")
   if(MLIR_INCLUDE_TESTS)
     add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")