Add instrumentation for Google Sanitizers (openxla#1959)

Google Sanitiziers (https://github.com/google/sanitizers) have been
integrated into Clang & GNU.

Add support for ASan + UBSan if a CMake option is set. This should help
catch some of the same errors that can be provided typicaly within
internal Google development.

TODO: Add support for TSan which needs to run individually since it
cannot be instrumented alongside ASan.

---------

Co-authored-by: mlevesquedion <mlevesquedion@google.com>

.github/workflows/buildAndTestCMake.yml

@@ -74,6 +74,15 @@ jobs:
           CMAKE_BUILD_TYPE: Release
           MLIR_ENABLE_BINDINGS_PYTHON: ON
 
+    - name: Build and Test StableHLO (with AddressSanitizer)
+      shell: bash
+      run: |
+        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
+      env:
+          CMAKE_BUILD_TYPE: Release
+          STABLEHLO_ENABLE_BINDINGS_PYTHON: OFF
+          STABLEHLO_ENABLE_SANITIZER: address
+
     - name: Build and Test StableHLO (with Python bindings)
       shell: bash
       run: |

CMakeLists.txt

@@ -34,12 +34,18 @@ if(POLICY CMP0116)
   cmake_policy(SET CMP0116 OLD)
 endif()
 
+# Support for return(PROPAGATE ...) in functions.
+if (POLICY CMP0140)
+  cmake_policy(SET CMP0140 NEW)
+endif()
+
 #-------------------------------------------------------------------------------
 # Options and settings
 #-------------------------------------------------------------------------------
 option(STABLEHLO_BUILD_EMBEDDED "Build StableHLO as part of another project" OFF)
 option(STABLEHLO_ENABLE_BINDINGS_PYTHON "Enables StableHLO Python bindings" OFF)
 option(STABLEHLO_ENABLE_STRICT_BUILD "Build StableHLO with strict warnings and warnings as errors" OFF)
+option(STABLEHLO_ENABLE_SANITIZER "Enable a sanitizer [OFF, address]" OFF)
 
 #-------------------------------------------------------------------------------
 # Project setup and globals
@@ -115,6 +121,9 @@ else()
   message(STATUS "Building StableHLO embedded in another project")
 endif()
 
+# Add the CMake modules specific to StableHLO
+list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
+
 if(LLVM_ENABLE_ZLIB)
   find_package(ZLIB)
 endif()
@@ -130,6 +139,14 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR})
 link_directories(${LLVM_BUILD_LIBRARY_DIR})
 add_definitions(${LLVM_DEFINITIONS})
 
+
+#-------------------------------------------------------------------------------
+# Sanitizer configuration
+#-------------------------------------------------------------------------------
+
+include(SetupSanitizers)
+setup_sanitizers()
+
 #-------------------------------------------------------------------------------
 # Python configuration
 #-------------------------------------------------------------------------------

build_tools/github_actions/ci_build_cmake.sh

@@ -31,6 +31,9 @@ CMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE:-RelWithDebInfo}"
 
 # Turn on building Python bindings
 STABLEHLO_ENABLE_BINDINGS_PYTHON="${STABLEHLO_ENABLE_BINDINGS_PYTHON:-OFF}"
+# Turn on building Sanitizers
+# Note: This is not congruent with building python bindings
+STABLEHLO_ENABLE_SANITIZER="${STABLEHLO_ENABLE_SANITIZER:-OFF}"
 
 # Configure StableHLO
 # CMAKE_PLATFORM_NO_VERSIONED_SONAME Disables generation of "version soname"
@@ -48,6 +51,7 @@ cmake -GNinja \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache \
   -DSTABLEHLO_ENABLE_STRICT_BUILD=ON \
   -DCMAKE_PLATFORM_NO_VERSIONED_SONAME:BOOL=ON \
+  -DSTABLEHLO_ENABLE_SANITIZER="$STABLEHLO_ENABLE_SANITIZER" \
   -DSTABLEHLO_ENABLE_BINDINGS_PYTHON="$STABLEHLO_ENABLE_BINDINGS_PYTHON"
 
 # Build and Test StableHLO

cmake/SetupSanitizers.cmake

@@ -0,0 +1,56 @@
+# Copyright 2024 The StableHLO Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# setup_sanitizers: sets up the compile and link options necessary to use
+#   various sanitizers.
+# At the moment we disallow setting up the sanitizers for Python bindings
+# since this seems especially broken when building with Clang.
+# As of 1/29/2024 here are some of the issues encountered with Python:
+# * Clang only links in the sanitizer runtime for the main executable, not
+#    for the shared libraries. Since Python is un-sanitized, this is
+#    problematic.
+# * Linking to the Clang Sanitizer is in the directory
+#    /usr/lib/llvm-16/lib/clang/16/lib/linux which needs to be added to RPATH
+# * We are using LLVM with RTTI enabled, but installed LLVM has RTTI disabled.
+#    This causes a linker error when linking to the sanitizer runtime.
+#
+# Other MLIR based projects (IREE) seem to have similar issues and just skip
+# the Python bindings for sanitization builds.
+# TODO(fzakaria): Revisit a better way or if this can be solved.
+# @see: https://github.com/google/sanitizers/wiki
+function(setup_sanitizers)
+    if (NOT STABLEHLO_ENABLE_SANITIZER)
+        return()
+    endif ()
+
+    string(TOLOWER "${STABLEHLO_ENABLE_SANITIZER}" STABLEHLO_ENABLE_SANITIZER_LOWERCASE)
+    if (STABLEHLO_ENABLE_SANITIZER_LOWERCASE STREQUAL "off")
+        return()
+    endif ()
+
+    if (STABLEHLO_ENABLE_BINDINGS_PYTHON)
+        message(FATAL_ERROR "STABLEHLO_ENABLE_SANITIZER must be set to OFF when building Python bindings")
+        return()
+    endif ()
+
+    if (STABLEHLO_ENABLE_SANITIZER_LOWERCASE STREQUAL "address")
+        add_compile_options(-fsanitize=address -fsanitize=undefined -fsanitize=leak -fno-omit-frame-pointer)
+        link_libraries(-fsanitize=address -fsanitize=undefined -fsanitize=leak)
+    else ()
+        message(FATAL_ERROR "Unknown sanitizer type: ${STABLEHLO_ENABLE_SANITIZER}")
+    endif ()
+endfunction()
+
+