From 9ad87fd165fc27f97560bbd96b604917e2588200 Mon Sep 17 00:00:00 2001 From: Frederick Stein <43850145+fstein93@users.noreply.github.com> Date: Thu, 10 Nov 2022 14:02:31 +0100 Subject: [PATCH] Improve numerical stability of quartically-scaling SOS-MP2/RPA gradients (#2390) * 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 --- src/common/kahan_sum.F | 81 +++- src/input_cp2k_mp2.F | 11 + src/mp2_setup.F | 1 + src/mp2_types.F | 1 + src/rpa_grad.F | 416 ++++++++++++------ .../RI_RPA_grad_MP2_H2O.inp | 5 + 6 files changed, 381 insertions(+), 134 deletions(-) diff --git a/src/common/kahan_sum.F b/src/common/kahan_sum.F index 366900624a..c55251e534 100644 --- a/src/common/kahan_sum.F +++ b/src/common/kahan_sum.F @@ -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 diff --git a/src/input_cp2k_mp2.F b/src/input_cp2k_mp2.F index cd7d63e286..7ade01aabf 100644 --- a/src/input_cp2k_mp2.F +++ b/src/input_cp2k_mp2.F @@ -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 ! ************************************************************************************************** diff --git a/src/mp2_setup.F b/src/mp2_setup.F index fe71dbfbc4..dfe9e6d9e9 100644 --- a/src/mp2_setup.F +++ b/src/mp2_setup.F @@ -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) diff --git a/src/mp2_types.F b/src/mp2_types.F index 697bce85f0..71a93325b1 100644 --- a/src/mp2_types.F +++ b/src/mp2_types.F @@ -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 diff --git a/src/rpa_grad.F b/src/rpa_grad.F index 74e59f6927..3950de28cb 100644 --- a/src/rpa_grad.F +++ b/src/rpa_grad.F @@ -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) & diff --git a/tests/QS/regtest-ri-rpa-grad/RI_RPA_grad_MP2_H2O.inp b/tests/QS/regtest-ri-rpa-grad/RI_RPA_grad_MP2_H2O.inp index ca242e0563..f40d5a3010 100644 --- a/tests/QS/regtest-ri-rpa-grad/RI_RPA_grad_MP2_H2O.inp +++ b/tests/QS/regtest-ri-rpa-grad/RI_RPA_grad_MP2_H2O.inp @@ -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