From 1018edadf9bbbf513ed4d6fcb8167c994bd1e1fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ji=C5=99=C3=AD=20Vysko=C4=8Dil?= Date: Wed, 6 Nov 2024 12:21:57 +0100 Subject: [PATCH] Use cusolverMpSygvd for generalized eigensolver --- .gitignore | 4 + src/fm/cp_fm_cusolver.c | 157 ++++++++++++++++++++++++++++++++++++ src/fm/cp_fm_cusolver_api.F | 81 ++++++++++++++++++- src/fm/cp_fm_diag.F | 35 ++++---- 4 files changed, 262 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index 77ed28e608..8a78ba763d 100644 --- a/.gitignore +++ b/.gitignore @@ -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-* diff --git a/src/fm/cp_fm_cusolver.c b/src/fm/cp_fm_cusolver.c index 97e439914f..80a69f101d 100644 --- a/src/fm/cp_fm_cusolver.c +++ b/src/fm/cp_fm_cusolver.c @@ -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 diff --git a/src/fm/cp_fm_cusolver_api.F b/src/fm/cp_fm_cusolver_api.F index 5efac02d41..cfaa3207c9 100644 --- a/src/fm/cp_fm_cusolver_api.F +++ b/src/fm/cp_fm_cusolver_api.F @@ -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 diff --git a/src/fm/cp_fm_diag.F b/src/fm/cp_fm_diag.F index 504e369114..0490f227ac 100644 --- a/src/fm/cp_fm_diag.F +++ b/src/fm/cp_fm_diag.F @@ -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)