Add NCCL support for CuSolverMP

This commit is contained in:
Jiří Vyskočil 2026-02-02 12:29:35 +01:00 committed by Ole Schütt
parent e84522140f
commit e7b8a776ce
6 changed files with 195 additions and 43 deletions

View file

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

View file

@ -0,0 +1,31 @@
#!-------------------------------------------------------------------------------------------------!
#! CP2K: A general program to perform molecular dynamics simulations !
#! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
#! !
#! 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)

View file

@ -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

View file

@ -1807,6 +1807,7 @@ target_compile_definitions(
$<$<STREQUAL:"${CP2K_BLAS_VENDOR}","MKL">:__MKL>
$<$<STREQUAL:"${CP2K_BLAS_VENDOR}","Apple">:__ACCELERATE>
$<$<BOOL:${CP2K_USE_CUSOLVER_MP}>:__CUSOLVERMP>
$<$<BOOL:${CP2K_CUSOLVERMP_USE_NCCL}>:__CUSOLVERMP_NCCL>
$<$<BOOL:${CP2K_USE_CRAY_PM_ENERGY}>:__CRAY_PM_ENERGY>
$<$<BOOL:${CP2K_USE_CRAY_PM_ACCEL_ENERGY}>:__CRAY_PM_ACCEL_ENERGY>
${CP2K_GPU_DFLAGS})

View file

@ -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

View file

@ -9,12 +9,17 @@
#include "../offload/offload_library.h"
#include <assert.h>
#include <cal.h>
#include <cuda_runtime.h>
#include <cusolverMp.h>
#include <mpi.h>
#include <stdlib.h>
#if defined(__CUSOLVERMP_NCCL)
#include <nccl.h>
#else
#include <cal.h>
#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
}