Coriolis: Update CAu and CAv

This commit is probably larger then it should but, but it appears to
correctly compute CAu and CAv on the GPU.

Array updates from GPU have been removed where no longer needed.

src/core/MOM_CoriolisAdv.F90

@@ -296,20 +296,22 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
   !$omp target enter data map(alloc: hArea_u, hArea_v)
   !$omp target enter data map(alloc: rel_vort, abs_vort, q, Ih_q)
   !$omp target enter data map(alloc: a, b, c, d, ep_u, ep_v)
+  !$omp target enter data map(alloc: KE, KEx, KEy)
   ! TODO: These Stokes_VF fields seem associated with diagnostics
   !$omp target enter data map(alloc: dvSdx, duSdy, stk_vort, qS) if (Stokes_VF)
+  !$omp target enter data map(alloc: CAuS, CAvS) if (Stokes_VF)
   !$omp target enter data map(alloc: uh_center, vh_center) if (CS%Coriolis_En_Dis)
   ! TODO: May also need SADOURNEY75_ENERGY
   !$omp target enter data map(alloc: uh_min, vh_min) if (CS%Coriolis_En_Dis)
   !$omp target enter data map(alloc: uh_max, vh_max) if (CS%Coriolis_En_Dis)
 
-  !$omp target enter data map(alloc: KE, KEx, KEy)
-
   ! Diagnostics
   !$omp target enter data map(alloc: RV) if (CS%id_RV > 0)
   !$omp target enter data map(alloc: PV) if (CS%id_PV > 0)
   !$omp target enter data map(alloc: q2) &
   !$omp   if(associated(AD%rv_x_u) .or. associated(AD%rv_x_v))
+  !$omp target enter data map(alloc: AD%gradKEu) if (associated(AD%gradKEu))
+  !$omp target enter data map(alloc: AD%gradKEv) if (associated(AD%gradKEv))
 
   ! TODO: Do this outside of the function
   !$omp target enter data map(to: u, v, h, uh, vh)
@@ -715,8 +717,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
 
     ! Calculate KE and the gradient of KE
     call gradKE(u, v, h, KE, KEx, KEy, k, OBC, G, GV, US, CS)
-
-    !$omp target update from(KE, KEx, KEy)
+    ! TODO: Can KE be removed from this function?
 
     !$omp target
     ! Calculate the tendencies of zonal velocity due to the Coriolis
@@ -799,30 +800,21 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
         endif
       enddo ; enddo
     endif
-    !$omp end target
-    !$omp target update from(CAu(:,:,k))
-
-    !$omp target update from(abs_vort, q)
-    !$omp target update from(a, b, c, d, ep_u, ep_v)
-    !$omp target update from(uh_min, vh_min) if (CS%Coriolis_En_Dis)
-    !$omp target update from(uh_max, vh_max) if (CS%Coriolis_En_Dis)
-
-    ! Diagnostics
-    !$omp target update from(RV) if (CS%id_RV > 0)
-    !$omp target update from(PV) if (CS%id_PV > 0)
-    !$omp target update from(q2) &
-    !$omp     if(associated(AD%rv_x_u) .or. associated(AD%rv_x_v))
 
     ! Add in the additional terms with Arakawa & Lamb.
     if ((CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. &
-        (CS%Coriolis_Scheme == AL_BLEND)) then ; do j=js,je ; do I=Isq,Ieq
-      CAu(I,j,k) = CAu(I,j,k) + &
-            ((ep_u(i,j)*uh(I-1,j,k)) - (ep_u(i+1,j)*uh(I+1,j,k))) * G%IdxCu(I,j)
-    enddo ; enddo ; endif
+        (CS%Coriolis_Scheme == AL_BLEND)) then
+      !$omp parallel loop collapse(2)
+      do j=js,je ; do I=Isq,Ieq
+        CAu(I,j,k) = CAu(I,j,k) + &
+              ((ep_u(i,j)*uh(I-1,j,k)) - (ep_u(i+1,j)*uh(I+1,j,k))) * G%IdxCu(I,j)
+      enddo ; enddo
+    endif
 
     if (Stokes_VF) then
       if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then
         ! Computing the diagnostic Stokes contribution to CAu
+        !$omp parallel loop collapse(2)
         do j=js,je ; do I=Isq,Ieq
           CAuS(I,j,k) = 0.25 * &
                 ((qS(I,J) * (vh(i+1,J,k) + vh(i,J,k))) + &
@@ -832,6 +824,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     endif
 
     if (CS%bound_Coriolis) then
+      !$omp parallel loop collapse(2)
       do j=js,je ; do I=Isq,Ieq
         fv1 = abs_vort(I,J) * v(i+1,J,k)
         fv2 = abs_vort(I,J) * v(i,J,k)
@@ -847,11 +840,13 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     endif
 
     ! Term - d(KE)/dx.
+    !$omp parallel loop collapse(2)
     do j=js,je ; do I=Isq,Ieq
       CAu(I,j,k) = CAu(I,j,k) - KEx(I,j)
     enddo ; enddo
 
     if (associated(AD%gradKEu)) then
+      !$omp parallel loop collapse(2)
       do j=js,je ; do I=Isq,Ieq
         AD%gradKEu(I,j,k) = -KEx(I,j)
       enddo ; enddo
@@ -863,6 +858,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     if (CS%Coriolis_Scheme == SADOURNY75_ENERGY) then
       if (CS%Coriolis_En_Dis) then
         ! Energy dissipating biased scheme, Hallberg 200x
+        !$omp parallel loop collapse(2)
         do J=Jsq,Jeq ; do i=is,ie
           if (q(I-1,J)*v(i,J,k) == 0.0) then
             temp1 = q(I-1,J) * ( (uh_max(i-1,j)+uh_max(i-1,j+1)) &
@@ -884,13 +880,15 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
         enddo ; enddo
       else
         ! Energy conserving scheme, Sadourny 1975
+        !$omp parallel loop collapse(2)
         do J=Jsq,Jeq ; do i=is,ie
           CAv(i,J,k) = - 0.25* &
               ((q(I-1,J)*(uh(I-1,j,k) + uh(I-1,j+1,k))) + &
                (q(I,J)*(uh(I,j,k) + uh(I,j+1,k)))) * G%IdyCv(i,J)
         enddo ; enddo
       endif
     elseif (CS%Coriolis_Scheme == SADOURNY75_ENSTRO) then
+      !$omp parallel loop collapse(2)
       do J=Jsq,Jeq ; do i=is,ie
         CAv(i,J,k) = -0.125 * (G%IdyCv(i,J) * (q(I-1,J) + q(I,J))) * &
                      ((uh(I-1,j,k) + uh(I-1,j+1,k)) + (uh(I,j,k) + uh(I,j+1,k)))
@@ -899,6 +897,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
             (CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. &
             (CS%Coriolis_Scheme == AL_BLEND)) then
       ! (Global) Energy and (Local) Enstrophy conserving, Arakawa & Hsu 1990
+      !$omp parallel loop collapse(2)
       do J=Jsq,Jeq ; do i=is,ie
         CAv(i,J,k) = - (((a(I-1,j)   * uh(I-1,j,k)) + &
                          (c(I,j+1)   * uh(I,j+1,k)))  &
@@ -909,6 +908,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
       ! An enstrophy conserving scheme robust to vanishing layers
       ! Note: Heffs are in lieu of h_at_u that should be returned by the
       !       continuity solver. AJA
+      !$omp parallel loop collapse(2)
       do J=Jsq,Jeq ; do i=is,ie
         Heff1 = abs(uh(I,j,k) * G%IdyCu(I,j)) / (eps_vel+abs(u(I,j,k)))
         Heff1 = max(Heff1, min(h(i,j,k),h(i+1,j,k)))
@@ -939,14 +939,18 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     endif
     ! Add in the additonal terms with Arakawa & Lamb.
     if ((CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. &
-        (CS%Coriolis_Scheme == AL_BLEND)) then ; do J=Jsq,Jeq ; do i=is,ie
-      CAv(i,J,k) = CAv(i,J,k) + &
-            ((ep_v(i,j)*vh(i,J-1,k)) - (ep_v(i,j+1)*vh(i,J+1,k))) * G%IdyCv(i,J)
-    enddo ; enddo ; endif
+        (CS%Coriolis_Scheme == AL_BLEND)) then
+      !$omp parallel loop collapse(2)
+      do J=Jsq,Jeq ; do i=is,ie
+        CAv(i,J,k) = CAv(i,J,k) + &
+              ((ep_v(i,j)*vh(i,J-1,k)) - (ep_v(i,j+1)*vh(i,J+1,k))) * G%IdyCv(i,J)
+      enddo ; enddo
+    endif
 
     if (Stokes_VF) then
       if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then
         ! Computing the diagnostic Stokes contribution to CAv
+        !$omp parallel loop collapse(2)
         do J=Jsq,Jeq ; do i=is,ie
           CAvS(I,j,k) = 0.25 * &
                 ((qS(I,J) * (uh(I,j+1,k) + uh(I,j,k))) + &
@@ -956,6 +960,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     endif
 
     if (CS%bound_Coriolis) then
+      !$omp parallel loop collapse(2)
       do J=Jsq,Jeq ; do i=is,ie
         fu1 = -abs_vort(I,J) * u(I,j+1,k)
         fu2 = -abs_vort(I,J) * u(I,j,k)
@@ -971,14 +976,24 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
     endif
 
     ! Term - d(KE)/dy.
+    !$omp parallel loop collapse(2)
     do J=Jsq,Jeq ; do i=is,ie
       CAv(i,J,k) = CAv(i,J,k) - KEy(i,J)
     enddo ; enddo
     if (associated(AD%gradKEv)) then
+      !$omp parallel loop collapse(2)
       do J=Jsq,Jeq ; do i=is,ie
         AD%gradKEv(i,J,k) = -KEy(i,J)
       enddo ; enddo
     endif
+    !$omp end target
+    !$omp target update from(CAu(:,:,k), CAv(:,:,k))
+
+    ! Diagnostics
+    !$omp target update from(RV(:,:,k)) if (CS%id_RV > 0)
+    !$omp target update from(PV(:,:,k)) if (CS%id_PV > 0)
+    !$omp target update from(q2) &
+    !$omp     if(associated(AD%rv_x_u) .or. associated(AD%rv_x_v))
 
     if (associated(AD%rv_x_u) .or. associated(AD%rv_x_v)) then
       ! Calculate the Coriolis-like acceleration due to relative vorticity.
@@ -1028,12 +1043,15 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav
   !$omp target exit data map(delete: hArea_u, hArea_v)
   !$omp target exit data map(delete: rel_vort, abs_vort, q, Ih_q)
   !$omp target exit data map(delete: a, b, c, d, ep_u, ep_v)
+  !$omp target exit data map(delete: KE, KEx, KEy)
   !$omp target exit data map(delete: dvSdx, duSdy, stk_vort, qS) if (Stokes_VF)
+  !$omp target exit data map(delete: CAuS, CAvS) if (Stokes_VF)
   !$omp target exit data map(delete: uh_center, vh_center) if (CS%Coriolis_En_Dis)
   !$omp target exit data map(delete: uh_min, vh_min) if (CS%Coriolis_En_Dis)
   !$omp target exit data map(delete: uh_max, vh_max) if (CS%Coriolis_En_Dis)
+  !$omp target exit data map(delete: AD%gradKEu) if (associated(AD%gradKEu))
+  !$omp target exit data map(delete: AD%gradKEv) if (associated(AD%gradKEv))
 
-  !$omp target exit data map(delete: KE, KEx, KEy)
 
   ! TODO: Move outside function
   !$omp target exit data map(delete: u, v, h, uh, vh)