mirror of
https://github.com/cp2k/cp2k.git
synced 2026-07-28 14:15:19 -04:00
Use cusolverMpSygvd for generalized eigensolver
This commit is contained in:
Jiří Vyskočil
Ole Schütt
parent
e5c64e915b
commit
1018edadf9
4 changed files with 262 additions and 15 deletions
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -371,3 +371,7 @@ tags
|
|||
# End of https://www.gitignore.io/api/python,fortran,c,c++,cuda,emacs,vim,cmake
|
||||
|
||||
!tools/fedora/cp2k.spec
|
||||
|
||||
# CLion directories
|
||||
/.idea/
|
||||
/cmake-build-*
|
||||
|
|
|
|||
|
|
@ -316,6 +316,163 @@ void cp_fm_diag_cusolver(const int fortran_comm, const int matrix_desc[9],
|
|||
MPI_Barrier(comm);
|
||||
}
|
||||
|
||||
/*******************************************************************************
|
||||
* \brief Driver routine to solve A*x = lambda*B*x with cuSOLVERMp sygvd.
|
||||
* \author Jiri Vyskocil
|
||||
******************************************************************************/
|
||||
void cp_fm_diag_cusolver_sygvd(const int fortran_comm,
|
||||
const int a_matrix_desc[9],
|
||||
const int b_matrix_desc[9], const int nprow,
|
||||
const int npcol, const int myprow,
|
||||
const int mypcol, const int n,
|
||||
const double *aMatrix, const double *bMatrix,
|
||||
double *eigenvectors, double *eigenvalues) {
|
||||
|
||||
offload_activate_chosen_device();
|
||||
const int local_device = offload_get_chosen_device();
|
||||
|
||||
MPI_Comm comm = MPI_Comm_f2c(fortran_comm);
|
||||
int rank, nranks;
|
||||
MPI_Comm_rank(comm, &rank);
|
||||
MPI_Comm_size(comm, &nranks);
|
||||
|
||||
// Create CAL communicator
|
||||
cal_comm_t cal_comm = NULL;
|
||||
cal_comm_create_params_t params;
|
||||
params.allgather = &allgather;
|
||||
params.req_test = &req_test;
|
||||
params.req_free = &req_free;
|
||||
params.data = &comm;
|
||||
params.rank = rank;
|
||||
params.nranks = nranks;
|
||||
params.local_device = local_device;
|
||||
CAL_CHECK(cal_comm_create(params, &cal_comm));
|
||||
|
||||
// Create CUDA stream and cuSOLVER handle
|
||||
cudaStream_t stream = NULL;
|
||||
CUDA_CHECK(cudaStreamCreate(&stream));
|
||||
|
||||
cusolverMpHandle_t cusolvermp_handle = NULL;
|
||||
CUSOLVER_CHECK(cusolverMpCreate(&cusolvermp_handle, local_device, stream));
|
||||
|
||||
// Define grid for device computation
|
||||
cusolverMpGrid_t grid = NULL;
|
||||
CUSOLVER_CHECK(cusolverMpCreateDeviceGrid(cusolvermp_handle, &grid, cal_comm,
|
||||
nprow, npcol,
|
||||
CUSOLVERMP_GRID_MAPPING_ROW_MAJOR));
|
||||
|
||||
// Matrix descriptors for A, B, and Z
|
||||
const int mb_a = a_matrix_desc[4];
|
||||
const int nb_a = a_matrix_desc[5];
|
||||
const int rsrc_a = a_matrix_desc[6];
|
||||
const int csrc_a = a_matrix_desc[7];
|
||||
const int ldA = a_matrix_desc[8];
|
||||
|
||||
const int mb_b = b_matrix_desc[4];
|
||||
const int nb_b = b_matrix_desc[5];
|
||||
const int rsrc_b = b_matrix_desc[6];
|
||||
const int csrc_b = b_matrix_desc[7];
|
||||
const int ldB = b_matrix_desc[8];
|
||||
|
||||
// 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);
|
||||
|
||||
const int np_a = cusolverMpNUMROC(n, mb_a, myprow, rsrc_a, nprow);
|
||||
const int nq_a = cusolverMpNUMROC(n, nb_a, mypcol, csrc_a, npcol);
|
||||
|
||||
const cublasFillMode_t uplo = CUBLAS_FILL_MODE_LOWER;
|
||||
const cusolverEigType_t itype = CUSOLVER_EIG_TYPE_1;
|
||||
const cusolverEigMode_t jobz = CUSOLVER_EIG_MODE_VECTOR;
|
||||
const cudaDataType_t data_type = CUDA_R_64F;
|
||||
|
||||
// Create matrix descriptors
|
||||
cusolverMpMatrixDescriptor_t descrA = NULL;
|
||||
cusolverMpMatrixDescriptor_t descrB = NULL;
|
||||
cusolverMpMatrixDescriptor_t descrZ = NULL;
|
||||
|
||||
CUSOLVER_CHECK(cusolverMpCreateMatrixDesc(&descrA, grid, data_type, n, n,
|
||||
mb_a, nb_a, rsrc_a, csrc_a, np_a));
|
||||
CUSOLVER_CHECK(cusolverMpCreateMatrixDesc(&descrB, grid, data_type, n, n,
|
||||
mb_b, nb_b, rsrc_b, csrc_b, np_a));
|
||||
CUSOLVER_CHECK(cusolverMpCreateMatrixDesc(&descrZ, grid, data_type, n, n,
|
||||
mb_a, nb_a, rsrc_a, csrc_a, np_a));
|
||||
|
||||
// Allocate device memory for matrices
|
||||
double *dev_A = NULL, *dev_B = NULL;
|
||||
size_t matrix_local_size = ldA * nq_a * sizeof(double);
|
||||
CUDA_CHECK(cudaMalloc((void **)&dev_A, matrix_local_size));
|
||||
CUDA_CHECK(cudaMalloc((void **)&dev_B, matrix_local_size));
|
||||
|
||||
// Copy matrices from host to device
|
||||
CUDA_CHECK(cudaMemcpyAsync(dev_A, aMatrix, matrix_local_size,
|
||||
cudaMemcpyHostToDevice, stream));
|
||||
CUDA_CHECK(cudaMemcpyAsync(dev_B, bMatrix, matrix_local_size,
|
||||
cudaMemcpyHostToDevice, stream));
|
||||
|
||||
// Allocate device memory for eigenvalues and eigenvectors
|
||||
double *dev_Z = NULL, *eigenvalues_dev = NULL;
|
||||
CUDA_CHECK(cudaMalloc((void **)&dev_Z, matrix_local_size));
|
||||
CUDA_CHECK(cudaMalloc((void **)&eigenvalues_dev, n * sizeof(double)));
|
||||
|
||||
// Allocate workspace
|
||||
size_t work_dev_size = 0, work_host_size = 0;
|
||||
CUSOLVER_CHECK(cusolverMpSygvd_bufferSize(
|
||||
cusolvermp_handle, itype, jobz, uplo, n, 1, 1, descrA, 1, 1, descrB, 1, 1,
|
||||
descrZ, data_type, &work_dev_size, &work_host_size));
|
||||
|
||||
void *work_dev = NULL, *work_host = NULL;
|
||||
CUDA_CHECK(cudaMalloc(&work_dev, work_dev_size));
|
||||
CUDA_CHECK(cudaMallocHost(&work_host, work_host_size));
|
||||
|
||||
// Allocate device memory for info
|
||||
int *info_dev = NULL;
|
||||
CUDA_CHECK(cudaMalloc((void **)&info_dev, sizeof(int)));
|
||||
|
||||
// Call cusolverMpSygvd
|
||||
CUSOLVER_CHECK(cusolverMpSygvd(
|
||||
cusolvermp_handle, itype, jobz, uplo, n, dev_A, 1, 1, descrA, dev_B, 1, 1,
|
||||
descrB, eigenvalues_dev, dev_Z, 1, 1, descrZ, data_type, work_dev,
|
||||
work_dev_size, work_host, work_host_size, info_dev));
|
||||
|
||||
// Wait for computation to finish
|
||||
CUDA_CHECK(cudaStreamSynchronize(stream));
|
||||
|
||||
// Check info
|
||||
int info;
|
||||
CUDA_CHECK(cudaMemcpy(&info, info_dev, sizeof(int), cudaMemcpyDeviceToHost));
|
||||
if (info != 0) {
|
||||
fprintf(stderr, "ERROR: cusolverMpSygvd failed with info = %d\n", info);
|
||||
abort();
|
||||
}
|
||||
|
||||
// Copy results back to host
|
||||
CUDA_CHECK(cudaMemcpyAsync(eigenvectors, dev_Z, matrix_local_size,
|
||||
cudaMemcpyDeviceToHost, stream));
|
||||
CUDA_CHECK(cudaMemcpyAsync(eigenvalues, eigenvalues_dev, n * sizeof(double),
|
||||
cudaMemcpyDeviceToHost, stream));
|
||||
|
||||
// Wait for copy to finish
|
||||
CUDA_CHECK(cudaStreamSynchronize(stream));
|
||||
|
||||
// Clean up resources
|
||||
CUDA_CHECK(cudaFree(dev_A));
|
||||
CUDA_CHECK(cudaFree(dev_B));
|
||||
CUDA_CHECK(cudaFree(dev_Z));
|
||||
CUDA_CHECK(cudaFree(eigenvalues_dev));
|
||||
CUDA_CHECK(cudaFree(info_dev));
|
||||
CUDA_CHECK(cudaFree(work_dev));
|
||||
CUDA_CHECK(cudaFreeHost(work_host));
|
||||
CUSOLVER_CHECK(cusolverMpDestroyMatrixDesc(descrA));
|
||||
CUSOLVER_CHECK(cusolverMpDestroyMatrixDesc(descrB));
|
||||
CUSOLVER_CHECK(cusolverMpDestroyMatrixDesc(descrZ));
|
||||
CUSOLVER_CHECK(cusolverMpDestroyGrid(grid));
|
||||
CUSOLVER_CHECK(cusolverMpDestroy(cusolvermp_handle));
|
||||
CUDA_CHECK(cudaStreamDestroy(stream));
|
||||
CAL_CHECK(cal_comm_destroy(cal_comm));
|
||||
|
||||
MPI_Barrier(comm); // Synchronize MPI ranks
|
||||
}
|
||||
#endif
|
||||
|
||||
// EOF
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ MODULE cp_fm_cusolver_api
|
|||
PRIVATE
|
||||
|
||||
PUBLIC :: cp_fm_diag_cusolver
|
||||
PUBLIC :: cp_fm_general_cusolver
|
||||
|
||||
CONTAINS
|
||||
|
||||
|
|
@ -98,5 +99,83 @@ CONTAINS
|
|||
CALL timestop(handle)
|
||||
END SUBROUTINE cp_fm_diag_cusolver
|
||||
|
||||
END MODULE cp_fm_cusolver_api
|
||||
! **************************************************************************************************
|
||||
!> \brief Driver routine to solve generalized eigenvalue problem A*x = lambda*B*x with cuSOLVERMp.
|
||||
!> \param aMatrix the first matrix for the generalized eigenvalue problem
|
||||
!> \param bMatrix the second matrix for the generalized eigenvalue problem
|
||||
!> \param eigenvectors eigenvectors of the input matrix
|
||||
!> \param eigenvalues eigenvalues of the input matrix
|
||||
! **************************************************************************************************
|
||||
SUBROUTINE cp_fm_general_cusolver(aMatrix, bMatrix, eigenvectors, eigenvalues)
|
||||
USE ISO_C_BINDING, ONLY: C_INT, C_DOUBLE
|
||||
TYPE(cp_fm_type), INTENT(IN) :: aMatrix, bMatrix, eigenvectors
|
||||
REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: eigenvalues
|
||||
|
||||
CHARACTER(len=*), PARAMETER :: routineN = 'cp_fm_general_cusolver'
|
||||
|
||||
INTEGER(kind=C_INT) :: handle, n, nmo
|
||||
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: eigenvalues_buffer
|
||||
TYPE(cp_blacs_env_type), POINTER :: context
|
||||
INTERFACE
|
||||
SUBROUTINE cp_fm_general_cusolver_c(fortran_comm, a_matrix_desc, b_matrix_desc, &
|
||||
nprow, npcol, myprow, mypcol, &
|
||||
n, aMatrix, bMatrix, eigenvectors, eigenvalues) &
|
||||
BIND(C, name="cp_fm_diag_cusolver_sygvd")
|
||||
IMPORT :: C_INT, C_DOUBLE
|
||||
INTEGER(kind=C_INT), VALUE :: fortran_comm
|
||||
INTEGER(kind=C_INT), DIMENSION(*) :: a_matrix_desc, b_matrix_desc
|
||||
INTEGER(kind=C_INT), VALUE :: nprow
|
||||
INTEGER(kind=C_INT), VALUE :: npcol
|
||||
INTEGER(kind=C_INT), VALUE :: myprow
|
||||
INTEGER(kind=C_INT), VALUE :: mypcol
|
||||
INTEGER(kind=C_INT), VALUE :: n
|
||||
REAL(kind=C_DOUBLE), DIMENSION(*) :: aMatrix
|
||||
REAL(kind=C_DOUBLE), DIMENSION(*) :: bMatrix
|
||||
REAL(kind=C_DOUBLE), DIMENSION(*) :: eigenvectors
|
||||
REAL(kind=C_DOUBLE), DIMENSION(*) :: eigenvalues
|
||||
END SUBROUTINE cp_fm_general_cusolver_c
|
||||
END INTERFACE
|
||||
|
||||
CALL timeset(routineN, handle)
|
||||
|
||||
#if defined(__CUSOLVERMP)
|
||||
n = INT(aMatrix%matrix_struct%nrow_global, C_INT)
|
||||
context => aMatrix%matrix_struct%context
|
||||
|
||||
! Allocate eigenvalues_buffer
|
||||
ALLOCATE (eigenvalues_buffer(n))
|
||||
|
||||
CALL cp_fm_general_cusolver_c( &
|
||||
fortran_comm=INT(aMatrix%matrix_struct%para_env%get_handle(), C_INT), &
|
||||
a_matrix_desc=INT(aMatrix%matrix_struct%descriptor, C_INT), &
|
||||
b_matrix_desc=INT(bMatrix%matrix_struct%descriptor, C_INT), &
|
||||
nprow=INT(context%num_pe(1), C_INT), &
|
||||
npcol=INT(context%num_pe(2), C_INT), &
|
||||
myprow=INT(context%mepos(1), C_INT), &
|
||||
mypcol=INT(context%mepos(2), C_INT), &
|
||||
n=n, &
|
||||
aMatrix=aMatrix%local_data, &
|
||||
bMatrix=bMatrix%local_data, &
|
||||
eigenvectors=eigenvectors%local_data, &
|
||||
eigenvalues=eigenvalues_buffer)
|
||||
|
||||
nmo = SIZE(eigenvalues)
|
||||
eigenvalues(1:nmo) = eigenvalues_buffer(1:nmo)
|
||||
|
||||
DEALLOCATE (eigenvalues_buffer)
|
||||
#else
|
||||
MARK_USED(aMatrix)
|
||||
MARK_USED(bMatrix)
|
||||
MARK_USED(eigenvectors)
|
||||
eigenvalues = 0.0_dp
|
||||
MARK_USED(n)
|
||||
MARK_USED(nmo)
|
||||
MARK_USED(eigenvalues_buffer)
|
||||
MARK_USED(context)
|
||||
CPABORT("CP2K compiled without the cuSOLVERMp library.")
|
||||
#endif
|
||||
|
||||
CALL timestop(handle)
|
||||
END SUBROUTINE cp_fm_general_cusolver
|
||||
|
||||
END MODULE cp_fm_cusolver_api
|
||||
|
|
|
|||
|
|
@ -34,7 +34,8 @@ MODULE cp_fm_diag
|
|||
set_elpa_kernel, &
|
||||
set_elpa_print, &
|
||||
set_elpa_qr
|
||||
USE cp_fm_cusolver_api, ONLY: cp_fm_diag_cusolver
|
||||
USE cp_fm_cusolver_api, ONLY: cp_fm_diag_cusolver, &
|
||||
cp_fm_general_cusolver
|
||||
#if defined(__DLAF)
|
||||
USE cp_fm_dlaf_api, ONLY: cp_fm_diag_dlaf
|
||||
#endif
|
||||
|
|
@ -1413,19 +1414,25 @@ CONTAINS
|
|||
|
||||
CALL cp_fm_get_info(amatrix, nrow_global=nao)
|
||||
nmo = SIZE(eigenvalues)
|
||||
! Cholesky decompose S=U(T)U
|
||||
CALL cp_fm_cholesky_decompose(bmatrix)
|
||||
! Invert to get U^(-1)
|
||||
CALL cp_fm_triangular_invert(bmatrix)
|
||||
! Reduce to get U^(-T) * H * U^(-1)
|
||||
CALL cp_fm_triangular_multiply(bmatrix, amatrix, side="R")
|
||||
CALL cp_fm_triangular_multiply(bmatrix, amatrix, transpose_tr=.TRUE.)
|
||||
! Diagonalize
|
||||
CALL choose_eigv_solver(matrix=amatrix, eigenvectors=work, &
|
||||
eigenvalues=eigenvalues)
|
||||
! Restore vectors C = U^(-1) * C*
|
||||
CALL cp_fm_triangular_multiply(bmatrix, work)
|
||||
CALL cp_fm_to_fm(work, eigenvectors, nmo)
|
||||
|
||||
IF (diag_type == FM_DIAG_TYPE_CUSOLVER .AND. nao >= 64) THEN
|
||||
! Use cuSolverMP generalized eigenvalue solver for large matrices
|
||||
CALL cp_fm_general_cusolver(amatrix, bmatrix, eigenvectors, eigenvalues)
|
||||
ELSE
|
||||
! Cholesky decompose S=U(T)U
|
||||
CALL cp_fm_cholesky_decompose(bmatrix)
|
||||
! Invert to get U^(-1)
|
||||
CALL cp_fm_triangular_invert(bmatrix)
|
||||
! Reduce to get U^(-T) * H * U^(-1)
|
||||
CALL cp_fm_triangular_multiply(bmatrix, amatrix, side="R")
|
||||
CALL cp_fm_triangular_multiply(bmatrix, amatrix, transpose_tr=.TRUE.)
|
||||
! Diagonalize
|
||||
CALL choose_eigv_solver(matrix=amatrix, eigenvectors=work, &
|
||||
eigenvalues=eigenvalues)
|
||||
! Restore vectors C = U^(-1) * C*
|
||||
CALL cp_fm_triangular_multiply(bmatrix, work)
|
||||
CALL cp_fm_to_fm(work, eigenvectors, nmo)
|
||||
END IF
|
||||
|
||||
CALL timestop(handle)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue