K1 on GPU

This commit is contained in:
Anthony Scemama 2023-08-03 01:21:16 +02:00
parent 2763954c19
commit eec4af61e2
6 changed files with 155 additions and 140 deletions

View File

@ -120,7 +120,7 @@ subroutine run_ccsd_space_orb
gpu_data = gpu_init(nO, nV, cholesky_mo_num, &
cc_space_v_oo_chol, cc_space_v_ov_chol, cc_space_v_vo_chol, cc_space_v_vv_chol, &
cc_space_v_oooo, cc_space_v_vooo, cc_space_v_oovv, cc_space_v_vvoo, &
cc_space_v_oovo, cc_space_v_ovvo, cc_space_v_ovoo, cc_space_f_vo)
cc_space_v_oovo, cc_space_v_ovvo, cc_space_v_ovov, cc_space_v_ovoo, cc_space_f_vo)
do while (not_converged)

View File

@ -492,11 +492,12 @@ subroutine compute_r2_space_chol(nO,nV,t1,t2,tau,H_oo,H_vv,H_vo,r2,max_r2,gpu_da
double precision, allocatable :: J1(:,:,:,:)
allocate(J1(nO,nV,nV,nO))
J1 = 0.d0
double precision, allocatable :: K1(:,:,:,:)
allocate(K1(nO,nV,nO,nV))
call compute_r2_space_chol_gpu(nO,nV,cholesky_mo_num,gpu_data,t1,t2,tau, &
H_vv, g_occ, J1, r2)
H_vv, g_occ, J1, K1, r2)
!---
double precision, allocatable :: X_oovv(:,:,:,:)
@ -506,10 +507,6 @@ subroutine compute_r2_space_chol(nO,nV,t1,t2,tau,H_oo,H_vv,H_vo,r2,max_r2,gpu_da
double precision, allocatable :: X_vovo(:,:,:,:), Y_oovo(:,:,:,:)
double precision, allocatable :: K1(:,:,:,:)
allocate(K1(nO,nV,nO,nV))
call compute_K1_chol(nO,nV,t1,t2,cc_space_v_ovoo,cc_space_v_vvoo, &
cc_space_v_ovov,K1)
allocate(X_ovvo(nO,nV,nV,nO))
!$omp parallel &
@ -721,106 +718,3 @@ subroutine compute_r2_space_chol(nO,nV,t1,t2,tau,H_oo,H_vv,H_vo,r2,max_r2,gpu_da
end
! K1
subroutine compute_K1_chol(nO,nV,t1,t2,v_ovoo,v_vvoo,v_ovov,K1)
implicit none
integer, intent(in) :: nO,nV
double precision, intent(in) :: t1(nO, nV)
double precision, intent(in) :: t2(nO, nO, nV, nV)
double precision, intent(in) :: v_vvoo(nV,nV,nO,nO), v_ovov(nO,nV,nO,nV)
double precision, intent(in) :: v_ovoo(nO,nV,nO,nO)
double precision, intent(out) :: K1(nO, nV, nO, nV)
double precision, allocatable :: X(:,:,:,:), Y(:,:,:,:), Z(:,:,:,:)
integer :: a,tmp_a,b,k,l,c,d,tmp_c,tmp_d,i,j,u,v, beta, gam
allocate(X(nV,nO,nV,nO),Y(nO,nV,nV,nO),Z(nO,nV,nV,nO))
!$omp parallel &
!$omp shared(nO,nV,K1,X,Y,v_vvoo,v_ovov,t1,t2) &
!$omp private(i,beta,a,u,j,b) &
!$omp default(none)
!$omp do
do beta = 1, nV
do i = 1, nO
do a = 1, nV
do u = 1, nO
K1(u,a,i,beta) = v_ovov(u,a,i,beta)
enddo
enddo
enddo
enddo
!$omp end do nowait
do i = 1, nO
!$omp do
do a = 1, nV
do j = 1, nO
do b = 1, nV
X(b,j,a,i) = - v_vvoo(b,a,i,j)
enddo
enddo
enddo
!$omp end do nowait
enddo
do j = 1, nO
!$omp do
do b = 1, nV
do beta = 1, nV
do u = 1, nO
Y(u,beta,b,j) = 0.5d0 * t2(u,j,b,beta) + t1(u,b) * t1(j,beta)
enddo
enddo
enddo
!$omp end do
enddo
!$omp end parallel
call dgemm('N','N',nO*nV*nO,nV,nO, &
-1d0, v_ovoo, size(v_ovoo,1) * size(v_ovoo,2) * size(v_ovoo,3), &
t1 , size(t1,1), &
1d0, K1 , size(K1,1) * size(K1,2) * size(K1,3))
double precision, allocatable :: K1tmp(:,:,:,:), t1v(:,:,:)
allocate(K1tmp(nO,nO,nV,nV), t1v(cholesky_mo_num,nO,nO))
call dgemm('N','T', cholesky_mo_num*nO, nO, nV, 1.d0, &
cc_space_v_ov_chol, cholesky_mo_num*nO, t1, nO, 0.d0, &
t1v, cholesky_mo_num*nO)
call dgemm('T','N', nO*nO, nV*nV, cholesky_mo_num, 1.d0, &
t1v, cholesky_mo_num, cc_space_v_vv_chol, cholesky_mo_num, 0.d0, &
K1tmp, nO*nO)
deallocate(t1v)
! Y(u,beta,b,j) * X(b,j,a,i) = Z(u,beta,a,i)
call dgemm('N','N',nV*nO,nO*nV,nV*nO, &
1d0, Y, size(Y,1) * size(Y,2), &
X, size(X,1) * size(X,2), &
0d0, Z, size(Z,1) * size(Z,2))
!$omp parallel &
!$omp shared(nO,nV,K1,Z,K1tmp) &
!$omp private(i,beta,a,u) &
!$omp default(none)
!$omp do
do beta = 1, nV
do i = 1, nO
do a = 1, nV
do u = 1, nO
K1(u,a,i,beta) = K1(u,a,i,beta) + K1tmp(u,i,a,beta) + Z(u,beta,a,i)
enddo
enddo
enddo
enddo
!$omp end do
!$omp end parallel
deallocate(K1tmp,X,Y,Z)
end

View File

@ -17,6 +17,7 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
double* H_vv,
double* g_occ,
double* J1,
double* K1,
double* r2)
{
@ -50,6 +51,7 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
double* d_cc_space_v_vvoo = data[igpu].cc_space_v_vvoo;
double* d_cc_space_v_oovo = data[igpu].cc_space_v_oovo;
double* d_cc_space_v_ovvo = data[igpu].cc_space_v_ovvo;
double* d_cc_space_v_ovov = data[igpu].cc_space_v_ovov;
double* d_cc_space_v_ovoo = data[igpu].cc_space_v_ovoo;
double* d_cc_space_f_vo = data[igpu].cc_space_f_vo;
@ -76,7 +78,6 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cudaMalloc((void **)&d_t2, nO*nO*nV*nV * sizeof(double));
cublasSetMatrix(nO*nO, nV*nV, sizeof(double), t2, lda, d_t2, lda);
#pragma omp sections
{
@ -533,19 +534,12 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nV ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
double* d_X_vovo;
cudaMalloc((void **)&d_X_vovo, nV*nO*nV*nO * sizeof(double));
alpha = 0.0;
beta = 1.0;
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
for (int i=0 ; i<nO ; ++i) {
for (int gam=0 ; gam<nV ; ++gam) {
cublasSetStream(handle, stream[gam]);
@ -624,7 +618,6 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
#pragma omp section
{
double* d_J1;
cudaMalloc((void **)&d_J1, nO*nV*nV*nO * sizeof(double));
@ -792,17 +785,10 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nV ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
double* d_Y_vovo;
cudaMalloc((void **)&d_Y_vovo, nV*nO*nV*nO * sizeof(double));
for (int i=0 ; i<nO ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha = 1.0;
beta = -0.5;
for (int j=0 ; j<nO ; ++j) {
@ -814,17 +800,10 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_T, nV, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nO ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
double* d_X_ovvo;
cudaMalloc((void **)&d_X_ovvo, nO*nV*nV*nO * sizeof(double));
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha = 1.0;
beta = 0.0;
for (int j=0 ; j<nO ; ++j) {
@ -878,13 +857,144 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
}
// #pragma omp section
// {
// }
#pragma omp section
{
double* d_K1;
cudaMalloc((void **)&d_K1, nO*nV*nO*nV * sizeof(double));
alpha = 1.0;
beta = 0.0;
A = d_cc_space_v_ovov; lda = nO*nV;
B = d_cc_space_v_ovov; ldb = nO*nV;
C = d_K1; ldc = nO*nV;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO*nV, nO*nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
alpha = -1.0;
beta = 1.0;
m=nO*nV*nO; n=nV; k=nO;
A=d_cc_space_v_ovoo; lda=nO*nV*nO;
B=d_t1; ldb=nO;
C=d_K1; ldc=nO*nV*nO;
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
double* d_X;
cudaMalloc((void **)&d_X, nV*nO*nV*nO * sizeof(double));
double* d_Y;
cudaMalloc((void **)&d_Y, nO*nV*nV*nO * sizeof(double));
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha =-1.0;
beta = 0.0;
for (int j=0 ; j<nO ; ++j) {
for (int i=0 ; i<nO ; ++i) {
cublasSetStream(handle, stream[i]);
A = &(d_cc_space_v_vvoo[nV*nV*(i+nO*j)]); lda = nV;
B = &(d_cc_space_v_vvoo[nV*nV*(i+nO*j)]); ldb = nV;
C = &(d_X[nV*(j+nO*nV*i)]); ldc = nV*nO;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nV, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
alpha = 0.5;
for (int j=0 ; j<nO ; ++j) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
beta = t1[j+b*nO];
A = &(d_t2[nO*(j+nO*nV*b)]); lda = nO*nO;
B = d_t1; ldb = nO;
C = &(d_Y[nO*(b+nV*nV*j)]); ldc = nO*nV;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nV ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
double* d_Z;
cudaMalloc((void **)&d_Z, nO*nV*nV*nO * sizeof(double));
alpha = 1.0;
beta = 0.0;
m=nV*nO; n=nO*nV; k=nV*nO;
A=d_Y; lda=nO*nV;
B=d_X; ldb=nV*nO;
C=d_Z; ldc=nO*nV;
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
cudaFree(d_X);
cudaFree(d_Y);
double* d_t1v;
cudaMalloc((void **)&d_t1v, cholesky_mo_num*nO*nO * sizeof(double));
alpha = 1.0;
beta = 0.0;
m=cholesky_mo_num*nO; n=nO; k=nV;
A=d_cc_space_v_ov_chol; lda=cholesky_mo_num*nO;
B=d_t1; ldb=nO;
C=d_t1v; ldc=cholesky_mo_num*nO;
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
double* d_K1tmp;
cudaMalloc((void **)&d_K1tmp, nO*nO*nV*nV * sizeof(double));
alpha = 1.0;
beta = 0.0;
m=nO*nO; n=nV*nV; k=cholesky_mo_num;
A=d_t1v; lda=cholesky_mo_num;
B=d_cc_space_v_vv_chol; ldb=cholesky_mo_num;
C=d_K1tmp; ldc=nO*nO;
cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
cudaFree(d_t1v);
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha = 1.0;
beta = 1.0;
for (int i=0 ; i<nO ; ++i) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
A = &(d_K1[nO*nV*(i+nO*b)]); lda = nO;
B = &(d_K1tmp[nO*(i+nO*nV*b)]); ldb = nO*nO;
C = &(d_K1[nO*nV*(i+nO*b)]); ldc = nO;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nO ; ++i) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
A = &(d_K1[nO*nV*(i+nO*b)]); lda = nO;
B = &(d_Z[nO*(b+nV*nV*i)]); ldb = nO*nV;
C = &(d_K1[nO*nV*(i+nO*b)]); ldc = nO;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nV ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
cudaFree(d_K1tmp);
// cudaFree(d_Z);
lda = nO*nV;
cublasGetMatrix(nO*nV, nO*nV, sizeof(double), d_K1, lda, K1, lda);
cudaFree(d_K1);
}
} // end sections
lda = cholesky_mo_num * nV;
cudaMalloc((void **)&d_tmp_cc, lda * nV * sizeof(double));

View File

@ -9,6 +9,7 @@ typedef struct {
double* cc_space_v_vvoo;
double* cc_space_v_oovo;
double* cc_space_v_ovvo;
double* cc_space_v_ovov;
double* cc_space_v_ovoo;
double* cc_space_f_vo;
} gpu_data;

View File

@ -15,7 +15,8 @@ gpu_data* gpu_init(
double* cc_space_v_oooo, double* cc_space_v_vooo,
double* cc_space_v_oovv, double* cc_space_v_vvoo,
double* cc_space_v_oovo, double* cc_space_v_ovvo,
double* cc_space_v_ovoo, double* cc_space_f_vo)
double* cc_space_v_ovov, double* cc_space_v_ovoo,
double* cc_space_f_vo)
{
int ngpus = 1;
cudaGetDeviceCount(&ngpus);
@ -77,6 +78,11 @@ gpu_data* gpu_init(
cudaMalloc((void **)&d_cc_space_v_ovvo, nO*nV*nV*nO * sizeof(double));
cublasSetMatrix(lda, nV*nO, sizeof(double), cc_space_v_ovvo, lda, d_cc_space_v_ovvo, lda);
double* d_cc_space_v_ovov;
lda = nO*nV;
cudaMalloc((void **)&d_cc_space_v_ovov, nO*nV*nV*nO * sizeof(double));
cublasSetMatrix(lda, nV*nO, sizeof(double), cc_space_v_ovov, lda, d_cc_space_v_ovov, lda);
double* d_cc_space_v_ovoo;
lda = nO*nV;
cudaMalloc((void **)&d_cc_space_v_ovoo, nO*nV*nO*nO * sizeof(double));
@ -96,6 +102,7 @@ gpu_data* gpu_init(
data[igpu].cc_space_v_vvoo = d_cc_space_v_vvoo;
data[igpu].cc_space_v_oovo = d_cc_space_v_oovo;
data[igpu].cc_space_v_ovvo = d_cc_space_v_ovvo;
data[igpu].cc_space_v_ovov = d_cc_space_v_ovov;
data[igpu].cc_space_v_ovoo = d_cc_space_v_ovoo;
data[igpu].cc_space_f_vo = d_cc_space_f_vo;
}

View File

@ -7,7 +7,8 @@ module gpu_module
type(c_ptr) function gpu_init(nO, nV, cholesky_mo_num, &
cc_space_v_oo_chol, cc_space_v_ov_chol, cc_space_v_vo_chol, cc_space_v_vv_chol, &
cc_space_v_oooo, cc_space_v_vooo, cc_space_v_oovv, cc_space_v_vvoo, &
cc_space_v_oovo, cc_space_v_ovvo, cc_space_v_ovoo, cc_space_f_vo) bind(C)
cc_space_v_oovo, cc_space_v_ovvo, cc_space_v_ovov, cc_space_v_ovoo, &
cc_space_f_vo) bind(C)
import c_int, c_double, c_ptr
integer(c_int), intent(in), value :: nO, nV, cholesky_mo_num
real(c_double), intent(in) :: cc_space_v_oo_chol(cholesky_mo_num,nO,nO)
@ -20,12 +21,13 @@ module gpu_module
real(c_double), intent(in) :: cc_space_v_vvoo(nV,nV,nO,nO)
real(c_double), intent(in) :: cc_space_v_oovo(nO,nO,nV,nO)
real(c_double), intent(in) :: cc_space_v_ovvo(nO,nV,nV,nO)
real(c_double), intent(in) :: cc_space_v_ovov(nO,nV,nO,nV)
real(c_double), intent(in) :: cc_space_v_ovoo(nO,nV,nO,nO)
real(c_double), intent(in) :: cc_space_f_vo(nV,nO)
end function
subroutine compute_r2_space_chol_gpu(nO,nV,cholesky_mo_num, gpu_data, t1, t2, tau,&
H_vv, g_occ, J1, r2) bind(C)
H_vv, g_occ, J1, K1, r2) bind(C)
import c_int, c_double, c_ptr
integer(c_int), intent(in), value :: nO, nV, cholesky_mo_num
type(c_ptr), value :: gpu_data
@ -35,6 +37,7 @@ module gpu_module
real(c_double), intent(in) :: H_vv(nV,nV)
real(c_double), intent(in) :: g_occ(nO,nO)
real(c_double) :: J1(nO,nV,nV,nO)
real(c_double) :: K1(nO,nV,nO,nV)
real(c_double), intent(out) :: r2(nO,nO,nV,nV)
end subroutine