r2 on GPU

This commit is contained in:
Anthony Scemama 2023-08-03 03:21:32 +02:00
parent 30e2e086a0
commit a541abd04a
3 changed files with 91 additions and 86 deletions

View File

@ -490,82 +490,8 @@ subroutine compute_r2_space_chol(nO,nV,t1,t2,tau,H_oo,H_vv,H_vo,r2,max_r2,gpu_da
!$omp end do
!$omp end parallel
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, K1, r2)
!---
double precision, allocatable :: X_oovv(:,:,:,:)
double precision, allocatable :: X_vovv(:,:,:,:)
double precision, allocatable :: X_ovvo(:,:,:,:)
double precision, allocatable :: tcc(:,:,:), tcc2(:,:,:)
double precision, allocatable :: X_vovo(:,:,:,:), Y_oovo(:,:,:,:)
double precision, allocatable :: Y_voov(:,:,:,:)
double precision, allocatable :: Z_ovov(:,:,:,:)
double precision, allocatable :: Y_ovov(:,:,:,:), X_ovov(:,:,:,:)
allocate(X_ovov(nO,nV,nO,nV),Y_ovov(nO,nV,nO,nV))
!$omp parallel &
!$omp shared(nO,nV,K1,X_ovov,Y_ovov,t2) &
!$omp private(u,v,gam,beta,i,a) &
!$omp default(none)
!$omp do
do a = 1, nV
do i = 1, nO
do gam = 1, nV
do u = 1, nO
X_ovov(u,gam,i,a) = K1(u,a,i,gam)
enddo
enddo
enddo
enddo
!$omp end do nowait
!$omp do
do beta = 1, nV
do v = 1, nO
do a = 1, nV
do i = 1, nO
Y_ovov(i,a,v,beta) = t2(i,v,beta,a)
enddo
enddo
enddo
enddo
!$omp end do
!$omp end parallel
deallocate(K1)
allocate(Z_ovov(nO,nV,nO,nV))
call dgemm('N','N',nO*nV,nO*nV,nO*nV, &
1d0, X_ovov, size(X_ovov,1) * size(X_ovov,2), &
Y_ovov, size(Y_ovov,1) * size(Y_ovov,2), &
0d0, Z_ovov, size(Y_ovov,1) * size(Y_ovov,2))
deallocate(X_ovov,Y_ovov)
!$omp parallel &
!$omp shared(nO,nV,r2,Z_ovov) &
!$omp private(u,v,gam,beta) &
!$omp default(none)
!$omp do
do gam = 1, nV
do beta = 1, nV
do v = 1, nO
do u = 1, nO
r2(u,v,beta,gam) = r2(u,v,beta,gam) - Z_ovov(u,gam,v,beta) - Z_ovov(v,beta,u,gam)
enddo
enddo
enddo
enddo
!$omp end do
!$omp end parallel
deallocate(Z_ovov)
H_vv, g_occ, r2)
! Change the sign for consistency with the code in spin orbitals

View File

@ -16,7 +16,6 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
double* tau,
double* H_vv,
double* g_occ,
double* K1,
double* r2)
{
@ -24,6 +23,7 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cudaGetDeviceCount(&ngpus);
double* J1 = malloc(nO*nV*nV*nO*sizeof(double));
double* K1 = malloc(nO*nV*nV*nO*sizeof(double));
#pragma omp parallel num_threads(ngpus)
{
@ -629,13 +629,6 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasSetStream(handle, NULL);
cudaFree(d_Y_oovv);
}
} // end sections
lda = nO*nV;
cublasSetMatrix(lda, nO*nV, sizeof(double), K1, lda, d_K1, lda);
#pragma omp sections
{
// g_occ
#pragma omp section
@ -693,6 +686,14 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cudaFree(d_X_oovv);
}
} // end sections
lda = nO*nV;
cublasSetMatrix(lda, nO*nV, sizeof(double), K1, lda, d_K1, lda);
#pragma omp sections
{
#pragma omp section
{
double* d_X_vovv;
@ -1112,10 +1113,87 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
#pragma omp section
{
double* d_X_ovov;
cudaMalloc((void **)&d_X_ovov, nO*nV*nO*nV * sizeof(double));
double* d_Y_ovov;
cudaMalloc((void **)&d_Y_ovov, nO*nV*nO*nV * sizeof(double));
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha = 1.0;
beta = 0.0;
for (int a=0 ; a<nV ; ++a) {
for (int g=0 ; g<nV ; ++g) {
cublasSetStream(handle, stream[g]);
A = &(d_K1[nO*(a+nV*nO*g)]); lda = nO*nV;
B = &(d_K1[nO*(a+nV*nO*g)]); ldb = nO*nV;
C = &(d_X_ovov[nO*(g+nV*nO*a)]); ldc = nO*nV;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
alpha = 1.0;
for (int v=0 ; v<nO ; ++v) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
A = &(d_t2[nO*(v+nO*b)]); lda = nO*nO*nV;
B = &(d_t2[nO*(v+nO*b)]); ldb = nO*nO*nV;
C = &(d_Y_ovov[nO*nV*(v+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);
double* d_Z_ovov;
cudaMalloc((void **)&d_Z_ovov, nO*nV*nO*nV * sizeof(double));
alpha = 1.0;
beta = 0.0;
m=nO*nV; n=nO*nV; k=nO*nV;
A=d_X_ovov; lda=nO*nV;
B=d_Y_ovov; ldb=nO*nV;
C=d_Z_ovov; 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_ovov);
cudaFree(d_Y_ovov);
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
alpha = 1.0;
beta = -1.0;
for (int b=0 ; b<nV ; ++b) {
for (int g=0 ; g<nV ; ++g) {
cublasSetStream(handle, stream[g]);
A = &(d_r2[nO*nO*(b+nV*g)]); lda = nO;
B = &(d_Z_ovov[nO*(g+nV*nO*b)]); ldb = nO*nV;
C = &(d_r2[nO*nO*(b+nV*g)]); ldc = nO;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
for (int g=0 ; g<nV ; ++g) {
cublasSetStream(handle, stream[g]);
A = &(d_r2[nO*nO*(b+nV*g)]); lda = nO;
B = &(d_Z_ovov[nO*(b+nV*nO*g)]); ldb = nO*nV;
C = &(d_r2[nO*nO*(b+nV*g)]); ldc = nO;
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);
cudaFree(d_Z_ovov);
}
} // end sections
cudaFree(d_K1);
lda = cholesky_mo_num * nV;
cudaMalloc((void **)&d_tmp_cc, lda * nV * sizeof(double));
@ -1209,6 +1287,8 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasDestroy(handle);
}
free(K1);
free(J1);
}

View File

@ -27,7 +27,7 @@ module gpu_module
end function
subroutine compute_r2_space_chol_gpu(nO,nV,cholesky_mo_num, gpu_data, t1, t2, tau,&
H_vv, g_occ, K1, r2) bind(C)
H_vv, g_occ, 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
@ -36,7 +36,6 @@ module gpu_module
real(c_double), intent(in) :: tau(nO,nO,nV,nV)
real(c_double), intent(in) :: H_vv(nV,nV)
real(c_double), intent(in) :: g_occ(nO,nO)
real(c_double) :: K1(nO,nV,nO,nV)
real(c_double), intent(out) :: r2(nO,nO,nV,nV)
end subroutine