Address OMP issue for large systems (#45)

* Remove OMP reduction for multidimensional arrays.

Signed-off-by: thfroitzheim <[email protected]>

* Change to private allocation

* Update notation

---------

Signed-off-by: thfroitzheim <[email protected]>

Author: thfroitzheim

src/multicharge/model.F90

@@ -159,11 +159,16 @@ subroutine get_amat_0d(self, mol, amat)
    integer :: iat, jat, izp, jzp
    real(wp) :: vec(3), r2, gam, tmp
 
+   ! Thread-private array for reduction
+   real(wp), allocatable :: amat_local(:, :)
+
    amat(:, :) = 0.0_wp
 
-   !$omp parallel do default(none) schedule(runtime) &
-   !$omp reduction(+:amat) shared(mol, self) &
-   !$omp private(iat, izp, jat, jzp, gam, vec, r2, tmp)
+   !$omp parallel default(none) &
+   !$omp shared(amat, mol, self) &
+   !$omp private(iat, izp, jat, jzp, gam, vec, r2, tmp, amat_local)
+   allocate(amat_local, source=amat)
+   !$omp do schedule(runtime) 
    do iat = 1, mol%nat
       izp = mol%id(iat)
       do jat = 1, iat-1
@@ -172,12 +177,18 @@ subroutine get_amat_0d(self, mol, amat)
          r2 = vec(1)**2 + vec(2)**2 + vec(3)**2
          gam = 1.0_wp / (self%rad(izp)**2 + self%rad(jzp)**2)
          tmp = erf(sqrt(r2*gam))/sqrt(r2)
-         amat(jat, iat) = amat(jat, iat) + tmp
-         amat(iat, jat) = amat(iat, jat) + tmp
+         amat_local(jat, iat) = amat_local(jat, iat) + tmp
+         amat_local(iat, jat) = amat_local(iat, jat) + tmp
       end do
       tmp = self%eta(izp) + sqrt2pi / self%rad(izp)
-      amat(iat, iat) = amat(iat, iat) + tmp
+      amat_local(iat, iat) = amat_local(iat, iat) + tmp
    end do
+   !$omp end do
+   !$omp critical (get_amat_0d_)
+   amat(:, :) = amat(:, :) + amat_local(:, :)
+   !$omp end critical (get_amat_0d_)
+   deallocate(amat_local)
+   !$omp end parallel
 
    amat(mol%nat+1, 1:mol%nat+1) = 1.0_wp
    amat(1:mol%nat+1, mol%nat+1) = 1.0_wp
@@ -196,15 +207,20 @@ subroutine get_amat_3d(self, mol, wsc, alpha, amat)
    real(wp) :: vec(3), gam, wsw, dtmp, rtmp, vol
    real(wp), allocatable :: dtrans(:, :), rtrans(:, :)
 
+   ! Thread-private array for reduction
+   real(wp), allocatable :: amat_local(:, :)
+
    amat(:, :) = 0.0_wp
 
    vol = abs(matdet_3x3(mol%lattice))
    call get_dir_trans(mol%lattice, dtrans)
    call get_rec_trans(mol%lattice, rtrans)
 
-   !$omp parallel do default(none) schedule(runtime) &
-   !$omp reduction(+:amat) shared(mol, self, wsc, dtrans, rtrans, alpha, vol) &
-   !$omp private(iat, izp, jat, jzp, gam, wsw, vec, dtmp, rtmp)
+   !$omp parallel default(none) &
+   !$omp shared(amat, mol, self, wsc, dtrans, rtrans, alpha, vol) &
+   !$omp private(iat, izp, jat, jzp, gam, wsw, vec, dtmp, rtmp, amat_local)
+   allocate(amat_local, source=amat)
+   !$omp do schedule(runtime) 
    do iat = 1, mol%nat
       izp = mol%id(iat)
       do jat = 1, iat-1
@@ -215,8 +231,8 @@ subroutine get_amat_3d(self, mol, wsc, alpha, amat)
             vec = mol%xyz(:, iat) - mol%xyz(:, jat) - wsc%trans(:, wsc%tridx(img, jat, iat))
             call get_amat_dir_3d(vec, gam, alpha, dtrans, dtmp)
             call get_amat_rec_3d(vec, vol, alpha, rtrans, rtmp)
-            amat(jat, iat) = amat(jat, iat) + (dtmp + rtmp) * wsw
-            amat(iat, jat) = amat(iat, jat) + (dtmp + rtmp) * wsw
+            amat_local(jat, iat) = amat_local(jat, iat) + (dtmp + rtmp) * wsw
+            amat_local(iat, jat) = amat_local(iat, jat) + (dtmp + rtmp) * wsw
          end do
       end do
 
@@ -226,12 +242,18 @@ subroutine get_amat_3d(self, mol, wsc, alpha, amat)
          vec = wsc%trans(:, wsc%tridx(img, iat, iat))
          call get_amat_dir_3d(vec, gam, alpha, dtrans, dtmp)
          call get_amat_rec_3d(vec, vol, alpha, rtrans, rtmp)
-         amat(iat, iat) = amat(iat, iat) + (dtmp + rtmp) * wsw
+         amat_local(iat, iat) = amat_local(iat, iat) + (dtmp + rtmp) * wsw
       end do
 
       dtmp = self%eta(izp) + sqrt2pi / self%rad(izp) - 2 * alpha / sqrtpi
-      amat(iat, iat) = amat(iat, iat) + dtmp
+      amat_local(iat, iat) = amat_local(iat, iat) + dtmp
    end do
+   !$omp end do
+   !$omp critical (get_amat_3d_)
+   amat(:, :) = amat(:, :) + amat_local(:, :)
+   !$omp end critical (get_amat_3d_)
+   deallocate(amat_local)
+   !$omp end parallel
 
    amat(mol%nat+1, 1:mol%nat+1) = 1.0_wp
    amat(1:mol%nat+1, mol%nat+1) = 1.0_wp
@@ -295,13 +317,22 @@ subroutine get_damat_0d(self, mol, qvec, dadr, dadL, atrace)
    integer :: iat, jat, izp, jzp
    real(wp) :: vec(3), r2, gam, arg, dtmp, dG(3), dS(3, 3)
 
+   ! Thread-private arrays for reduction
+   real(wp), allocatable :: atrace_local(:, :)
+   real(wp), allocatable :: dadr_local(:, :, :), dadL_local(:, :, :)
+
    atrace(:, :) = 0.0_wp
    dadr(:, :, :) = 0.0_wp
    dadL(:, :, :) = 0.0_wp
 
-   !$omp parallel do default(none) schedule(runtime) &
-   !$omp reduction(+:atrace, dadr, dadL) shared(mol, self, qvec) &
-   !$omp private(iat, izp, jat, jzp, gam, r2, vec, dG, dS, dtmp, arg)
+   !$omp parallel default(none) &
+   !$omp shared(atrace, dadr, dadL, mol, self, qvec) &
+   !$omp private(iat, izp, jat, jzp, gam, r2, vec, dG, dS, dtmp, arg) &
+   !$omp private(atrace_local, dadr_local, dadL_local)
+   allocate(atrace_local, source=atrace)
+   allocate(dadr_local, source=dadr)
+   allocate(dadL_local, source=dadL)
+   !$omp do schedule(runtime)
    do iat = 1, mol%nat
       izp = mol%id(iat)
       do jat = 1, iat-1
@@ -313,14 +344,22 @@ subroutine get_damat_0d(self, mol, qvec, dadr, dadL, atrace)
          dtmp = 2.0_wp*gam*exp(-arg)/(sqrtpi*r2)-erf(sqrt(arg))/(r2*sqrt(r2))
          dG = dtmp*vec
          dS = spread(dG, 1, 3) * spread(vec, 2, 3)
-         atrace(:, iat) = +dG*qvec(jat) + atrace(:, iat)
-         atrace(:, jat) = -dG*qvec(iat) + atrace(:, jat)
-         dadr(:, iat, jat) = +dG*qvec(iat)
-         dadr(:, jat, iat) = -dG*qvec(jat)
-         dadL(:, :, jat) = +dS*qvec(iat) + dadL(:, :, jat)
-         dadL(:, :, iat) = +dS*qvec(jat) + dadL(:, :, iat)
+         atrace_local(:, iat) = +dG*qvec(jat) + atrace_local(:, iat)
+         atrace_local(:, jat) = -dG*qvec(iat) + atrace_local(:, jat)
+         dadr_local(:, iat, jat) = +dG*qvec(iat)
+         dadr_local(:, jat, iat) = -dG*qvec(jat)
+         dadL_local(:, :, jat) = +dS*qvec(iat) + dadL_local(:, :, jat)
+         dadL_local(:, :, iat) = +dS*qvec(jat) + dadL_local(:, :, iat)
       end do
    end do
+   !$omp end do
+   !$omp critical (get_damat_0d_)
+   atrace(:, :) = atrace(:, :) + atrace_local(:, :)
+   dadr(:, :, :) = dadr(:, :, :) + dadr_local(:, :, :)
+   dadL(:, :, :) = dadL(:, :, :) + dadL_local(:, :, :)
+   !$omp end critical (get_damat_0d_)
+   deallocate(dadL_local, dadr_local, atrace_local)
+   !$omp end parallel
 
 end subroutine get_damat_0d
 
@@ -339,6 +378,10 @@ subroutine get_damat_3d(self, mol, wsc, alpha, qvec, dadr, dadL, atrace)
    real(wp) :: dGd(3), dSd(3, 3), dGr(3), dSr(3, 3)
    real(wp), allocatable :: dtrans(:, :), rtrans(:, :)
 
+   ! Thread-private arrays for reduction
+   real(wp), allocatable :: atrace_local(:, :)
+   real(wp), allocatable :: dadr_local(:, :, :), dadL_local(:, :, :)
+
    atrace(:, :) = 0.0_wp
    dadr(:, :, :) = 0.0_wp
    dadL(:, :, :) = 0.0_wp
@@ -347,11 +390,15 @@ subroutine get_damat_3d(self, mol, wsc, alpha, qvec, dadr, dadL, atrace)
    call get_dir_trans(mol%lattice, dtrans)
    call get_rec_trans(mol%lattice, rtrans)
 
-   !$omp parallel do default(none) schedule(runtime) &
-   !$omp reduction(+:atrace, dadr, dadL) &
+   !$omp parallel default(none) &
    !$omp shared(mol, self, wsc, alpha, vol, dtrans, rtrans, qvec) &
-   !$omp private(iat, izp, jat, jzp, img, gam, wsw, vec, dG, dS, &
-   !$omp& dGr, dSr, dGd, dSd)
+   !$omp shared(atrace, dadr, dadL) &
+   !$omp private(iat, izp, jat, jzp, img, gam, wsw, vec, dG, dS) & 
+   !$omp private(dGr, dSr, dGd, dSd, atrace_local, dadr_local, dadL_local)
+   allocate(atrace_local, source=atrace)
+   allocate(dadr_local, source=dadr)
+   allocate(dadL_local, source=dadL)
+   !$omp do schedule(runtime)
    do iat = 1, mol%nat
       izp = mol%id(iat)
       do jat = 1, iat-1
@@ -367,12 +414,12 @@ subroutine get_damat_3d(self, mol, wsc, alpha, qvec, dadr, dadL, atrace)
             dG = dG + (dGd + dGr) * wsw
             dS = dS + (dSd + dSr) * wsw
          end do
-         atrace(:, iat) = +dG*qvec(jat) + atrace(:, iat)
-         atrace(:, jat) = -dG*qvec(iat) + atrace(:, jat)
-         dadr(:, iat, jat) = +dG*qvec(iat) + dadr(:, iat, jat)
-         dadr(:, jat, iat) = -dG*qvec(jat) + dadr(:, jat, iat)
-         dadL(:, :, jat) = +dS*qvec(iat) + dadL(:, :, jat)
-         dadL(:, :, iat) = +dS*qvec(jat) + dadL(:, :, iat)
+         atrace_local(:, iat) = +dG*qvec(jat) + atrace_local(:, iat)
+         atrace_local(:, jat) = -dG*qvec(iat) + atrace_local(:, jat)
+         dadr_local(:, iat, jat) = +dG*qvec(iat) + dadr_local(:, iat, jat)
+         dadr_local(:, jat, iat) = -dG*qvec(jat) + dadr_local(:, jat, iat)
+         dadL_local(:, :, jat) = +dS*qvec(iat) + dadL_local(:, :, jat)
+         dadL_local(:, :, iat) = +dS*qvec(jat) + dadL_local(:, :, iat)
       end do
 
       dS(:, :) = 0.0_wp
@@ -384,8 +431,16 @@ subroutine get_damat_3d(self, mol, wsc, alpha, qvec, dadr, dadL, atrace)
          call get_damat_rec_3d(vec, vol, alpha, rtrans, dGr, dSr)
          dS = dS + (dSd + dSr) * wsw
       end do
-      dadL(:, :, iat) = +dS*qvec(iat) + dadL(:, :, iat)
+      dadL_local(:, :, iat) = +dS*qvec(iat) + dadL_local(:, :, iat)
    end do
+   !$omp end do
+   !$omp critical (get_damat_3d_)
+   atrace(:, :) = atrace(:, :) + atrace_local(:, :)
+   dadr(:, :, :) = dadr(:, :, :) + dadr_local(:, :, :)
+   dadL(:, :, :) = dadL(:, :, :) + dadL_local(:, :, :)
+   !$omp end critical (get_damat_3d_)
+   deallocate(dadL_local, dadr_local, atrace_local)
+   !$omp end parallel
 
 end subroutine get_damat_3d