From 6dc54fe8d670a3ff15f6ba49929deb0202e93948 Mon Sep 17 00:00:00 2001
From: Nikita Shulga <2453524+malfet@users.noreply.github.com>
Date: Sat, 14 Oct 2023 20:35:04 +0000
Subject: [PATCH] [BE] Compile FBGEMM with ASAN (#111266)

If `USE_ASAN` is set, compile FBGEMM with ASAN as well, by setting `USE_SANITIZER` to `address,undefined`

This fixes regression in sanitizer coverage introduced by https://github.com/pytorch/pytorch/pull/93147  that change effects of sanitizer from the entire project to just torch libraries, and finally allows one to reliably catch regression reported in https://github.com/pytorch/pytorch/issues/111189

Pull Request resolved: https://github.com/pytorch/pytorch/pull/111266
Approved by: https://github.com/huydhn
---
 caffe2/CMakeLists.txt    | 8 ++++++++
 cmake/Dependencies.cmake | 3 +++
 2 files changed, 11 insertions(+)

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index b3888ebde6737..0511b8fd0f63c 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -1688,6 +1688,14 @@ if(BUILD_TEST)
           if(USE_FBGEMM)
             target_link_libraries(${test_name}_${CPU_CAPABILITY} fbgemm)
           endif()
+          if(USE_ASAN)
+            if(TARGET Sanitizer::address)
+              target_link_libraries(${test_name}_${CPU_CAPABILITY} Sanitizer::address)
+            endif()
+            if(TARGET Sanitizer::undefined)
+              target_link_libraries(${test_name}_${CPU_CAPABILITY} Sanitizer::undefined)
+            endif()
+          endif()
         else()
           add_executable(${test_name}_${CPU_CAPABILITY} "${test_src}")
           target_link_libraries(${test_name}_${CPU_CAPABILITY} torch_library gtest_main)
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
index 503895aa4b152..0602d534dc4c1 100644
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@@ -815,6 +815,9 @@ if(USE_FBGEMM)
     else()
       set(FBGEMM_LIBRARY_TYPE "static" CACHE STRING "")
     endif()
+    if(USE_ASAN)
+      set(USE_SANITIZER "address,undefined" CACHE STRING "-fsanitize options for FBGEMM")
+    endif()
     add_subdirectory("${FBGEMM_SOURCE_DIR}")
     set_property(TARGET fbgemm_generic PROPERTY POSITION_INDEPENDENT_CODE ON)
     set_property(TARGET fbgemm_avx2 PROPERTY POSITION_INDEPENDENT_CODE ON)