mirror of
https://github.com/cp2k/cp2k.git
synced 2026-07-21 14:35:15 -04:00
Add NCCL support for CuSolverMP
This commit is contained in:
parent
e84522140f
commit
e7b8a776ce
6 changed files with 195 additions and 43 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
31
cmake/modules/FindNccl.cmake
Normal file
31
cmake/modules/FindNccl.cmake
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue