From e7b8a776cee0fb6ee947e70ab190ff10ba8f4801 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ji=C5=99=C3=AD=20Vysko=C4=8Dil?= Date: Mon, 2 Feb 2026 12:29:35 +0100 Subject: [PATCH] Add NCCL support for CuSolverMP --- cmake/modules/FindCuSolverMP.cmake | 46 +++++- cmake/modules/FindNccl.cmake | 31 ++++ docs/technologies/eigensolvers/cusolvermp.md | 12 ++ src/CMakeLists.txt | 1 + src/cp2k_info.F | 3 + src/fm/cp_fm_cusolver.c | 145 ++++++++++++++----- 6 files changed, 195 insertions(+), 43 deletions(-) create mode 100644 cmake/modules/FindNccl.cmake diff --git a/cmake/modules/FindCuSolverMP.cmake b/cmake/modules/FindCuSolverMP.cmake index 3ea733fc5e..ce6b3c13ff 100644 --- a/cmake/modules/FindCuSolverMP.cmake +++ b/cmake/modules/FindCuSolverMP.cmake @@ -13,24 +13,60 @@ include(FindPackageHandleStandardArgs) include(cp2k_utils) find_package(ucc REQUIRED) -find_package(Cal REQUIRED) +# First, find CuSolverMP library and headers cp2k_set_default_paths(CUSOLVER_MP "CUSOLVER_MP") - cp2k_find_libraries(CUSOLVER_MP "cusolverMp") cp2k_include_dirs(CUSOLVER_MP "cusolverMp.h") +# CuSolverMP 0.7+ uses NCCL for communication, older versions use Cal. We need +# to detect the version to require the correct communication library. +if(CP2K_CUSOLVER_MP_INCLUDE_DIRS) + file(STRINGS "${CP2K_CUSOLVER_MP_INCLUDE_DIRS}/cusolverMp.h" _ver_major_line + REGEX "^#define CUSOLVERMP_VER_MAJOR") + file(STRINGS "${CP2K_CUSOLVER_MP_INCLUDE_DIRS}/cusolverMp.h" _ver_minor_line + REGEX "^#define CUSOLVERMP_VER_MINOR") + + string(REGEX MATCH "[0-9]+" _ver_major "${_ver_major_line}") + string(REGEX MATCH "[0-9]+" _ver_minor "${_ver_minor_line}") + + if(_ver_major STREQUAL "" OR _ver_minor STREQUAL "") + message(FATAL_ERROR "Could not determine CuSolverMP version from header") + endif() + + message(STATUS "Found CuSolverMP version: ${_ver_major}.${_ver_minor}") + + # CuSolverMP 0.7+ uses NCCL, older versions use Cal + if(_ver_major GREATER 0 OR _ver_minor GREATER_EQUAL 7) + find_package(Nccl REQUIRED) + set(CP2K_CUSOLVERMP_USE_NCCL + ON + CACHE BOOL "CuSolverMP uses NCCL for communication") + else() + find_package(Cal REQUIRED) + set(CP2K_CUSOLVERMP_USE_NCCL + OFF + CACHE BOOL "CuSolverMP uses Cal for communication") + endif() +endif() + find_package_handle_standard_args( CuSolverMP DEFAULT_MSG CP2K_CUSOLVER_MP_LINK_LIBRARIES CP2K_CUSOLVER_MP_INCLUDE_DIRS) if(CP2K_CUSOLVER_MP_FOUND AND NOT TARGET cp2k::CUSOLVER_MP::cusolver_mp) add_library(cp2k::CUSOLVER_MP::cusolver_mp INTERFACE IMPORTED) + + if(CP2K_CUSOLVERMP_USE_NCCL) + set(_comm_lib "cp2k::NCCL::nccl") + else() + set(_comm_lib "cp2k::CAL::cal") + endif() + set_target_properties( cp2k::CUSOLVER_MP::cusolver_mp - PROPERTIES - INTERFACE_LINK_LIBRARIES - "${CP2K_CUSOLVER_MP_LINK_LIBRARIES};cp2k::CAL::cal;cp2k::UCC::ucc") + PROPERTIES INTERFACE_LINK_LIBRARIES + "${CP2K_CUSOLVER_MP_LINK_LIBRARIES};${_comm_lib};cp2k::UCC::ucc") set_target_properties( cp2k::CUSOLVER_MP::cusolver_mp PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${CP2K_CUSOLVER_MP_INCLUDE_DIRS}") diff --git a/cmake/modules/FindNccl.cmake b/cmake/modules/FindNccl.cmake new file mode 100644 index 0000000000..9e9cbfc74d --- /dev/null +++ b/cmake/modules/FindNccl.cmake @@ -0,0 +1,31 @@ +#!-------------------------------------------------------------------------------------------------! +#! CP2K: A general program to perform molecular dynamics simulations ! +#! Copyright 2000-2026 CP2K developers group ! +#! ! +#! SPDX-License-Identifier: GPL-2.0-or-later ! +#!-------------------------------------------------------------------------------------------------! + +include(FindPackageHandleStandardArgs) +include(cp2k_utils) + +cp2k_set_default_paths(NCCL "Nccl") + +cp2k_find_libraries(NCCL "nccl") +cp2k_include_dirs(NCCL "nccl.h") + +find_package_handle_standard_args(Nccl DEFAULT_MSG CP2K_NCCL_LINK_LIBRARIES + CP2K_NCCL_INCLUDE_DIRS) + +if(CP2K_NCCL_FOUND AND NOT TARGET cp2k::NCCL::nccl) + add_library(cp2k::NCCL::nccl INTERFACE IMPORTED) + set_target_properties( + cp2k::NCCL::nccl PROPERTIES INTERFACE_LINK_LIBRARIES + "${CP2K_NCCL_LINK_LIBRARIES}") + set_target_properties( + cp2k::NCCL::nccl PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + "${CP2K_NCCL_INCLUDE_DIRS}") +endif() + +mark_as_advanced(CP2K_NCCL_LINK_LIBRARIES) +mark_as_advanced(CP2K_NCCL_INCLUDE_DIRS) +mark_as_advanced(CP2K_NCCL_FOUND) diff --git a/docs/technologies/eigensolvers/cusolvermp.md b/docs/technologies/eigensolvers/cusolvermp.md index 9295cc3ffb..ecd86a4f7c 100644 --- a/docs/technologies/eigensolvers/cusolvermp.md +++ b/docs/technologies/eigensolvers/cusolvermp.md @@ -5,9 +5,16 @@ tools for the solution of dense linear systems and eigenvalue problems. ## Dependencies +[cuSOLVERmp] \< 0.7 + - [CAL]: requires `libcal.\*` in the `$PATH` - [UCC]: requires `libucc.\*` and `libucs.\*` in the `$PATH` +[cuSOLVERmp] >= 0.7 + +- [NCCL]: requires `libnccl.\*` in the `$PATH` +- [UCC]: requires `libucc.\*` and `libucs.\*` in the `$PATH` + ## CMake [cuSOLVERMp] is enabled with the following CMake option: @@ -16,6 +23,11 @@ tools for the solution of dense linear systems and eigenvalue problems. -DCP2K_USE_CUSOLVER_MP=ON ``` +The `FindCuSolverMP.cmake` module tries to automatically deduce the [cuSOLVERmp] version from the +`cusolvermp.h` header file, and enables the `CP2K_CUSOLVERMP_USE_NCCL` CMake option in case it finds +cuSOLVERmp >= 0.7. + [cal]: https://developer.download.nvidia.com/compute/cublasmp/redist/libcal/ [cusolvermp]: https://docs.nvidia.com/cuda/cusolvermp/ +[nccl]: https://developer.nvidia.com/nccl [ucc]: https://github.com/openucx/ucc diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7a910aa362..4dbe3e72d7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1807,6 +1807,7 @@ target_compile_definitions( $<$:__MKL> $<$:__ACCELERATE> $<$:__CUSOLVERMP> + $<$:__CUSOLVERMP_NCCL> $<$:__CRAY_PM_ENERGY> $<$:__CRAY_PM_ACCEL_ENERGY> ${CP2K_GPU_DFLAGS}) diff --git a/src/cp2k_info.F b/src/cp2k_info.F index 02a7e02f68..7dc813cc1f 100644 --- a/src/cp2k_info.F +++ b/src/cp2k_info.F @@ -255,6 +255,9 @@ CONTAINS #if defined(__CUSOLVERMP) flags = TRIM(flags)//" cusolvermp" #endif +#if defined(__CUSOLVERMP_NCCL) + flags = TRIM(flags)//" cusolvermp_nccl" +#endif #if defined(__DLAF) flags = TRIM(flags)//" dlaf" #endif diff --git a/src/fm/cp_fm_cusolver.c b/src/fm/cp_fm_cusolver.c index 01f5dcfe49..1206717ffb 100644 --- a/src/fm/cp_fm_cusolver.c +++ b/src/fm/cp_fm_cusolver.c @@ -9,12 +9,17 @@ #include "../offload/offload_library.h" #include -#include #include #include #include #include +#if defined(__CUSOLVERMP_NCCL) +#include +#else +#include +#endif + /******************************************************************************* * \brief Run given CUDA command and upon failure abort with a nice message. * \author Ole Schuett @@ -29,6 +34,22 @@ } \ } while (0) +#if defined(__CUSOLVERMP_NCCL) +/******************************************************************************* + * \brief Run given NCCL command and upon failure abort with a nice message. + * \author Jiri Vyskocil + ******************************************************************************/ +#define NCCL_CHECK(cmd) \ + do { \ + ncclResult_t status = cmd; \ + if (status != ncclSuccess) { \ + fprintf(stderr, "ERROR: %s %s %d\n", ncclGetErrorString(status), \ + __FILE__, __LINE__); \ + abort(); \ + } \ + } while (0) + +#else /******************************************************************************* * \brief Decode given cal error. * \author Ole Schuett @@ -70,6 +91,44 @@ static char *calGetErrorString(calError_t status) { } \ } while (0) +/******************************************************************************* + * \brief Callback for cal library to initiate an allgather operation. + * \author Ole Schuett + ******************************************************************************/ +static calError_t allgather(void *src_buf, void *recv_buf, size_t size, + void *data, void **req) { + const MPI_Comm comm = *(MPI_Comm *)data; + MPI_Request *request = malloc(sizeof(MPI_Request)); + *req = request; + const int status = MPI_Iallgather(src_buf, size, MPI_BYTE, recv_buf, size, + MPI_BYTE, comm, request); + return (status == MPI_SUCCESS) ? CAL_OK : CAL_ERROR; +} + +/******************************************************************************* + * \brief Callback for cal library to test if a request has completed. + * \author Ole Schuett + ******************************************************************************/ +static calError_t req_test(void *req) { + MPI_Request *request = (MPI_Request *)(req); + int completed; + const int status = MPI_Test(request, &completed, MPI_STATUS_IGNORE); + if (status != MPI_SUCCESS) { + return CAL_ERROR; + } + return completed ? CAL_OK : CAL_ERROR_INPROGRESS; +} + +/******************************************************************************* + * \brief Callback for cal library to free a request. + * \author Ole Schuett + ******************************************************************************/ +static calError_t req_free(void *req) { + free(req); + return CAL_OK; +} +#endif /* __CUSOLVERMP_NCCL */ + /******************************************************************************* * \brief Decode given cusolver error. * \author Ole Schuett @@ -145,43 +204,6 @@ static char *cusolverGetErrorString(cusolverStatus_t status) { } \ } while (0) -/******************************************************************************* - * \brief Callback for cal library to initiate an allgather operation. - * \author Ole Schuett - ******************************************************************************/ -static calError_t allgather(void *src_buf, void *recv_buf, size_t size, - void *data, void **req) { - const MPI_Comm comm = *(MPI_Comm *)data; - MPI_Request *request = malloc(sizeof(MPI_Request)); - *req = request; - const int status = MPI_Iallgather(src_buf, size, MPI_BYTE, recv_buf, size, - MPI_BYTE, comm, request); - return (status == MPI_SUCCESS) ? CAL_OK : CAL_ERROR; -} - -/******************************************************************************* - * \brief Callback for cal library to test if a request has completed. - * \author Ole Schuett - ******************************************************************************/ -static calError_t req_test(void *req) { - MPI_Request *request = (MPI_Request *)(req); - int completed; - const int status = MPI_Test(request, &completed, MPI_STATUS_IGNORE); - if (status != MPI_SUCCESS) { - return CAL_ERROR; - } - return completed ? CAL_OK : CAL_ERROR_INPROGRESS; -} - -/******************************************************************************* - * \brief Callback for cal library to free a request. - * \author Ole Schuett - ******************************************************************************/ -static calError_t req_free(void *req) { - free(req); - return CAL_OK; -} - /******************************************************************************* * \brief Driver routine to diagonalize a matrix with the cuSOLVERMp library. * \author Ole Schuett @@ -199,6 +221,17 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9], MPI_Comm_rank(comm, &rank); MPI_Comm_size(comm, &nranks); +#if defined(__CUSOLVERMP_NCCL) + // Create NCCL communicator. + ncclUniqueId nccl_id; + if (rank == 0) { + NCCL_CHECK(ncclGetUniqueId(&nccl_id)); + } + MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, comm); + + ncclComm_t nccl_comm; + NCCL_CHECK(ncclCommInitRank(&nccl_comm, nranks, nccl_id, rank)); +#else // Create CAL communicator. cal_comm_t cal_comm = NULL; cal_comm_create_params_t params; @@ -210,6 +243,7 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9], params.nranks = nranks; params.local_device = local_device; CAL_CHECK(cal_comm_create(params, &cal_comm)); +#endif // Create various handles. cudaStream_t stream = NULL; @@ -219,9 +253,15 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9], CUSOLVER_CHECK(cusolverMpCreate(&cusolvermp_handle, local_device, stream)); cusolverMpGrid_t grid = NULL; +#if defined(__CUSOLVERMP_NCCL) + CUSOLVER_CHECK(cusolverMpCreateDeviceGrid(cusolvermp_handle, &grid, nccl_comm, + nprow, npcol, + CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); +#else CUSOLVER_CHECK(cusolverMpCreateDeviceGrid(cusolvermp_handle, &grid, cal_comm, nprow, npcol, CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); +#endif const int mb = matrix_desc[4]; const int nb = matrix_desc[5]; const int rsrc = matrix_desc[6]; @@ -281,7 +321,9 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9], // Wait for solver to finish. CUDA_CHECK(cudaStreamSynchronize(stream)); +#if !defined(__CUSOLVERMP_NCCL) CAL_CHECK(cal_stream_sync(cal_comm, stream)); +#endif // Check info. int info = -1; @@ -310,7 +352,11 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9], CUSOLVER_CHECK(cusolverMpDestroyGrid(grid)); CUSOLVER_CHECK(cusolverMpDestroy(cusolvermp_handle)); CUDA_CHECK(cudaStreamDestroy(stream)); +#if defined(__CUSOLVERMP_NCCL) + NCCL_CHECK(ncclCommDestroy(nccl_comm)); +#else CAL_CHECK(cal_comm_destroy(cal_comm)); +#endif // Sync MPI ranks to include load imbalance in total timings. MPI_Barrier(comm); @@ -336,6 +382,17 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm, MPI_Comm_rank(comm, &rank); MPI_Comm_size(comm, &nranks); +#if defined(__CUSOLVERMP_NCCL) + // Create NCCL communicator. + ncclUniqueId nccl_id; + if (rank == 0) { + NCCL_CHECK(ncclGetUniqueId(&nccl_id)); + } + MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, comm); + + ncclComm_t nccl_comm; + NCCL_CHECK(ncclCommInitRank(&nccl_comm, nranks, nccl_id, rank)); +#else // Create CAL communicator cal_comm_t cal_comm = NULL; cal_comm_create_params_t params; @@ -347,6 +404,7 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm, params.nranks = nranks; params.local_device = local_device; CAL_CHECK(cal_comm_create(params, &cal_comm)); +#endif // Create CUDA stream and cuSOLVER handle cudaStream_t stream = NULL; @@ -357,9 +415,15 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm, // Define grid for device computation cusolverMpGrid_t grid = NULL; +#if defined(__CUSOLVERMP_NCCL) + CUSOLVER_CHECK(cusolverMpCreateDeviceGrid(cusolvermp_handle, &grid, nccl_comm, + nprow, npcol, + CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); +#else CUSOLVER_CHECK(cusolverMpCreateDeviceGrid(cusolvermp_handle, &grid, cal_comm, nprow, npcol, CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); +#endif // Matrix descriptors for A, B, and Z const int mb_a = a_matrix_desc[4]; @@ -377,6 +441,7 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm, // Ensure consistency in block sizes and sources assert(mb_a == mb_b && nb_a == nb_b); assert(rsrc_a == rsrc_b && csrc_a == csrc_b); + (void)ldB; // Suppress unused variable warning const int np_a = cusolverMpNUMROC(n, mb_a, myprow, rsrc_a, nprow); const int nq_a = cusolverMpNUMROC(n, nb_a, mypcol, csrc_a, npcol); @@ -469,7 +534,11 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm, CUSOLVER_CHECK(cusolverMpDestroyGrid(grid)); CUSOLVER_CHECK(cusolverMpDestroy(cusolvermp_handle)); CUDA_CHECK(cudaStreamDestroy(stream)); +#if defined(__CUSOLVERMP_NCCL) + NCCL_CHECK(ncclCommDestroy(nccl_comm)); +#else CAL_CHECK(cal_comm_destroy(cal_comm)); +#endif MPI_Barrier(comm); // Synchronize MPI ranks }