Skip to content

Commit c86993b

Browse files
Flying-dragon-boxingFlying-dragon-boxing
authored
Feature: Enable NCCL library for GPU-Direct collective communications in plane wave (BPCG, stochastic KG, etc.) (#7301)
* Harden GPU MPI staging helpers * Add NCCL collectives for parallel_device * Fix NCCL headers in parallel_device * Route PGemm collectives through device wrappers * Tighten NCCL collective correctness * Relax NCCL discovery for existing environments * Decouple NCCL parallel_device from CUDA-aware MPI * Propagate NCCL headers to subdirectory targets * Fix: narrow CPU staging guards in para_gemm to respect NCCL collectives isend_dev has no NCCL path — keep guard as #ifndef __CUDA_MPI. reduce_dev / gatherv_dev have NCCL early-returns — exclude CPU staging when __NCCL_PARALLEL_DEVICE is defined (&& !defined).
1 parent 78e16d9 commit c86993b

6 files changed

Lines changed: 401 additions & 53 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ option(ENABLE_GOOGLEBENCH "Enable GOOGLE-benchmark usage" OFF)
5151
option(ENABLE_RAPIDJSON "Enable rapid-json usage" OFF)
5252
option(ENABLE_CNPY "Enable cnpy usage" OFF)
5353
option(ENABLE_CUSOLVERMP "Enable cusolvermp" OFF)
54+
option(ENABLE_NCCL_PARALLEL_DEVICE "Enable NCCL-backed collectives in parallel_device" OFF)
5455

5556
if(NOT DEFINED NVHPC_ROOT_DIR AND DEFINED ENV{NVHPC_ROOT})
5657
set(NVHPC_ROOT_DIR
@@ -451,6 +452,15 @@ if(USE_CUDA)
451452
if (USE_OPENMP AND OpenMP_CXX_FOUND)
452453
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler=${OpenMP_CXX_FLAGS}" CACHE STRING "CUDA flags" FORCE)
453454
endif()
455+
if (ENABLE_NCCL_PARALLEL_DEVICE)
456+
if (NOT ENABLE_MPI)
457+
message(FATAL_ERROR
458+
"ENABLE_NCCL_PARALLEL_DEVICE requires ENABLE_MPI=ON.")
459+
endif()
460+
add_compile_definitions(__NCCL_PARALLEL_DEVICE)
461+
include(cmake/SetupNccl.cmake)
462+
abacus_setup_nccl(${ABACUS_BIN_NAME})
463+
endif()
454464
if (ENABLE_CUSOLVERMP)
455465
# Keep cuSOLVERMp discovery/linking logic in a dedicated module.
456466
include(cmake/SetupCuSolverMp.cmake)

‎cmake/SetupNccl.cmake‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
include_guard(GLOBAL)
2+
3+
include(CheckIncludeFileCXX)
4+
5+
function(abacus_setup_nccl target_name)
6+
find_library(NCCL_LIBRARY NAMES nccl
7+
HINTS ${NCCL_PATH} ${NVHPC_ROOT_DIR}
8+
PATH_SUFFIXES lib lib64 comm_libs/nccl/lib)
9+
find_path(NCCL_INCLUDE_DIR NAMES nccl.h
10+
HINTS ${NCCL_PATH} ${NVHPC_ROOT_DIR}
11+
PATHS ${CUDAToolkit_ROOT}
12+
PATH_SUFFIXES include comm_libs/nccl/include)
13+
14+
check_include_file_cxx("nccl.h" HAVE_NCCL_HEADER)
15+
16+
if(NOT NCCL_LIBRARY)
17+
set(NCCL_LIBRARY nccl)
18+
endif()
19+
20+
if(NOT NCCL_INCLUDE_DIR AND NOT HAVE_NCCL_HEADER)
21+
message(FATAL_ERROR
22+
"NCCL not found. Set NCCL_PATH or NVHPC_ROOT_DIR.")
23+
endif()
24+
25+
if(NCCL_INCLUDE_DIR)
26+
message(STATUS "Found NCCL for parallel_device: ${NCCL_LIBRARY}")
27+
else()
28+
message(STATUS "Using default compiler/linker search paths for NCCL: ${NCCL_LIBRARY}")
29+
endif()
30+
if(NOT TARGET NCCL::NCCL)
31+
add_library(NCCL::NCCL IMPORTED INTERFACE)
32+
if(NCCL_INCLUDE_DIR)
33+
set_target_properties(NCCL::NCCL PROPERTIES
34+
INTERFACE_LINK_LIBRARIES "${NCCL_LIBRARY}"
35+
INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}")
36+
else()
37+
set_target_properties(NCCL::NCCL PROPERTIES
38+
INTERFACE_LINK_LIBRARIES "${NCCL_LIBRARY}")
39+
endif()
40+
endif()
41+
42+
if(NCCL_INCLUDE_DIR)
43+
# `parallel_device.cpp` is compiled inside the later `base` OBJECT library,
44+
# so the header path must also be visible to targets created in subdirs.
45+
include_directories(${NCCL_INCLUDE_DIR})
46+
target_include_directories(${target_name} PRIVATE ${NCCL_INCLUDE_DIR})
47+
endif()
48+
target_link_libraries(${target_name} NCCL::NCCL)
49+
endfunction()

‎source/source_base/module_device/device_check.h‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,23 @@ static const char* _cufftGetErrorString(cufftResult_t error)
219219
#define CHECK_CUDA_SYNC() do {} while (0)
220220
#endif
221221

222+
// NCCL check macro: shared by cuSOLVER MP (non-CAL path) and parallel device
223+
#if (defined(__CUSOLVERMP) && !defined(__USE_CAL)) || defined(__NCCL_PARALLEL_DEVICE)
224+
#include <nccl.h>
225+
226+
#define CHECK_NCCL(func) \
227+
do \
228+
{ \
229+
ncclResult_t status = (func); \
230+
if (status != ncclSuccess) \
231+
{ \
232+
fprintf(stderr, "In File %s : NCCL API failed at line %d with error: %s (%d)\n", __FILE__, __LINE__, \
233+
ncclGetErrorString(status), status); \
234+
exit(EXIT_FAILURE); \
235+
} \
236+
} while (0)
237+
#endif
238+
222239
// cuSOLVER MP support
223240
#ifdef __CUSOLVERMP
224241
#include <cusolverMp.h>
@@ -262,20 +279,6 @@ static const char* _calGetErrorString(calError_t error)
262279
exit(EXIT_FAILURE); \
263280
} \
264281
} while (0)
265-
#else // !__USE_CAL (use NCCL)
266-
#include <nccl.h>
267-
268-
#define CHECK_NCCL(func) \
269-
do \
270-
{ \
271-
ncclResult_t status = (func); \
272-
if (status != ncclSuccess) \
273-
{ \
274-
fprintf(stderr, "In File %s : NCCL API failed at line %d with error: %s (%d)\n", __FILE__, __LINE__, \
275-
ncclGetErrorString(status), status); \
276-
exit(EXIT_FAILURE); \
277-
} \
278-
} while (0)
279282
#endif // __USE_CAL
280283

281284
#endif // __CUSOLVERMP

‎source/source_base/para_gemm.cpp‎

Lines changed: 16 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ void PGemmCN<T, Device>::set_dimension(
133133
if (std::is_same<Device, base_device::DEVICE_GPU>::value)
134134
{
135135
resmem_dev_op()(C_local_tmp_, size_C_local);
136-
#ifndef __CUDA_MPI
136+
#if !defined(__CUDA_MPI) && !defined(__NCCL_PARALLEL_DEVICE)
137137
C_global_tmp_.resize(size_C_global);
138138
#endif
139139
}
@@ -277,38 +277,27 @@ void PGemmCN<T, Device>::multiply_col(const T alpha, const T* A, const T* B, con
277277

278278
if (this->gatherC)
279279
{
280-
#ifdef __CUDA_MPI
281-
T* Clocal_mpi = C_local;
282-
T* Cglobal_mpi = C;
283-
#else
284-
T* Clocal_mpi = C_tmp_.data();
285-
T* Cglobal_mpi = nullptr;
280+
T* reduce_tmp = nullptr;
281+
T* gather_tmp = nullptr;
282+
#if !defined(__CUDA_MPI) && !defined(__NCCL_PARALLEL_DEVICE)
286283
if (std::is_same<Device, base_device::DEVICE_GPU>::value)
287284
{
288-
syncmem_d2h_op()(Clocal_mpi, C_local, size_C_local);
289-
Cglobal_mpi = C_global_tmp_.data();
290-
}
291-
else
292-
{
293-
Cglobal_mpi = C;
285+
reduce_tmp = C_tmp_.data();
286+
gather_tmp = C_global_tmp_.data();
294287
}
295288
#endif
296289
if (this->row_nproc > 1)
297290
{
298-
Parallel_Common::reduce_data(Clocal_mpi, size_C_local, row_world);
291+
Parallel_Common::reduce_dev<T, Device>(C_local, size_C_local, row_world, reduce_tmp);
299292
}
300-
Parallel_Common::gatherv_data(Clocal_mpi,
301-
size_C_local,
302-
Cglobal_mpi,
303-
recv_counts.data(),
304-
displs.data(),
305-
col_world);
306-
#ifndef __CUDA_MPI
307-
if (std::is_same<Device, base_device::DEVICE_GPU>::value)
308-
{
309-
syncmem_h2d_op()(C, Cglobal_mpi, size_C_global);
310-
}
311-
#endif
293+
Parallel_Common::gatherv_dev<T, Device>(C_local,
294+
size_C_local,
295+
C,
296+
recv_counts.data(),
297+
displs.data(),
298+
col_world,
299+
reduce_tmp,
300+
gather_tmp);
312301
}
313302
else
314303
{
@@ -409,4 +398,4 @@ template class PGemmCN<std::complex<double>, base_device::DEVICE_GPU>;
409398
template class PGemmCN<std::complex<float>, base_device::DEVICE_GPU>;
410399
#endif
411400

412-
} // namespace ModuleBase
401+
} // namespace ModuleBase

0 commit comments

Comments
 (0)