feat: register ENABLE_MULTI_DEVICE and ENABLE_UCX as CMake options (#3343)

No change of default value (still ON).
These were hidden cmake vars before that patch.
Fix issue #3289

Signed-off-by: William Tambellini <wtambellini@sdl.com>
Co-authored-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>
This commit is contained in:
William Tambellini 2025-04-13 19:30:23 -07:00 committed by GitHub
parent 75e13f4f88
commit af67bf00a8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 21 additions and 10 deletions

View File

@ -39,6 +39,10 @@ option(FAST_BUILD "Skip compiling some kernels to accelerate compiling" OFF)
option(FAST_MATH "Compiling in fast math mode" OFF)
option(INDEX_RANGE_CHECK "Compiling with index range checks" OFF)
option(COMPRESS_FATBIN "Compress everything in fatbin" ON)
option(ENABLE_MULTI_DEVICE
"Enable building with multi device support (requires NCCL, MPI,...)" ON)
option(ENABLE_UCX "Enable building with UCX (Uniform Communication X) support"
ON)
# Always use static NVRTC for IP protection reasons.
set(USE_SHARED_NVRTC OFF)
@ -347,7 +351,6 @@ set(CMAKE_CUDA_RUNTIME_LIBRARY Static)
find_library(RT_LIB rt)
set_ifndef(ENABLE_MULTI_DEVICE 1)
if(ENABLE_MULTI_DEVICE)
# NCCL dependencies
set_ifndef(NCCL_LIB_DIR /usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu/)
@ -438,9 +441,15 @@ endif()
# set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -G")
set(CMAKE_CXX_FLAGS
"${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss -DENABLE_MULTI_DEVICE=${ENABLE_MULTI_DEVICE}"
)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DBUILD_SYSTEM=cmake_oss ")
# note: cmake expr generation $<BOOL:${ENABLE_MULTI_DEVICE}> is a build time
# evaluation so hard to debug at cmake time
if(ENABLE_MULTI_DEVICE)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_MULTI_DEVICE=1")
else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_MULTI_DEVICE=0")
endif()
# Fix linking issue with TRT 10, the detailed description about `--mcmodel` can
# be found in
@ -653,7 +662,6 @@ endif()
# Defer UCX/UCXX setup until after USE_CXX11_ABI is well defined, as UCXX will
# need to be built to have aligned symbols
set_ifndef(ENABLE_UCX 1)
if(ENABLE_UCX)
# Only enable UCX related features if the system has UCX library
find_package(ucx)
@ -680,7 +688,11 @@ if(ENABLE_UCX)
NO_DEFAULT_PATH)
endif()
endif()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=${ENABLE_UCX}")
if(ENABLE_UCX)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=1")
else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_UCX=0")
endif()
file(STRINGS "${TRT_INCLUDE_DIR}/NvInferVersion.h" VERSION_STRINGS
REGEX "#define NV_TENSORRT_.*")

View File

@ -59,7 +59,7 @@ function(filter_cuda_archs ARCH SOURCES_VAR)
endif()
endfunction()
if(ENABLE_MULTI_DEVICE EQUAL 0)
if(NOT ENABLE_MULTI_DEVICE)
list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$")
endif()

View File

@ -166,9 +166,8 @@ function(process_target target_name enable_hopper enable_blackwell)
${target_name} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-Wno-psabi>)
endif()
target_compile_options(${target_name}
PRIVATE "-DENABLE_MULTI_DEVICE=${ENABLE_MULTI_DEVICE}")
if(ENABLE_MULTI_DEVICE)
target_compile_options(${target_name} PRIVATE "-DENABLE_MULTI_DEVICE=1")
target_link_libraries(${target_name} PRIVATE ${MPI_C_LIBRARIES})
endif()

View File

@ -14,7 +14,7 @@
# License for the specific language governing permissions and limitations under
# the License.
#
if(ENABLE_MULTI_DEVICE EQUAL 0)
if(NOT ENABLE_MULTI_DEVICE)
add_library(userbuffers_src OBJECT ub_interface.cpp userbuffersManager.cpp)
else()
file(GLOB_RECURSE SRC_CPP *.cpp)