comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-29 06:35:28 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve numerical stability of quartically-scaling SOS-MP2/RPA gradients (#2390)

Browse source 
* Use accurate_dot_product with SOS-MP2 gradients

* Add keyword for extended precision

* Attempt to implement more accurate contraction

* Unroll loops of Kahan summation

We gain a factor of 2 in performance, important with RPA.

* OpenMP parallelization

* Add timing regions

* Block dot product

* Add keyword to set blksize of dot_products

* Use ddot instead of dot_product

* Use kahan summation for RPA, virt block

* Add loop for Pab

* Use Kahan summation scheme for the calculation of Pij

* Block P direction

* Cleanup and acceleration

* Adjust one of the tests to use the second branch

* Reshuffle the tests
This commit is contained in:
Frederick Stein 2022-11-10 14:02:31 +01:00 • committed by GitHub
parent a213519a15
commit 9ad87fd165
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 381 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

81
src/common/kahan_sum.F
View file

@ -30,12 +30,13 @@ MODULE kahan_sum
IMPLICIT NONE
PRIVATE
PUBLIC :: accurate_dot_product, accurate_sum
PUBLIC :: accurate_dot_product, accurate_sum, accurate_dot_product_2
INTEGER, PARAMETER :: sp = KIND(0.0), dp = KIND(0.0D0)
REAL(KIND=sp), PARAMETER :: szero = 0.0_sp
REAL(KIND=dp), PARAMETER :: dzero = 0.0_dp
COMPLEX(KIND=sp), PARAMETER :: czero = (0.0_sp, 0.0_sp)
COMPLEX(KIND=dp), PARAMETER :: zzero = (0.0_dp, 0.0_dp)
INTEGER, PARAMETER :: dblksize = 8
INTERFACE accurate_sum
MODULE PROCEDURE &
@ -54,6 +55,11 @@ MODULE kahan_sum
kahan_dot_product_s2, kahan_dot_product_d2, kahan_dot_product_z2, &
kahan_dot_product_d3, kahan_dot_product_masked_d3
END INTERFACE accurate_dot_product
INTERFACE accurate_dot_product_2
MODULE PROCEDURE kahan_blocked_dot_product_d1
END INTERFACE
CONTAINS
! **************************************************************************************************
!> \brief ...
@ -136,17 +142,30 @@ CONTAINS
REAL(KIND=dp) :: ks
INTEGER :: i, n
REAL(KIND=dp) :: c, t, y
REAL(KIND=dp), DIMENSION(dblksize) :: c, ks_local, t, y
ks = dzero; t = dzero; y = dzero; c = dzero
t = dzero; y = dzero; c = dzero; ks_local = dzero
n = SIZE(array1)
DO i = 1, n
y = array1(i)*array2(i) - c
t = ks + y
c = (t - ks) - y
ks = t
DO i = 1, MOD(n, dblksize)
y(1) = array1(i)*array2(i) - c(1)
t(1) = ks_local(1) + y(1)
c(1) = (t(1) - ks_local(1)) - y(1)
ks_local(1) = t(1)
END DO
DO i = MOD(n, dblksize) + 1, n, dblksize
y = array1(i:i + (dblksize - 1))*array2(i:i + (dblksize - 1)) - c
t = ks_local + y
c = (t - ks_local) - y
ks_local = t
END DO
DO i = 2, dblksize
y(1) = ks_local(i) - (c(1) + c(i))
t(1) = ks_local(1) + y(1)
c(1) = (t(1) - ks_local(1)) - y(1)
ks_local(1) = t(1)
END DO
ks = ks_local(1)
END FUNCTION kahan_dot_product_d1
! **************************************************************************************************
@ -1570,4 +1589,50 @@ CONTAINS
END IF
END FUNCTION kahan_sum_z7
! **************************************************************************************************
!> \brief computes the accurate sum of blocks of regular dot products
!> \param array1 array of real numbers
!> \param array2 another array of real numbers
!> \param blksize ...
!> \return dot product
! **************************************************************************************************
FUNCTION kahan_blocked_dot_product_d1(array1, array2, blksize) RESULT(ks)
REAL(KIND=dp), DIMENSION(:), INTENT(in) :: array1, array2
INTEGER, INTENT(IN), OPTIONAL :: blksize
REAL(KIND=dp) :: ks
INTEGER :: my_blksize
REAL(KIND=dp) :: DDOT
my_blksize = 32
IF (PRESENT(blksize)) my_blksize = blksize
IF (my_blksize <= 1) THEN
! The original should be faster
ks = accurate_dot_product(array1, array2)
ELSE IF (my_blksize >= SIZE(array1)) THEN
! Just use standard dot product from BLAS for performance
ks = DDOT(SIZE(array1), array1(1), 1, array2(1), 1)
ELSE
ks = 0.0_dp
BLOCK
INTEGER :: i, n, stripesize
REAL(KIND=dp) :: c, dotproduct, t, y
t = dzero; y = dzero; c = dzero
n = SIZE(array1)
DO i = 1, n, my_blksize
! Remove 1 to save an operation in the length
stripesize = MIN(my_blksize, n - i + 1)
! Perform dot product on the given stripe
dotproduct = DDOT(stripesize, array1(i), 1, array2(i), 1)
y = dotproduct - c
t = ks + y
c = (t - ks) - y
ks = t
END DO
END BLOCK
END IF
END FUNCTION kahan_blocked_dot_product_d1
END MODULE kahan_sum

11
src/input_cp2k_mp2.F
View file

@ -464,6 +464,17 @@ CONTAINS
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create( &
keyword, __LOCATION__, &
name="DOT_PRODUCT_BLKSIZE", &
description="Dot products for the calculation of the RPA/SOS-MP2 density matrices "// &
"are calculated in batches of the size given by this keyword. Larger block sizes "// &
"improve the performance but reduce the numerical accuracy. Recommended block sizes are multiples of the number of "// &
"doubles per cache line (usually 8). Ignored with MP2 gradients. Set it to -1 to prevent blocking.", &
default_i_val=32)
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
END SUBROUTINE create_canonical_gradients
! **************************************************************************************************

1
src/mp2_setup.F
View file

@ -374,6 +374,7 @@ CONTAINS
CALL section_vals_val_get(mp2_section, "CANONICAL_GRADIENTS%FREE_HFX_BUFFER", l_val=mp2_env%ri_grad%free_hfx_buffer)
CALL section_vals_val_get(mp2_section, "CANONICAL_GRADIENTS%EPS_CANONICAL", r_val=mp2_env%ri_grad%eps_canonical)
CALL section_vals_val_get(mp2_section, "CANONICAL_GRADIENTS%USE_OLD_GRADIENT_CODE", l_val=mp2_env%ri_grad%use_old_grad)
CALL section_vals_val_get(mp2_section, "CANONICAL_GRADIENTS%DOT_PRODUCT_BLKSIZE", i_val=mp2_env%ri_grad%dot_blksize)
cphf_section => section_vals_get_subs_vals(mp2_section, "CANONICAL_GRADIENTS%CPHF")
IF (ASSOCIATED(cphf_section)) THEN
CALL section_vals_val_get(cphf_section, "MAX_ITER", i_val=mp2_env%ri_grad%cphf_max_num_iter)

1
src/mp2_types.F
View file

@ -235,6 +235,7 @@ MODULE mp2_types
REAL(dp) :: eps_canonical
LOGICAL :: free_hfx_buffer
LOGICAL :: use_old_grad
INTEGER :: dot_blksize
END TYPE
TYPE mp2_type

416
src/rpa_grad.F
View file

@ -38,6 +38,8 @@ MODULE rpa_grad
group_dist_proc,&
maxsize,&
release_group_dist
USE kahan_sum, ONLY: accurate_dot_product,&
accurate_dot_product_2
USE kinds, ONLY: dp,&
int_8
USE libint_2c_3c, ONLY: compare_potential_types
@ -759,6 +761,8 @@ CONTAINS
CALL cp_fm_get_info(fm_mat_Q(1)%matrix, nrow_global=dimen_RI, nrow_local=nrow_local, ncol_local=ncol_local, &
col_indices=col_indices, row_indices=row_indices)
IF (mp2_env%ri_grad%dot_blksize == -1) mp2_env%ri_grad%dot_blksize = nrow_local
IF (.NOT. do_ri_sos_laplace_mp2) THEN
CALL cp_fm_create(fm_work_PQ, fm_mat_Q(1)%matrix%matrix_struct)
@ -831,7 +835,7 @@ CONTAINS
IF (do_ri_sos_laplace_mp2) THEN
CALL calc_P_sos_mp2(homo(ispin), fm_mat_S(ispin)%matrix, fm_work_iaP, &
rpa_grad%sos_mp2_work_occ(ispin), rpa_grad%sos_mp2_work_virt(ispin), &
omega, weight, virtual(ispin), Eigenval(:, ispin))
omega, weight, virtual(ispin), Eigenval(:, ispin), mp2_env%ri_grad%dot_blksize)
CALL calc_fm_mat_S_laplace(fm_work_iaP, homo(ispin), virtual(ispin), Eigenval(:, ispin), omega)
@ -884,14 +888,17 @@ CONTAINS
!> \param weight ...
!> \param virtual ...
!> \param Eigenval ...
!> \param dot_blksize ...
! **************************************************************************************************
SUBROUTINE calc_P_sos_mp2(homo, fm_mat_S, fm_work_iaP, sos_mp2_work_occ, sos_mp2_work_virt, omega, weight, virtual, Eigenval)
SUBROUTINE calc_P_sos_mp2(homo, fm_mat_S, fm_work_iaP, sos_mp2_work_occ, sos_mp2_work_virt, &
omega, weight, virtual, Eigenval, dot_blksize)
INTEGER, INTENT(IN) :: homo
TYPE(cp_fm_type), INTENT(IN) :: fm_mat_S, fm_work_iaP
TYPE(sos_mp2_grad_work_type), INTENT(INOUT) :: sos_mp2_work_occ, sos_mp2_work_virt
REAL(KIND=dp), INTENT(IN) :: omega, weight
INTEGER, INTENT(IN) :: virtual
REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: Eigenval
INTEGER, INTENT(IN) :: dot_blksize
CHARACTER(LEN=*), PARAMETER :: routineN = 'calc_P_sos_mp2'
@ -953,12 +960,12 @@ CONTAINS
IF (num_ij_pairs > 0) THEN
CALL calc_Pij_degen(fm_work_iaP, fm_mat_S, sos_mp2_work_occ%pair_list, &
virtual, sos_mp2_work_occ%P(homo + 1:), Eigenval(:homo), omega, weight, &
sos_mp2_work_occ%index2send, sos_mp2_work_occ%index2recv)
sos_mp2_work_occ%index2send, sos_mp2_work_occ%index2recv, dot_blksize)
END IF
IF (num_ab_pairs > 0) THEN
CALL calc_Pab_degen(fm_work_iaP, fm_mat_S, sos_mp2_work_virt%pair_list, &
virtual, sos_mp2_work_virt%P(virtual + 1:), Eigenval(homo + 1:), omega, weight, &
sos_mp2_work_virt%index2send, sos_mp2_work_virt%index2recv)
sos_mp2_work_virt%index2send, sos_mp2_work_virt%index2recv, dot_blksize)
END IF
CALL timestop(handle)
@ -997,28 +1004,33 @@ CONTAINS
CHARACTER(LEN=*), PARAMETER :: routineN = 'calc_P_rpa'
INTEGER :: handle, handle2, my_a, my_a_size, my_a_start, my_i, my_i_size, my_i_start, &
my_P_size, my_prow, proc_a_recv, proc_a_send, proc_i_recv, proc_i_send, proc_recv, &
proc_send, proc_shift, recv_a_end, recv_a_size, recv_a_start, recv_i_end, recv_i_size, &
recv_i_start, tag
INTEGER :: handle, handle2, my_a, my_a_size, my_a_start, my_b, my_i, my_i_size, my_i_start, &
my_j, my_P_size, my_prow, P_end, P_start, proc_a_recv, proc_a_send, proc_i_recv, &
proc_i_send, proc_recv, proc_send, proc_shift, recv_a, recv_a_end, recv_a_size, &
recv_a_start, recv_i, recv_i_end, recv_i_size, recv_i_start, stripesize, tag
INTEGER, DIMENSION(:, :), POINTER :: blacs2mpi
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), TARGET :: buffer_1D
REAL(KIND=dp), DIMENSION(:, :), POINTER :: buffer_2D, mat_S_2D
REAL(KIND=dp) :: my_compens, my_pab, my_pij, s
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), TARGET :: buffer_1D, buffer_compens_1D
REAL(KIND=dp), DIMENSION(:, :), POINTER :: buffer_2D, buffer_compens_2D
REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: buffer_3D, mat_S_3D
TYPE(cp_para_env_type), POINTER :: para_env
CALL timeset(routineN, handle)
! We allocate it at every step to reduce potential memory conflicts with COSMA
IF (.NOT. c_associated(mp2_env%local_gemm_ctx)) THEN
CALL local_gemm_create(mp2_env%local_gemm_ctx, LOCAL_GEMM_PU_GPU)
CALL local_gemm_set_op_threshold_gpu(mp2_env%local_gemm_ctx, 128*128*128*2)
IF (mp2_env%ri_grad%dot_blksize >= 8) THEN
IF (.NOT. c_associated(mp2_env%local_gemm_ctx)) THEN
CALL local_gemm_create(mp2_env%local_gemm_ctx, LOCAL_GEMM_PU_GPU)
CALL local_gemm_set_op_threshold_gpu(mp2_env%local_gemm_ctx, 128*128*128*2)
END IF
END IF
tag = 47
my_P_size = SIZE(mat_work_iaP_3D, 1)
IF (mp2_env%ri_grad%dot_blksize == -1) mp2_env%ri_grad%dot_blksize = my_P_size
CALL cp_fm_struct_get(fm_struct_S, para_env=para_env)
CALL get_blacs_info(fm_struct_S%context, my_process_row=my_prow, blacs2mpi=blacs2mpi, para_env=para_env)
@ -1028,12 +1040,16 @@ CONTAINS
! We have to remap the indices because mp_sendrecv requires a 3D array (because of mat_work_iaP_3D)
! and dgemm requires 2D arrays
! Fortran 2008 does allow pointer remapping independently of the ranks but GCC 7 does not properly support it
mat_S_2D(1:my_P_size*my_a_size, 1:my_i_size) => mat_S_1D(1:INT(my_P_size, int_8)*my_a_size*my_i_size)
mat_S_3D(1:my_P_size, 1:my_a_size, 1:my_i_size) => mat_S_1D(1:INT(my_P_size, int_8)*my_a_size*my_i_size)
ALLOCATE (buffer_1D(MAX(INT(maxsize(gd_homo), KIND=int_8)*my_a_size, &
INT(maxsize(gd_virtual), KIND=int_8)*my_i_size)*my_P_size))
! Allocate buffers for vector version of kahan summation
IF (mp2_env%ri_grad%dot_blksize >= 8) THEN
ALLOCATE (buffer_compens_1D(MAX(my_a_size*maxsize(gd_virtual), my_i_size*maxsize(gd_homo))))
END IF
DO proc_shift = 0, grid(1) - 1
proc_a_send = MODULO(mepos(1) + proc_shift, grid(1))
proc_a_recv = MODULO(mepos(1) - proc_shift, grid(1))
@ -1051,53 +1067,117 @@ CONTAINS
para_env%group, tag)
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_a", handle2)
! This loop imitates the actual tensor contraction
DO my_i = 1, my_i_size
CALL local_gemm("T", "N", my_a_size, recv_a_size, my_P_size, &
-weight, mat_S_3D(:, :, my_i), my_P_size, &
buffer_3D(:, :, my_i), my_P_size, &
1.0_dp, P_ab(:, recv_a_start:recv_a_end), my_a_size, mp2_env%local_gemm_ctx)
END DO
CALL timestop(handle2)
IF (mp2_env%ri_grad%dot_blksize >= 8) THEN
buffer_compens_2D(1:my_a_size, 1:recv_a_size) => buffer_compens_1D(1:my_a_size*recv_a_size)
CALL timeset(routineN//"_scale_a", handle2)
DO my_i = 1, my_i_size
! Scale both matrices
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
DO my_a = 1, recv_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + recv_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
END DO
CALL timestop(handle2)
buffer_compens_2D = 0.0_dp
CALL timeset(routineN//"_dgemm_a", handle2)
DO my_i = 1, my_i_size
! Repeat contraction
CALL local_gemm("T", "N", my_a_size, recv_a_size, my_P_size, &
weight, mat_S_3D(:, :, my_i), my_P_size, &
buffer_3D(:, :, my_i), my_P_size, &
1.0_dp, P_ab(:, recv_a_start:recv_a_end), my_a_size, mp2_env%local_gemm_ctx)
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_a", handle2)
! This loop imitates the actual tensor contraction
DO my_i = 1, my_i_size
DO P_start = 1, my_P_size, mp2_env%ri_grad%dot_blksize
stripesize = MIN(mp2_env%ri_grad%dot_blksize, my_P_size - P_start + 1)
P_end = P_start + stripesize - 1
CALL timeset(routineN//"_scale_a", handle2)
! Remove scaling of mat_S for later
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1))/omega)
CALL local_gemm("T", "N", my_a_size, recv_a_size, stripesize, &
-weight, mat_S_3D(P_start:P_end, :, my_i), stripesize, &
buffer_3D(P_start:P_end, :, my_i), stripesize, &
-1.0_dp, buffer_compens_2D, my_a_size, mp2_env%local_gemm_ctx)
!$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(2) SHARED(recv_a_size,my_a_size,P_ab,recv_a_start,buffer_compens_2D) &
!$OMP PRIVATE(my_a,my_b,my_pab,my_compens,s)
DO my_a = 1, recv_a_size
DO my_b = 1, my_a_size
my_pab = P_ab(my_b, recv_a_start - 1 + my_a)
my_compens = buffer_compens_2D(my_b, my_a)
s = my_pab + my_compens
buffer_compens_2D(my_b, my_a) = (s - my_pab) - my_compens
P_ab(my_b, recv_a_start - 1 + my_a) = s
END DO
END DO
END DO
END DO
DO my_a = 1, recv_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + recv_a_start - 1) - Eigenval(my_i + my_i_start - 1))/omega)
CALL timestop(handle2)
CALL timeset(routineN//"_scale_a", handle2)
DO my_i = 1, my_i_size
! Scale both matrices
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
DO my_a = 1, recv_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + recv_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
END DO
END DO
CALL timestop(handle2)
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_a", handle2)
DO my_i = 1, my_i_size
DO P_start = 1, my_P_size, mp2_env%ri_grad%dot_blksize
stripesize = MIN(mp2_env%ri_grad%dot_blksize, my_P_size - P_start + 1)
P_end = P_start + stripesize - 1
! Repeat contraction
CALL local_gemm("T", "N", my_a_size, recv_a_size, stripesize, &
weight, mat_S_3D(P_start:P_end, :, my_i), stripesize, &
buffer_3D(P_start:P_end, :, my_i), stripesize, &
-1.0_dp, buffer_compens_2D, my_a_size, mp2_env%local_gemm_ctx)
!$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(2) SHARED(recv_a_size,my_a_size,P_ab,recv_a_start,buffer_compens_2D) &
!$OMP PRIVATE(my_a,my_b,my_pab,my_compens,s)
DO my_a = 1, recv_a_size
DO my_b = 1, my_a_size
my_pab = P_ab(my_b, recv_a_start - 1 + my_a)
my_compens = buffer_compens_2D(my_b, my_a)
s = my_pab + my_compens
buffer_compens_2D(my_b, my_a) = (s - my_pab) - my_compens
P_ab(my_b, recv_a_start - 1 + my_a) = s
END DO
END DO
END DO
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_scale_a", handle2)
! Remove scaling of mat_S for later
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1))/omega)
END DO
DO my_a = 1, recv_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + recv_a_start - 1) - Eigenval(my_i + my_i_start - 1))/omega)
END DO
END DO
CALL timestop(handle2)
ELSE
BLOCK
REAL(KIND=dp) :: tmp, e_i, e_a, e_b, omega2
CALL timeset(routineN//"_accurate_a", handle2)
omega2 = -omega**2
!$OMP PARALLEL DO COLLAPSE(2) DEFAULT(NONE)&
!$OMP SHARED(my_a_size,recv_a_size,my_i_size,mat_S_3D,homo,Eigenval,omega2,my_a_start,&
!$OMP recv_a_start,P_ab,weight,buffer_3D,my_i_start,mp2_env)&
!$OMP PRIVATE(tmp,my_a,recv_a,my_i,e_a,e_b,e_i)
DO my_a = 1, my_a_size
DO recv_a = 1, recv_a_size
e_a = Eigenval(homo + my_a_start - 1 + my_a)
e_b = Eigenval(homo + recv_a_start - 1 + recv_a)
tmp = 0.0_dp
DO my_i = 1, my_i_size
e_i = -Eigenval(my_i_start - 1 + my_i)
tmp = tmp + accurate_dot_product(mat_S_3D(:, my_a, my_i), buffer_3D(:, recv_a, my_i)) &
*(1.0_dp + omega2/((e_a + e_i)*(e_b + e_i)))
END DO
P_ab(my_a, recv_a_start - 1 + recv_a) = P_ab(my_a, recv_a_start - 1 + recv_a) - weight*tmp
END DO
END DO
CALL timestop(handle2)
END BLOCK
END IF
END DO
DO proc_shift = 0, grid(2) - 1
@ -1119,52 +1199,122 @@ CONTAINS
para_env%group, tag)
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_i", handle2)
! This loop imitates the actual tensor contraction
CALL local_gemm("T", "N", my_i_size, recv_i_size, my_a_size*my_P_size, &
weight, mat_S_2D, my_P_size*my_a_size, &
buffer_2D, my_P_size*my_a_size, &
1.0_dp, P_ij(:, recv_i_start:recv_i_end), my_i_size, mp2_env%local_gemm_ctx)
CALL timestop(handle2)
IF (mp2_env%ri_grad%dot_blksize >= 8) THEN
buffer_compens_2D(1:my_i_size, 1:recv_i_size) => buffer_compens_1D(1:my_i_size*recv_i_size)
buffer_compens_2D = 0.0_dp
CALL timeset(routineN//"_scale_i", handle2)
! Scale both matrices
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
END DO
DO my_i = 1, recv_i_size
DO my_a = 1, my_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + recv_i_start - 1)))
END DO
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_i", handle2)
! This loop imitates the actual tensor contraction
DO my_a = 1, my_a_size
DO P_start = 1, my_P_size, mp2_env%ri_grad%dot_blksize
stripesize = MIN(mp2_env%ri_grad%dot_blksize, my_P_size - P_start + 1)
P_end = P_start + stripesize - 1
CALL timeset(routineN//"_dgemm_i", handle2)
! Repeat contraction
CALL local_gemm("T", "N", my_i_size, recv_i_size, my_a_size*my_P_size, &
-weight, mat_S_2D, my_P_size*my_a_size, &
buffer_2D, my_P_size*my_a_size, &
1.0_dp, P_ij(:, recv_i_start:recv_i_end), my_i_size, mp2_env%local_gemm_ctx)
CALL timestop(handle2)
CALL local_gemm("T", "N", my_i_size, recv_i_size, stripesize, &
weight, mat_S_3D(P_start:P_end, my_a, :), stripesize, &
buffer_3D(P_start:P_end, my_a, :), stripesize, &
-1.0_dp, buffer_compens_2D(:, 1:recv_i_start), my_i_size, mp2_env%local_gemm_ctx)
CALL timeset(routineN//"_scale_i", handle2)
! Remove scaling factor of mat_S for later
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + recv_i_start - 1))/omega)
END DO
END DO
CALL timestop(handle2)
!$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(2) SHARED(recv_i_size,my_i_size,P_ij,recv_i_start,buffer_compens_2D) &
!$OMP PRIVATE(my_i,my_j,my_pij,my_compens,s)
DO my_i = 1, recv_i_size
DO my_j = 1, my_i_size
my_pij = P_ij(my_j, recv_i_start - 1 + my_i)
my_compens = buffer_compens_2D(my_j, my_i)
s = my_pij + my_compens
buffer_compens_2D(my_j, my_i) = (s - my_pij) - my_compens
P_ij(my_j, recv_i_start - 1 + my_i) = s
END DO
END DO
END DO
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_scale_i", handle2)
! Scale both matrices
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + my_i_start - 1)))
END DO
END DO
DO my_i = 1, recv_i_size
DO my_a = 1, my_a_size
buffer_3D(:, my_a, my_i) = buffer_3D(:, my_a, my_i)* &
(omega/(Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + recv_i_start - 1)))
END DO
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_dgemm_i", handle2)
DO my_a = 1, my_a_size
DO P_start = 1, my_P_size, mp2_env%ri_grad%dot_blksize
stripesize = MIN(mp2_env%ri_grad%dot_blksize, my_P_size - P_start + 1)
P_end = P_start + stripesize - 1
! Repeat contraction
CALL local_gemm("T", "N", my_i_size, recv_i_size, stripesize, &
-weight, mat_S_3D(P_start:P_end, my_a, :), stripesize, &
buffer_3D(P_start:P_end, my_a, :), stripesize, &
-1.0_dp, buffer_compens_2D(:, 1:recv_i_start), my_i_size, mp2_env%local_gemm_ctx)
!$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(2) SHARED(recv_i_size,my_i_size,P_ij,recv_i_start,buffer_compens_2D) &
!$OMP PRIVATE(my_i,my_j,my_pij,my_compens,s)
DO my_i = 1, recv_i_size
DO my_j = 1, my_i_size
my_pij = P_ij(my_j, recv_i_start - 1 + my_i)
my_compens = buffer_compens_2D(my_j, my_i)
s = my_pij + my_compens
buffer_compens_2D(my_j, my_i) = (s - my_pij) - my_compens
P_ij(my_j, recv_i_start - 1 + my_i) = s
END DO
END DO
END DO
END DO
CALL timestop(handle2)
CALL timeset(routineN//"_scale_i", handle2)
! Remove scaling factor of mat_S for later
DO my_i = 1, my_i_size
DO my_a = 1, my_a_size
mat_S_3D(:, my_a, my_i) = mat_S_3D(:, my_a, my_i)* &
((Eigenval(homo + my_a + my_a_start - 1) - Eigenval(my_i + recv_i_start - 1))/omega)
END DO
END DO
CALL timestop(handle2)
ELSE
BLOCK
REAL(KIND=dp) :: tmp, e_i, e_a, e_j, omega2
CALL timeset(routineN//"_accurate_i", handle2)
omega2 = -omega**2
!$OMP PARALLEL DO COLLAPSE(2) DEFAULT(NONE)&
!$OMP SHARED(my_a_size,recv_i_size,my_i_size,mat_S_3D,homo,Eigenval,omega2,my_a_start,&
!$OMP recv_i_start,P_ij,weight,buffer_3D,my_i_start)&
!$OMP PRIVATE(tmp,my_a,recv_i,my_i,e_i,e_j,e_a)
DO my_i = 1, my_i_size
DO recv_i = 1, recv_i_size
e_i = Eigenval(my_i_start - 1 + my_i)
e_j = Eigenval(recv_i_start - 1 + recv_i)
tmp = 0.0_dp
DO my_a = 1, my_a_size
e_a = Eigenval(homo + my_a_start - 1 + my_a)
tmp = tmp + accurate_dot_product(mat_S_3D(:, my_a, my_i), buffer_3D(:, my_a, recv_i)) &
*(1.0_dp + omega2/((e_a - e_i)*(e_a - e_j)))
END DO
P_ij(my_i, recv_i_start - 1 + recv_i) = P_ij(my_i, recv_i_start - 1 + recv_i) + weight*tmp
END DO
END DO
CALL timestop(handle2)
END BLOCK
END IF
END DO
! release memory allocated by local_gemm when run on GPU. local_gemm_ctx is null on cpu only runs
CALL local_gemm_destroy(mp2_env%local_gemm_ctx)
mp2_env%local_gemm_ctx = C_NULL_PTR
IF (mp2_env%ri_grad%dot_blksize >= 8) THEN
! release memory allocated by local_gemm when run on GPU. local_gemm_ctx is null on cpu only runs
CALL local_gemm_destroy(mp2_env%local_gemm_ctx)
mp2_env%local_gemm_ctx = C_NULL_PTR
DEALLOCATE (buffer_compens_1D)
END IF
CALL timestop(handle)
@ -1201,8 +1351,10 @@ CONTAINS
!> \param weight ...
!> \param index2send ...
!> \param index2recv ...
!> \param dot_blksize ...
! **************************************************************************************************
SUBROUTINE calc_Pij_degen(fm_work_iaP, fm_mat_S, pair_list, virtual, P_ij, Eigenval, omega, weight, index2send, index2recv)
SUBROUTINE calc_Pij_degen(fm_work_iaP, fm_mat_S, pair_list, virtual, P_ij, Eigenval, &
omega, weight, index2send, index2recv, dot_blksize)
TYPE(cp_fm_type), INTENT(IN) :: fm_work_iaP, fm_mat_S
INTEGER, DIMENSION(:, :), INTENT(IN) :: pair_list
INTEGER, INTENT(IN) :: virtual
@ -1210,6 +1362,7 @@ CONTAINS
REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: Eigenval
REAL(KIND=dp), INTENT(IN) :: omega, weight
TYPE(one_dim_int_array), DIMENSION(0:), INTENT(IN) :: index2send, index2recv
INTEGER, INTENT(IN) :: dot_blksize
CHARACTER(LEN=*), PARAMETER :: routineN = 'calc_Pij_degen'
@ -1219,7 +1372,7 @@ CONTAINS
recv_size, send_size, size_recv_buffer, size_send_buffer, tag
INTEGER, DIMENSION(:), POINTER :: col_indices, ncol_locals
INTEGER, DIMENSION(:, :), POINTER :: blacs2mpi
REAL(KIND=dp) :: ddot, trace
REAL(KIND=dp) :: trace
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: buffer_recv, buffer_send
TYPE(cp_blacs_env_type), POINTER :: context
TYPE(cp_para_env_type), POINTER :: para_env
@ -1237,16 +1390,18 @@ CONTAINS
tag = 42
DO col_local = 1, ncol_local
col_global = col_indices(col_local)
DO ij_counter = 1, num_ij_pairs
iocc = MAX(1, col_global - 1)/virtual + 1
avirt = col_global - (iocc - 1)*virtual
my_i = pair_list(1, ij_counter)
my_j = pair_list(2, ij_counter)
DO ij_counter = 1, num_ij_pairs
trace = 0.0_dp
my_i = pair_list(1, ij_counter)
my_j = pair_list(2, ij_counter)
DO col_local = 1, ncol_local
col_global = col_indices(col_local)
iocc = MAX(1, col_global - 1)/virtual + 1
avirt = col_global - (iocc - 1)*virtual
IF (iocc /= my_j) CYCLE
pcol = cp_fm_indxg2p((my_i - 1)*virtual + avirt, ncol_block, 0, first_p_pos_col, num_pe_col)
@ -1254,12 +1409,12 @@ CONTAINS
my_col_local = cp_fm_indxg2l((my_i - 1)*virtual + avirt, ncol_block, 0, first_p_pos_col, num_pe_col)
trace = ddot(nrow_local, fm_mat_S%local_data(:, my_col_local), 1, fm_work_iaP%local_data(:, col_local), 1)
P_ij(ij_counter) = P_ij(ij_counter) &
- trace*sinh_over_x(0.5_dp*(Eigenval(my_i) - Eigenval(my_j))*omega)*omega*weight
trace = trace + accurate_dot_product_2(fm_mat_S%local_data(:, my_col_local), fm_work_iaP%local_data(:, col_local), &
dot_blksize)
END DO
P_ij(ij_counter) = P_ij(ij_counter) - trace*sinh_over_x(0.5_dp*(Eigenval(my_i) - Eigenval(my_j))*omega)*omega*weight
END DO
IF (num_pe_col > 1) THEN
@ -1322,7 +1477,8 @@ CONTAINS
my_col_local = cp_fm_indxg2l((my_i - 1)*virtual + avirt, ncol_block, 0, first_p_pos_col, num_pe_col)
trace = trace + ddot(nrow_local, fm_mat_S%local_data(:, my_col_local), 1, buffer_recv(:, col_local), 1)
trace = trace + accurate_dot_product_2(fm_mat_S%local_data(:, my_col_local), buffer_recv(:, col_local), &
dot_blksize)
END DO
P_ij(ij_counter) = P_ij(ij_counter) &
@ -1354,8 +1510,10 @@ CONTAINS
!> \param weight ...
!> \param index2send ...
!> \param index2recv ...
!> \param dot_blksize ...
! **************************************************************************************************
SUBROUTINE calc_Pab_degen(fm_work_iaP, fm_mat_S, pair_list, virtual, P_ab, Eigenval, omega, weight, index2send, index2recv)
SUBROUTINE calc_Pab_degen(fm_work_iaP, fm_mat_S, pair_list, virtual, P_ab, Eigenval, &
omega, weight, index2send, index2recv, dot_blksize)
TYPE(cp_fm_type), INTENT(IN) :: fm_work_iaP, fm_mat_S
INTEGER, DIMENSION(:, :), INTENT(IN) :: pair_list
INTEGER, INTENT(IN) :: virtual
@ -1363,6 +1521,7 @@ CONTAINS
REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: Eigenval
REAL(KIND=dp), INTENT(IN) :: omega, weight
TYPE(one_dim_int_array), DIMENSION(0:), INTENT(IN) :: index2send, index2recv
INTEGER, INTENT(IN) :: dot_blksize
CHARACTER(LEN=*), PARAMETER :: routineN = 'calc_Pab_degen'
@ -1372,7 +1531,7 @@ CONTAINS
proc_shift, recv_size, send_size, size_recv_buffer, size_send_buffer, tag
INTEGER, DIMENSION(:), POINTER :: col_indices, ncol_locals
INTEGER, DIMENSION(:, :), POINTER :: blacs2mpi
REAL(KIND=dp) :: ddot, trace
REAL(KIND=dp) :: trace
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: buffer_recv, buffer_send
TYPE(cp_blacs_env_type), POINTER :: context
TYPE(cp_para_env_type), POINTER :: para_env
@ -1390,28 +1549,32 @@ CONTAINS
tag = 43
DO col_local = 1, ncol_local
col_global = col_indices(col_local)
DO ab_counter = 1, num_ab_pairs
iocc = MAX(1, col_global - 1)/virtual + 1
avirt = col_global - (iocc - 1)*virtual
my_a = pair_list(1, ab_counter)
my_b = pair_list(2, ab_counter)
DO ab_counter = 1, num_ab_pairs
trace = 0.0_dp
my_a = pair_list(1, ab_counter)
my_b = pair_list(2, ab_counter)
DO col_local = 1, ncol_local
col_global = col_indices(col_local)
iocc = MAX(1, col_global - 1)/virtual + 1
avirt = col_global - (iocc - 1)*virtual
IF (avirt /= my_b) CYCLE
pcol = cp_fm_indxg2p((iocc - 1)*virtual + my_a, ncol_block, 0, first_p_pos_col, num_pe_col)
IF (pcol /= my_pcol) CYCLE
my_col_local = cp_fm_indxg2l((iocc - 1)*virtual + my_a, ncol_block, 0, first_p_pos_col, num_pe_col)
trace = ddot(nrow_local, fm_mat_S%local_data(:, my_col_local), 1, fm_work_iaP%local_data(:, col_local), 1)
P_ab(ab_counter) = P_ab(ab_counter) &
+ trace*sinh_over_x(0.5_dp*(Eigenval(my_a) - Eigenval(my_b))*omega)*omega*weight
trace = trace + accurate_dot_product_2(fm_mat_S%local_data(:, my_col_local), fm_work_iaP%local_data(:, col_local), &
dot_blksize)
END DO
P_ab(ab_counter) = P_ab(ab_counter) &
+ trace*sinh_over_x(0.5_dp*(Eigenval(my_a) - Eigenval(my_b))*omega)*omega*weight
END DO
IF (num_pe_col > 1) THEN
@ -1474,7 +1637,8 @@ CONTAINS
my_col_local = cp_fm_indxg2l((iocc - 1)*virtual + my_a, ncol_block, 0, first_p_pos_col, num_pe_col)
trace = trace + ddot(nrow_local, fm_mat_S%local_data(:, my_col_local), 1, buffer_recv(:, col_local), 1)
trace = trace + accurate_dot_product_2(fm_mat_S%local_data(:, my_col_local), buffer_recv(:, col_local), &
dot_blksize)
END DO
P_ab(ab_counter) = P_ab(ab_counter) &

5
tests/QS/regtest-ri-rpa-grad/RI_RPA_grad_MP2_H2O.inp
View file

@ -61,6 +61,11 @@
REL_CUTOFF 20
&END WFC_GPW
&END INTEGRALS
&CANONICAL_GRADIENTS
EPS_CANONICAL 0.1
# Just for testing, use it only if you need the accuracy
DOT_PRODUCT_BLKSIZE 1
&END
MEMORY 200.
NUMBER_PROC 1
&END