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)}")