From be179fce7a34dd67730ce9037802720c8db659b2 Mon Sep 17 00:00:00 2001
From: Muhammad Asif Manzoor <mmanzoor@tenstorrent.com>
Date: Thu, 12 Dec 2024 08:48:48 -0500
Subject: [PATCH] Add checks for address sanitizer build (#1562)

* Address sanitization build causes errors for python bindings and shared libs.
Other MLIR based projects (IREE, Stablehlo) seem to have similar issues and just
skip building these components for sanitization builds.
* Address sanitizer build can be enabled by setting '-DCMAKE_BUILD_TYPE=Asan
-DTTMLIR_ENABLE_SHARED_LIB=OFF -DTTMLIR_ENABLE_BINDINGS_PYTHON=OFF' options with
cmake command.
---
 CMakeLists.txt     | 16 ++++++++++++++++
 lib/CMakeLists.txt |  2 +-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index bebff7a0fd..bb6bcda75f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -9,6 +9,11 @@ option(TT_RUNTIME_ENABLE_PERF_TRACE "Enable performance mode" OFF)
 option(TTMLIR_ENABLE_RUNTIME "Enable runtime" OFF)
 option(TTMLIR_ENABLE_STABLEHLO "Enable StableHLO support" OFF)
 option(TTMLIR_ENABLE_OP_MODEL "Enable OpModel support" OFF)
+option(TTMLIR_ENABLE_SHARED_LIB "Enable Shared lib building" ON)
+
+if (NOT TTMLIR_ENABLE_RUNTIME)
+  set (TTMLIR_ENABLE_SHARED_LIB OFF)
+endif()
 
 if (TTMLIR_ENABLE_STABLEHLO)
   add_compile_definitions(TTMLIR_ENABLE_STABLEHLO)
@@ -50,6 +55,17 @@ set(Python3_EXECUTABLE $ENV{TTMLIR_VENV_DIR}/bin/python3)
 include(FindMLIR)
 include(TTMLIRVersion)
 
+if (CMAKE_BUILD_TYPE STREQUAL "Asan")
+  # Address sanitization does not work with python bindings.
+  if(MLIR_ENABLE_BINDINGS_PYTHON AND TTMLIR_ENABLE_BINDINGS_PYTHON)
+    message(FATAL_ERROR "TTMLIR_ENABLE_BINDINGS_PYTHON must be set to OFF when building TTMLIR with address sanitization")
+  endif()
+  # Clang only links in the sanitizer runtime for the main executable, not for the shared libraries.
+  if (TTMLIR_ENABLE_SHARED_LIB)
+    message(FATAL_ERROR "TTMLIR_ENABLE_SHARED_LIB must be set to OFF when building TTMLIR with address sanitization")
+  endif()
+endif()
+
 set(TTMLIR_TOOLCHAIN_DIR $ENV{TTMLIR_TOOLCHAIN_DIR})
 set(TTMLIR_SOURCE_DIR ${PROJECT_SOURCE_DIR})
 set(TTMLIR_BINARY_DIR ${PROJECT_BINARY_DIR})
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
index 881d6545dc..76c6b215a0 100644
--- a/lib/CMakeLists.txt
+++ b/lib/CMakeLists.txt
@@ -10,7 +10,7 @@ add_subdirectory(Scheduler)
 
 # Shared library will include runtime code
 # so we only build it if runtime is enabled
-if (TTMLIR_ENABLE_RUNTIME)
+if (TTMLIR_ENABLE_SHARED_LIB)
     add_subdirectory(SharedLib)
 endif()