coradv_split_impl Subroutine

   subroutine coradv_split_impl(h, u, v, fx, fy, kec, idxCu, idyCv, &
                                nx, ny, nz, nghost, nx_phys, ny_phys, &
                                rim_w, east_w, dt, &
                                wpv_full, wpv_rim, wpv_east, &
                                wke_full, wke_rim, wke_east)
      !! Device pass: per-face PV/∇KE work split, three regions.
      !! Faces attributed to their west/south cell's region mask.
      integer, intent(in) :: nx, ny, nz, nghost, nx_phys, ny_phys
      integer, intent(in) :: rim_w, east_w
      real(wp), intent(in) :: h(nx, ny, nz)
      real(wp), intent(in) :: u(nx + 1, ny, nz)
      real(wp), intent(in) :: v(nx, ny + 1, nz)
      real(wp), intent(in) :: fx(nx + 1, ny, nz)
      real(wp), intent(in) :: fy(nx, ny + 1, nz)
      real(wp), intent(in) :: kec(nx, ny, nz)
      real(wp), intent(in) :: idxCu(nx + 1, ny), idyCv(nx, ny + 1)
      real(wp), intent(in) :: dt
      real(wp), intent(out) :: wpv_full, wpv_rim, wpv_east
      real(wp), intent(out) :: wke_full, wke_rim, wke_east

      integer :: i, j, k, i_lo, i_hi, j_lo, j_hi
      real(wp) :: h_face, f_tot, f_ke, f_pv, u_mid, w_pv, w_ke
      logical :: in_rim, in_east

      i_lo = nghost + 1
      i_hi = nghost + nx_phys
      j_lo = nghost + 1
      j_hi = nghost + ny_phys

      wpv_full = 0.0_wp
      wpv_rim = 0.0_wp
      wpv_east = 0.0_wp
      wke_full = 0.0_wp
      wke_rim = 0.0_wp
      wke_east = 0.0_wp

      ! u-faces: interior faces i in [i_lo+1, i_hi]; region from the
      ! west cell (i-1 ≥ i_lo).
      do concurrent(k=1:nz, j=j_lo:j_hi, i=i_lo + 1:i_hi) &
         reduce(+:wpv_full, wpv_rim, wpv_east, wke_full, wke_rim, wke_east)
         h_face = 0.5_wp*(h(i - 1, j, k) + h(i, j, k))
         f_tot = fx(i, j, k)
         f_ke = -(kec(i, j, k) - kec(i - 1, j, k))*idxCu(i, j)
         f_pv = f_tot - f_ke
         u_mid = u(i, j, k) - 0.5_wp*dt*f_tot
         w_pv = h_face*u_mid*dt*f_pv
         w_ke = h_face*u_mid*dt*f_ke
         wpv_full = wpv_full + w_pv
         wke_full = wke_full + w_ke
         in_rim = (i - 1 - i_lo < rim_w .or. i_hi - (i - 1) < rim_w .or. &
                   j - j_lo < rim_w .or. j_hi - j < rim_w)
         in_east = (i_hi - (i - 1) < east_w)
         if (in_rim) then
            wpv_rim = wpv_rim + w_pv
            wke_rim = wke_rim + w_ke
         end if
         if (in_east) then
            wpv_east = wpv_east + w_pv
            wke_east = wke_east + w_ke
         end if
      end do

      ! v-faces: interior faces j in [j_lo+1, j_hi]; region from the
      ! south cell.
      do concurrent(k=1:nz, j=j_lo + 1:j_hi, i=i_lo:i_hi) &
         reduce(+:wpv_full, wpv_rim, wpv_east, wke_full, wke_rim, wke_east)
         h_face = 0.5_wp*(h(i, j - 1, k) + h(i, j, k))
         f_tot = fy(i, j, k)
         f_ke = -(kec(i, j, k) - kec(i, j - 1, k))*idyCv(i, j)
         f_pv = f_tot - f_ke
         u_mid = v(i, j, k) - 0.5_wp*dt*f_tot
         w_pv = h_face*u_mid*dt*f_pv
         w_ke = h_face*u_mid*dt*f_ke
         wpv_full = wpv_full + w_pv
         wke_full = wke_full + w_ke
         in_rim = (i - i_lo < rim_w .or. i_hi - i < rim_w .or. &
                   j - 1 - j_lo < rim_w .or. j_hi - (j - 1) < rim_w)
         in_east = (i_hi - i < east_w)
         if (in_rim) then
            wpv_rim = wpv_rim + w_pv
            wke_rim = wke_rim + w_ke
         end if
         if (in_east) then
            wpv_east = wpv_east + w_pv
            wke_east = wke_east + w_ke
         end if
      end do
   end subroutine coradv_split_impl