Parallelize the calculation of 2c integrals for RI with GPW using OpenMP

cp2k · May 5, 2021 · 8de0492 · 8de0492
1 parent 2703552
commit 8de0492
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/src/mp2_eri_gpw.F b/src/mp2_eri_gpw.F
@@ -52,6 +52,7 @@ MODULE mp2_eri_gpw
    USE qs_integrate_potential,          ONLY: integrate_pgf_product,&
                                               integrate_v_rspace
    USE qs_kind_types,                   ONLY: get_qs_kind,&
+                                              get_qs_kind_set,&
                                               qs_kind_type
    USE qs_ks_types,                     ONLY: qs_ks_env_type
    USE qs_neighbor_list_types,          ONLY: neighbor_list_set_p_type
@@ -179,7 +180,8 @@ SUBROUTINE mp2_eri_2c_integrate_gpw(qs_env, para_env_sub, dimen_RI, mo_coeff, my
       CHARACTER(LEN=*), PARAMETER :: routineN = 'mp2_eri_2c_integrate_gpw'
 
       INTEGER :: dir, handle, handle2, i, i_counter, iatom, igrid_level, ikind, ipgf, iset, lb(3), &
-         LLL, location(3), na1, na2, ncoa, nseta, offset, sgfa, tp(3), ub(3)
+         LLL, location(3), max_nseta, na1, na2, ncoa, nseta, offset, sgfa, tp(3), ub(3)
+      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: offset_2d
       INTEGER, DIMENSION(:), POINTER                     :: la_max, la_min, npgfa, nsgfa
       INTEGER, DIMENSION(:, :), POINTER                  :: first_sgfa
       LOGICAL                                            :: map_it_here
@@ -243,8 +245,32 @@ SUBROUTINE mp2_eri_2c_integrate_gpw(qs_env, para_env_sub, dimen_RI, mo_coeff, my
 
          CALL timestop(handle2)
 
-         ! integrate the little bastards
          offset = 0
+         ! Prepare offset ahead of OMP parallel loop
+         CALL get_qs_kind_set(qs_kind_set=qs_kind_set, maxnset=max_nseta, basis_type="RI_AUX")
+         ALLOCATE (offset_2d(max_nseta, natom))
+
+         DO iatom = 1, natom
+            ikind = kind_of(iatom)
+            CALL get_qs_kind(qs_kind=qs_kind_set(ikind), basis_set=basis_set_a, basis_type="RI_AUX")
+            nseta = basis_set_a%nset
+            nsgfa => basis_set_a%nsgf_set
+            DO iset = 1, nseta
+               offset = offset + nsgfa(iset)
+               offset_2d(iset, iatom) = offset
+            END DO
+         END DO
+
+         ! integrate the little bastards
+         !$OMP PARALLEL DO DEFAULT(NONE) &
+         !$OMP SHARED(natom, particle_set, cell, pw_env_sub, rs_v, offset_2d, &
+         !$OMP        qs_kind_set, ncoset, para_env_sub, dft_control, i_counter, &
+         !$OMP        kind_of, l_local_col) &
+         !$OMP PRIVATE(iatom, ikind, basis_set_a, first_sgfa, la_max, la_min, npgfa, &
+         !$OMP         nseta, nsgfa, rpgfa, set_radius_a, sphi_a, zeta, &
+         !$OMP         ra, iset, ncoa, I_tmp2, I_ab, igrid_level, &
+         !$OMP         map_it_here, dir, tp, lb, ub, location, ipgf, &
+         !$OMP         sgfa, na1, na2, radius, offset)
          DO iatom = 1, natom
             ikind = kind_of(iatom)
             CALL get_qs_kind(qs_kind=qs_kind_set(ikind), basis_set=basis_set_a, basis_type="RI_AUX")
@@ -271,6 +297,8 @@ SUBROUTINE mp2_eri_2c_integrate_gpw(qs_env, para_env_sub, dimen_RI, mo_coeff, my
                ALLOCATE (I_ab(nsgfa(iset), 1))
                I_ab = 0.0_dp
 
+               offset = offset_2d(iset, iatom)
+
                igrid_level = gaussian_gridlevel(pw_env_sub%gridlevel_info, MINVAL(zeta(:, iset)))
                map_it_here = .FALSE.
                IF (.NOT. ALL(rs_v(igrid_level)%rs_grid%desc%perd == 1)) THEN
@@ -298,8 +326,6 @@ SUBROUTINE mp2_eri_2c_integrate_gpw(qs_env, para_env_sub, dimen_RI, mo_coeff, my
                   IF (MODULO(offset, para_env_sub%num_pe) == para_env_sub%mepos) map_it_here = .TRUE.
                ENDIF
 
-               offset = offset + nsgfa(iset)
-
                IF (map_it_here) THEN
                   DO ipgf = 1, npgfa(iset)
                      sgfa = first_sgfa(1, iset)
@@ -340,6 +366,8 @@ SUBROUTINE mp2_eri_2c_integrate_gpw(qs_env, para_env_sub, dimen_RI, mo_coeff, my
 
             END DO
          END DO
+         !$OMP END PARALLEL DO
+         DEALLOCATE (offset_2d)
 
          DO i = 1, SIZE(rs_v)
             CALL rs_grid_release(rs_v(i)%rs_grid)

diff --git a/src/pw_env/gaussian_gridlevels.F b/src/pw_env/gaussian_gridlevels.F
@@ -152,8 +152,10 @@ FUNCTION gaussian_gridlevel(gridlevel_info, exponent) RESULT(gridlevel)
             gridlevel = i
          ENDIF
       ENDDO
+      !$OMP CRITICAL(add_counts)
       gridlevel_info%total_count = gridlevel_info%total_count + 1
       gridlevel_info%count(gridlevel) = gridlevel_info%count(gridlevel) + 1
+      !$OMP END CRITICAL(add_counts)
 
    END FUNCTION gaussian_gridlevel