More dgemm

This commit is contained in:
Anthony Scemama 2023-08-02 19:51:12 +02:00
parent b376fe685c
commit 9d5df395aa
2 changed files with 74 additions and 49 deletions

View File

@ -745,63 +745,18 @@ subroutine compute_J1_chol(nO,nV,t1,t2,v_ovvo,v_ovoo,v_vvoo,J1)
double precision, allocatable :: tmp_cc(:,:,:), J1_tmp(:,:,:,:)
!TODO: I am here
!TODO: I am here
!- cc_space_v_vvoo(a,b,i,j) * (0.5d0 * t2(u,j,b,beta) + t1(u,b) * t1(j,beta)) &
double precision, allocatable :: X_voov(:,:,:,:), Z_ovvo(:,:,:,:)
allocate(X_voov(nV,nO,nO,nV), Z_ovvo(nO,nV,nV,nO))
!$omp parallel &
!$omp shared(nO,nV,t2,t1,Y_ovov,X_voov,v_vvoo) &
!$omp private(i,beta,a,u,b,j) &
!$omp default(none)
!$omp do
do b = 1, nV
do j = 1, nO
do beta = 1, nV
do u = 1, nO
Y_ovov(u,beta,j,b) = 0.5d0 * t2(u,j,b,beta) + t1(u,b) * t1(j,beta)
enddo
enddo
enddo
enddo
!$omp end do nowait
!$omp do
do b = 1, nV
do j = 1, nO
do i = 1, nO
do a = 1, nV
X_voov(a,i,j,b) = v_vvoo(a,b,i,j)
enddo
enddo
enddo
enddo
!$omp end do
!$omp end parallel
call dgemm('N','T',nO*nV,nV*nO,nO*nV, &
-1d0, Y_ovov, size(Y_ovov,1) * size(Y_ovov,2), &
X_voov, size(X_voov,1) * size(X_voov,2), &
0d0, Z_ovvo, size(Z_ovvo,1) * size(Z_ovvo,2))
deallocate(X_voov)
double precision, allocatable :: X_ovvo(:,:,:,:), Y_vovo(:,:,:,:)
allocate(X_ovvo(nO,nV,nV,nO),Y_vovo(nV,nO,nV,nO))
allocate(X_ovvo(nO,nV,nV,nO))
allocate(Y_vovo(nV,nO,nV,nO))
!$omp parallel &
!$omp shared(nO,nV,J1,Z_ovvo,t2,Y_vovo,v_vvoo,X_ovvo) &
!$omp private(i,beta,a,u,j,b) &
!$omp default(none)
do i = 1, nO
!$omp do
do beta = 1, nV
do a = 1, nV
do u = 1, nO
J1(u,a,beta,i) = J1(u,a,beta,i) + Z_ovvo(u,beta,a,i)
enddo
enddo
enddo
!$omp end do nowait
enddo
!+ 0.5d0 * (2d0 * cc_space_v_vvoo(a,b,i,j) - cc_space_v_vvoo(b,a,i,j)) * t2(u,j,beta,b)
do j = 1, nO
@ -851,7 +806,7 @@ subroutine compute_J1_chol(nO,nV,t1,t2,v_ovvo,v_ovoo,v_vvoo,J1)
enddo
!$omp end parallel
deallocate(X_ovvo,Z_ovvo,Y_ovov)
deallocate(X_ovvo,Z_ovvo)
end

View File

@ -680,6 +680,7 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
B=d_t1; ldb=nO;
C=d_Y_ovov; ldc=nO*nV*nO;
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
cudaFree(d_X_ovoo);
alpha = 1.0;
beta = -1.0;
@ -744,6 +745,75 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
cublasSetStream(handle, NULL);
cudaFree(d_J1_tmp);
//---
double* d_X_voov;
lda = nV*nO;
cudaMalloc((void **)&d_X_voov, nV*nO*nO*nV * sizeof(double));
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
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_ovov[nO*(b+nV*j)]); ldc = nO*nV*nO;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nO, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
alpha = 1.0;
beta = 0.0;
for (int j=0 ; j<nO ; ++j) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
A = &(d_cc_space_v_vvoo[nV*(b+nV*nO*j)]); lda = nV*nV;
B = &(d_cc_space_v_vvoo[nV*(b+nV*nO*j)]); ldb = nV*nV;
C = &(d_X_voov[nV*nO*(j+nO*b)]); ldc = nV;
cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nV, nO, &alpha, A, lda, &beta, B, ldb, C, ldc);
}
}
for (int i=0 ; i<nV ; ++i) {
cudaStreamDestroy(stream[i]);
}
cublasSetStream(handle, NULL);
double* d_Z_ovvo;
lda = nO*nV;
cudaMalloc((void **)&d_Z_ovvo, nO*nV*nV*nO * sizeof(double));
alpha = -1.0;
beta = 0.0;
m=nO*nV; n=nV*nO; k=nO*nV;
A=d_Y_ovov; lda=nO*nV;
B=d_X_voov; ldb=nV*nO;
C=d_Z_ovvo; ldc=nO*nV;
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
cudaFree(d_X_voov);
cudaFree(d_Y_ovov);
alpha = 1.0;
beta = 1.0;
for (int i=0 ; i<nV ; ++i) {
cudaStreamCreate(&(stream[i]));
}
for (int i=0 ; i<nO ; ++i) {
for (int b=0 ; b<nV ; ++b) {
cublasSetStream(handle, stream[b]);
A = &(d_J1[nO*nV*(b+nV*i)]); lda = nO;
B = &(d_Z_ovvo[nO*(b+nV*nV*i)]); ldb=nO*nV;
C = &(d_J1[nO*nV*(b+nV*i)]); 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_Z_ovvo);
lda = nO*nV;
cublasGetMatrix(nO*nV, nV*nO, sizeof(double), d_J1, lda, J1, lda);
cudaFree(d_J1);