cuSOLVERMp: allow scheduler GPU binding (#5227)

Co-authored-by: Thomas D. Kuehne <tkuehne@cp2k.org>
This commit is contained in:
Dynamics of Condensed Matter 2026-05-17 18:29:39 +02:00 committed by GitHub
parent ef9d0990c2
commit 8039b896f5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 40 deletions

View file

@ -34,8 +34,9 @@ diagonalization paths that avoid a CP2K-side Cholesky reduction where the select
supports them. With [cuSOLVERMp], this currently covers real symmetric and complex Hermitian
generalized eigenproblems through `cusolverMpSygvd`.
When [cuSOLVERMp] uses the NCCL backend, CP2K requires at most one local MPI rank per visible GPU.
Use one MPI rank per GPU or reduce the number of local MPI ranks when only one GPU is visible.
When [cuSOLVERMp] uses the NCCL backend, CP2K relies on the launcher or scheduler to bind MPI ranks
to GPUs. For example, assigning one GPU per rank with `CUDA_VISIBLE_DEVICES` is supported. CP2K only
checks that the CUDA device selected for each MPI rank is visible to that rank.
[cal]: https://developer.download.nvidia.com/compute/cublasmp/redist/libcal/
[cusolvermp]: https://docs.nvidia.com/cuda/cusolvermp/

View file

@ -54,57 +54,31 @@
} while (0)
/*******************************************************************************
* \brief Check that local MPI ranks can be mapped one-to-one to GPUs.
* \brief Check that the selected CUDA device is visible to this MPI rank.
******************************************************************************/
static void check_nccl_rank_device_mapping(MPI_Comm comm,
const int local_device) {
static void check_nccl_device_visibility(MPI_Comm comm,
const int local_device) {
int device_count = 0;
CUDA_CHECK(cudaGetDeviceCount(&device_count));
MPI_Comm node_comm;
MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &node_comm);
int rank, local_rank, local_nranks;
int rank;
MPI_Comm_rank(comm, &rank);
MPI_Comm_rank(node_comm, &local_rank);
MPI_Comm_size(node_comm, &local_nranks);
if (device_count <= 0) {
if (local_rank == 0) {
fprintf(stderr,
"ERROR: cuSOLVERMp with NCCL requires at least one visible GPU "
"on every node. No CUDA device is visible on the node containing "
"MPI rank %d.\n",
rank);
}
MPI_Comm_free(&node_comm);
fprintf(stderr,
"ERROR: cuSOLVERMp with NCCL requires at least one visible GPU. "
"No CUDA device is visible to MPI rank %d.\n",
rank);
MPI_Abort(comm, EXIT_FAILURE);
}
if (local_device < 0 || local_device >= device_count) {
fprintf(stderr,
"ERROR: cuSOLVERMp with NCCL selected invalid CUDA device %d on "
"MPI rank %d; %d CUDA device(s) are visible on this node.\n",
"MPI rank %d; %d CUDA device(s) are visible to this rank.\n",
local_device, rank, device_count);
MPI_Comm_free(&node_comm);
MPI_Abort(comm, EXIT_FAILURE);
}
if (local_nranks > device_count) {
if (local_rank == 0) {
fprintf(stderr,
"ERROR: cuSOLVERMp with NCCL in CP2K requires at most one local "
"MPI rank per visible GPU. This node has %d local MPI ranks and "
"%d visible GPU(s). Use fewer MPI ranks per node, expose more "
"GPUs with CUDA_VISIBLE_DEVICES, or choose another "
"PREFERRED_DIAG_LIBRARY.\n",
local_nranks, device_count);
}
MPI_Comm_free(&node_comm);
MPI_Abort(comm, EXIT_FAILURE);
}
MPI_Comm_free(&node_comm);
}
#else
@ -279,7 +253,7 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9],
MPI_Comm_size(comm, &nranks);
#if defined(__CUSOLVERMP_NCCL)
check_nccl_rank_device_mapping(comm, local_device);
check_nccl_device_visibility(comm, local_device);
#endif
offload_activate_chosen_device();
@ -444,7 +418,7 @@ void cp_fm_diag_cusolver_sygvd(const int fortran_comm,
MPI_Comm_size(comm, &nranks);
#if defined(__CUSOLVERMP_NCCL)
check_nccl_rank_device_mapping(comm, local_device);
check_nccl_device_visibility(comm, local_device);
#endif
offload_activate_chosen_device();
@ -646,7 +620,7 @@ void cp_cfm_diag_cusolver_hegvd(
MPI_Comm_size(comm, &nranks);
#if defined(__CUSOLVERMP_NCCL)
check_nccl_rank_device_mapping(comm, local_device);
check_nccl_device_visibility(comm, local_device);
#endif
offload_activate_chosen_device();