From 5169d129cea2dfe4cca1314cec885f2170b0b25b Mon Sep 17 00:00:00 2001
From: Jingyue Wu <wujingyue@gmail.com>
Date: Sat, 23 Sep 2023 11:18:05 -0700
Subject: [PATCH] Add an option to build with asan. (#933)

Tested:

```
$ pip install -e .
$ nm -an bin/nvfuser_tests | grep asan
<empty>
$ python setup.py develop --build-with-asan
$ nm -an bin/nvfuser_tests | grep asan
...
00000000024f3840 B __odr_asan._ZN7nvfuser38NVFuserTest_FusionDefinition_CUDA_Test10test_info_E
00000000024f38e0 B __odr_asan._ZN7nvfuser35NVFuserTest_PyFusionCache_CUDA_Test10test_info_E
00000000024f3980 B __odr_asan._ZN7nvfuser43NVFuserTest_RecordFunctorEquality_CUDA_Test10test_info_E
0000000002528fa0 B __odr_asan._ZN5torch8autograd47_GLOBAL__N__551c3bf8_14_rng_kernels_cu_8fdfe6fb13graph_task_idE
```
---
 CMakeLists.txt | 7 +++++++
 setup.py       | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index d7013b042f4..d8cd5d85fd7 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -18,6 +18,7 @@ set(NVFUSER_SRCS_DIR "${NVFUSER_ROOT}/csrc")
 set(NVFUSER_THIRD_PARTY_DIR "${NVFUSER_ROOT}/third_party")
 
 option(NVFUSER_STANDALONE_BUILD_WITH_UCC "" OFF)
+option(NVFUSER_BUILD_WITH_ASAN "Build nvFuser with asan" OFF)
 
 if(PROJECT_IS_TOP_LEVEL)
   find_package(Torch REQUIRED)
@@ -246,6 +247,11 @@ target_include_directories(${NVFUSER_CODEGEN} PUBLIC
 )
 set_property(TARGET ${NVFUSER_CODEGEN} PROPERTY CXX_STANDARD 17)
 
+if(NVFUSER_BUILD_WITH_ASAN)
+  add_compile_options(-fsanitize=address)
+  add_link_options(-fsanitize=address)
+endif()
+
 if(PROJECT_IS_TOP_LEVEL)
   target_link_libraries(${NVFUSER_CODEGEN} PRIVATE torch ${TORCH_LIBRARIES})
 
@@ -632,6 +638,7 @@ message(STATUS "")
 message(STATUS "******** Nvfuser configuration summary ********")
 message(STATUS "  UCC_FOUND: ${UCC_FOUND}")
 message(STATUS "  NVFUSER_STANDALONE_BUILD_WITH_UCC  : ${NVFUSER_STANDALONE_BUILD_WITH_UCC}")
+message(STATUS "  NVFUSER_BUILD_WITH_ASAN            : ${NVFUSER_BUILD_WITH_ASAN}")
 
 if(NVFUSER_STANDALONE_BUILD_WITH_UCC)
   message(STATUS "    UCC_HOME: $ENV{UCC_HOME}")
diff --git a/setup.py b/setup.py
index c6cf00fb2dc..7244a4b6bc1 100644
--- a/setup.py
+++ b/setup.py
@@ -61,6 +61,7 @@
 NO_BENCHMARK = False
 NO_NINJA = False
 BUILD_WITH_UCC = False
+BUILD_WITH_ASAN = False
 PATCH_NVFUSER = True
 OVERWRITE_VERSION = False
 VERSION_TAG = None
@@ -87,6 +88,9 @@
     if arg == "--build-with-ucc":
         BUILD_WITH_UCC = True
         continue
+    if arg == "--build-with-asan":
+        BUILD_WITH_ASAN = True
+        continue
     if arg == "--debug":
         BUILD_TYPE = "Debug"
         continue
@@ -296,6 +300,8 @@ def cmake(build_dir: str = "", install_prefix: str = "./nvfuser"):
         cmd_str.append(f"-DPython_EXECUTABLE={sys.executable}")
     if not NO_BENCHMARK:
         cmd_str.append("-DBUILD_NVFUSER_BENCHMARK=ON")
+    if BUILD_WITH_ASAN:
+        cmd_str.append("-DNVFUSER_BUILD_WITH_ASAN=ON")
     cmd_str.append(".")
 
     print(f"Configuring CMake with {' '.join(cmd_str)}")