diff --git a/config_src/drivers/unit_tests/test_reproducing_sum.F90 b/config_src/drivers/unit_tests/test_reproducing_sum.F90 index 2a9af42538..712c7e1e72 100644 --- a/config_src/drivers/unit_tests/test_reproducing_sum.F90 +++ b/config_src/drivers/unit_tests/test_reproducing_sum.F90 @@ -5,7 +5,7 @@ program test_reproducing_sum use MOM_coms, only : PE_here, root_PE, num_PEs, reproducing_sum -use MOM_coms, only : sum_across_PEs, max_across_PEs, max_count_prec +use MOM_coms, only : sum_across_PEs, max_across_PEs use MOM_domains, only : MOM_domain_type, create_MOM_domain, MOM_infra_init, MOM_infra_end use MOM_domains, only : MOM_define_layout use MOM_error_handler, only : MOM_error, MOM_mesg, FATAL, MOM_set_verbosity @@ -95,13 +95,9 @@ program test_reproducing_sum endif ! tot_fastR and tot_R should be identical unless too many values are summed if (abs(tot_fastR - tot_R) > 0.) then - if (n < max_count_prec) then - write(mesg,'("Mismatch between reproducing and fast reproducing sums.",4ES13.5)') & - tot_fastR, tot_R, tot_fastR - tot_R, ( tot_fastR - tot_R ) / tot_R - tests_failed = tests_failed .or. .true. - else - write(mesg,'("Too many values were summed for the fast reproducing sum to work.")') - endif + write(mesg,'("Mismatch between reproducing and fast reproducing sums.",4ES13.5)') & + tot_fastR, tot_R, tot_fastR - tot_R, ( tot_fastR - tot_R ) / tot_R + tests_failed = tests_failed .or. .true. call MOM_mesg(mesg) endif diff --git a/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 b/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 index aeca65b863..f76c38fe10 100644 --- a/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 +++ b/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 @@ -5,6 +5,12 @@ !> Wraps the MPP cpu clock functions !! !! The functions and constants should be accessed via mom_cpu_clock +!! +!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu +!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline +!! with no call-site changes. It requires nvfortran and the NVTX library +!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles +!! exactly as before and links no extra library. module MOM_cpu_clock_infra ! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA +#ifdef MOM_USE_NVTX +use nvtx, only : nvtxStartRange, nvtxEndRange +#endif implicit none ; private +#ifdef MOM_USE_NVTX +!> The largest clock handle for which an NVTX range name is retained. +integer, parameter :: MAX_NVTX_CLOCKS = 4096 +!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry +!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test +!! the same condition, so starts and ends stay balanced for handles that were never named +!! or that fall outside the table. +character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = "" +#endif + ! Public entities public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra subroutine cpu_clock_begin(id) integer, intent(in) :: id !< Handle for clock +#ifdef MOM_USE_NVTX + ! Opened before, and closed after, the mpp clock so the NVTX range encloses it. + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id))) + endif +#endif call mpp_clock_begin(id) end subroutine cpu_clock_begin @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id) integer, intent(in) :: id !< Handle for clock call mpp_clock_end(id) +#ifdef MOM_USE_NVTX + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange + endif +#endif end subroutine cpu_clock_end @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain) endif cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain) +#ifdef MOM_USE_NVTX + if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) & + nvtx_clock_names(cpu_clock_id) = name +#endif end function cpu_clock_id end module MOM_cpu_clock_infra diff --git a/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 b/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 index aeca65b863..f76c38fe10 100644 --- a/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 +++ b/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 @@ -5,6 +5,12 @@ !> Wraps the MPP cpu clock functions !! !! The functions and constants should be accessed via mom_cpu_clock +!! +!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu +!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline +!! with no call-site changes. It requires nvfortran and the NVTX library +!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles +!! exactly as before and links no extra library. module MOM_cpu_clock_infra ! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA +#ifdef MOM_USE_NVTX +use nvtx, only : nvtxStartRange, nvtxEndRange +#endif implicit none ; private +#ifdef MOM_USE_NVTX +!> The largest clock handle for which an NVTX range name is retained. +integer, parameter :: MAX_NVTX_CLOCKS = 4096 +!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry +!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test +!! the same condition, so starts and ends stay balanced for handles that were never named +!! or that fall outside the table. +character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = "" +#endif + ! Public entities public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra subroutine cpu_clock_begin(id) integer, intent(in) :: id !< Handle for clock +#ifdef MOM_USE_NVTX + ! Opened before, and closed after, the mpp clock so the NVTX range encloses it. + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id))) + endif +#endif call mpp_clock_begin(id) end subroutine cpu_clock_begin @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id) integer, intent(in) :: id !< Handle for clock call mpp_clock_end(id) +#ifdef MOM_USE_NVTX + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange + endif +#endif end subroutine cpu_clock_end @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain) endif cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain) +#ifdef MOM_USE_NVTX + if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) & + nvtx_clock_names(cpu_clock_id) = name +#endif end function cpu_clock_id end module MOM_cpu_clock_infra diff --git a/src/core/MOM_CoriolisAdv.F90 b/src/core/MOM_CoriolisAdv.F90 index 71794b16af..b54d3ff3e1 100644 --- a/src/core/MOM_CoriolisAdv.F90 +++ b/src/core/MOM_CoriolisAdv.F90 @@ -2,6 +2,8 @@ ! See the LICENSE file for licensing information. ! SPDX-License-Identifier: Apache-2.0 +#include "do_concurrent_compat.h" + !> Accelerations due to the Coriolis force and momentum advection module MOM_CoriolisAdv @@ -54,6 +56,7 @@ module MOM_CoriolisAdv !! Valid values are: !! - PV_ADV_CENTERED - centered (aka Sadourny, 75) !! - PV_ADV_UPWIND1 - upwind, first order + integer :: nkblock !< The k block size used in Coriolis/advection calculations [nondim]. real :: F_eff_max_blend !< The factor by which the maximum effective Coriolis !! acceleration from any point can be increased when !! blending different discretizations with the @@ -162,48 +165,47 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! Local variables real, dimension(SZIB_(G),SZJB_(G)) :: & + Area_q ! The sum of the ocean areas at the 4 adjacent thickness points [L2 ~> m2]. + real, dimension(SZIB_(G),SZJB_(G),merge(GV%ke,CS%nkblock,CS%nkblock==0)) :: & q, & ! Layer potential vorticity [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. qS, & ! Layer Stokes vorticity [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. Ih_q, & ! The inverse of thickness interpolated to q points [H-1 ~> m-1 or m2 kg-1]. h_q, & ! The thickness interpolated to q points [H-1 ~> m-1 or m2 kg-1]. - Area_q ! The sum of the ocean areas at the 4 adjacent thickness points [L2 ~> m2]. - - real, dimension(SZIB_(G),SZJ_(G)) :: & - a, b, c, d ! a, b, c, & d are combinations of the potential vorticities - ! surrounding an h grid point. At small scales, a = q/4, - ! b = q/4, etc. All are in [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1], - ! and use the indexing of the corresponding u point. + dvdx, dudy, & ! Contributions to the circulation around q-points [L2 T-1 ~> m2 s-1] + dvSdx, duSdy, & ! idem. for Stokes drift [L2 T-1 ~> m2 s-1] + rel_vort, & ! Relative vorticity at q-points [T-1 ~> s-1]. + abs_vort, & ! Absolute vorticity at q-points [T-1 ~> s-1]. + stk_vort, & ! Stokes vorticity at q-points [T-1 ~> s-1]. + q2 ! Relative vorticity over thickness [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. - real, dimension(SZI_(G),SZJ_(G)) :: & - Area_h, & ! The ocean area at h points [L2 ~> m2]. Area_h is used to find the - ! average thickness in the denominator of q. 0 for land points. - KE ! Kinetic energy per unit mass [L2 T-2 ~> m2 s-2], KE = (u^2 + v^2)/2. - real, dimension(SZIB_(G),SZJ_(G)) :: & + real, dimension(SZIB_(G),SZJ_(G),merge(GV%ke,CS%nkblock,CS%nkblock==0)) :: & + a, b, c, d, & ! a, b, c, & d are combinations of the potential vorticities + ! surrounding an h grid point. At small scales, a = q/4, + ! b = q/4, etc. All are in [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1], + ! and use the indexing of the corresponding u point. hArea_u, & ! The cell area weighted thickness interpolated to u points ! times the effective areas [H L2 ~> m3 or kg]. KEx, & ! The zonal gradient of Kinetic energy per unit mass [L T-2 ~> m s-2], ! KEx = d/dx KE. uh_center ! Transport based on arithmetic mean h at u-points [H L2 T-1 ~> m3 s-1 or kg s-1] - real, dimension(SZI_(G),SZJB_(G)) :: & - hArea_v, & ! The cell area weighted thickness interpolated to v points - ! times the effective areas [H L2 ~> m3 or kg]. - KEy, & ! The meridional gradient of Kinetic energy per unit mass [L T-2 ~> m s-2], - ! KEy = d/dy KE. - vh_center ! Transport based on arithmetic mean h at v-points [H L2 T-1 ~> m3 s-1 or kg s-1] + real, dimension(SZI_(G),SZJ_(G)) :: & + Area_h ! The ocean area at h points [L2 ~> m2]. Area_h is used to find the + ! average thickness in the denominator of q. 0 for land points. + real, dimension(SZI_(G),SZJ_(G),merge(GV%ke,CS%nkblock,CS%nkblock==0)) :: & + KE, & ! Kinetic energy per unit mass [L2 T-2 ~> m2 s-2], KE = (u^2 + v^2)/2. uh_min, uh_max, & ! The smallest and largest estimates of the zonal volume fluxes through ! the faces (i.e. u*h*dy) [H L2 T-1 ~> m3 s-1 or kg s-1] vh_min, vh_max, & ! The smallest and largest estimates of the meridional volume fluxes through ! the faces (i.e. v*h*dx) [H L2 T-1 ~> m3 s-1 or kg s-1] ep_u, ep_v ! Additional pseudo-Coriolis terms in the Arakawa and Lamb ! discretization [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. - real, dimension(SZIB_(G),SZJB_(G)) :: & - dvdx, dudy, & ! Contributions to the circulation around q-points [L2 T-1 ~> m2 s-1] - dvSdx, duSdy, & ! idem. for Stokes drift [L2 T-1 ~> m2 s-1] - rel_vort, & ! Relative vorticity at q-points [T-1 ~> s-1]. - abs_vort, & ! Absolute vorticity at q-points [T-1 ~> s-1]. - stk_vort, & ! Stokes vorticity at q-points [T-1 ~> s-1]. - q2 ! Relative vorticity over thickness [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. + real, dimension(SZI_(G),SZJB_(G),merge(GV%ke,CS%nkblock,CS%nkblock==0)) :: & + hArea_v, & ! The cell area weighted thickness interpolated to v points + ! times the effective areas [H L2 ~> m3 or kg]. + KEy, & ! The meridional gradient of Kinetic energy per unit mass [L T-2 ~> m s-2], + ! KEy = d/dy KE. + vh_center ! Transport based on arithmetic mean h at v-points [H L2 T-1 ~> m3 s-1 or kg s-1] real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)) :: & PV, & ! A diagnostic array of the potential vorticities [H-1 T-1 ~> m-1 s-1 or m2 kg-1 s-1]. RV ! A diagnostic array of the relative vorticities [T-1 ~> s-1]. @@ -245,7 +247,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav real :: h_tiny ! A very small thickness [H ~> m or kg m-2]. real :: UHeff, VHeff ! More temporary variables [H L2 T-1 ~> m3 s-1 or kg s-1]. real :: QUHeff,QVHeff ! More temporary variables [H L2 T-2 ~> m3 s-2 or kg s-2]. - integer :: i, j, k, n, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz + integer :: i, j, k, n, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, nkblock, k_start, k_end, kmax, kk integer :: Is_q, Ie_q, Js_q, Je_q ! The scheme-dependent range of values at which vorticity is set. logical :: Stokes_VF real :: u_v, v_u ! u_v is the u velocity at v point, v_u is the v velocity at u point [L T-1 ~> m s-1] @@ -270,6 +272,7 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB ; nz = GV%ke + nkblock = merge(GV%ke, CS%nkblock, CS%nkblock==0) vol_neglect = GV%H_subroundoff * (1e-4 * US%m_to_L)**2 area_neglect = (1e-4 * US%m_to_L)**2 eps_vel = 1.0e-10*US%m_s_to_L_T @@ -359,7 +362,17 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav !$omp target enter data map(to: pbv, pbv%por_face_areaU, pbv%por_face_areaV) & !$omp if (CS%Coriolis_En_Dis) - do k=1,nz + ! Hoist AL_BLEND k-independent scalars out of the block loop. + Fe_m2 = 0.0 ; rat_lin = 0.0 + if (CS%Coriolis_Scheme == AL_BLEND) then + Fe_m2 = CS%F_eff_max_blend - 2.0 + rat_lin = 1.5 * Fe_m2 / max(CS%wt_lin_blend, 1.0e-16) + if (CS%F_eff_max_blend <= 2.0) then ; Fe_m2 = -1. ; rat_lin = -1.0 ; endif + endif + + do k_start=1,nz,nkblock + k_end = min(k_start+nkblock-1, nz) + kmax = k_end - k_start + 1 ! Here the second order accurate layer potential vorticities, q, ! are calculated. hq is second order accurate in space. Relative ! vorticity is second order accurate everywhere with free slip b.c.s, @@ -367,48 +380,56 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! First calculate the contributions to the circulation around the q-point. if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - dvSdx(I,J) = (-Waves%us_y(i+1,J,k)*G%dyCv(i+1,J)) - & - (-Waves%us_y(i,J,k)*G%dyCv(i,J)) - duSdy(I,J) = (-Waves%us_x(I,j+1,k)*G%dxCu(I,j+1)) - & - (-Waves%us_x(I,j,k)*G%dxCu(I,j)) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + dvSdx(I,J,kk) = (-Waves%us_y(i+1,J,k)*G%dyCv(i+1,J)) - & + (-Waves%us_y(i,J,k)*G%dyCv(i,J)) + duSdy(I,J,kk) = (-Waves%us_x(I,j+1,k)*G%dxCu(I,j+1)) - & + (-Waves%us_x(I,j,k)*G%dxCu(I,j)) enddo endif if (.not. Waves%Passive_Stokes_VF) then - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - dvdx(I,J) = ((v(i+1,J,k)-Waves%us_y(i+1,J,k))*G%dyCv(i+1,J)) - & - ((v(i,J,k)-Waves%us_y(i,J,k))*G%dyCv(i,J)) - dudy(I,J) = ((u(I,j+1,k)-Waves%us_x(I,j+1,k))*G%dxCu(I,j+1)) - & - ((u(I,j,k)-Waves%us_x(I,j,k))*G%dxCu(I,j)) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + dvdx(I,J,kk) = ((v(i+1,J,k)-Waves%us_y(i+1,J,k))*G%dyCv(i+1,J)) - & + ((v(i,J,k)-Waves%us_y(i,J,k))*G%dyCv(i,J)) + dudy(I,J,kk) = ((u(I,j+1,k)-Waves%us_x(I,j+1,k))*G%dxCu(I,j+1)) - & + ((u(I,j,k)-Waves%us_x(I,j,k))*G%dxCu(I,j)) enddo else - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - dvdx(I,J) = (v(i+1,J,k)*G%dyCv(i+1,J)) - (v(i,J,k)*G%dyCv(i,J)) - dudy(I,J) = (u(I,j+1,k)*G%dxCu(I,j+1)) - (u(I,j,k)*G%dxCu(I,j)) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + dvdx(I,J,kk) = (v(i+1,J,k)*G%dyCv(i+1,J)) - (v(i,J,k)*G%dyCv(i,J)) + dudy(I,J,kk) = (u(I,j+1,k)*G%dxCu(I,j+1)) - (u(I,j,k)*G%dxCu(I,j)) enddo endif else - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - dvdx(I,J) = (v(i+1,J,k)*G%dyCv(i+1,J)) - (v(i,J,k)*G%dyCv(i,J)) - dudy(I,J) = (u(I,j+1,k)*G%dxCu(I,j+1)) - (u(I,j,k)*G%dxCu(I,j)) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + dvdx(I,J,kk) = (v(i+1,J,k)*G%dyCv(i+1,J)) - (v(i,J,k)*G%dyCv(i,J)) + dudy(I,J,kk) = (u(I,j+1,k)*G%dxCu(I,j+1)) - (u(I,j,k)*G%dxCu(I,j)) enddo endif - do concurrent (J=Js_q:Je_q, i=Is_q:Ie_q+1) - hArea_v(i,J) = 0.5*((Area_h(i,j) * h(i,j,k)) + (Area_h(i,j+1) * h(i,j+1,k))) + do concurrent (kk=1:kmax, J=Js_q:Je_q, i=Is_q:Ie_q+1) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + hArea_v(i,J,kk) = 0.5*((Area_h(i,j) * h(i,j,k)) + (Area_h(i,j+1) * h(i,j+1,k))) enddo - do concurrent (j=Js_q:Je_q+1, I=Is_q:Ie_q) - hArea_u(I,j) = 0.5*((Area_h(i,j) * h(i,j,k)) + (Area_h(i+1,j) * h(i+1,j,k))) + do concurrent (kk=1:kmax, j=Js_q:Je_q+1, I=Is_q:Ie_q) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + hArea_u(I,j,kk) = 0.5*((Area_h(i,j) * h(i,j,k)) + (Area_h(i+1,j) * h(i+1,j,k))) enddo if (CS%Coriolis_En_Dis) then - do concurrent (J=Jsq:Jeq+1, I=is-1:ie) - uh_center(I,j) = 0.5 * ((G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k)) * (h(i,j,k) + h(i+1,j,k)) + do concurrent (kk=1:kmax, J=Jsq:Jeq+1, I=is-1:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + uh_center(I,j,kk) = 0.5 * ((G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k)) * (h(i,j,k) + h(i+1,j,k)) enddo - do concurrent (J=js-1:je, i=Isq:Ieq+1) - vh_center(i,J) = 0.5 * ((G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k)) * (h(i,j,k) + h(i,j+1,k)) + do concurrent (kk=1:kmax, J=js-1:je, i=Isq:Ieq+1) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + vh_center(i,J,kk) = 0.5 * ((G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k)) * (h(i,j,k) + h(i,j+1,k)) enddo endif @@ -416,222 +437,225 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! velocity points on open boundaries. if (associated(OBC)) then !$omp target update from(Area_h) - !$omp target update from(dvdx, dudy) - !$omp target update from(hArea_u, hArea_v) - !$omp target update from(uh_center, vh_center) - - do n=1,OBC%number_of_segments - if (.not. OBC%segment(n)%on_pe) cycle - - I = OBC%segment(n)%HI%IsdB ; J = OBC%segment(n)%HI%JsdB - - if (OBC%segment(n)%is_N_or_S .and. (J >= Js_q) .and. (J <= Je_q)) then - select case (OBC%vorticity_config) - case (OBC_VORTICITY_ZERO) - do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB - dvdx(I,J) = 0. ; dudy(I,J) = 0. - enddo - case (OBC_VORTICITY_FREESLIP) - do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB - dudy(I,J) = 0. - enddo - case (OBC_VORTICITY_COMPUTED) - do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB + do k=k_start,k_end ! TODO: port (OBC GPU path not yet implemented) + kk = k - k_start + 1 + !$omp target update from(dvdx(:,:,kk), dudy(:,:,kk)) + !$omp target update from(hArea_u(:,:,kk), hArea_v(:,:,kk)) + !$omp target update from(uh_center(:,:,kk), vh_center(:,:,kk)) + + do n=1,OBC%number_of_segments + if (.not. OBC%segment(n)%on_pe) cycle + + I = OBC%segment(n)%HI%IsdB ; J = OBC%segment(n)%HI%JsdB + + if (OBC%segment(n)%is_N_or_S .and. (J >= Js_q) .and. (J <= Je_q)) then + select case (OBC%vorticity_config) + case (OBC_VORTICITY_ZERO) + do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB + dvdx(I,J,kk) = 0. ; dudy(I,J,kk) = 0. + enddo + case (OBC_VORTICITY_FREESLIP) + do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB + dudy(I,J,kk) = 0. + enddo + case (OBC_VORTICITY_COMPUTED) + do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB + if (OBC%segment(n)%direction == OBC_DIRECTION_N) then + dudy(I,J,kk) = 2.0*(OBC%segment(n)%tangential_vel(I,J,k) - u(I,j,k))*G%dxCu(I,j) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) + dudy(I,J,kk) = 2.0*(u(I,j+1,k) - OBC%segment(n)%tangential_vel(I,J,k))*G%dxCu(I,j+1) + endif + enddo + case (OBC_VORTICITY_SPECIFIED) + do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB + if (OBC%segment(n)%direction == OBC_DIRECTION_N) then + dudy(I,J,kk) = OBC%segment(n)%tangential_grad(I,J,k)*G%dxCu(I,j)*G%dyBu(I,J) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) + dudy(I,J,kk) = OBC%segment(n)%tangential_grad(I,J,k)*G%dxCu(I,j+1)*G%dyBu(I,J) + endif + enddo + end select + + ! Project thicknesses across OBC points with a no-gradient condition. + do i=max(Is_q,OBC%segment(n)%HI%isd), min(Ie_q+1,OBC%segment(n)%HI%ied) if (OBC%segment(n)%direction == OBC_DIRECTION_N) then - dudy(I,J) = 2.0*(OBC%segment(n)%tangential_vel(I,J,k) - u(I,j,k))*G%dxCu(I,j) + hArea_v(i,J,kk) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j,k) else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) - dudy(I,J) = 2.0*(u(I,j+1,k) - OBC%segment(n)%tangential_vel(I,J,k))*G%dxCu(I,j+1) + hArea_v(i,J,kk) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j+1,k) endif enddo - case (OBC_VORTICITY_SPECIFIED) - do I=OBC%segment(n)%HI%IsdB,OBC%segment(n)%HI%IedB - if (OBC%segment(n)%direction == OBC_DIRECTION_N) then - dudy(I,J) = OBC%segment(n)%tangential_grad(I,J,k)*G%dxCu(I,j)*G%dyBu(I,J) - else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) - dudy(I,J) = OBC%segment(n)%tangential_grad(I,J,k)*G%dxCu(I,j+1)*G%dyBu(I,J) + + if (CS%Coriolis_En_Dis) then + do i=max(Isq,OBC%segment(n)%HI%isd), min(Ieq+1,OBC%segment(n)%HI%ied) + if (OBC%segment(n)%direction == OBC_DIRECTION_N) then + vh_center(i,J,kk) = (G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k) * h(i,j,k) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) + vh_center(i,J,kk) = (G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k) * h(i,j+1,k) + endif + enddo + endif + elseif (OBC%segment(n)%is_E_or_W .and. (I >= Is_q) .and. (I <= Ie_q)) then + select case (OBC%vorticity_config) + case (OBC_VORTICITY_ZERO) + do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB + dvdx(I,J,kk) = 0. ; dudy(I,J,kk) = 0. + enddo + case (OBC_VORTICITY_FREESLIP) + do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB + dvdx(I,J,kk) = 0. + enddo + case (OBC_VORTICITY_COMPUTED) + do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB + if (OBC%segment(n)%direction == OBC_DIRECTION_E) then + dvdx(I,J,kk) = 2.0*(OBC%segment(n)%tangential_vel(I,J,k) - v(i,J,k))*G%dyCv(i,J) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) + dvdx(I,J,kk) = 2.0*(v(i+1,J,k) - OBC%segment(n)%tangential_vel(I,J,k))*G%dyCv(i+1,J) + endif + enddo + case (OBC_VORTICITY_SPECIFIED) + do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB + if (OBC%segment(n)%direction == OBC_DIRECTION_E) then + dvdx(I,J,kk) = OBC%segment(n)%tangential_grad(I,J,k)*G%dyCv(i,J)*G%dxBu(I,J) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) + dvdx(I,J,kk) = OBC%segment(n)%tangential_grad(I,J,k)*G%dyCv(i+1,J)*G%dxBu(I,J) + endif + enddo + end select + + ! Project thicknesses across OBC points with a no-gradient condition. + do j=max(Js_q,OBC%segment(n)%HI%jsd), min(Je_q+1,OBC%segment(n)%HI%jed) + if (OBC%segment(n)%direction == OBC_DIRECTION_E) then + hArea_u(I,j,kk) = 0.5*(Area_h(i,j) + Area_h(i+1,j)) * h(i,j,k) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) + hArea_u(I,j,kk) = 0.5*(Area_h(i,j) + Area_h(i+1,j)) * h(i+1,j,k) endif enddo - end select - - ! Project thicknesses across OBC points with a no-gradient condition. - do i=max(Is_q,OBC%segment(n)%HI%isd), min(Ie_q+1,OBC%segment(n)%HI%ied) - if (OBC%segment(n)%direction == OBC_DIRECTION_N) then - hArea_v(i,J) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j,k) - else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) - hArea_v(i,J) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j+1,k) + + if (CS%Coriolis_En_Dis) then + do j=max(Jsq,OBC%segment(n)%HI%jsd), min(Jeq+1,OBC%segment(n)%HI%jed) + if (OBC%segment(n)%direction == OBC_DIRECTION_E) then + uh_center(I,j,kk) = (G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k) * h(i,j,k) + else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) + uh_center(I,j,kk) = (G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k) * h(i+1,j,k) + endif + enddo endif - enddo + endif + enddo - if (CS%Coriolis_En_Dis) then - do i=max(Isq,OBC%segment(n)%HI%isd), min(Ieq+1,OBC%segment(n)%HI%ied) + ! Now project thicknesses across cell-corner points in the OBCs. The two + ! projections have to occur in sequence and can not be combined easily. + do n=1,OBC%number_of_segments + if (.not. OBC%segment(n)%on_pe) cycle + I = OBC%segment(n)%HI%IsdB ; J = OBC%segment(n)%HI%JsdB + if (OBC%segment(n)%is_N_or_S .and. (J >= Js_q) .and. (J <= Je_q)) then + do I = max(Is_q,OBC%segment(n)%HI%IsdB), min(Ie_q,OBC%segment(n)%HI%IedB) if (OBC%segment(n)%direction == OBC_DIRECTION_N) then - vh_center(i,J) = (G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k) * h(i,j,k) + if (Area_h(i,j) + Area_h(i+1,j) > 0.0) then + hArea_u(I,j+1,kk) = hArea_u(I,j,kk) * ((Area_h(i,j+1) + Area_h(i+1,j+1)) / & + (Area_h(i,j) + Area_h(i+1,j))) + else ; hArea_u(I,j+1,kk) = 0.0 ; endif else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) - vh_center(i,J) = (G%dx_Cv(i,J)*pbv%por_face_areaV(i,J,k)) * v(i,J,k) * h(i,j+1,k) + if (Area_h(i,j+1) + Area_h(i+1,j+1) > 0.0) then + hArea_u(I,j,kk) = hArea_u(I,j+1,kk) * ((Area_h(i,j) + Area_h(i+1,j)) / & + (Area_h(i,j+1) + Area_h(i+1,j+1))) + else ; hArea_u(I,j,kk) = 0.0 ; endif endif enddo - endif - elseif (OBC%segment(n)%is_E_or_W .and. (I >= Is_q) .and. (I <= Ie_q)) then - select case (OBC%vorticity_config) - case (OBC_VORTICITY_ZERO) - do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB - dvdx(I,J) = 0. ; dudy(I,J) = 0. - enddo - case (OBC_VORTICITY_FREESLIP) - do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB - dvdx(I,J) = 0. - enddo - case (OBC_VORTICITY_COMPUTED) - do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB + elseif (OBC%segment(n)%is_E_or_W .and. (I >= Is_q) .and. (I <= Ie_q)) then + do J = max(Js_q,OBC%segment(n)%HI%JsdB), min(Je_q,OBC%segment(n)%HI%JedB) if (OBC%segment(n)%direction == OBC_DIRECTION_E) then - dvdx(I,J) = 2.0*(OBC%segment(n)%tangential_vel(I,J,k) - v(i,J,k))*G%dyCv(i,J) + if (Area_h(i,j) + Area_h(i,j+1) > 0.0) then + hArea_v(i+1,J,kk) = hArea_v(i,J,kk) * ((Area_h(i+1,j) + Area_h(i+1,j+1)) / & + (Area_h(i,j) + Area_h(i,j+1))) + else ; hArea_v(i+1,J,kk) = 0.0 ; endif else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) - dvdx(I,J) = 2.0*(v(i+1,J,k) - OBC%segment(n)%tangential_vel(I,J,k))*G%dyCv(i+1,J) - endif - enddo - case (OBC_VORTICITY_SPECIFIED) - do J=OBC%segment(n)%HI%JsdB,OBC%segment(n)%HI%JedB - if (OBC%segment(n)%direction == OBC_DIRECTION_E) then - dvdx(I,J) = OBC%segment(n)%tangential_grad(I,J,k)*G%dyCv(i,J)*G%dxBu(I,J) - else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) - dvdx(I,J) = OBC%segment(n)%tangential_grad(I,J,k)*G%dyCv(i+1,J)*G%dxBu(I,J) - endif - enddo - end select - - ! Project thicknesses across OBC points with a no-gradient condition. - do j=max(Js_q,OBC%segment(n)%HI%jsd), min(Je_q+1,OBC%segment(n)%HI%jed) - if (OBC%segment(n)%direction == OBC_DIRECTION_E) then - hArea_u(I,j) = 0.5*(Area_h(i,j) + Area_h(i+1,j)) * h(i,j,k) - else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) - hArea_u(I,j) = 0.5*(Area_h(i,j) + Area_h(i+1,j)) * h(i+1,j,k) - endif - enddo - - if (CS%Coriolis_En_Dis) then - do j=max(Jsq,OBC%segment(n)%HI%jsd), min(Jeq+1,OBC%segment(n)%HI%jed) - if (OBC%segment(n)%direction == OBC_DIRECTION_E) then - uh_center(I,j) = (G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k) * h(i,j,k) - else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) - uh_center(I,j) = (G%dy_Cu(I,j)*pbv%por_face_areaU(I,j,k)) * u(I,j,k) * h(i+1,j,k) + hArea_v(i,J,kk) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j+1,k) + if (Area_h(i+1,j) + Area_h(i+1,j+1) > 0.0) then + hArea_v(i,J,kk) = hArea_v(i+1,J,kk) * ((Area_h(i,j) + Area_h(i,j+1)) / & + (Area_h(i+1,j) + Area_h(i+1,j+1))) + else ; hArea_v(i,J,kk) = 0.0 ; endif endif enddo endif - endif - enddo - endif - - if (associated(OBC)) then - do n=1,OBC%number_of_segments - if (.not. OBC%segment(n)%on_pe) cycle - ! Now project thicknesses across cell-corner points in the OBCs. The two - ! projections have to occur in sequence and can not be combined easily. - I = OBC%segment(n)%HI%IsdB ; J = OBC%segment(n)%HI%JsdB - if (OBC%segment(n)%is_N_or_S .and. (J >= Js_q) .and. (J <= Je_q)) then - do I = max(Is_q,OBC%segment(n)%HI%IsdB), min(Ie_q,OBC%segment(n)%HI%IedB) - if (OBC%segment(n)%direction == OBC_DIRECTION_N) then - if (Area_h(i,j) + Area_h(i+1,j) > 0.0) then - hArea_u(I,j+1) = hArea_u(I,j) * ((Area_h(i,j+1) + Area_h(i+1,j+1)) / & - (Area_h(i,j) + Area_h(i+1,j))) - else ; hArea_u(I,j+1) = 0.0 ; endif - else ! (OBC%segment(n)%direction == OBC_DIRECTION_S) - if (Area_h(i,j+1) + Area_h(i+1,j+1) > 0.0) then - hArea_u(I,j) = hArea_u(I,j+1) * ((Area_h(i,j) + Area_h(i+1,j)) / & - (Area_h(i,j+1) + Area_h(i+1,j+1))) - else ; hArea_u(I,j) = 0.0 ; endif - endif - enddo - elseif (OBC%segment(n)%is_E_or_W .and. (I >= Is_q) .and. (I <= Ie_q)) then - do J = max(Js_q,OBC%segment(n)%HI%JsdB), min(Je_q,OBC%segment(n)%HI%JedB) - if (OBC%segment(n)%direction == OBC_DIRECTION_E) then - if (Area_h(i,j) + Area_h(i,j+1) > 0.0) then - hArea_v(i+1,J) = hArea_v(i,J) * ((Area_h(i+1,j) + Area_h(i+1,j+1)) / & - (Area_h(i,j) + Area_h(i,j+1))) - else ; hArea_v(i+1,J) = 0.0 ; endif - else ! (OBC%segment(n)%direction == OBC_DIRECTION_W) - hArea_v(i,J) = 0.5 * (Area_h(i,j) + Area_h(i,j+1)) * h(i,j+1,k) - if (Area_h(i+1,j) + Area_h(i+1,j+1) > 0.0) then - hArea_v(i,J) = hArea_v(i+1,J) * ((Area_h(i,j) + Area_h(i,j+1)) / & - (Area_h(i+1,j) + Area_h(i+1,j+1))) - else ; hArea_v(i,J) = 0.0 ; endif - endif - enddo - endif - enddo + enddo - !$omp target update to(dvdx, dudy) - !$omp target update to(hArea_u, hArea_v) - !$omp target update to(uh_center, vh_center) + !$omp target update to(dvdx(:,:,kk), dudy(:,:,kk)) + !$omp target update to(hArea_u(:,:,kk), hArea_v(:,:,kk)) + !$omp target update to(uh_center(:,:,kk), vh_center(:,:,kk)) + enddo ! k=k_start,k_end OBC endif if (CS%no_slip) then - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - rel_vort(I,J) = (2.0 - G%mask2dBu(I,J)) * (dvdx(I,J) - dudy(I,J)) * G%IareaBu(I,J) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) + rel_vort(I,J,kk) = (2.0 - G%mask2dBu(I,J)) * (dvdx(I,J,kk) - dudy(I,J,kk)) * G%IareaBu(I,J) enddo if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then - do concurrent (J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) - stk_vort(I,J) = (2.0 - G%mask2dBu(I,J)) * (dvSdx(I,J) - duSdy(I,J)) * G%IareaBu(I,J) + do concurrent (kk=1:kmax, J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) + stk_vort(I,J,kk) = (2.0 - G%mask2dBu(I,J)) * (dvSdx(I,J,kk) - duSdy(I,J,kk)) * G%IareaBu(I,J) enddo endif endif else - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - rel_vort(I,J) = G%mask2dBu(I,J) * (dvdx(I,J) - dudy(I,J)) * G%IareaBu(I,J) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) + rel_vort(I,J,kk) = G%mask2dBu(I,J) * (dvdx(I,J,kk) - dudy(I,J,kk)) * G%IareaBu(I,J) enddo if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then - do concurrent (J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) - stk_vort(I,J) = (2.0 - G%mask2dBu(I,J)) * (dvSdx(I,J) - duSdy(I,J)) * G%IareaBu(I,J) + do concurrent (kk=1:kmax, J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) + stk_vort(I,J,kk) = (2.0 - G%mask2dBu(I,J)) * (dvSdx(I,J,kk) - duSdy(I,J,kk)) * G%IareaBu(I,J) enddo endif endif endif - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - abs_vort(I,J) = G%CoriolisBu(I,J) + rel_vort(I,J) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) + abs_vort(I,J,kk) = G%CoriolisBu(I,J) + rel_vort(I,J,kk) enddo - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - hArea_q = (hArea_u(I,j) + hArea_u(I,j+1)) + (hArea_v(i,J) + hArea_v(i+1,J)) - Ih_q(I,J) = Area_q(I,J) / (hArea_q + vol_neglect) - q(I,J) = abs_vort(I,J) * Ih_q(I,J) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(hArea_q)) + hArea_q = (hArea_u(I,j,kk) + hArea_u(I,j+1,kk)) + (hArea_v(i,J,kk) + hArea_v(i+1,J,kk)) + Ih_q(I,J,kk) = Area_q(I,J) / (hArea_q + vol_neglect) + q(I,J,kk) = abs_vort(I,J,kk) * Ih_q(I,J,kk) enddo ! NOTE: `h_q` is only used by WENO and was pulled out of the above loop to ! improve GPU performance, but it may need to be moved back. if (use_weno) then - do concurrent (J=Js_q:Je_q, I=Is_q:Ie_q) - hArea_q = (hArea_u(I,j) + hArea_u(I,j+1)) + (hArea_v(i,J) + hArea_v(i+1,J)) - h_q(I,J) = hArea_q / max(Area_q(I,J), area_neglect) + do concurrent (kk=1:kmax, J=Js_q:Je_q, I=Is_q:Ie_q) DO_LOCALITY(local(hArea_q)) + hArea_q = (hArea_u(I,j,kk) + hArea_u(I,j+1,kk)) + (hArea_v(i,J,kk) + hArea_v(i+1,J,kk)) + h_q(I,J,kk) = hArea_q / max(Area_q(I,J), area_neglect) enddo endif if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then - do concurrent (J=js-1:Jeq, I=is-1:Ieq) - qS(I,J) = stk_vort(I,J) * Ih_q(I,J) + do concurrent (kk=1:kmax, J=js-1:Jeq, I=is-1:Ieq) + qS(I,J,kk) = stk_vort(I,J,kk) * Ih_q(I,J,kk) enddo endif endif if (CS%id_rv > 0) then - do concurrent (J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) - RV(I,J,k) = rel_vort(I,J) + do concurrent (kk=1:kmax, J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + RV(I,J,k) = rel_vort(I,J,kk) enddo endif if (CS%id_PV > 0) then - do concurrent (J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) - PV(I,J,k) = q(I,J) + do concurrent (kk=1:kmax, J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + PV(I,J,k) = q(I,J,kk) enddo endif if (associated(AD%rv_x_v) .or. associated(AD%rv_x_u)) then - do concurrent (J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) - q2(I,J) = rel_vort(I,J) * Ih_q(I,J) + do concurrent (kk=1:kmax, J=Jsq-1:Jeq+1, I=Isq-1:Ieq+1) + q2(I,J,kk) = rel_vort(I,J,kk) * Ih_q(I,J,kk) enddo endif @@ -640,34 +664,30 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! scheme. All are defined at u grid points. if (CS%Coriolis_Scheme == ARAKAWA_HSU90) then - do concurrent (j=Jsq:Jeq+1, I=is-1:Ieq) - a(I,j) = (q(I,J) + (q(I+1,J) + q(I,J-1))) * C1_12 - d(I,j) = ((q(I,J) + q(I+1,J-1)) + q(I,J-1)) * C1_12 + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, I=is-1:Ieq) + a(I,j,kk) = (q(I,J,kk) + (q(I+1,J,kk) + q(I,J-1,kk))) * C1_12 + d(I,j,kk) = ((q(I,J,kk) + q(I+1,J-1,kk)) + q(I,J-1,kk)) * C1_12 enddo - do concurrent (j=Jsq:Jeq+1, I=Isq:Ieq) - b(I,j) = (q(I,J) + (q(I-1,J) + q(I,J-1))) * C1_12 - c(I,j) = ((q(I,J) + q(I-1,J-1)) + q(I,J-1)) * C1_12 + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, I=Isq:Ieq) + b(I,j,kk) = (q(I,J,kk) + (q(I-1,J,kk) + q(I,J-1,kk))) * C1_12 + c(I,j,kk) = ((q(I,J,kk) + q(I-1,J-1,kk)) + q(I,J-1,kk)) * C1_12 enddo elseif (CS%Coriolis_Scheme == ARAKAWA_LAMB81) then - do concurrent (j=Jsq:Jeq+1, I=Isq:Ieq+1) - a(I-1,j) = (2.0*(q(I,J) + q(I-1,J-1)) + (q(I-1,J) + q(I,J-1))) * C1_24 - d(I-1,j) = ((q(I,j) + q(I-1,J-1)) + 2.0*(q(I-1,J) + q(I,J-1))) * C1_24 - b(I,j) = ((q(I,J) + q(I-1,J-1)) + 2.0*(q(I-1,J) + q(I,J-1))) * C1_24 - c(I,j) = (2.0*(q(I,J) + q(I-1,J-1)) + (q(I-1,J) + q(I,J-1))) * C1_24 - ep_u(i,j) = ((q(I,J) - q(I-1,J-1)) + (q(I-1,J) - q(I,J-1))) * C1_24 - ep_v(i,j) = (-(q(I,J) - q(I-1,J-1)) + (q(I-1,J) - q(I,J-1))) * C1_24 + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, I=Isq:Ieq+1) + a(I-1,j,kk) = (2.0*(q(I,J,kk) + q(I-1,J-1,kk)) + (q(I-1,J,kk) + q(I,J-1,kk))) * C1_24 + d(I-1,j,kk) = ((q(I,j,kk) + q(I-1,J-1,kk)) + 2.0*(q(I-1,J,kk) + q(I,J-1,kk))) * C1_24 + b(I,j,kk) = ((q(I,J,kk) + q(I-1,J-1,kk)) + 2.0*(q(I-1,J,kk) + q(I,J-1,kk))) * C1_24 + c(I,j,kk) = (2.0*(q(I,J,kk) + q(I-1,J-1,kk)) + (q(I-1,J,kk) + q(I,J-1,kk))) * C1_24 + ep_u(i,j,kk) = ((q(I,J,kk) - q(I-1,J-1,kk)) + (q(I-1,J,kk) - q(I,J-1,kk))) * C1_24 + ep_v(i,j,kk) = (-(q(I,J,kk) - q(I-1,J-1,kk)) + (q(I-1,J,kk) - q(I,J-1,kk))) * C1_24 enddo elseif (CS%Coriolis_Scheme == AL_BLEND) then - Fe_m2 = CS%F_eff_max_blend - 2.0 - rat_lin = 1.5 * Fe_m2 / max(CS%wt_lin_blend, 1.0e-16) - - ! This allows the code to always give Sadourny Energy - if (CS%F_eff_max_blend <= 2.0) then ; Fe_m2 = -1. ; rat_lin = -1.0 ; endif - - do concurrent (j=Jsq:Jeq+1, I=Isq:Ieq+1) - min_Ihq = MIN(Ih_q(I-1,J-1), Ih_q(I,J-1), Ih_q(I-1,J), Ih_q(I,J)) - max_Ihq = MAX(Ih_q(I-1,J-1), Ih_q(I,J-1), Ih_q(I-1,J), Ih_q(I,J)) + ! Fe_m2 and rat_lin are k-independent; computed before the block loop. + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, I=Isq:Ieq+1) & + DO_LOCALITY(local(min_Ihq, max_Ihq, rat_m1, AL_wt, Sad_wt)) + min_Ihq = MIN(Ih_q(I-1,J-1,kk), Ih_q(I,J-1,kk), Ih_q(I-1,J,kk), Ih_q(I,J,kk)) + max_Ihq = MAX(Ih_q(I-1,J-1,kk), Ih_q(I,J-1,kk), Ih_q(I-1,J,kk), Ih_q(I,J,kk)) rat_m1 = 1.0e15 if (max_Ihq < 1.0e15*min_Ihq) rat_m1 = max_Ihq / min_Ihq - 1.0 ! The weights used here are designed to keep the effective Coriolis @@ -688,20 +708,20 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav Sad_wt = 1.0 - (CS%wt_lin_blend / rat_lin) * (rat_m1 - 2.0*rat_lin) else ; Sad_wt = 1.0 ; endif - a(I-1,j) = Sad_wt * 0.25 * q(I-1,J) + (1.0 - Sad_wt) * & - ( ((2.0-AL_wt)* q(I-1,J) + AL_wt*q(I,J-1)) + & - 2.0 * (q(I,J) + q(I-1,J-1)) ) * C1_24 - d(I-1,j) = Sad_wt * 0.25 * q(I-1,J-1) + (1.0 - Sad_wt) * & - ( ((2.0-AL_wt)* q(I-1,J-1) + AL_wt*q(I,J)) + & - 2.0 * (q(I-1,J) + q(I,J-1)) ) * C1_24 - b(I,j) = Sad_wt * 0.25 * q(I,J) + (1.0 - Sad_wt) * & - ( ((2.0-AL_wt)* q(I,J) + AL_wt*q(I-1,J-1)) + & - 2.0 * (q(I-1,J) + q(I,J-1)) ) * C1_24 - c(I,j) = Sad_wt * 0.25 * q(I,J-1) + (1.0 - Sad_wt) * & - ( ((2.0-AL_wt)* q(I,J-1) + AL_wt*q(I-1,J)) + & - 2.0 * (q(I,J) + q(I-1,J-1)) ) * C1_24 - ep_u(i,j) = AL_wt * ((q(I,J) - q(I-1,J-1)) + (q(I-1,J) - q(I,J-1))) * C1_24 - ep_v(i,j) = AL_wt * (-(q(I,J) - q(I-1,J-1)) + (q(I-1,J) - q(I,J-1))) * C1_24 + a(I-1,j,kk) = Sad_wt * 0.25 * q(I-1,J,kk) + (1.0 - Sad_wt) * & + ( ((2.0-AL_wt)* q(I-1,J,kk) + AL_wt*q(I,J-1,kk)) + & + 2.0 * (q(I,J,kk) + q(I-1,J-1,kk)) ) * C1_24 + d(I-1,j,kk) = Sad_wt * 0.25 * q(I-1,J-1,kk) + (1.0 - Sad_wt) * & + ( ((2.0-AL_wt)* q(I-1,J-1,kk) + AL_wt*q(I,J,kk)) + & + 2.0 * (q(I-1,J,kk) + q(I,J-1,kk)) ) * C1_24 + b(I,j,kk) = Sad_wt * 0.25 * q(I,J,kk) + (1.0 - Sad_wt) * & + ( ((2.0-AL_wt)* q(I,J,kk) + AL_wt*q(I-1,J-1,kk)) + & + 2.0 * (q(I-1,J,kk) + q(I,J-1,kk)) ) * C1_24 + c(I,j,kk) = Sad_wt * 0.25 * q(I,J-1,kk) + (1.0 - Sad_wt) * & + ( ((2.0-AL_wt)* q(I,J-1,kk) + AL_wt*q(I-1,J,kk)) + & + 2.0 * (q(I,J,kk) + q(I-1,J-1,kk)) ) * C1_24 + ep_u(i,j,kk) = AL_wt * ((q(I,J,kk) - q(I-1,J-1,kk)) + (q(I-1,J,kk) - q(I,J-1,kk))) * C1_24 + ep_v(i,j,kk) = AL_wt * (-(q(I,J,kk) - q(I-1,J-1,kk)) + (q(I-1,J,kk) - q(I,J-1,kk))) * C1_24 enddo endif @@ -710,8 +730,9 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! c1 = 1.0-1.5*RANGE ; c2 = 1.0-RANGE ; c3 = 2.0 ; slope = 0.5 c1 = 1.0-1.5*0.5 ; c2 = 1.0-0.5 ; c3 = 2.0 ; slope = 0.5 - do concurrent (j=Jsq:Jeq+1, I=is-1:ie) - uhc = uh_center(I,j) + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, I=is-1:ie) DO_LOCALITY(local(k, uhc, uhm)) + k = k_start + kk - 1 + uhc = uh_center(I,j,kk) uhm = uh(I,j,k) ! This sometimes matters with some types of open boundary conditions. if (G%dy_Cu(I,j) == 0.0) uhc = uhm @@ -726,14 +747,15 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav endif if (uhc > uhm) then - uh_min(I,j) = uhm ; uh_max(I,j) = uhc + uh_min(I,j,kk) = uhm ; uh_max(I,j,kk) = uhc else - uh_max(I,j) = uhm ; uh_min(I,j) = uhc + uh_max(I,j,kk) = uhm ; uh_min(I,j,kk) = uhc endif enddo - do concurrent (J=js-1:je, i=Isq:Ieq+1) - vhc = vh_center(i,J) + do concurrent (kk=1:kmax, J=js-1:je, i=Isq:Ieq+1) DO_LOCALITY(local(k, vhc, vhm)) + k = k_start + kk - 1 + vhc = vh_center(i,J,kk) vhm = vh(i,J,k) ! This sometimes matters with some types of open boundary conditions. if (G%dx_Cv(i,J) == 0.0) vhc = vhm @@ -748,15 +770,15 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav endif if (vhc > vhm) then - vh_min(i,J) = vhm ; vh_max(i,J) = vhc + vh_min(i,J,kk) = vhm ; vh_max(i,J,kk) = vhc else - vh_max(i,J) = vhm ; vh_min(i,J) = vhc + vh_max(i,J,kk) = vhm ; vh_min(i,J,kk) = vhc endif enddo endif ! Calculate KE and the gradient of KE - call gradKE(u(:,:,k), v(:,:,k), h(:,:,k), KE, KEx, KEy, G, GV, US, CS) + call gradKE(u, v, h, KE, KEx, KEy, k_start, k_end, nkblock, G, GV, US, CS) ! TODO: Can KE be removed from this function? ! Calculate the tendencies of zonal velocity due to the Coriolis @@ -765,51 +787,57 @@ 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 - do concurrent (j=js:je, I=Isq:Ieq) - if (q(I,J)*u(I,j,k) == 0.0) then - temp1 = q(I,J) * ( (vh_max(i,j)+vh_max(i+1,j)) & - + (vh_min(i,j)+vh_min(i+1,j)) )*0.5 - elseif (q(I,J)*u(I,j,k) < 0.0) then - temp1 = q(I,J) * (vh_max(i,j)+vh_max(i+1,j)) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k, temp1, temp2)) + k = k_start + kk - 1 + if (q(I,J,kk)*u(I,j,k) == 0.0) then + temp1 = q(I,J,kk) * ( (vh_max(i,j,kk)+vh_max(i+1,j,kk)) & + + (vh_min(i,j,kk)+vh_min(i+1,j,kk)) )*0.5 + elseif (q(I,J,kk)*u(I,j,k) < 0.0) then + temp1 = q(I,J,kk) * (vh_max(i,j,kk)+vh_max(i+1,j,kk)) else - temp1 = q(I,J) * (vh_min(i,j)+vh_min(i+1,j)) + temp1 = q(I,J,kk) * (vh_min(i,j,kk)+vh_min(i+1,j,kk)) endif - if (q(I,J-1)*u(I,j,k) == 0.0) then - temp2 = q(I,J-1) * ( (vh_max(i,j-1)+vh_max(i+1,j-1)) & - + (vh_min(i,j-1)+vh_min(i+1,j-1)) )*0.5 - elseif (q(I,J-1)*u(I,j,k) < 0.0) then - temp2 = q(I,J-1) * (vh_max(i,j-1)+vh_max(i+1,j-1)) + if (q(I,J-1,kk)*u(I,j,k) == 0.0) then + temp2 = q(I,J-1,kk) * ( (vh_max(i,j-1,kk)+vh_max(i+1,j-1,kk)) & + + (vh_min(i,j-1,kk)+vh_min(i+1,j-1,kk)) )*0.5 + elseif (q(I,J-1,kk)*u(I,j,k) < 0.0) then + temp2 = q(I,J-1,kk) * (vh_max(i,j-1,kk)+vh_max(i+1,j-1,kk)) else - temp2 = q(I,J-1) * (vh_min(i,j-1)+vh_min(i+1,j-1)) + temp2 = q(I,J-1,kk) * (vh_min(i,j-1,kk)+vh_min(i+1,j-1,kk)) endif CAu(I,j,k) = 0.25 * G%IdxCu(I,j) * (temp1 + temp2) enddo else ! Energy conserving scheme, Sadourny 1975 - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 CAu(I,j,k) = 0.25 * & - ((q(I,J) * (vh(i+1,J,k) + vh(i,J,k))) + & - (q(I,J-1) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) + ((q(I,J,kk) * (vh(i+1,J,k) + vh(i,J,k))) + & + (q(I,J-1,kk) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) enddo endif elseif (CS%Coriolis_Scheme == SADOURNY75_ENSTRO) then - do concurrent (j=js:je, I=Isq:Ieq) - CAu(I,j,k) = 0.125 * (G%IdxCu(I,j) * (q(I,J) + q(I,J-1))) * & + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAu(I,j,k) = 0.125 * (G%IdxCu(I,j) * (q(I,J,kk) + q(I,J-1,kk))) * & ((vh(i+1,J,k) + vh(i,J,k)) + (vh(i,J-1,k) + vh(i+1,J-1,k))) enddo elseif ((CS%Coriolis_Scheme == ARAKAWA_HSU90) .or. & (CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. & (CS%Coriolis_Scheme == AL_BLEND)) then ! (Global) Energy and (Local) Enstrophy conserving, Arakawa & Hsu 1990 - do concurrent (j=js:je, I=Isq:Ieq) - CAu(I,j,k) = (((a(I,j) * vh(i+1,J,k)) + (c(I,j) * vh(i,J-1,k))) + & - ((b(I,j) * vh(i,J,k)) + (d(I,j) * vh(i+1,J-1,k)))) * G%IdxCu(I,j) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAu(I,j,k) = (((a(I,j,kk) * vh(i+1,J,k)) + (c(I,j,kk) * vh(i,J-1,k))) + & + ((b(I,j,kk) * vh(i,J,k)) + (d(I,j,kk) * vh(i+1,J-1,k)))) * G%IdxCu(I,j) enddo elseif (CS%Coriolis_Scheme == ROBUST_ENSTRO) then ! An enstrophy conserving scheme robust to vanishing layers ! Note: Heffs are in lieu of h_at_v that should be returned by the ! continuity solver. AJA - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) & + DO_LOCALITY(local(k, Heff1, Heff2, Heff3, Heff4, VHeff, QVHeff)) + k = k_start + kk - 1 Heff1 = abs(vh(i,J,k) * G%IdxCv(i,J)) / (eps_vel+abs(v(i,J,k))) Heff1 = max(Heff1, min(h(i,j,k),h(i,j+1,k))) Heff1 = min(Heff1, max(h(i,j,k),h(i,j+1,k))) @@ -823,18 +851,20 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav Heff4 = max(Heff4, min(h(i+1,j-1,k),h(i+1,j,k))) Heff4 = min(Heff4, max(h(i+1,j-1,k),h(i+1,j,k))) if (CS%PV_Adv_Scheme == PV_ADV_CENTERED) then - CAu(I,j,k) = 0.5*(abs_vort(I,J)+abs_vort(I,J-1)) * & + CAu(I,j,k) = 0.5*(abs_vort(I,J,kk)+abs_vort(I,J-1,kk)) * & ((vh(i,J,k) + vh(i+1,J-1,k)) + (vh(i,J-1,k) + vh(i+1,J,k)) ) / & (h_tiny + ((Heff1+Heff4) + (Heff2+Heff3)) ) * G%IdxCu(I,j) elseif (CS%PV_Adv_Scheme == PV_ADV_UPWIND1) then VHeff = ((vh(i,J,k) + vh(i+1,J-1,k)) + (vh(i,J-1,k) + vh(i+1,J,k)) ) - QVHeff = 0.5*( ((abs_vort(I,J)+abs_vort(I,J-1))*VHeff) & - - ((abs_vort(I,J)-abs_vort(I,J-1))*abs(VHeff)) ) + QVHeff = 0.5*( ((abs_vort(I,J,kk)+abs_vort(I,J-1,kk))*VHeff) & + - ((abs_vort(I,J,kk)-abs_vort(I,J-1,kk))*abs(VHeff)) ) CAu(I,j,k) = (QVHeff / ( h_tiny + ((Heff1+Heff4) + (Heff2+Heff3)) ) ) * G%IdxCu(I,j) endif enddo elseif (CS%Coriolis_Scheme == wenovi7th_PV_ENSTRO) then - !$omp target update from(u, vh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(u(:,:,k), vh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do j=js,je ; do I=Isq,Ieq v_u = 0.25*G%IdxCu(I,j)*((vh(i+1,J,k) + vh(i,J,k)) + (vh(i,J-1,k) + vh(i+1,J-1,k))) ! check whether there is masked land points in the stencil @@ -849,8 +879,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (seventh_order == 1) then ! all values are valid, we use seventh order reconstruction u_q8(:) = (u(I,j-4:j+3,k) + u(I,j-3:j+4,k)) * 0.5 - call weno_seven_h_weight_reconstruction(abs_vort(I,J-4:J+3), & - h_q(I,J-4:J+3), & + call weno_seven_h_weight_reconstruction(abs_vort(I,J-4:J+3,kk), & + h_q(I,J-4:J+3,kk), & u_q8, & GV%H_subroundoff, v_u, q_u, cs%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) @@ -858,8 +888,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (fifth_order == 1) then ! all values are valid, we use fifth order reconstruction u_q6(:) = (u(I,j-3:j+2,k) + u(I,j-2:j+3,k)) * 0.5 - call weno_five_h_weight_reconstruction(abs_vort(I,J-3:J+2), & - h_q(I,J-3:J+2), & + call weno_five_h_weight_reconstruction(abs_vort(I,J-3:J+2,kk), & + h_q(I,J-3:J+2,kk), & u_q6, & GV%H_subroundoff, v_u, q_u, CS%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) @@ -867,24 +897,27 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (third_order == 1) then ! only the middle values are valid, we use third order reconstruction u_q4(:) = (u(I,j-2:j+1,k) + u(I,j-1:j+2,k)) * 0.5 - call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1), & - h_q(I,J-2:J+1), & + call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1,kk), & + h_q(I,J-2:J+1,kk), & u_q4, & GV%H_subroundoff, v_u, q_u, CS%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) else ! Upwind first order if (v_u>0.) then - q_u = q(I,J-1) + q_u = q(I,J-1,kk) else - q_u = q(I,J) + q_u = q(I,J,kk) endif CAu(I,j,k) = (q_u * v_u) endif enddo ; enddo - !$omp target update to(CAu) + !$omp target update to(CAu(:,:,k)) + enddo elseif (CS%Coriolis_Scheme == wenovi5th_PV_ENSTRO) then - !$omp target update from(u, vh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(u(:,:,k), vh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do j=js,je ; do I=Isq,Ieq v_u = 0.25*G%IdxCu(I,j)*((vh(i+1,J,k) + vh(i,J,k)) + (vh(i,J-1,k) + vh(i+1,J-1,k))) third_order = (G%mask2dCu(I,j-2) * G%mask2dCu(I,j-1) * G%mask2dCu(I,j) * & @@ -895,8 +928,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (fifth_order == 1) then ! all values are valid, we use fifth order reconstruction u_q6(:) = (u(I,j-3:j+2,k) + u(I,j-2:j+3,k)) * 0.5 - call weno_five_h_weight_reconstruction(abs_vort(I,J-3:J+2), & - h_q(I,J-3:J+2), & + call weno_five_h_weight_reconstruction(abs_vort(I,J-3:J+2,kk), & + h_q(I,J-3:J+2,kk), & u_q6, & GV%H_subroundoff, v_u, q_u, CS%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) @@ -904,23 +937,27 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (third_order == 1) then ! only the middle values are valid, we use third order reconstruction u_q4(:) = (u(I,j-2:j+1,k) + u(I,j-1:j+2,k)) * 0.5 - call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1), & - h_q(I,J-2:J+1), & + call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1,kk), & + h_q(I,J-2:J+1,kk), & u_q4, & GV%H_subroundoff, v_u, q_u, CS%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) else ! Upwind first order if (v_u>0.) then - q_u = q(I,J-1) + q_u = q(I,J-1,kk) else - q_u = q(I,J) + q_u = q(I,J,kk) endif CAu(I,j,k) = (q_u * v_u) endif enddo ; enddo + !$omp target update to(CAu(:,:,k)) + enddo elseif (CS%Coriolis_Scheme == wenovi3rd_PV_ENSTRO) then - !$omp target update from(u, vh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(u(:,:,k), vh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do j=js,je ; do I=Isq,Ieq v_u = 0.25*G%IdxCu(I,j)*((vh(i+1,J,k) + vh(i,J,k)) + (vh(i,J-1,k) + vh(i+1,J-1,k))) third_order = (G%mask2dCu(I,j-2) * G%mask2dCu(I,j-1) * G%mask2dCu(I,j) * & @@ -930,50 +967,54 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (third_order == 1) then ! only the middle values are valid, we use third order reconstruction u_q4(:) = (u(I,j-2:j+1,k) + u(I,j-1:j+2,k)) * 0.5 - call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1), & - h_q(I,J-2:J+1), & + call weno_three_h_weight_reconstruction(abs_vort(I,J-2:J+1,kk), & + h_q(I,J-2:J+1,kk), & u_q4, & GV%H_subroundoff, v_u, q_u, CS%weno_velocity_smooth) CAu(I,j,k) = (q_u * v_u) else ! Upwind first order if (v_u>0.) then - q_u = q(I,J-1) + q_u = q(I,J-1,kk) else - q_u = q(I,J) + q_u = q(I,J,kk) endif CAu(I,j,k) = (q_u * v_u) endif enddo ; enddo - !$omp target update to(CAu) + !$omp target update to(CAu(:,:,k)) + enddo endif ! Add in the additional terms with Arakawa & Lamb. if ((CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. & (CS%Coriolis_Scheme == AL_BLEND)) then - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 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) + ((ep_u(i,j,kk)*uh(I-1,j,k)) - (ep_u(i+1,j,kk)*uh(I+1,j,k))) * G%IdxCu(I,j) enddo endif if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then ! Computing the diagnostic Stokes contribution to CAu - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 CAuS(I,j,k) = 0.25 * & - ((qS(I,J) * (vh(i+1,J,k) + vh(i,J,k))) + & - (qS(I,J-1) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) + ((qS(I,J,kk) * (vh(i+1,J,k) + vh(i,J,k))) + & + (qS(I,J-1,kk) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) enddo endif endif if (CS%bound_Coriolis) then - do concurrent (j=js:je, I=Isq:Ieq) - fv1 = abs_vort(I,J) * v(i+1,J,k) - fv2 = abs_vort(I,J) * v(i,J,k) - fv3 = abs_vort(I,J-1) * v(i+1,J-1,k) - fv4 = abs_vort(I,J-1) * v(i,J-1,k) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k, fv1, fv2, fv3, fv4, max_fv, min_fv)) + k = k_start + kk - 1 + fv1 = abs_vort(I,J,kk) * v(i+1,J,k) + fv2 = abs_vort(I,J,kk) * v(i,J,k) + fv3 = abs_vort(I,J-1,kk) * v(i+1,J-1,k) + fv4 = abs_vort(I,J-1,kk) * v(i,J-1,k) max_fv = max(fv1, fv2, fv3, fv4) min_fv = min(fv1, fv2, fv3, fv4) @@ -984,13 +1025,15 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav endif ! Term - d(KE)/dx. - do concurrent (j=js:je, I=Isq:Ieq) - CAu(I,j,k) = CAu(I,j,k) - KEx(I,j) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAu(I,j,k) = CAu(I,j,k) - KEx(I,j,kk) enddo if (associated(AD%gradKEu)) then - do concurrent (j=js:je, I=Isq:Ieq) - AD%gradKEu(I,j,k) = -KEx(I,j) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + AD%gradKEu(I,j,k) = -KEx(I,j,kk) enddo endif @@ -1000,53 +1043,59 @@ 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 - do concurrent (J=Jsq:Jeq, 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)) & - + (uh_min(i-1,j)+uh_min(i-1,j+1)) )*0.5 - elseif (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)) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k, temp1, temp2)) + k = k_start + kk - 1 + if (q(I-1,J,kk)*v(i,J,k) == 0.0) then + temp1 = q(I-1,J,kk) * ( (uh_max(i-1,j,kk)+uh_max(i-1,j+1,kk)) & + + (uh_min(i-1,j,kk)+uh_min(i-1,j+1,kk)) )*0.5 + elseif (q(I-1,J,kk)*v(i,J,k) > 0.0) then + temp1 = q(I-1,J,kk) * (uh_max(i-1,j,kk)+uh_max(i-1,j+1,kk)) else - temp1 = q(I-1,J) * (uh_min(i-1,j)+uh_min(i-1,j+1)) + temp1 = q(I-1,J,kk) * (uh_min(i-1,j,kk)+uh_min(i-1,j+1,kk)) endif - if (q(I,J)*v(i,J,k) == 0.0) then - temp2 = q(I,J) * ( (uh_max(i,j)+uh_max(i,j+1)) & - + (uh_min(i,j)+uh_min(i,j+1)) )*0.5 - elseif (q(I,J)*v(i,J,k) > 0.0) then - temp2 = q(I,J) * (uh_max(i,j)+uh_max(i,j+1)) + if (q(I,J,kk)*v(i,J,k) == 0.0) then + temp2 = q(I,J,kk) * ( (uh_max(i,j,kk)+uh_max(i,j+1,kk)) & + + (uh_min(i,j,kk)+uh_min(i,j+1,kk)) )*0.5 + elseif (q(I,J,kk)*v(i,J,k) > 0.0) then + temp2 = q(I,J,kk) * (uh_max(i,j,kk)+uh_max(i,j+1,kk)) else - temp2 = q(I,J) * (uh_min(i,j)+uh_min(i,j+1)) + temp2 = q(I,J,kk) * (uh_min(i,j,kk)+uh_min(i,j+1,kk)) endif CAv(i,J,k) = -0.25 * G%IdyCv(i,J) * (temp1 + temp2) enddo else ! Energy conserving scheme, Sadourny 1975 - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 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) + ((q(I-1,J,kk)*(uh(I-1,j,k) + uh(I-1,j+1,k))) + & + (q(I,J,kk)*(uh(I,j,k) + uh(I,j+1,k)))) * G%IdyCv(i,J) enddo endif elseif (CS%Coriolis_Scheme == SADOURNY75_ENSTRO) then - do concurrent (J=Jsq:Jeq, i=is:ie) - CAv(i,J,k) = -0.125 * (G%IdyCv(i,J) * (q(I-1,J) + q(I,J))) * & + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAv(i,J,k) = -0.125 * (G%IdyCv(i,J) * (q(I-1,J,kk) + q(I,J,kk))) * & ((uh(I-1,j,k) + uh(I-1,j+1,k)) + (uh(I,j,k) + uh(I,j+1,k))) enddo elseif ((CS%Coriolis_Scheme == ARAKAWA_HSU90) .or. & (CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. & (CS%Coriolis_Scheme == AL_BLEND)) then ! (Global) Energy and (Local) Enstrophy conserving, Arakawa & Hsu 1990 - do concurrent (J=Jsq:Jeq, 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))) & - + ((b(I,j) * uh(I,j,k)) + & - (d(I-1,j+1) * uh(I-1,j+1,k)))) * G%IdyCv(i,J) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAv(i,J,k) = - (((a(I-1,j,kk) * uh(I-1,j,k)) + & + (c(I,j+1,kk) * uh(I,j+1,k))) & + + ((b(I,j,kk) * uh(I,j,k)) + & + (d(I-1,j+1,kk) * uh(I-1,j+1,k)))) * G%IdyCv(i,J) enddo elseif (CS%Coriolis_Scheme == ROBUST_ENSTRO) then ! 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 - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) & + DO_LOCALITY(local(k, Heff1, Heff2, Heff3, Heff4, UHeff, QUHeff)) + k = k_start + kk - 1 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))) Heff1 = min(Heff1, max(h(i,j,k),h(i+1,j,k))) @@ -1060,21 +1109,23 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav Heff4 = max(Heff4, min(h(i-1,j+1,k),h(i,j+1,k))) Heff4 = min(Heff4, max(h(i-1,j+1,k),h(i,j+1,k))) if (CS%PV_Adv_Scheme == PV_ADV_CENTERED) then - CAv(i,J,k) = - 0.5*(abs_vort(I,J)+abs_vort(I-1,J)) * & + CAv(i,J,k) = - 0.5*(abs_vort(I,J,kk)+abs_vort(I-1,J,kk)) * & ((uh(I ,j ,k)+uh(I-1,j+1,k)) + & (uh(I-1,j ,k)+uh(I ,j+1,k)) ) / & (h_tiny + ((Heff1+Heff4) +(Heff2+Heff3)) ) * G%IdyCv(i,J) elseif (CS%PV_Adv_Scheme == PV_ADV_UPWIND1) then UHeff = ((uh(I ,j ,k)+uh(I-1,j+1,k)) + & (uh(I-1,j ,k)+uh(I ,j+1,k)) ) - QUHeff = 0.5*( ((abs_vort(I,J)+abs_vort(I-1,J))*UHeff) & - - ((abs_vort(I,J)-abs_vort(I-1,J))*abs(UHeff)) ) + QUHeff = 0.5*( ((abs_vort(I,J,kk)+abs_vort(I-1,J,kk))*UHeff) & + - ((abs_vort(I,J,kk)-abs_vort(I-1,J,kk))*abs(UHeff)) ) CAv(i,J,k) = - QUHeff / & (h_tiny + ((Heff1+Heff4) +(Heff2+Heff3)) ) * G%IdyCv(i,J) endif enddo elseif (CS%Coriolis_Scheme == wenovi7th_PV_ENSTRO) then - !$omp target update from(v, uh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(v(:,:,k), uh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do J=Jsq,Jeq ; do i=is,ie u_v = 0.25*G%IdyCv(i,J)*((uh(I-1,j,k) + uh(I-1,j+1,k)) + (uh(I,j,k) + uh(I,j+1,k))) @@ -1090,8 +1141,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (seventh_order == 1) then v_q8(:) = (v(i-4:i+3,J,k) + v(i-3:i+4,J,k)) * 0.5 ! all values are valid, we use seventh order reconstruction - call weno_seven_h_weight_reconstruction(abs_vort(I-4:I+3,J), & - h_q(I-4:I+3,J), & + call weno_seven_h_weight_reconstruction(abs_vort(I-4:I+3,J,kk), & + h_q(I-4:I+3,J,kk), & v_q8, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) @@ -1099,8 +1150,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (fifth_order == 1) then v_q6(:) = (v(i-3:i+2,J,k) + v(i-2:i+3,J,k)) * 0.5 ! all values are valid, we use fifth order reconstruction - call weno_five_h_weight_reconstruction(abs_vort(I-3:I+2,J), & - h_q(I-3:I+2,J), & + call weno_five_h_weight_reconstruction(abs_vort(I-3:I+2,J,kk), & + h_q(I-3:I+2,J,kk), & v_q6, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) @@ -1108,24 +1159,27 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (third_order == 1) then v_q4(:) = (v(i-2:i+1,J,k) + v(i-1:i+2,J,k)) * 0.5 ! ! only the middle values are valid, we use third order reconstruction - call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J), & - h_q(I-2:I+1,J), & + call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J,kk), & + h_q(I-2:I+1,J,kk), & v_q4, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) else ! Upwind first order! if (u_v>0.) then - q_v = q(I-1,J) + q_v = q(I-1,J,kk) else - q_v = q(I,J) + q_v = q(I,J,kk) endif CAv(i,J,k) = - (q_v * u_v) endif enddo ; enddo - !$omp target update to(CAv) + !$omp target update to(CAv(:,:,k)) + enddo elseif (CS%Coriolis_Scheme == wenovi5th_PV_ENSTRO) then - !$omp target update from(v, uh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(v(:,:,k), uh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do J=Jsq,Jeq ; do i=is,ie u_v = 0.25*G%IdyCv(i,J)*((uh(I-1,j,k) + uh(I-1,j+1,k)) + (uh(I,j,k) + uh(I,j+1,k))) @@ -1138,8 +1192,8 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (fifth_order == 1) then v_q6(:) = (v(i-3:i+2,J,k) + v(i-2:i+3,J,k)) * 0.5 ! all values are valid, we use fifth order reconstruction - call weno_five_h_weight_reconstruction(abs_vort(I-3:I+2,J), & - h_q(I-3:I+2,J), & + call weno_five_h_weight_reconstruction(abs_vort(I-3:I+2,J,kk), & + h_q(I-3:I+2,J,kk), & v_q6, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) @@ -1147,25 +1201,28 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav elseif (third_order == 1) then v_q4(:) = (v(i-2:i+1,J,k) + v(i-1:i+2,J,k)) * 0.5 ! ! only the middle values are valid, we use third order reconstruction - call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J), & - h_q(I-2:I+1,J), & + call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J,kk), & + h_q(I-2:I+1,J,kk), & v_q4, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) else if (u_v>0.) then - q_v = q(I-1,J) + q_v = q(I-1,J,kk) else - q_v = q(I,J) + q_v = q(I,J,kk) endif CAv(i,J,k) = - (q_v * u_v) endif enddo ; enddo - !$omp target update to(CAv) + !$omp target update to(CAv(:,:,k)) + enddo elseif (CS%Coriolis_Scheme == wenovi3rd_PV_ENSTRO) then - !$omp target update from(v, uh, abs_vort, h_q, q) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + !$omp target update from(v(:,:,k), uh(:,:,k), abs_vort(:,:,kk), h_q(:,:,kk), q(:,:,kk)) do J=Jsq,Jeq ; do i=is,ie u_v = 0.25*G%IdyCv(i,J)*((uh(I-1,j,k) + uh(I-1,j+1,k)) + (uh(I,j,k) + uh(I,j+1,k))) @@ -1177,50 +1234,54 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav if (third_order == 1) then v_q4(:) = (v(i-2:i+1,J,k) + v(i-1:i+2,J,k)) * 0.5 ! ! only the middle values are valid, we use third order reconstruction - call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J), & - h_q(I-2:I+1,J), & + call weno_three_h_weight_reconstruction(abs_vort(I-2:I+1,J,kk), & + h_q(I-2:I+1,J,kk), & v_q4, & GV%H_subroundoff, u_v, q_v, CS%weno_velocity_smooth) CAv(i,J,k) = - (q_v * u_v) else if (u_v>0.) then - q_v = q(I-1,J) + q_v = q(I-1,J,kk) else - q_v = q(I,J) + q_v = q(I,J,kk) endif CAv(i,J,k) = - (q_v * u_v) endif enddo ; enddo - !$omp target update to(CAv) + !$omp target update to(CAv(:,:,k)) + enddo endif ! Add in the additonal terms with Arakawa & Lamb. if ((CS%Coriolis_Scheme == ARAKAWA_LAMB81) .or. & (CS%Coriolis_Scheme == AL_BLEND)) then - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 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) + ((ep_v(i,j,kk)*vh(i,J-1,k)) - (ep_v(i,j+1,kk)*vh(i,J+1,k))) * G%IdyCv(i,J) enddo endif if (Stokes_VF) then if (CS%id_CAuS>0 .or. CS%id_CAvS>0) then ! Computing the diagnostic Stokes contribution to CAv - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 CAvS(i,J,k) = 0.25 * & - ((qS(I,J) * (uh(I,j+1,k) + uh(I,j,k))) + & - (qS(I-1,J) * (uh(I-1,j,k) + uh(I-1,j+1,k)))) * G%IdyCv(i,J) + ((qS(I,J,kk) * (uh(I,j+1,k) + uh(I,j,k))) + & + (qS(I-1,J,kk) * (uh(I-1,j,k) + uh(I-1,j+1,k)))) * G%IdyCv(i,J) enddo endif endif if (CS%bound_Coriolis) then - do concurrent (J=Jsq:Jeq, i=is:ie) - fu1 = -abs_vort(I,J) * u(I,j+1,k) - fu2 = -abs_vort(I,J) * u(I,j,k) - fu3 = -abs_vort(I-1,J) * u(I-1,j+1,k) - fu4 = -abs_vort(I-1,J) * u(I-1,j,k) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k, fu1, fu2, fu3, fu4, max_fu, min_fu)) + k = k_start + kk - 1 + fu1 = -abs_vort(I,J,kk) * u(I,j+1,k) + fu2 = -abs_vort(I,J,kk) * u(I,j,k) + fu3 = -abs_vort(I-1,J,kk) * u(I-1,j+1,k) + fu4 = -abs_vort(I-1,J,kk) * u(I-1,j,k) max_fu = max(fu1, fu2, fu3, fu4) min_fu = min(fu1, fu2, fu3, fu4) @@ -1231,12 +1292,14 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav endif ! Term - d(KE)/dy. - do concurrent (J=Jsq:Jeq, i=is:ie) - CAv(i,J,k) = CAv(i,J,k) - KEy(i,J) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + CAv(i,J,k) = CAv(i,J,k) - KEy(i,J,kk) enddo if (associated(AD%gradKEv)) then - do concurrent (J=Jsq:Jeq, i=is:ie) - AD%gradKEv(i,J,k) = -KEy(i,J) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + AD%gradKEv(i,J,k) = -KEy(i,J,kk) enddo endif @@ -1244,43 +1307,47 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav ! Calculate the Coriolis-like acceleration due to relative vorticity. if (CS%Coriolis_Scheme == SADOURNY75_ENERGY) then if (associated(AD%rv_x_u)) then - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 AD%rv_x_u(i,J,k) = - 0.25* & - ((q2(I-1,j)*(uh(I-1,j,k) + uh(I-1,j+1,k))) + & - (q2(I,j)*(uh(I,j,k) + uh(I,j+1,k)))) * G%IdyCv(i,J) + ((q2(I-1,j,kk)*(uh(I-1,j,k) + uh(I-1,j+1,k))) + & + (q2(I,j,kk)*(uh(I,j,k) + uh(I,j+1,k)))) * G%IdyCv(i,J) enddo endif if (associated(AD%rv_x_v)) then - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 AD%rv_x_v(I,j,k) = 0.25 * & - ((q2(I,j) * (vh(i+1,J,k) + vh(i,J,k))) + & - (q2(I,j-1) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) + ((q2(I,j,kk) * (vh(i+1,J,k) + vh(i,J,k))) + & + (q2(I,j-1,kk) * (vh(i,J-1,k) + vh(i+1,J-1,k)))) * G%IdxCu(I,j) enddo endif else if (associated(AD%rv_x_u)) then - do concurrent (J=Jsq:Jeq, i=is:ie) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) DO_LOCALITY(local(k)) + k = k_start + kk - 1 AD%rv_x_u(i,J,k) = -G%IdyCv(i,J) * C1_12 * & - (((((q2(I,J) + q2(I-1,J-1)) + q2(I-1,J)) * uh(I-1,j,k)) + & - (((q2(I-1,J) + q2(I,J+1)) + q2(I,J)) * uh(I,j+1,k))) + & - ((((q2(I-1,J) + q2(I,J-1)) + q2(I,J)) * uh(I,j,k))+ & - (((q2(I,J) + q2(I-1,J+1)) + q2(I-1,J)) * uh(I-1,j+1,k)))) + (((((q2(I,J,kk) + q2(I-1,J-1,kk)) + q2(I-1,J,kk)) * uh(I-1,j,k)) + & + (((q2(I-1,J,kk) + q2(I,J+1,kk)) + q2(I,J,kk)) * uh(I,j+1,k))) + & + ((((q2(I-1,J,kk) + q2(I,J-1,kk)) + q2(I,J,kk)) * uh(I,j,k))+ & + (((q2(I,J,kk) + q2(I-1,J+1,kk)) + q2(I-1,J,kk)) * uh(I-1,j+1,k)))) enddo endif if (associated(AD%rv_x_v)) then - do concurrent (j=js:je, I=Isq:Ieq) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) DO_LOCALITY(local(k)) + k = k_start + kk - 1 AD%rv_x_v(I,j,k) = G%IdxCu(I,j) * C1_12 * & - (((((q2(I+1,J) + q2(I,J-1)) + q2(I,J)) * vh(i+1,J,k)) + & - (((q2(I-1,J-1) + q2(I,J)) + q2(I,J-1)) * vh(i,J-1,k))) + & - ((((q2(I-1,J) + q2(I,J-1)) + q2(I,J)) * vh(i,J,k)) + & - (((q2(I+1,J-1) + q2(I,J)) + q2(I,J-1)) * vh(i+1,J-1,k)))) + (((((q2(I+1,J,kk) + q2(I,J-1,kk)) + q2(I,J,kk)) * vh(i+1,J,k)) + & + (((q2(I-1,J-1,kk) + q2(I,J,kk)) + q2(I,J-1,kk)) * vh(i,J-1,k))) + & + ((((q2(I-1,J,kk) + q2(I,J-1,kk)) + q2(I,J,kk)) * vh(i,J,k)) + & + (((q2(I+1,J-1,kk) + q2(I,J,kk)) + q2(I,J-1,kk)) * vh(i+1,J-1,k)))) enddo endif endif endif - enddo ! k-loop. + enddo ! end of k_start block loop. !$omp target exit data map(delete: Area_h, Area_q) !$omp target exit data map(delete: dvdx, dudy) @@ -1350,154 +1417,166 @@ subroutine CorAdCalc(u, v, h, uh, vh, CAu, CAv, OBC, AD, G, GV, US, CS, pbv, Wav end subroutine CorAdCalc -!> Calculates the acceleration due to the gradient of kinetic energy in one layer. -subroutine gradKE(u, v, h, KE, KEx, KEy, G, GV, US, CS) - type(ocean_grid_type), intent(in) :: G !< Ocean grid structure - type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure - real, dimension(SZIB_(G),SZJ_(G)), intent(in) :: u !< Zonal velocity [L T-1 ~> m s-1] - real, dimension(SZI_(G),SZJB_(G)), intent(in) :: v !< Meridional velocity [L T-1 ~> m s-1] - real, dimension(SZI_(G),SZJ_(G)), intent(in) :: h !< Layer thickness [H ~> m or kg m-2] - real, dimension(SZI_(G),SZJ_(G)), intent(out) :: KE !< Kinetic energy per unit mass [L2 T-2 ~> m2 s-2] - real, dimension(SZIB_(G),SZJ_(G)), intent(out) :: KEx !< Zonal acceleration due to kinetic - !! energy gradient [L T-2 ~> m s-2] - real, dimension(SZI_(G),SZJB_(G)), intent(out) :: KEy !< Meridional acceleration due to kinetic - !! energy gradient [L T-2 ~> m s-2] - type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type - type(CoriolisAdv_CS), intent(in) :: CS !< Control structure for MOM_CoriolisAdv +!> Calculates the acceleration due to the gradient of kinetic energy for a k-block. +subroutine gradKE(u, v, h, KE, KEx, KEy, k_start, k_end, nkblock, G, GV, US, CS) + type(ocean_grid_type), intent(in) :: G !< Ocean grid structure + type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure + integer, intent(in) :: k_start !< First layer in the k-block [nondim]. + integer, intent(in) :: k_end !< Last layer in the k-block [nondim]. + integer, intent(in) :: nkblock !< Size of the k-block [nondim]. + real, dimension(SZIB_(G),SZJ_(G),SZK_(GV)), intent(in) :: u !< Zonal velocity [L T-1 ~> m s-1] + real, dimension(SZI_(G),SZJB_(G),SZK_(GV)), intent(in) :: v !< Meridional velocity [L T-1 ~> m s-1] + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h !< Layer thickness [H ~> m or kg m-2] + real, dimension(SZI_(G),SZJ_(G),nkblock), intent(out) :: KE !< Kinetic energy per unit mass [L2 T-2 ~> m2 s-2] + real, dimension(SZIB_(G),SZJ_(G),nkblock), intent(out) :: KEx !< Zonal acceleration due to kinetic + !! energy gradient [L T-2 ~> m s-2] + real, dimension(SZI_(G),SZJB_(G),nkblock), intent(out) :: KEy !< Meridional acceleration due to kinetic + !! energy gradient [L T-2 ~> m s-2] + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + type(CoriolisAdv_CS), intent(in) :: CS !< Control structure for MOM_CoriolisAdv ! Local variables real :: um, up, vm, vp ! Temporary variables [L T-1 ~> m s-1]. real :: um2, up2, vm2, vp2 ! Temporary variables [L2 T-2 ~> m2 s-2]. real :: um2a, up2a, vm2a, vp2a ! Temporary variables [L4 T-2 ~> m4 s-2]. real :: third_order_u, third_order_v ! Product of mask values to determine the boundary - integer :: i, j, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, n + integer :: i, j, k, kk, kmax, is, ie, js, je, Isq, Ieq, Jsq, Jeq real, parameter :: C1_12 = 1.0/12.0 ! The ratio of 1/12 [nondim] - is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke + is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec Isq = G%IscB ; Ieq = G%IecB ; Jsq = G%JscB ; Jeq = G%JecB + kmax = k_end - k_start + 1 ! Calculate KE (Kinetic energy for use in the -grad(KE) acceleration term). if (CS%KE_Scheme == KE_ARAKAWA) then ! The following calculation of Kinetic energy includes the metric terms ! identified in Arakawa & Lamb 1982 as important for KE conservation. It ! also includes the possibility of partially-blocked tracer cell faces. - do concurrent (j=Jsq:Jeq+1, i=Isq:Ieq+1) - KE(i,j) = ( ( (G%areaCu( I ,j)*(u( I ,j)*u( I ,j))) + & - (G%areaCu(I-1,j)*(u(I-1,j)*u(I-1,j))) ) + & - ( (G%areaCv(i, J )*(v(i, J )*v(i, J ))) + & - (G%areaCv(i,J-1)*(v(i,J-1)*v(i,J-1))) ) )*0.25*G%IareaT(i,j) + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, i=Isq:Ieq+1) DO_LOCALITY(local(k)) + k = k_start + kk - 1 + KE(i,j,kk) = ( ( (G%areaCu( I ,j)*(u( I ,j,k)*u( I ,j,k))) + & + (G%areaCu(I-1,j)*(u(I-1,j,k)*u(I-1,j,k))) ) + & + ( (G%areaCv(i, J )*(v(i, J ,k)*v(i, J ,k))) + & + (G%areaCv(i,J-1)*(v(i,J-1,k)*v(i,J-1,k))) ) )*0.25*G%IareaT(i,j) enddo elseif (CS%KE_Scheme == KE_SIMPLE_GUDONOV) then ! The following discretization of KE is based on the one-dimensional Gudonov ! scheme which does not take into account any geometric factors - do concurrent (j=Jsq:Jeq+1, i=Isq:Ieq+1) - up = 0.5*( u(I-1,j) + ABS( u(I-1,j) ) ) ; up2 = up*up - um = 0.5*( u( I ,j) - ABS( u( I ,j) ) ) ; um2 = um*um - vp = 0.5*( v(i,J-1) + ABS( v(i,J-1) ) ) ; vp2 = vp*vp - vm = 0.5*( v(i, J ) - ABS( v(i, J ) ) ) ; vm2 = vm*vm - KE(i,j) = ( max(up2,um2) + max(vp2,vm2) ) *0.5 + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, i=Isq:Ieq+1) & + DO_LOCALITY(local(k, up, um, vp, vm, up2, um2, vp2, vm2)) + k = k_start + kk - 1 + up = 0.5*( u(I-1,j,k) + ABS( u(I-1,j,k) ) ) ; up2 = up*up + um = 0.5*( u( I ,j,k) - ABS( u( I ,j,k) ) ) ; um2 = um*um + vp = 0.5*( v(i,J-1,k) + ABS( v(i,J-1,k) ) ) ; vp2 = vp*vp + vm = 0.5*( v(i, J ,k) - ABS( v(i, J ,k) ) ) ; vm2 = vm*vm + KE(i,j,kk) = ( max(up2,um2) + max(vp2,vm2) ) *0.5 enddo elseif (CS%KE_Scheme == KE_GUDONOV) then ! The following discretization of KE is based on the one-dimensional Gudonov ! scheme but has been adapted to take horizontal grid factors into account - do concurrent (j=Jsq:Jeq+1, i=Isq:Ieq+1) - up = 0.5*( u(I-1,j) + ABS( u(I-1,j) ) ) ; up2a = up*up*G%areaCu(I-1,j) - um = 0.5*( u( I ,j) - ABS( u( I ,j) ) ) ; um2a = um*um*G%areaCu( I ,j) - vp = 0.5*( v(i,J-1) + ABS( v(i,J-1) ) ) ; vp2a = vp*vp*G%areaCv(i,J-1) - vm = 0.5*( v(i, J ) - ABS( v(i, J ) ) ) ; vm2a = vm*vm*G%areaCv(i, J ) - KE(i,j) = ( max(um2a,up2a) + max(vm2a,vp2a) )*0.5*G%IareaT(i,j) + do concurrent (kk=1:kmax, j=Jsq:Jeq+1, i=Isq:Ieq+1) & + DO_LOCALITY(local(k, up, um, vp, vm, up2a, um2a, vp2a, vm2a)) + k = k_start + kk - 1 + up = 0.5*( u(I-1,j,k) + ABS( u(I-1,j,k) ) ) ; up2a = up*up*G%areaCu(I-1,j) + um = 0.5*( u( I ,j,k) - ABS( u( I ,j,k) ) ) ; um2a = um*um*G%areaCu( I ,j) + vp = 0.5*( v(i,J-1,k) + ABS( v(i,J-1,k) ) ) ; vp2a = vp*vp*G%areaCv(i,J-1) + vm = 0.5*( v(i, J ,k) - ABS( v(i, J ,k) ) ) ; vm2a = vm*vm*G%areaCv(i, J ) + KE(i,j,kk) = ( max(um2a,up2a) + max(vm2a,vp2a) )*0.5*G%IareaT(i,j) enddo elseif (CS%KE_Scheme == KE_UP3) then ! The following discretization of KE is based on the one-dimensional third-order ! upwind scheme which does not take horizontal grid factors into account - ! TODO: GPU data tansfers? - if (CS%KE_use_limiter) then - do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 - ! compute the masking to make sure that inland values are not used - third_order_u = (G%mask2dCu(I-2,j) * G%mask2dCu(I-1,j)* & - G%mask2dCu(I,j) * G%mask2dCu(I+1,j)) - - if (third_order_u == 1) then - up = (7.0 * (u(I-1,j) + u(I,j)) - (u(I-2,j) + u(I+1,j))) * C1_12 - call UP3_Koren_limiter_reconstruction(u(I-2:I+1,j), up, um) - else - up = (u(I-1,j) + u(I,j))*0.5 - if (up>0.) then - um = u(I-1,j) - elseif (up<0.) then - um = u(I,j) + do k=k_start,k_end ! TODO: port + kk = k - k_start + 1 + ! TODO: GPU data tansfers? + if (CS%KE_use_limiter) then + do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 + ! compute the masking to make sure that inland values are not used + third_order_u = (G%mask2dCu(I-2,j) * G%mask2dCu(I-1,j)* & + G%mask2dCu(I,j) * G%mask2dCu(I+1,j)) + + if (third_order_u == 1) then + up = (7.0 * (u(I-1,j,k) + u(I,j,k)) - (u(I-2,j,k) + u(I+1,j,k))) * C1_12 + call UP3_Koren_limiter_reconstruction(u(I-2:I+1,j,k), up, um) else - um = up + up = (u(I-1,j,k) + u(I,j,k))*0.5 + if (up>0.) then + um = u(I-1,j,k) + elseif (up<0.) then + um = u(I,j,k) + else + um = up + endif endif - endif - third_order_v = (G%mask2dCv(i,J-2) * G%mask2dCv(i,J-1)* & - G%mask2dCv(i,J) * G%mask2dCv(i,J+1)) - if (third_order_v ==1) then - vp = (7.0 * (v(i,J-1) + v(i,J)) - (v(i,J-2) + v(i,J+1))) * C1_12 - call UP3_Koren_limiter_reconstruction(v(i,J-2:J+1), vp, vm) - else - vp = (v(i,J-1) + v(i,J))*0.5 - if (vp>0.) then - vm = v(i,J-1) - elseif (vp<0.) then - vm = v(i,J) + third_order_v = (G%mask2dCv(i,J-2) * G%mask2dCv(i,J-1)* & + G%mask2dCv(i,J) * G%mask2dCv(i,J+1)) + if (third_order_v ==1) then + vp = (7.0 * (v(i,J-1,k) + v(i,J,k)) - (v(i,J-2,k) + v(i,J+1,k))) * C1_12 + call UP3_Koren_limiter_reconstruction(v(i,J-2:J+1,k), vp, vm) else - vm = vp + vp = (v(i,J-1,k) + v(i,J,k))*0.5 + if (vp>0.) then + vm = v(i,J-1,k) + elseif (vp<0.) then + vm = v(i,J,k) + else + vm = vp + endif endif - endif - KE(i,j) = ( (um*um) + (vm*vm) )*0.5 - enddo ; enddo - else - do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 - ! compute the masking to make sure that inland values are not used - third_order_u = (G%mask2dCu(I-2,j) * G%mask2dCu(I-1,j)* & - G%mask2dCu(I,j) * G%mask2dCu(I+1,j)) - - if (third_order_u == 1) then - up = (7.0 * (u(I-1,j) + u(I,j)) - (u(I-2,j) + u(I+1,j))) * C1_12 - call UP3_reconstruction(u(I-2:I+1,j), up, um) - else - up = (u(I-1,j) + u(I,j))*0.5 - if (up>0.) then - um = u(I-1,j) - elseif (up<0.) then - um = u(I,j) + KE(i,j,kk) = ( (um*um) + (vm*vm) )*0.5 + enddo ; enddo + else + do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 + ! compute the masking to make sure that inland values are not used + third_order_u = (G%mask2dCu(I-2,j) * G%mask2dCu(I-1,j)* & + G%mask2dCu(I,j) * G%mask2dCu(I+1,j)) + + if (third_order_u == 1) then + up = (7.0 * (u(I-1,j,k) + u(I,j,k)) - (u(I-2,j,k) + u(I+1,j,k))) * C1_12 + call UP3_reconstruction(u(I-2:I+1,j,k), up, um) else - um = up + up = (u(I-1,j,k) + u(I,j,k))*0.5 + if (up>0.) then + um = u(I-1,j,k) + elseif (up<0.) then + um = u(I,j,k) + else + um = up + endif endif - endif - third_order_v = (G%mask2dCv(i,J-2) * G%mask2dCv(i,J-1)* & - G%mask2dCv(i,J) * G%mask2dCv(i,J+1)) - if (third_order_v ==1) then - vp = (7.0 * (v(i,J-1) + v(i,J)) - (v(i,J-2) + v(i,J+1))) * C1_12 - call UP3_reconstruction(v(i,J-2:J+1), vp, vm) - else - vp = (v(i,J-1) + v(i,J))*0.5 - if (vp>0.) then - vm = v(i,J-1) - elseif (vp<0.) then - vm = v(i,J) + third_order_v = (G%mask2dCv(i,J-2) * G%mask2dCv(i,J-1)* & + G%mask2dCv(i,J) * G%mask2dCv(i,J+1)) + if (third_order_v ==1) then + vp = (7.0 * (v(i,J-1,k) + v(i,J,k)) - (v(i,J-2,k) + v(i,J+1,k))) * C1_12 + call UP3_reconstruction(v(i,J-2:J+1,k), vp, vm) else - vm = vp + vp = (v(i,J-1,k) + v(i,J,k))*0.5 + if (vp>0.) then + vm = v(i,J-1,k) + elseif (vp<0.) then + vm = v(i,J,k) + else + vm = vp + endif endif - endif - KE(i,j) = ( (um*um) + (vm*vm) )*0.5 - enddo ; enddo - endif + KE(i,j,kk) = ( (um*um) + (vm*vm) )*0.5 + enddo ; enddo + endif + enddo endif ! Term - d(KE)/dx. - do concurrent (j=js:je, I=Isq:Ieq) - KEx(I,j) = (KE(i+1,j) - KE(i,j)) * G%IdxCu_OBCmask(I,j) + do concurrent (kk=1:kmax, j=js:je, I=Isq:Ieq) + KEx(I,j,kk) = (KE(i+1,j,kk) - KE(i,j,kk)) * G%IdxCu_OBCmask(I,j) enddo ! Term - d(KE)/dy. - do concurrent (J=Jsq:Jeq, i=is:ie) - KEy(i,J) = (KE(i,j+1) - KE(i,j)) * G%IdyCv_OBCmask(i,J) + do concurrent (kk=1:kmax, J=Jsq:Jeq, i=is:ie) + KEy(i,J,kk) = (KE(i,j+1,kk) - KE(i,j,kk)) * G%IdyCv_OBCmask(i,J) enddo end subroutine gradKE @@ -2012,6 +2091,11 @@ subroutine CoriolisAdv_init(Time, G, GV, US, param_file, diag, AD, CS) character(len=400) :: mesg logical :: use_weno integer :: isd, ied, jsd, jed, IsdB, IedB, JsdB, JedB, nz +#ifdef __NVCOMPILER_OPENMP_GPU + integer, parameter :: default_nkblock = 0 +#else + integer, parameter :: default_nkblock = 1 +#endif isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed ; nz = GV%ke IsdB = G%IsdB ; IedB = G%IedB ; JsdB = G%JsdB ; JedB = G%JedB @@ -2021,6 +2105,11 @@ subroutine CoriolisAdv_init(Time, G, GV, US, param_file, diag, AD, CS) ! Read all relevant parameters and write them to the model log. call log_version(param_file, mdl, version, "") + call get_param(param_file, mdl, "CORIOLIS_ADV_NKBLOCK", CS%nkblock, & + "The k-direction block size used in Coriolis and momentum advection "//& + "calculations. The default 0 setting dynamically uses the full vertical column.", & + default=default_nkblock, layoutParam=.true.) + if (CS%nkblock < 0) call MOM_error(FATAL, "CORIOLIS_ADV_NKBLOCK must be >= 0.") call get_param(param_file, mdl, "NOSLIP", CS%no_slip, & "If true, no slip boundary conditions are used; otherwise "//& "free slip boundary conditions are assumed. The "//& diff --git a/src/core/MOM_continuity_PPM.F90 b/src/core/MOM_continuity_PPM.F90 index 616649fe64..5ea7deda00 100644 --- a/src/core/MOM_continuity_PPM.F90 +++ b/src/core/MOM_continuity_PPM.F90 @@ -33,10 +33,6 @@ module MOM_continuity_PPM public zonal_BT_mass_flux, meridional_BT_mass_flux public set_continuity_loop_bounds -! These were found to give best performance in limited tests. -integer, parameter :: default_niblock = 32 !< Default i block size for array calculations [nondim]. -integer, parameter :: default_njblock = 4 !< Default j block size for array calculations [nondim]. - !>@{ CPU time clock IDs integer :: id_clock_reconstruct, id_clock_update, id_clock_correct !>@} @@ -79,6 +75,7 @@ module MOM_continuity_PPM !! averaged areas. integer :: niblock !< The i block size used in array calculations [nondim]. integer :: njblock !< The j block size used in array calculations [nondim]. + integer :: nkblock !< The k block size used in reconstruction calculations [nondim]. end type continuity_PPM_CS @@ -169,8 +166,6 @@ subroutine continuity_PPM(u, v, hin, h, uh, vh, dt, G, GV, US, CS, OBC, pbv, uhb niblock = CS%niblock njblock = CS%njblock - if (niblock == 0) niblock = default_niblock - if (njblock == 0) njblock = default_njblock h_min = GV%Angstrom_H @@ -189,11 +184,9 @@ subroutine continuity_PPM(u, v, hin, h, uh, vh, dt, G, GV, US, CS, OBC, pbv, uhb ! First advect zonally, with loop bounds that accomodate the subsequent meridional advection. LB = set_continuity_loop_bounds(G, CS, i_stencil=.false., j_stencil=.true.) - ! set default block sizes for OpenMP offload - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = LB%ieh-LB%ish+2 - !$ if (CS%njblock == 0) njblock = LB%jeh-LB%jsh+1 - !$ endif + ! set whole-domain block sizes when ni/jblock is 0 + if (niblock == 0) niblock = LB%ieh-LB%ish+2 + if (njblock == 0) njblock = LB%jeh-LB%jsh+1 call zonal_edge_thickness(hin, h_W, h_E, G, GV, US, CS, OBC, LB) call zonal_mass_flux(u, hin, h_W, h_E, uh, dt, G, GV, US, CS, OBC, pbv%por_face_areaU, & @@ -203,10 +196,8 @@ subroutine continuity_PPM(u, v, hin, h, uh, vh, dt, G, GV, US, CS, OBC, pbv, uhb ! Now advect meridionally, using the updated thicknesses to determine the fluxes. LB = set_continuity_loop_bounds(G, CS, i_stencil=.false., j_stencil=.false.) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = LB%ieh-LB%ish+1 - !$ if (CS%njblock == 0) njblock = LB%jeh-LB%jsh+2 - !$ endif + if (niblock == 0) niblock = LB%ieh-LB%ish+1 + if (njblock == 0) njblock = LB%jeh-LB%jsh+2 call meridional_edge_thickness(h, h_S, h_N, G, GV, US, CS, OBC, LB) call meridional_mass_flux(v, h, h_S, h_N, vh, dt, G, GV, US, CS, OBC, pbv%por_face_areaV, & @@ -217,10 +208,8 @@ subroutine continuity_PPM(u, v, hin, h, uh, vh, dt, G, GV, US, CS, OBC, pbv, uhb ! First advect meridionally, with loop bounds that accomodate the subsequent zonal advection. LB = set_continuity_loop_bounds(G, CS, i_stencil=.true., j_stencil=.false.) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = LB%ieh-LB%ish+1 - !$ if (CS%njblock == 0) njblock = LB%jeh-LB%jsh+2 - !$ endif + if (niblock == 0) niblock = LB%ieh-LB%ish+1 + if (njblock == 0) njblock = LB%jeh-LB%jsh+2 call meridional_edge_thickness(hin, h_S, h_N, G, GV, US, CS, OBC, LB) call meridional_mass_flux(v, hin, h_S, h_N, vh, dt, G, GV, US, CS, OBC, pbv%por_face_areaV, & niblock, njblock, LB, vhbt, visc_rem_v, v_cor, BT_cont, dv_cor) @@ -228,10 +217,8 @@ subroutine continuity_PPM(u, v, hin, h, uh, vh, dt, G, GV, US, CS, OBC, pbv, uhb ! Now advect zonally, using the updated thicknesses to determine the fluxes. LB = set_continuity_loop_bounds(G, CS, i_stencil=.false., j_stencil=.false.) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = LB%ieh-LB%ish+2 - !$ if (CS%njblock == 0) njblock = LB%jeh-LB%jsh+1 - !$ endif + if (niblock == 0) niblock = LB%ieh-LB%ish+2 + if (njblock == 0) njblock = LB%jeh-LB%jsh+1 call zonal_edge_thickness(h, h_W, h_E, G, GV, US, CS, OBC, LB) call zonal_mass_flux(u, h, h_W, h_E, uh, dt, G, GV, US, CS, OBC, pbv%por_face_areaU, & niblock, njblock, LB, uhbt, visc_rem_u, u_cor, BT_cont, du_cor) @@ -277,22 +264,16 @@ subroutine continuity_3d_fluxes(u, v, h, uh, vh, dt, G, GV, US, CS, OBC, pbv) niblock = CS%niblock njblock = CS%njblock - if (niblock == 0) niblock = default_niblock - if (njblock == 0) njblock = default_njblock call zonal_edge_thickness(h, h_W, h_E, G, GV, US, CS, OBC) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = G%iec-G%isc+2 - !$ if (CS%njblock == 0) njblock = G%jec-G%jsc+1 - !$ endif + if (niblock == 0) niblock = G%iec-G%isc+2 + if (njblock == 0) njblock = G%jec-G%jsc+1 call zonal_mass_flux(u, h, h_W, h_E, uh, dt, G, GV, US, CS, OBC, pbv%por_face_areaU, & niblock=niblock, njblock=njblock) call meridional_edge_thickness(h, h_S, h_N, G, GV, US, CS, OBC) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = G%iec-G%isc+1 - !$ if (CS%njblock == 0) njblock = G%jec-G%jsc+2 - !$ endif + if (niblock == 0) niblock = G%iec-G%isc+1 + if (njblock == 0) njblock = G%jec-G%jsc+2 call meridional_mass_flux(v, h, h_S, h_N, vh, dt, G, GV, US, CS, OBC, pbv%por_face_areaV, & niblock=niblock, njblock=njblock) @@ -398,8 +379,6 @@ subroutine continuity_adjust_vel(u, v, h, dt, G, GV, US, CS, OBC, pbv, uhbt, vhb niblock = CS%niblock njblock = CS%njblock - if (niblock == 0) niblock = default_niblock - if (njblock == 0) njblock = default_njblock ! It might not be necessary to separate the input velocity array from the adjusted velocities, ! but it seems safer to do so, even if it might be less efficient. @@ -407,19 +386,15 @@ subroutine continuity_adjust_vel(u, v, h, dt, G, GV, US, CS, OBC, pbv, uhbt, vhb v_in(:,:,:) = v(:,:,:) call zonal_edge_thickness(h, h_W, h_E, G, GV, US, CS, OBC) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = G%iec-G%isc+2 - !$ if (CS%njblock == 0) njblock = G%jec-G%jsc+1 - !$ endif + if (niblock == 0) niblock = G%iec-G%isc+2 + if (njblock == 0) njblock = G%jec-G%jsc+1 call zonal_mass_flux(u_in, h, h_W, h_E, uh, dt, G, GV, US, CS, OBC, pbv%por_face_areaU, & niblock=niblock, njblock=njblock, & uhbt=uhbt, visc_rem_u=visc_rem_u, u_cor=u) call meridional_edge_thickness(h, h_S, h_N, G, GV, US, CS, OBC) - !$ if (omp_get_num_devices() > 0) then - !$ if (CS%niblock == 0) niblock = G%iec-G%isc+1 - !$ if (CS%njblock == 0) njblock = G%jec-G%jsc+2 - !$ endif + if (niblock == 0) niblock = G%iec-G%isc+1 + if (njblock == 0) njblock = G%jec-G%jsc+2 call meridional_mass_flux(v_in, h, h_S, h_N, vh, dt, G, GV, US, CS, OBC, pbv%por_face_areaV, & niblock=niblock, njblock=njblock, & vhbt=vhbt, visc_rem_v=visc_rem_v, v_cor=v) @@ -524,7 +499,7 @@ subroutine zonal_edge_thickness(h_in, h_W, h_E, G, GV, US, CS, OBC, LB_in) ! Local variables type(cont_loop_bounds_type) :: LB - integer :: i, j, k, ish, ieh, jsh, jeh, nz + integer :: i, j, k, ish, ieh, jsh, jeh, nz, nkblock call cpu_clock_begin(id_clock_reconstruct) @@ -534,6 +509,8 @@ subroutine zonal_edge_thickness(h_in, h_W, h_E, G, GV, US, CS, OBC, LB_in) LB%ish = G%isc ; LB%ieh = G%iec ; LB%jsh = G%jsc ; LB%jeh = G%jec endif ish = LB%ish ; ieh = LB%ieh ; jsh = LB%jsh ; jeh = LB%jeh ; nz = GV%ke + nkblock = CS%nkblock + if (nkblock == 0) nkblock = nz if (CS%upwind_1st) then do concurrent (k=1:nz, j=jsh:jeh, i=ish-1:ieh+1) @@ -541,7 +518,7 @@ subroutine zonal_edge_thickness(h_in, h_W, h_E, G, GV, US, CS, OBC, LB_in) enddo else call PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, & - 2.0*GV%Angstrom_H, CS%monotonic, CS%simple_2nd, OBC) + nkblock, 2.0*GV%Angstrom_H, CS%monotonic, CS%simple_2nd, OBC) endif call cpu_clock_end(id_clock_reconstruct) @@ -567,7 +544,7 @@ subroutine meridional_edge_thickness(h_in, h_S, h_N, G, GV, US, CS, OBC, LB_in) ! Local variables type(cont_loop_bounds_type) :: LB - integer :: i, j, k, ish, ieh, jsh, jeh, nz + integer :: i, j, k, ish, ieh, jsh, jeh, nz, nkblock call cpu_clock_begin(id_clock_reconstruct) @@ -577,6 +554,8 @@ subroutine meridional_edge_thickness(h_in, h_S, h_N, G, GV, US, CS, OBC, LB_in) LB%ish = G%isc ; LB%ieh = G%iec ; LB%jsh = G%jsc ; LB%jeh = G%jec endif ish = LB%ish ; ieh = LB%ieh ; jsh = LB%jsh ; jeh = LB%jeh ; nz = GV%ke + nkblock = CS%nkblock + if (nkblock == 0) nkblock = nz if (CS%upwind_1st) then ! untested @@ -585,7 +564,7 @@ subroutine meridional_edge_thickness(h_in, h_S, h_N, G, GV, US, CS, OBC, LB_in) enddo else call PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, & - 2.0*GV%Angstrom_H, CS%monotonic, CS%simple_2nd, OBC) + nkblock, 2.0*GV%Angstrom_H, CS%monotonic, CS%simple_2nd, OBC) endif call cpu_clock_end(id_clock_reconstruct) @@ -1104,7 +1083,7 @@ subroutine zonal_BT_mass_flux(u, h_in, h_W, h_E, uhbt, dt, G, GV, US, CS, OBC, p end subroutine zonal_BT_mass_flux !> Evaluates the zonal mass or volume fluxes in an element. -!NVF$ INLINE +!DIR$ ATTRIBUTES FORCEINLINE :: flux_elem elemental subroutine flux_elem(u, h, h_p1, h_L, h_L_p1, h_R, h_R_p1, uh, duhdu, visc_rem, & G_dy_Cu, G_IareaT, G_IareaT_p1, G_IdxT, G_IdxT_p1, dt, & vol_CFL, por_face_area) @@ -1167,7 +1146,7 @@ elemental subroutine flux_elem(u, h, h_p1, h_L, h_L_p1, h_R, h_R_p1, uh, duhdu, end subroutine flux_elem -!NVF$ INLINE +!DIR$ ATTRIBUTES FORCEINLINE :: flux_elem_OBC elemental subroutine flux_elem_OBC(u, h, h_p1, uh, duhdu, visc_rem, por_face_area, & G_dy_Cu, OBC, l_seg) real, intent(in) :: u !< Zonal/meridional velocity [L T-1 ~> m s-1]. @@ -2702,7 +2681,7 @@ subroutine set_merid_BT_cont(v, h_in, h_S, h_N, BT_cont, vh_tot_0, dvhdv_tot_0, end subroutine set_merid_BT_cont !> Calculates left/right edge values for PPM reconstruction. -subroutine PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, h_min, monotonic, simple_2nd, OBC) +subroutine PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, nkblock, h_min, monotonic, simple_2nd, OBC) type(ocean_grid_type), intent(in) :: G !< Ocean's grid structure. type(verticalGrid_type), intent(in) :: GV !< Ocean's vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h_in !< Layer thickness [H ~> m or kg m-2]. @@ -2711,6 +2690,7 @@ subroutine PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, h_min, monotonic, sim real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: h_E !< East edge thickness in the reconstruction, !! [H ~> m or kg m-2]. type(cont_loop_bounds_type), intent(in) :: LB !< Active loop bounds structure. + integer, intent(in) :: nkblock !< k block size for reconstruction calculations [nondim]. real, intent(in) :: h_min !< The minimum thickness !! that can be obtained by a concave parabolic fit [H ~> m or kg m-2] logical, intent(in) :: monotonic !< If true, use the @@ -2720,16 +2700,17 @@ subroutine PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, h_min, monotonic, sim !! arithmetic mean thicknesses as the default edge values !! for a simple 2nd order scheme. type(ocean_OBC_type), pointer :: OBC !< Open boundaries control structure. - integer :: k !< vertical grid index + integer :: k, kk !< vertical grid and k-block indices ! Local variables with useful mnemonic names. - real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: slp ! The slopes per grid point [H ~> m or kg m-2] + real, dimension(SZI_(G),SZJ_(G),max(1,nkblock)) :: & + slp ! The slopes per grid point in a k block [H ~> m or kg m-2] real, parameter :: oneSixth = 1./6. ! [nondim] real :: h_ip1, h_im1 ! Neighboring thicknesses or sensibly extrapolated values [H ~> m or kg m-2] real :: dMx, dMn ! The difference between the local thickness and the maximum (dMx) or ! minimum (dMn) of the surrounding values [H ~> m or kg m-2] character(len=256) :: mesg - integer :: i, j, isl, iel, jsl, jel, nz, n, stencil + integer :: i, j, isl, iel, jsl, jel, nz, n, stencil, ks, ke logical :: local_open_BC type(OBC_segment_type), pointer :: segment => NULL() @@ -2758,95 +2739,104 @@ subroutine PPM_reconstruction_x(h_in, h_W, h_E, G, GV, LB, h_min, monotonic, sim !$omp target enter data map(alloc: slp) - if (simple_2nd) then - ! untested - do concurrent (k =1:nz, j=jsl:jel, i=isl:iel) - h_im1 = G%mask2dT(i-1,j) * h_in(i-1,j,k) + (1.0-G%mask2dT(i-1,j)) * h_in(i,j,k) - h_ip1 = G%mask2dT(i+1,j) * h_in(i+1,j,k) + (1.0-G%mask2dT(i+1,j)) * h_in(i,j,k) - h_W(i,j,k) = 0.5*( h_im1 + h_in(i,j,k) ) - h_E(i,j,k) = 0.5*( h_ip1 + h_in(i,j,k) ) - enddo - else - do concurrent (k=1:nz, j=jsl:jel, i=isl-1:iel+1) - if ((G%mask2dT(i-1,j) * G%mask2dT(i,j) * G%mask2dT(i+1,j)) == 0.0) then - slp(i,j,k) = 0.0 - else - ! This uses a simple 2nd order slope. - slp(i,j,k) = 0.5 * (h_in(i+1,j,k) - h_in(i-1,j,k)) - ! Monotonic constraint, see Eq. B2 in Lin 1994, MWR (132) - dMx = max(h_in(i+1,j,k), h_in(i-1,j,k), h_in(i,j,k)) - h_in(i,j,k) - dMn = h_in(i,j,k) - min(h_in(i+1,j,k), h_in(i-1,j,k), h_in(i,j,k)) - slp(i,j,k) = sign(1.,slp(i,j,k)) * min(abs(slp(i,j,k)), 2. * min(dMx, dMn)) - ! * (G%mask2dT(i-1,j) * G%mask2dT(i,j) * G%mask2dT(i+1,j)) + do ks = 1, nz, nkblock + ke = min(ks + nkblock - 1, nz) + + if (simple_2nd) then + ! untested + do concurrent (k=ks:ke, j=jsl:jel, i=isl:iel) DO_LOCALITY(local(h_im1,h_ip1)) + h_im1 = G%mask2dT(i-1,j) * h_in(i-1,j,k) + (1.0-G%mask2dT(i-1,j)) * h_in(i,j,k) + h_ip1 = G%mask2dT(i+1,j) * h_in(i+1,j,k) + (1.0-G%mask2dT(i+1,j)) * h_in(i,j,k) + h_W(i,j,k) = 0.5*( h_im1 + h_in(i,j,k) ) + h_E(i,j,k) = 0.5*( h_ip1 + h_in(i,j,k) ) + enddo + else + do concurrent (k=ks:ke, j=jsl:jel, i=isl-1:iel+1) DO_LOCALITY(local(dMx,dMn,kk)) + kk = k - ks + 1 + if ((G%mask2dT(i-1,j) * G%mask2dT(i,j) * G%mask2dT(i+1,j)) == 0.0) then + slp(i,j,kk) = 0.0 + else + ! This uses a simple 2nd order slope. + slp(i,j,kk) = 0.5 * (h_in(i+1,j,k) - h_in(i-1,j,k)) + ! Monotonic constraint, see Eq. B2 in Lin 1994, MWR (132) + dMx = max(h_in(i+1,j,k), h_in(i-1,j,k), h_in(i,j,k)) - h_in(i,j,k) + dMn = h_in(i,j,k) - min(h_in(i+1,j,k), h_in(i-1,j,k), h_in(i,j,k)) + slp(i,j,kk) = sign(1.,slp(i,j,kk)) * min(abs(slp(i,j,kk)), 2. * min(dMx, dMn)) + ! * (G%mask2dT(i-1,j) * G%mask2dT(i,j) * G%mask2dT(i+1,j)) + endif + enddo + + if (local_open_BC) then + ! untested + do n=1, OBC%number_of_segments + segment => OBC%segment(n) + if (.not. segment%on_pe) cycle + if (segment%is_E_or_W) then + I=segment%HI%IsdB + do concurrent (k=ks:ke, j=segment%HI%jsd:segment%HI%jed) DO_LOCALITY(local(kk)) + kk = k - ks + 1 + slp(i+1,j,kk) = 0.0 + slp(i,j,kk) = 0.0 + enddo + endif + enddo endif - enddo + + do concurrent (k=ks:ke, j=jsl:jel, i=isl:iel) DO_LOCALITY(local(h_im1,h_ip1,kk)) + kk = k - ks + 1 + ! Neighboring values should take into account any boundaries. The 3 + ! following sets of expressions are equivalent. + ! h_im1 = h_in(i-1,j,k) ; if (G%mask2dT(i-1,j) < 0.5) h_im1 = h_in(i,j) + ! h_ip1 = h_in(i+1,j,k) ; if (G%mask2dT(i+1,j) < 0.5) h_ip1 = h_in(i,j) + h_im1 = G%mask2dT(i-1,j) * h_in(i-1,j,k) + (1.0-G%mask2dT(i-1,j)) * h_in(i,j,k) + h_ip1 = G%mask2dT(i+1,j) * h_in(i+1,j,k) + (1.0-G%mask2dT(i+1,j)) * h_in(i,j,k) + ! Left/right values following Eq. B2 in Lin 1994, MWR (132) + h_W(i,j,k) = 0.5*( h_im1 + h_in(i,j,k) ) + & + oneSixth*( slp(i-1,j,kk) - slp(i,j,kk) ) + h_E(i,j,k) = 0.5*( h_ip1 + h_in(i,j,k) ) + & + oneSixth*( slp(i,j,kk) - slp(i+1,j,kk) ) + enddo + endif if (local_open_BC) then ! untested do n=1, OBC%number_of_segments segment => OBC%segment(n) if (.not. segment%on_pe) cycle - if (segment%is_E_or_W) then + if (segment%direction == OBC_DIRECTION_E) then I=segment%HI%IsdB - do concurrent (k=1:nz, j=segment%HI%jsd:segment%HI%jed) - slp(i+1,j,k) = 0.0 - slp(i,j,k) = 0.0 + do concurrent (k=ks:ke, j=segment%HI%jsd:segment%HI%jed) + h_W(i+1,j,k) = h_in(i,j,k) + h_E(i+1,j,k) = h_in(i,j,k) + h_W(i,j,k) = h_in(i,j,k) + h_E(i,j,k) = h_in(i,j,k) + enddo + elseif (segment%direction == OBC_DIRECTION_W) then + I=segment%HI%IsdB + do concurrent (k=ks:ke, j=segment%HI%jsd:segment%HI%jed) + h_W(i,j,k) = h_in(i+1,j,k) + h_E(i,j,k) = h_in(i+1,j,k) + h_W(i+1,j,k) = h_in(i+1,j,k) + h_E(i+1,j,k) = h_in(i+1,j,k) enddo endif enddo endif - do concurrent (k=1:nz, j=jsl:jel, i=isl:iel) - ! Neighboring values should take into account any boundaries. The 3 - ! following sets of expressions are equivalent. - ! h_im1 = h_in(i-1,j,k) ; if (G%mask2dT(i-1,j) < 0.5) h_im1 = h_in(i,j) - ! h_ip1 = h_in(i+1,j,k) ; if (G%mask2dT(i+1,j) < 0.5) h_ip1 = h_in(i,j) - h_im1 = G%mask2dT(i-1,j) * h_in(i-1,j,k) + (1.0-G%mask2dT(i-1,j)) * h_in(i,j,k) - h_ip1 = G%mask2dT(i+1,j) * h_in(i+1,j,k) + (1.0-G%mask2dT(i+1,j)) * h_in(i,j,k) - ! Left/right values following Eq. B2 in Lin 1994, MWR (132) - h_W(i,j,k) = 0.5*( h_im1 + h_in(i,j,k) ) + oneSixth*( slp(i-1,j,k) - slp(i,j,k) ) - h_E(i,j,k) = 0.5*( h_ip1 + h_in(i,j,k) ) + oneSixth*( slp(i,j,k) - slp(i+1,j,k) ) - enddo - endif - - if (local_open_BC) then - ! untested - do n=1, OBC%number_of_segments - segment => OBC%segment(n) - if (.not. segment%on_pe) cycle - if (segment%direction == OBC_DIRECTION_E) then - I=segment%HI%IsdB - do concurrent (k=1:nz, j=segment%HI%jsd:segment%HI%jed) - h_W(i+1,j,k) = h_in(i,j,k) - h_E(i+1,j,k) = h_in(i,j,k) - h_W(i,j,k) = h_in(i,j,k) - h_E(i,j,k) = h_in(i,j,k) - enddo - elseif (segment%direction == OBC_DIRECTION_W) then - I=segment%HI%IsdB - do concurrent (k=1:nz, j=segment%HI%jsd:segment%HI%jed) - h_W(i,j,k) = h_in(i+1,j,k) - h_E(i,j,k) = h_in(i+1,j,k) - h_W(i+1,j,k) = h_in(i+1,j,k) - h_E(i+1,j,k) = h_in(i+1,j,k) - enddo - endif - enddo - endif - - if (monotonic) then - ! untested - call PPM_limit_CW84(h_in, h_W, h_E, G, GV, isl, iel, jsl, jel, nz) - else - call PPM_limit_pos(h_in, h_W, h_E, h_min, G, GV, isl, iel, jsl, jel, nz) - endif + if (monotonic) then + ! untested + call PPM_limit_CW84(h_in, h_W, h_E, G, GV, isl, iel, jsl, jel, ks, ke) + else + call PPM_limit_pos(h_in, h_W, h_E, h_min, G, GV, isl, iel, jsl, jel, ks, ke) + endif + enddo !$omp target exit data map(release: slp) end subroutine PPM_reconstruction_x !> Calculates left/right edge values for PPM reconstruction. -subroutine PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, h_min, monotonic, simple_2nd, OBC) +subroutine PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, nkblock, h_min, monotonic, simple_2nd, OBC) type(ocean_grid_type), intent(in) :: G !< Ocean's grid structure. type(verticalGrid_type), intent(in) :: GV !< Ocean's vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h_in !< Layer thickness [H ~> m or kg m-2]. @@ -2855,6 +2845,7 @@ subroutine PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, h_min, monotonic, sim real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(out) :: h_N !< North edge thickness in the reconstruction, !! [H ~> m or kg m-2]. type(cont_loop_bounds_type), intent(in) :: LB !< Active loop bounds structure. + integer, intent(in) :: nkblock !< k block size for reconstruction calculations [nondim]. real, intent(in) :: h_min !< The minimum thickness !! that can be obtained by a concave parabolic fit [H ~> m or kg m-2] logical, intent(in) :: monotonic !< If true, use the @@ -2864,16 +2855,17 @@ subroutine PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, h_min, monotonic, sim !! arithmetic mean thicknesses as the default edge values !! for a simple 2nd order scheme. type(ocean_OBC_type), pointer :: OBC !< Open boundaries control structure. - integer :: k !< vertical grid index + integer :: k, kk !< vertical grid and k-block indices ! Local variables with useful mnemonic names. - real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: slp ! The slopes per grid point [H ~> m or kg m-2] + real, dimension(SZI_(G),SZJ_(G),max(1,nkblock)) :: & + slp ! The slopes per grid point in a k block [H ~> m or kg m-2] real, parameter :: oneSixth = 1./6. ! [nondim] real :: h_jp1, h_jm1 ! Neighboring thicknesses or sensibly extrapolated values [H ~> m or kg m-2] real :: dMx, dMn ! The difference between the local thickness and the maximum (dMx) or ! minimum (dMn) of the surrounding values [H ~> m or kg m-2] character(len=256) :: mesg - integer :: i, j, isl, iel, jsl, jel, nz, n, stencil + integer :: i, j, isl, iel, jsl, jel, nz, n, stencil, ks, ke logical :: local_open_BC type(OBC_segment_type), pointer :: segment => NULL() @@ -2902,86 +2894,95 @@ subroutine PPM_reconstruction_y(h_in, h_S, h_N, G, GV, LB, h_min, monotonic, sim !$omp target enter data map(alloc: slp) - if (simple_2nd) then - ! untested - do concurrent (k=1:nz, j=jsl:jel, i=isl:iel) - h_jm1 = G%mask2dT(i,j-1) * h_in(i,j-1,k) + (1.0-G%mask2dT(i,j-1)) * h_in(i,j,k) - h_jp1 = G%mask2dT(i,j+1) * h_in(i,j+1,k) + (1.0-G%mask2dT(i,j+1)) * h_in(i,j,k) - h_S(i,j,k) = 0.5*( h_jm1 + h_in(i,j,k) ) - h_N(i,j,k) = 0.5*( h_jp1 + h_in(i,j,k) ) - enddo - else - do concurrent (k=1:nz, j=jsl-1:jel+1, i=isl:iel) - if ((G%mask2dT(i,j-1) * G%mask2dT(i,j) * G%mask2dT(i,j+1)) == 0.0) then - slp(i,j,k) = 0.0 - else - ! This uses a simple 2nd order slope. - slp(i,j,k) = 0.5 * (h_in(i,j+1,k) - h_in(i,j-1,k)) - ! Monotonic constraint, see Eq. B2 in Lin 1994, MWR (132) - dMx = max(h_in(i,j+1,k), h_in(i,j-1,k), h_in(i,j,k)) - h_in(i,j,k) - dMn = h_in(i,j,k) - min(h_in(i,j+1,k), h_in(i,j-1,k), h_in(i,j,k)) - slp(i,j,k) = sign(1.,slp(i,j,k)) * min(abs(slp(i,j,k)), 2. * min(dMx, dMn)) - ! * (G%mask2dT(i,j-1) * G%mask2dT(i,j) * G%mask2dT(i,j+1)) + do ks = 1, nz, nkblock + ke = min(ks + nkblock - 1, nz) + + if (simple_2nd) then + ! untested + do concurrent (k=ks:ke, j=jsl:jel, i=isl:iel) DO_LOCALITY(local(h_jm1,h_jp1)) + h_jm1 = G%mask2dT(i,j-1) * h_in(i,j-1,k) + (1.0-G%mask2dT(i,j-1)) * h_in(i,j,k) + h_jp1 = G%mask2dT(i,j+1) * h_in(i,j+1,k) + (1.0-G%mask2dT(i,j+1)) * h_in(i,j,k) + h_S(i,j,k) = 0.5*( h_jm1 + h_in(i,j,k) ) + h_N(i,j,k) = 0.5*( h_jp1 + h_in(i,j,k) ) + enddo + else + do concurrent (k=ks:ke, j=jsl-1:jel+1, i=isl:iel) DO_LOCALITY(local(dMx,dMn,kk)) + kk = k - ks + 1 + if ((G%mask2dT(i,j-1) * G%mask2dT(i,j) * G%mask2dT(i,j+1)) == 0.0) then + slp(i,j,kk) = 0.0 + else + ! This uses a simple 2nd order slope. + slp(i,j,kk) = 0.5 * (h_in(i,j+1,k) - h_in(i,j-1,k)) + ! Monotonic constraint, see Eq. B2 in Lin 1994, MWR (132) + dMx = max(h_in(i,j+1,k), h_in(i,j-1,k), h_in(i,j,k)) - h_in(i,j,k) + dMn = h_in(i,j,k) - min(h_in(i,j+1,k), h_in(i,j-1,k), h_in(i,j,k)) + slp(i,j,kk) = sign(1.,slp(i,j,kk)) * min(abs(slp(i,j,kk)), 2. * min(dMx, dMn)) + ! * (G%mask2dT(i,j-1) * G%mask2dT(i,j) * G%mask2dT(i,j+1)) + endif + enddo + + if (local_open_BC) then + ! untested + do n=1, OBC%number_of_segments + segment => OBC%segment(n) + if (.not. segment%on_pe) cycle + if (segment%is_N_or_S) then + J=segment%HI%JsdB + do concurrent (k=ks:ke, i=segment%HI%isd:segment%HI%ied) DO_LOCALITY(local(kk)) + kk = k - ks + 1 + slp(i,j+1,kk) = 0.0 + slp(i,j,kk) = 0.0 + enddo + endif + enddo endif - enddo + + do concurrent (k=ks:ke, j=jsl:jel, i=isl:iel) DO_LOCALITY(local(h_jm1,h_jp1,kk)) + kk = k - ks + 1 + ! Neighboring values should take into account any boundaries. The 3 + ! following sets of expressions are equivalent. + h_jm1 = G%mask2dT(i,j-1) * h_in(i,j-1,k) + (1.0-G%mask2dT(i,j-1)) * h_in(i,j,k) + h_jp1 = G%mask2dT(i,j+1) * h_in(i,j+1,k) + (1.0-G%mask2dT(i,j+1)) * h_in(i,j,k) + ! Left/right values following Eq. B2 in Lin 1994, MWR (132) + h_S(i,j,k) = 0.5*( h_jm1 + h_in(i,j,k) ) + & + oneSixth*( slp(i,j-1,kk) - slp(i,j,kk) ) + h_N(i,j,k) = 0.5*( h_jp1 + h_in(i,j,k) ) + & + oneSixth*( slp(i,j,kk) - slp(i,j+1,kk) ) + enddo + endif if (local_open_BC) then ! untested do n=1, OBC%number_of_segments segment => OBC%segment(n) if (.not. segment%on_pe) cycle - if (segment%is_N_or_S) then + if (segment%direction == OBC_DIRECTION_N) then + J=segment%HI%JsdB + do concurrent (k=ks:ke, i=segment%HI%isd:segment%HI%ied) + h_S(i,j+1,k) = h_in(i,j,k) + h_N(i,j+1,k) = h_in(i,j,k) + h_S(i,j,k) = h_in(i,j,k) + h_N(i,j,k) = h_in(i,j,k) + enddo + elseif (segment%direction == OBC_DIRECTION_S) then J=segment%HI%JsdB - do concurrent (k=1:nz, i=segment%HI%isd:segment%HI%ied) - slp(i,j+1,k) = 0.0 - slp(i,j,k) = 0.0 + do concurrent (k=ks:ke, i=segment%HI%isd:segment%HI%ied) + h_S(i,j,k) = h_in(i,j+1,k) + h_N(i,j,k) = h_in(i,j+1,k) + h_S(i,j+1,k) = h_in(i,j+1,k) + h_N(i,j+1,k) = h_in(i,j+1,k) enddo endif enddo endif - do concurrent (k=1:nz, j=jsl:jel, i=isl:iel) - ! Neighboring values should take into account any boundaries. The 3 - ! following sets of expressions are equivalent. - h_jm1 = G%mask2dT(i,j-1) * h_in(i,j-1,k) + (1.0-G%mask2dT(i,j-1)) * h_in(i,j,k) - h_jp1 = G%mask2dT(i,j+1) * h_in(i,j+1,k) + (1.0-G%mask2dT(i,j+1)) * h_in(i,j,k) - ! Left/right values following Eq. B2 in Lin 1994, MWR (132) - h_S(i,j,k) = 0.5*( h_jm1 + h_in(i,j,k) ) + oneSixth*( slp(i,j-1,k) - slp(i,j,k) ) - h_N(i,j,k) = 0.5*( h_jp1 + h_in(i,j,k) ) + oneSixth*( slp(i,j,k) - slp(i,j+1,k) ) - enddo - endif - - if (local_open_BC) then - ! untested - do n=1, OBC%number_of_segments - segment => OBC%segment(n) - if (.not. segment%on_pe) cycle - if (segment%direction == OBC_DIRECTION_N) then - J=segment%HI%JsdB - do concurrent (k=1:nz, i=segment%HI%isd:segment%HI%ied) - h_S(i,j+1,k) = h_in(i,j,k) - h_N(i,j+1,k) = h_in(i,j,k) - h_S(i,j,k) = h_in(i,j,k) - h_N(i,j,k) = h_in(i,j,k) - enddo - elseif (segment%direction == OBC_DIRECTION_S) then - J=segment%HI%JsdB - do concurrent (k=1:nz, i=segment%HI%isd:segment%HI%ied) - h_S(i,j,k) = h_in(i,j+1,k) - h_N(i,j,k) = h_in(i,j+1,k) - h_S(i,j+1,k) = h_in(i,j+1,k) - h_N(i,j+1,k) = h_in(i,j+1,k) - enddo - endif - enddo - endif - - if (monotonic) then - ! untested - call PPM_limit_CW84(h_in, h_S, h_N, G, GV, isl, iel, jsl, jel, nz) - else - call PPM_limit_pos(h_in, h_S, h_N, h_min, G, GV, isl, iel, jsl, jel, nz) - endif + if (monotonic) then + ! untested + call PPM_limit_CW84(h_in, h_S, h_N, G, GV, isl, iel, jsl, jel, ks, ke) + else + call PPM_limit_pos(h_in, h_S, h_N, h_min, G, GV, isl, iel, jsl, jel, ks, ke) + endif + enddo !$omp target exit data map(release: slp) @@ -2991,7 +2992,7 @@ end subroutine PPM_reconstruction_y !! to give a reconstruction that is positive-definite. Here this is !! reinterpreted as giving a constant thickness if the mean thickness is less !! than h_min, with a minimum of h_min otherwise. -subroutine PPM_limit_pos(h_in, h_L, h_R, h_min, G, GV, iis, iie, jis, jie, nz) +subroutine PPM_limit_pos(h_in, h_L, h_R, h_min, G, GV, iis, iie, jis, jie, ks, ke) type(ocean_grid_type), intent(in) :: G !< Ocean's grid structure. type(verticalGrid_type), intent(in) :: GV !< Ocean's vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), & @@ -3006,7 +3007,8 @@ subroutine PPM_limit_pos(h_in, h_L, h_R, h_min, G, GV, iis, iie, jis, jie, nz) integer, intent(in) :: iie !< End of i index range. integer, intent(in) :: jis !< Start of j index range. integer, intent(in) :: jie !< End of j index range. - integer, intent(in) :: nz !< End of k index range. + integer, intent(in) :: ks !< Start of k index range. + integer, intent(in) :: ke !< End of k index range. ! Local variables real :: curv ! The grid-normalized curvature of the three thicknesses [H ~> m or kg m-2] @@ -3014,7 +3016,7 @@ subroutine PPM_limit_pos(h_in, h_L, h_R, h_min, G, GV, iis, iie, jis, jie, nz) real :: scale ! A scaling factor to reduce the curvature of the fit [nondim] integer :: i,j,k - do concurrent (k=1:nz, j=jis:jie, i=iis:iie) + do concurrent (k=ks:ke, j=jis:jie, i=iis:iie) ! This limiter prevents undershooting minima within the domain with ! values less than h_min. curv = 3.0*((h_L(i,j,k) + h_R(i,j,k)) - 2.0*h_in(i,j,k)) @@ -3038,7 +3040,7 @@ end subroutine PPM_limit_pos !> This subroutine limits the left/right edge values of the PPM reconstruction !! according to the monotonic prescription of Colella and Woodward, 1984. -subroutine PPM_limit_CW84(h_in, h_L, h_R, G, GV, iis, iie, jis, jie, nz) +subroutine PPM_limit_CW84(h_in, h_L, h_R, G, GV, iis, iie, jis, jie, ks, ke) type(ocean_grid_type), intent(in) :: G !< Ocean's grid structure. type(verticalGrid_type), intent(in) :: GV !< Ocean's vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h_in !< Layer thickness [H ~> m or kg m-2]. @@ -3050,7 +3052,8 @@ subroutine PPM_limit_CW84(h_in, h_L, h_R, G, GV, iis, iie, jis, jie, nz) integer, intent(in) :: iie !< End of i index range. integer, intent(in) :: jis !< Start of j index range. integer, intent(in) :: jie !< End of j index range. - integer, intent(in) :: nz !< End of k index range. + integer, intent(in) :: ks !< Start of k index range. + integer, intent(in) :: ke !< End of k index range. ! Local variables real :: h_i ! A copy of the cell-average layer thickness [H ~> m or kg m-2] @@ -3061,7 +3064,7 @@ subroutine PPM_limit_CW84(h_in, h_L, h_R, G, GV, iis, iie, jis, jie, nz) integer :: i, j, k ! untested - do concurrent (k=1:nz, j=jis:jie, i=iis:iie) + do concurrent (k=ks:ke, j=jis:jie, i=iis:iie) DO_LOCALITY(local(h_i,RLdiff,RLdiff2,RLmean,FunFac)) ! This limiter monotonizes the parabola following ! Colella and Woodward, 1984, Eq. 1.10 h_i = h_in(i,j,k) @@ -3080,7 +3083,6 @@ subroutine PPM_limit_CW84(h_in, h_L, h_R, G, GV, iis, iie, jis, jie, nz) end subroutine PPM_limit_CW84 !> Return the maximum ratio of a/b or maxrat. -!NVF$ INLINE pure function ratio_max(a, b, maxrat) result(ratio) real, intent(in) :: a !< Numerator, in arbitrary units [A] real, intent(in) :: b !< Denominator, in arbitrary units [B] @@ -3114,7 +3116,17 @@ subroutine continuity_PPM_init(Time, G, GV, US, param_file, diag, CS, OBC) # include "version_variable.h" character(len=40) :: mdl = "MOM_continuity_PPM" ! This module's name. character(len=256) :: mesg - character(len=10) :: niblock_dflt_str, njblock_dflt_str + character(len=10) :: niblock_dflt_str, njblock_dflt_str, nkblock_dflt_str +#ifdef __NVCOMPILER_OPENMP_GPU + integer, parameter :: default_niblock = 0 !< Default i block size for array calculations [nondim]. + integer, parameter :: default_njblock = 0 !< Default j block size for array calculations [nondim]. + integer, parameter :: default_nkblock = 0 !< Default k block size for reconstruction calculations [nondim]. +#else + ! These were found to give best performance in limited tests. + integer, parameter :: default_niblock = 32 !< Default i block size for array calculations [nondim]. + integer, parameter :: default_njblock = 4 !< Default j block size for array calculations [nondim]. + integer, parameter :: default_nkblock = 1 !< Default k block size for reconstruction calculations [nondim]. +#endif CS%initialized = .true. @@ -3124,7 +3136,7 @@ subroutine continuity_PPM_init(Time, G, GV, US, param_file, diag, CS, OBC) endif ! Read all relevant parameters and write them to the model log. - call log_version(param_file, mdl, version, "") + call log_version(param_file, mdl, version, "", log_to_all=.true., layout=.true.) call get_param(param_file, mdl, "MONOTONIC_CONTINUITY", CS%monotonic, & "If true, CONTINUITY_PPM uses the Colella and Woodward "//& "monotonic limiter. The default (false) is to use a "//& @@ -3186,16 +3198,28 @@ subroutine continuity_PPM_init(Time, G, GV, US, param_file, diag, CS, OBC) "minimum is 0.", default=.false.) write(niblock_dflt_str, '(I0)') default_niblock write(njblock_dflt_str, '(I0)') default_njblock + write(nkblock_dflt_str, '(I0)') default_nkblock call get_param(param_file, mdl, "CONTINUITY_NIBLOCK", CS%niblock, & - "The i-direction block size used in the continuity solver. "//& - "If 0, defaults to "//trim(niblock_dflt_str)//", except when "//& - "running with OpenMP offload, in which case the full computational "//& - "domain width is used.", default=0) + "The i-direction block size used in the mass and volume flux calculations. "//& + "the default 0 setting is dynamic and fits the "//& + "full computational i-domain length.", default=default_niblock, layoutParam=.true.) call get_param(param_file, mdl, "CONTINUITY_NJBLOCK", CS%njblock, & - "The j-direction block size used in the continuity solver. "//& - "If 0, defaults to "//trim(njblock_dflt_str)//", except when "//& - "running with OpenMP offload, in which case the full computational "//& - "domain height is used.", default=0) + "The j-direction block size used in the mass and volume flux calculations. "//& + "the default 0 setting is dynamic and fits the "//& + "full computational j-domain length.", default=default_njblock, layoutParam=.true.) + call get_param(param_file, mdl, "CONTINUITY_NKBLOCK", CS%nkblock, & + "The k-direction block size used in PPM reconstruction edge value calculations. "//& + "the default 0 setting is dynamic and fits the "//& + "full vertical column.", default=default_nkblock, layoutParam=.true.) + if (CS%niblock < 0) & + call MOM_error(FATAL, "CONTINUITY_NIBLOCK must be nonnegative; "//& + "use 0 to select the default block size.") + if (CS%njblock < 0) & + call MOM_error(FATAL, "CONTINUITY_NJBLOCK must be nonnegative; "//& + "use 0 to select the default block size.") + if (CS%nkblock < 0) & + call MOM_error(FATAL, "CONTINUITY_NKBLOCK must be nonnegative; "//& + "use 0 to select the default block size.") CS%diag => diag !$omp target update to(CS) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 754f293b69..724a4ca28b 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -71,8 +71,10 @@ module MOM_EOS module procedure calculate_density_scalar module procedure calculate_density_1d module procedure calculate_density_2d + module procedure calculate_density_3d module procedure calculate_stanley_density_scalar module procedure calculate_stanley_density_1d + module procedure calculate_stanley_density_2d end interface calculate_density !> Calculates specific volume of sea water from T, S and P @@ -86,6 +88,7 @@ module MOM_EOS module procedure calculate_density_derivs_scalar, calculate_density_derivs_array module procedure calculate_density_derivs_1d module procedure calculate_density_derivs_2d + module procedure calculate_density_derivs_3d end interface calculate_density_derivs !> Calculate the derivatives of specific volume with temperature and salinity from T, S, and P @@ -97,6 +100,7 @@ module MOM_EOS !! salinity, and pressure from T, S and P interface calculate_density_second_derivs module procedure calculate_density_second_derivs_scalar, calculate_density_second_derivs_1d + module procedure calculate_density_second_derivs_2d end interface calculate_density_second_derivs !> Calculates the freezing point of sea water from T, S and P @@ -419,6 +423,68 @@ subroutine calculate_density_2d(T, S, pressure, rho, EOS, dom, rho_ref) end subroutine calculate_density_2d +!> Calls the appropriate subroutine to calculate density of sea water for 3-D array inputs. +subroutine calculate_density_3d(T, S, pressure, rho, EOS, dom, rho_ref) + real, intent(in) :: T(:,:,:) + !< Potential temperature referenced to the surface [C ~> degC] + real, intent(in) :: S(:,:,:) + !< Salinity [S ~> ppt] + real, intent(in) :: pressure(:,:,:) + !< Pressure [R L2 T-2 ~> Pa] + real, intent(inout) :: rho(:,:,:) + !< Density (in-situ if pressure is local) [R ~> kg m-3] + type(EOS_type), intent(in) :: EOS + !< Equation of state structure + integer, optional, intent(in) :: dom(3,2) + !< The domain of indices to work on, taking into account that arrays start + !! at 1. The first index is the rank (i, j, k) and the second is the bound + !! (1 = lower, 2 = upper). + real, optional, intent(in) :: rho_ref + !< A reference density [R ~> kg m-3] + + real, dimension(size(rho,1), size(rho,2), size(rho,3)) :: pres + ! Pressure converted to [Pa] + real, dimension(size(rho,1), size(rho,2), size(rho,3)) :: Ta + ! Temperature converted to [degC] + real, dimension(size(rho,1), size(rho,2), size(rho,3)) :: Sa + ! Salinity converted to [ppt] + integer :: is, ie, js, je, ks, ke + integer :: domain(3,2) + + if (present(dom)) then + domain(:,:) = dom(:,:) + else + domain(1,:) = [1, size(rho,1)] + domain(2,:) = [1, size(rho,2)] + domain(3,:) = [1, size(rho,3)] + endif + + is = domain(1,1) ; ie = domain(1,2) + js = domain(2,1) ; je = domain(2,2) + ks = domain(3,1) ; ke = domain(3,2) + + if ((EOS%RL2_T2_to_Pa == 1.0) .and. (EOS%R_to_kg_m3 == 1.0) .and. & + (EOS%C_to_degC == 1.0) .and. (EOS%S_to_ppt == 1.0)) then + call EOS%type%calculate_density_array_3d(T, S, pressure, rho, domain, & + rho_ref=rho_ref) + else ! This is the same as above, but with some extra work to rescale variables. + pres(is:ie, js:je, ks:ke) = EOS%RL2_T2_to_Pa * pressure(is:ie, js:je, ks:ke) + Ta(is:ie, js:je, ks:ke) = EOS%C_to_degC * T(is:ie, js:je, ks:ke) + Sa(is:ie, js:je, ks:ke) = EOS%S_to_ppt * S(is:ie, js:je, ks:ke) + + if (present(rho_ref)) then + call EOS%type%calculate_density_array_3d(Ta, Sa, pres, rho, domain, & + rho_ref=EOS%R_to_kg_m3*rho_ref) + else + call EOS%type%calculate_density_array_3d(Ta, Sa, pres, rho, domain) + endif + endif + + if (EOS%kg_m3_to_R /= 1.) & + rho(is:ie, js:je, ks:ke) = EOS%kg_m3_to_R * rho(is:ie, js:je, ks:ke) +end subroutine calculate_density_3d + + !> Calls the appropriate subroutine to calculate the density of sea water for 1-D array inputs !! including the variance of T, S and covariance of T-S, !! potentially limiting the domain of indices that are worked on. @@ -468,6 +534,39 @@ subroutine calculate_stanley_density_1d(T, S, pressure, Tvar, TScov, Svar, rho, end subroutine calculate_stanley_density_1d +!> Calls the Stanley density routine for each row of 2-D array inputs. +subroutine calculate_stanley_density_2d(T, S, pressure, Tvar, TScov, Svar, rho, EOS, dom, rho_ref, scale) + real, dimension(:,:), intent(in) :: T !< Potential temperature referenced to the surface [C ~> degC] + real, dimension(:,:), intent(in) :: S !< Salinity [S ~> ppt] + real, dimension(:,:), intent(in) :: pressure !< Pressure [R L2 T-2 ~> Pa] + real, dimension(:,:), intent(in) :: Tvar !< Variance of potential temperature [C2 ~> degC2] + real, dimension(:,:), intent(in) :: TScov !< Covariance of potential temperature and salinity [C S ~> degC ppt] + real, dimension(:,:), intent(in) :: Svar !< Variance of salinity [S2 ~> ppt2] + real, dimension(:,:), intent(inout) :: rho !< Density (in-situ if pressure is local) [R ~> kg m-3] + type(EOS_type), intent(in) :: EOS !< Equation of state structure + integer, dimension(2,2), optional, intent(in) :: dom !< The domain of indices to work on, taking + !! into account that arrays start at 1. + real, optional, intent(in) :: rho_ref !< A reference density [R ~> kg m-3] + real, optional, intent(in) :: scale !< A multiplicative factor by which to scale density + !! in combination with scaling stored in EOS [various] + integer :: j, js, je + integer, dimension(2) :: dom_row + + if (present(dom)) then + js = dom(2,1) ; je = dom(2,2) + dom_row = [dom(1,1), dom(1,2)] + else + js = 1 ; je = size(rho,2) + dom_row = [1, size(rho,1)] + endif + + do j=js,je + call calculate_stanley_density_1d(T(:,j), S(:,j), pressure(:,j), & + Tvar(:,j), TScov(:,j), Svar(:,j), & + rho(:,j), EOS, dom_row, rho_ref, scale) + enddo +end subroutine calculate_stanley_density_2d + !> Calls the appropriate subroutine to calculate the specific volume of sea water !! for 1-D array inputs. subroutine calculate_spec_vol_array(T, S, pressure, specvol, start, npts, EOS, spv_ref, scale) @@ -952,6 +1051,68 @@ subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, do end subroutine calculate_density_derivs_2d +!> Calls the appropriate subroutine to calculate density derivatives for 3-D array inputs. +subroutine calculate_density_derivs_3d(T, S, pressure, drho_dT, drho_dS, EOS, dom) + real, intent(in) :: T(:,:,:) + !< Potential temperature referenced to the surface [degC] + real, intent(in) :: S(:,:,:) + !< Salinity [ppt] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(inout) :: drho_dT(:,:,:) + !< The partial derivative of density with potential temperature + !! [kg m-3 degC-1] or other units determined by the optional scale argument + real, intent(inout) :: drho_dS(:,:,:) + !< The partial derivative of density with salinity, in [kg m-3 ppt-1] or + !! other units determined by the optional scale argument + type(EOS_type), intent(in) :: EOS + !< Equation of state structure + integer, optional, intent(in) :: dom(3,2) + !< The domain of indices to work on, taking into account that arrays start + !! at 1. The first index is the rank (i, j, k) and the second is the bound + !! (1 = lower, 2 = upper). + + ! Local variables + real :: Ta(size(T,1), size(T,2), size(T,3)) + ! Temperature converted to [degC] + real :: Sa(size(S,1), size(S,2), size(S,3)) + ! Salinity converted to [ppt] + real :: press(size(pressure,1), size(pressure,2), size(pressure,3)) + ! Pressure converted to [Pa] + integer :: is, ie, js, je, ks, ke + integer :: domain(3,2) + + if (present(dom)) then + domain(:,:) = dom(:,:) + else + domain(1,:) = [1, size(drho_dT, 1)] + domain(2,:) = [1, size(drho_dT, 2)] + domain(3,:) = [1, size(drho_dT, 3)] + endif + is = domain(1,1) ; ie = domain(1,2) + js = domain(2,1) ; je = domain(2,2) + ks = domain(3,1) ; ke = domain(3,2) + + if (.not. allocated(EOS%type)) call MOM_error(FATAL, & + "calculate_density_derivs_3d: EOS%form_of_EOS is not valid.") + + if (all([EOS%RL2_T2_to_Pa, EOS%C_to_degC, EOS%S_to_ppt] == 1.)) then + call EOS%type%calculate_density_derivs_3d(T, S, pressure, drho_dT, drho_dS, domain) + else + press(is:ie, js:je, ks:ke) = EOS%RL2_T2_to_Pa * pressure(is:ie, js:je, ks:ke) + Ta(is:ie, js:je, ks:ke) = EOS%C_to_degC * T(is:ie, js:je, ks:ke) + Sa(is:ie, js:je, ks:ke) = EOS%S_to_ppt * S(is:ie, js:je, ks:ke) + + call EOS%type%calculate_density_derivs_3d(Ta, Sa, press, drho_dT, drho_dS, domain) + endif + + if (EOS%kg_m3_to_R * EOS%C_to_degC /= 1.) & + drho_dT(is:ie, js:je, ks:ke) = EOS%kg_m3_to_R * EOS%C_to_degC * drho_dT(is:ie, js:je, ks:ke) + if (EOS%kg_m3_to_R * EOS%S_to_ppt /= 1.) & + drho_dS(is:ie, js:je, ks:ke) = EOS%kg_m3_to_R * EOS%S_to_ppt * drho_dS(is:ie, js:je, ks:ke) +end subroutine calculate_density_derivs_3d + + !> Calls the appropriate subroutines to calculate density derivatives by promoting a scalar !! to a one-element array subroutine calculate_density_derivs_scalar(T, S, pressure, drho_dT, drho_dS, EOS, scale) @@ -1071,6 +1232,86 @@ subroutine calculate_density_second_derivs_1d(T, S, pressure, drho_dS_dS, drho_d end subroutine calculate_density_second_derivs_1d +!> Calls the appropriate subroutine to calculate density second derivatives for 2D array inputs. +subroutine calculate_density_second_derivs_2d(T, S, pressure, drho_dS_dS, drho_dS_dT, drho_dT_dT, & + drho_dS_dP, drho_dT_dP, EOS, dom, scale) + real, intent(in) :: T(:,:) !< Potential temperature referenced to the surface [C ~> degC] + real, intent(in) :: S(:,:) !< Salinity [S ~> ppt] + real, intent(in) :: pressure(:,:) !< Pressure [R L2 T-2 ~> Pa] + real, intent(inout) :: drho_dS_dS(:,:) !< Partial derivative of beta with respect to S + !! [R S-2 ~> kg m-3 ppt-2] + real, intent(inout) :: drho_dS_dT(:,:) !< Partial derivative of beta with respect to T + !! [R S-1 C-1 ~> kg m-3 ppt-1 degC-1] + real, intent(inout) :: drho_dT_dT(:,:) !< Partial derivative of alpha with respect to T + !! [R C-2 ~> kg m-3 degC-2] + real, intent(inout) :: drho_dS_dP(:,:) !< Partial derivative of beta with respect to pressure + !! [T2 S-1 L-2 ~> kg m-3 ppt-1 Pa-1] + real, intent(inout) :: drho_dT_dP(:,:) !< Partial derivative of alpha with respect to pressure + !! [T2 C-1 L-2 ~> kg m-3 degC-1 Pa-1] + type(EOS_type), intent(in) :: EOS !< Equation of state structure + integer, optional, intent(in) :: dom(2,2) !< The domain of indices to work on + real, optional, intent(in) :: scale !< A multiplicative factor by which to scale density [various] + + ! Local variables + real :: Ta(size(T,1), size(T,2)) ! Temperature converted to [degC] + real :: Sa(size(S,1), size(S,2)) ! Salinity converted to [ppt] + real :: press(size(pressure,1), size(pressure,2)) ! Pressure converted to [Pa] + real :: rho_scale ! A factor to convert density from kg m-3 to the desired units [R m3 kg-1 ~> 1] + integer :: is, ie, js, je + integer :: domain(2,2) + + if (.not. allocated(EOS%type)) call MOM_error(FATAL, & + "calculate_density_second_derivs_2d: EOS%form_of_EOS is not valid.") + + if (present(dom)) then + domain(:,:) = dom(:,:) + else + domain(1,:) = [1, size(drho_dT_dT, 1)] + domain(2,:) = [1, size(drho_dT_dT, 2)] + endif + is = domain(1,1) ; ie = domain(1,2) + js = domain(2,1) ; je = domain(2,2) + + if (all([EOS%RL2_T2_to_Pa, EOS%C_to_degC, EOS%S_to_ppt] == 1.)) then + call EOS%type%calculate_density_second_derivs_2d(T, S, pressure, & + drho_dS_dS, drho_dS_dT, drho_dT_dT, drho_dS_dP, drho_dT_dP, domain) + else + press(is:ie, js:je) = EOS%RL2_T2_to_Pa * pressure(is:ie, js:je) + Ta(is:ie, js:je) = EOS%C_to_degC * T(is:ie, js:je) + Sa(is:ie, js:je) = EOS%S_to_ppt * S(is:ie, js:je) + call EOS%type%calculate_density_second_derivs_2d(Ta, Sa, press, & + drho_dS_dS, drho_dS_dT, drho_dT_dT, drho_dS_dP, drho_dT_dP, domain) + endif + + rho_scale = EOS%kg_m3_to_R + if (present(scale)) rho_scale = rho_scale * scale + if (rho_scale /= 1.0) then + drho_dS_dS(is:ie, js:je) = rho_scale * drho_dS_dS(is:ie, js:je) + drho_dS_dT(is:ie, js:je) = rho_scale * drho_dS_dT(is:ie, js:je) + drho_dT_dT(is:ie, js:je) = rho_scale * drho_dT_dT(is:ie, js:je) + drho_dS_dP(is:ie, js:je) = rho_scale * drho_dS_dP(is:ie, js:je) + drho_dT_dP(is:ie, js:je) = rho_scale * drho_dT_dP(is:ie, js:je) + endif + + if (EOS%RL2_T2_to_Pa /= 1.0) then + drho_dS_dP(is:ie, js:je) = EOS%RL2_T2_to_Pa * drho_dS_dP(is:ie, js:je) + drho_dT_dP(is:ie, js:je) = EOS%RL2_T2_to_Pa * drho_dT_dP(is:ie, js:je) + endif + + if (EOS%C_to_degC /= 1.0) then + drho_dS_dT(is:ie, js:je) = EOS%C_to_degC * drho_dS_dT(is:ie, js:je) + drho_dT_dT(is:ie, js:je) = EOS%C_to_degC**2 * drho_dT_dT(is:ie, js:je) + drho_dT_dP(is:ie, js:je) = EOS%C_to_degC * drho_dT_dP(is:ie, js:je) + endif + + if (EOS%S_to_ppt /= 1.0) then + drho_dS_dS(is:ie, js:je) = EOS%S_to_ppt**2 * drho_dS_dS(is:ie, js:je) + drho_dS_dT(is:ie, js:je) = EOS%S_to_ppt * drho_dS_dT(is:ie, js:je) + drho_dS_dP(is:ie, js:je) = EOS%S_to_ppt * drho_dS_dP(is:ie, js:je) + endif + +end subroutine calculate_density_second_derivs_2d + !> Calls the appropriate subroutine to calculate density second derivatives for scalar inputs. subroutine calculate_density_second_derivs_scalar(T, S, pressure, drho_dS_dS, drho_dS_dT, drho_dT_dT, & drho_dS_dP, drho_dT_dP, EOS, scale) diff --git a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 index 86110d6aa1..73ea02a400 100644 --- a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 +++ b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 @@ -183,10 +183,16 @@ module MOM_EOS_Roquet_rho procedure :: calculate_density_array => calculate_density_array_Roquet_rho !> Local implementation of generic calculate_density_array_2d for efficiency procedure :: calculate_density_array_2d => calculate_density_array_2d_Roquet_rho + !> Local implementation of generic calculate_density_array_3d for efficiency + procedure :: calculate_density_array_3d => calculate_density_array_3d_Roquet_rho !> Local implementation of generic calculate_spec_vol_array for efficiency procedure :: calculate_spec_vol_array => calculate_spec_vol_array_Roquet_rho !> Local implementation of generic calculate_density_derivs_2d for efficiency procedure :: calculate_density_derivs_2d => calculate_density_derivs_2d_Roquet_rho + !> Local implementation of generic calculate_density_derivs_3d for efficiency + procedure :: calculate_density_derivs_3d => calculate_density_derivs_3d_Roquet_rho + !> Local implementation of generic calculate_density_second_derivs_2d for efficiency + procedure :: calculate_density_second_derivs_2d => calculate_density_second_derivs_2d_Roquet_rho end type Roquet_rho_EOS @@ -455,23 +461,23 @@ elemental subroutine calculate_density_derivs_elem_Roquet_rho(this, T, S, pressu end subroutine calculate_density_derivs_elem_Roquet_rho -!> Second derivatives of density with respect to temperature, salinity, and pressure -elemental subroutine calculate_density_second_derivs_elem_Roquet_rho(this, T, S, pressure, & +!> Second derivatives of density with respect to temperature, salinity, and pressure. +!! Free-function form without "this" for use in do concurrent GPU regions with nvfortran. +elemental subroutine calculate_density_second_derivs_elem_Roquet_rho_loc(T, S, pressure, & drho_ds_ds, drho_ds_dt, drho_dt_dt, drho_ds_dp, drho_dt_dp) - class(Roquet_rho_EOS), intent(in) :: this !< This EOS - real, intent(in) :: T !< Conservative temperature [degC] - real, intent(in) :: S !< Absolute salinity [g kg-1] - real, intent(in) :: pressure !< Pressure [Pa] - real, intent(inout) :: drho_ds_ds !< Partial derivative of beta with respect - !! to S [kg m-3 ppt-2] - real, intent(inout) :: drho_ds_dt !< Partial derivative of beta with respect - !! to T [kg m-3 ppt-1 degC-1] - real, intent(inout) :: drho_dt_dt !< Partial derivative of alpha with respect - !! to T [kg m-3 degC-2] - real, intent(inout) :: drho_ds_dp !< Partial derivative of beta with respect - !! to pressure [kg m-3 ppt-1 Pa-1] = [s2 m-2 ppt-1] - real, intent(inout) :: drho_dt_dp !< Partial derivative of alpha with respect - !! to pressure [kg m-3 degC-1 Pa-1] = [s2 m-2 degC-1] + real, intent(in) :: T !< Conservative temperature [degC] + real, intent(in) :: S !< Absolute salinity [g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + real, intent(inout) :: drho_ds_ds !< Partial derivative of beta with respect + !! to S [kg m-3 ppt-2] + real, intent(inout) :: drho_ds_dt !< Partial derivative of beta with respect + !! to T [kg m-3 ppt-1 degC-1] + real, intent(inout) :: drho_dt_dt !< Partial derivative of alpha with respect + !! to T [kg m-3 degC-2] + real, intent(inout) :: drho_ds_dp !< Partial derivative of beta with respect + !! to pressure [kg m-3 ppt-1 Pa-1] = [s2 m-2 ppt-1] + real, intent(inout) :: drho_dt_dp !< Partial derivative of alpha with respect + !! to pressure [kg m-3 degC-1 Pa-1] = [s2 m-2 degC-1] ! Local variables real :: zp ! Pressure [Pa] @@ -545,6 +551,30 @@ elemental subroutine calculate_density_second_derivs_elem_Roquet_rho(this, T, S, + zt*(3.*EOS031 + (zs*(3.*EOS131) + zt*(4.*EOS041))) )) ) drho_dt_dp = (d2R_p0 + zp*(d2R_p1 + zp*d2R_p2)) +end subroutine calculate_density_second_derivs_elem_Roquet_rho_loc + +!> Wrapper for calculate_density_second_derivs_elem_Roquet_rho_loc created to preserve API +!! while calling without "this" variable that causes runtime errors on GPU with nvfortran. +elemental subroutine calculate_density_second_derivs_elem_Roquet_rho(this, T, S, pressure, & + drho_ds_ds, drho_ds_dt, drho_dt_dt, drho_ds_dp, drho_dt_dp) + class(Roquet_rho_EOS), intent(in) :: this !< This EOS + real, intent(in) :: T !< Conservative temperature [degC] + real, intent(in) :: S !< Absolute salinity [g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + real, intent(inout) :: drho_ds_ds !< Partial derivative of beta with respect + !! to S [kg m-3 ppt-2] + real, intent(inout) :: drho_ds_dt !< Partial derivative of beta with respect + !! to T [kg m-3 ppt-1 degC-1] + real, intent(inout) :: drho_dt_dt !< Partial derivative of alpha with respect + !! to T [kg m-3 degC-2] + real, intent(inout) :: drho_ds_dp !< Partial derivative of beta with respect + !! to pressure [kg m-3 ppt-1 Pa-1] = [s2 m-2 ppt-1] + real, intent(inout) :: drho_dt_dp !< Partial derivative of alpha with respect + !! to pressure [kg m-3 degC-1 Pa-1] = [s2 m-2 degC-1] + + call calculate_density_second_derivs_elem_Roquet_rho_loc(T, S, pressure, & + drho_ds_ds, drho_ds_dt, drho_dt_dt, drho_ds_dp, drho_dt_dp) + end subroutine calculate_density_second_derivs_elem_Roquet_rho !> Calculate the partial derivatives of specific volume with temperature and salinity @@ -721,6 +751,45 @@ subroutine calculate_density_array_2d_Roquet_rho(this, T, S, pressure, rho, dom, endif end subroutine calculate_density_array_2d_Roquet_rho +!> Calculate the in-situ density for 3D array inputs and outputs. +subroutine calculate_density_array_3d_Roquet_rho(this, T, S, pressure, rho, dom, rho_ref) + class(Roquet_rho_EOS), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:,:) + !< Conservative temperature [degC] + real, intent(in) :: S(:,:,:) + !< Absolute salinity [g kg-1] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: rho(:,:,:) + !< In situ density [kg m-3] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + real, optional, intent(in) :: rho_ref + !< A reference density [kg m-3] + + integer :: is, ie, js, je, ks, ke + integer :: i, j, k + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + ! The element functions are called via their free-function (_loc) forms rather than + ! through the polymorphic "this" binding, which causes runtime errors in do concurrent + ! regions offloaded to the GPU with nvfortran. + if (present(rho_ref)) then + do concurrent (k=ks:ke, j=js:je, i=is:ie) + rho(i,j,k) = density_anomaly_elem_Roquet_rho_loc(T(i,j,k), S(i,j,k), & + pressure(i,j,k), rho_ref) + enddo + else + do concurrent (k=ks:ke, j=js:je, i=is:ie) + rho(i,j,k) = density_elem_Roquet_rho_loc(T(i,j,k), S(i,j,k), pressure(i,j,k)) + enddo + endif +end subroutine calculate_density_array_3d_Roquet_rho + !> Calculate the in-situ density derivatives for 2D array inputs and outputs. subroutine calculate_density_derivs_2d_Roquet_rho(this, T, S, pressure, & drho_dT, drho_dS, dom) @@ -753,6 +822,79 @@ subroutine calculate_density_derivs_2d_Roquet_rho(this, T, S, pressure, & enddo end subroutine calculate_density_derivs_2d_Roquet_rho +!> Calculate the in-situ density derivatives for 3D array inputs and outputs. +subroutine calculate_density_derivs_3d_Roquet_rho(this, T, S, pressure, & + drho_dT, drho_dS, dom) + class(Roquet_rho_EOS), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:,:) + !< Conservative temperature [degC] + real, intent(in) :: S(:,:,:) + !< Absolute salinity [g kg-1] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: drho_dT(:,:,:) + !< Partial derivative of density with potential temperature [kg m-3 degC-1] + real, intent(out) :: drho_dS(:,:,:) + !< Partial derivative of density with salinity [kg m-3 ppt-1] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + + integer :: is, ie, js, je, ks, ke + integer :: i, j, k + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + ! The element subroutine is called via its free-function (_loc) form rather than + ! through the polymorphic "this" binding, which causes runtime errors in do concurrent + ! regions offloaded to the GPU with nvfortran. + do concurrent (k=ks:ke, j=js:je, i=is:ie) + call calculate_density_derivs_elem_Roquet_rho_loc(T(i,j,k), S(i,j,k), & + pressure(i,j,k), drho_dT(i,j,k), drho_dS(i,j,k)) + enddo +end subroutine calculate_density_derivs_3d_Roquet_rho + +!> Calculate the second derivatives of density for 2D array inputs and outputs. +subroutine calculate_density_second_derivs_2d_Roquet_rho(this, T, S, pressure, & + drho_dS_dS, drho_dS_dT, drho_dT_dT, drho_dS_dP, drho_dT_dP, dom) + class(Roquet_rho_EOS), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:) + !< Conservative temperature [degC] + real, intent(in) :: S(:,:) + !< Absolute salinity [g kg-1] + real, intent(in) :: pressure(:,:) + !< Pressure [Pa] + real, intent(inout) :: drho_dS_dS(:,:) + !< Partial derivative of beta with respect to S [kg m-3 ppt-2] + real, intent(inout) :: drho_dS_dT(:,:) + !< Partial derivative of beta with respect to T [kg m-3 ppt-1 degC-1] + real, intent(inout) :: drho_dT_dT(:,:) + !< Partial derivative of alpha with respect to T [kg m-3 degC-2] + real, intent(inout) :: drho_dS_dP(:,:) + !< Partial derivative of beta with respect to pressure [kg m-3 ppt-1 Pa-1] + real, intent(inout) :: drho_dT_dP(:,:) + !< Partial derivative of alpha with respect to pressure [kg m-3 degC-1 Pa-1] + integer, intent(in) :: dom(2,2) + !< Index bounds of domain. First index is rank, second is bounds + + integer :: is, ie, js, je + integer :: i, j + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + + ! The element subroutine is called via its free-function (_loc) form rather than + ! through the polymorphic "this" binding, which causes runtime errors in do concurrent + ! regions offloaded to the GPU with nvfortran. + do concurrent (j=js:je, i=is:ie) + call calculate_density_second_derivs_elem_Roquet_rho_loc(T(i,j), S(i,j), pressure(i,j), & + drho_dS_dS(i,j), drho_dS_dT(i,j), drho_dT_dT(i,j), drho_dS_dP(i,j), drho_dT_dP(i,j)) + enddo +end subroutine calculate_density_second_derivs_2d_Roquet_rho + !> Calculate the in-situ specific volume for 1D array inputs and outputs. subroutine calculate_spec_vol_array_Roquet_rho(this, T, S, pressure, specvol, start, npts, spv_ref) class(Roquet_rho_EOS), intent(in) :: this !< This EOS diff --git a/src/equation_of_state/MOM_EOS_Wright.F90 b/src/equation_of_state/MOM_EOS_Wright.F90 index c2861c451d..748ca12e69 100644 --- a/src/equation_of_state/MOM_EOS_Wright.F90 +++ b/src/equation_of_state/MOM_EOS_Wright.F90 @@ -71,10 +71,14 @@ module MOM_EOS_Wright procedure :: calculate_density_array => calculate_density_array_buggy_Wright !> Local implementation of generic calculate_density_array_2d for efficiency procedure :: calculate_density_array_2d => calculate_density_array_2d_buggy_Wright + !> Local implementation of generic calculate_density_array_3d for efficiency + procedure :: calculate_density_array_3d => calculate_density_array_3d_buggy_Wright !> Local implementation of generic calculate_spec_vol_array for efficiency procedure :: calculate_spec_vol_array => calculate_spec_vol_array_buggy_Wright !> Local implementation of generic calculate_density_derivs_2d for efficiency procedure :: calculate_density_derivs_2d => calculate_density_derivs_2d_buggy_Wright + !> Local implementation of generic calculate_density_derivs_3d for efficiency + procedure :: calculate_density_derivs_3d => calculate_density_derivs_3d_buggy_Wright end type buggy_Wright_EOS @@ -1016,6 +1020,46 @@ subroutine calculate_density_array_2d_buggy_Wright(this, T, S, pressure, rho, & endif end subroutine calculate_density_array_2d_buggy_Wright +!> Calculate the in-situ density for 3D array inputs and outputs. +subroutine calculate_density_array_3d_buggy_Wright(this, T, S, pressure, rho, & + dom, rho_ref) + class(buggy_Wright_EOS), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:,:) + !< Potential temperature relative to the surface [degC] + real, intent(in) :: S(:,:,:) + !< Salinity [PSU] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: rho(:,:,:) + !< In situ density [kg m-3] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + real, optional, intent(in) :: rho_ref + !< A reference density [kg m-3] + + integer :: is, ie, js, je, ks, ke + integer :: i, j, k + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + ! NOTE: There is an implicit copy of `this` which cannot yet be prevented. + ! Possibly because Nvidia cannot associate `this` with `EOS%type`. + + if (present(rho_ref)) then + do concurrent (k=ks:ke, j=js:je, i=is:ie) + rho(i,j,k) = density_anomaly_elem_buggy_Wright(this, T(i,j,k), S(i,j,k), & + pressure(i,j,k), rho_ref) + enddo + else + do concurrent (k=ks:ke, j=js:je, i=is:ie) + rho(i,j,k) = density_elem_buggy_Wright_loc( T(i,j,k), S(i,j,k), pressure(i,j,k)) + enddo + endif +end subroutine calculate_density_array_3d_buggy_Wright + !> Calculate the in-situ specific volume for 1D array inputs and outputs. subroutine calculate_spec_vol_array_buggy_Wright(this, T, S, pressure, specvol, start, npts, spv_ref) class(buggy_Wright_EOS), intent(in) :: this !< This EOS @@ -1075,6 +1119,39 @@ subroutine calculate_density_derivs_2d_buggy_Wright(this, T, S, pressure, & enddo end subroutine calculate_density_derivs_2d_buggy_Wright +!> Calculate the in-situ density derivatives for 3D array inputs and outputs. +subroutine calculate_density_derivs_3d_buggy_Wright(this, T, S, pressure, & + drho_dT, drho_dS, dom) + class(buggy_Wright_EOS), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:,:) + !< Potential temperature relative to the surface [degC] + real, intent(in) :: S(:,:,:) + !< Salinity [PSU] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: drho_dT(:,:,:) + !< Partial derivative of density with potential temperature [kg m-3 degC-1] + real, intent(out) :: drho_dS(:,:,:) + !< Partial derivative of density with salinity [kg m-3 PSU-1] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + + integer :: is, ie, js, je, ks, ke + integer :: i, j, k + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + ! NOTE: There is an implicit copy of `this` which cannot yet be prevented. + + do concurrent (k=ks:ke, j=js:je, i=is:ie) + call calculate_density_derivs_elem_buggy_Wright_loc( T(i,j,k), S(i,j,k), & + pressure(i,j,k), drho_dT(i,j,k), drho_dS(i,j,k)) + enddo +end subroutine calculate_density_derivs_3d_buggy_Wright + !> Set coefficients that can correct bugs un the buggy Wright equation of state. subroutine set_params_buggy_Wright(this, use_Wright_2nd_deriv_bug) class(buggy_Wright_EOS), intent(inout) :: this !< This EOS diff --git a/src/equation_of_state/MOM_EOS_base_type.F90 b/src/equation_of_state/MOM_EOS_base_type.F90 index 70d0c113dc..2b77af45e4 100644 --- a/src/equation_of_state/MOM_EOS_base_type.F90 +++ b/src/equation_of_state/MOM_EOS_base_type.F90 @@ -46,6 +46,8 @@ module MOM_EOS_base_type procedure :: calculate_density_array => a_calculate_density_array !> Calculates the in-situ density or density anomaly for 2d array inputs [m3 kg-1] procedure :: calculate_density_array_2d => a_calculate_density_array_2d + !> Calculates the in-situ density or density anomaly for 3d array inputs [m3 kg-1] + procedure :: calculate_density_array_3d => a_calculate_density_array_3d !> Calculates the in-situ specific volume or specific volume anomaly for scalar inputs [m3 kg-1] procedure :: calculate_spec_vol_scalar => a_calculate_spec_vol_scalar !> Calculates the in-situ specific volume or specific volume anomaly for array inputs [m3 kg-1] @@ -56,10 +58,14 @@ module MOM_EOS_base_type procedure :: calculate_density_derivs_array => a_calculate_density_derivs_array !> Calculates the derivatives of density for array inputs procedure :: calculate_density_derivs_2d => a_calculate_density_derivs_2d + !> Calculates the derivatives of density for 3d array inputs + procedure :: calculate_density_derivs_3d => a_calculate_density_derivs_3d !> Calculates the second derivatives of density for scalar inputs procedure :: calculate_density_second_derivs_scalar => a_calculate_density_second_derivs_scalar !> Calculates the second derivatives of density for array inputs procedure :: calculate_density_second_derivs_array => a_calculate_density_second_derivs_array + !> Calculates the second derivatives of density for 2d array inputs + procedure :: calculate_density_second_derivs_2d => a_calculate_density_second_derivs_2d !> Calculates the derivatives of specific volume for array inputs procedure :: calculate_specvol_derivs_array => a_calculate_specvol_derivs_array !> Calculates the compressibility for array inputs @@ -288,6 +294,37 @@ subroutine a_calculate_density_array_2d(this, T, S, pressure, rho, dom, rho_ref) endif end subroutine a_calculate_density_array_2d + !> Calculate the in-situ density for 3D array inputs and outputs. + subroutine a_calculate_density_array_3d(this, T, S, pressure, rho, dom, rho_ref) + class(EOS_base), intent(in) :: this !< This EOS + real, intent(in) :: T(:,:,:) + !< Potential temperature relative to the surface [degC] + real, intent(in) :: S(:,:,:) + !< Salinity [PSU] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: rho(:,:,:) + !< In situ density [kg m-3] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + real, optional, intent(in) :: rho_ref + !< A reference density [kg m-3] + + integer :: is, ie, js, je, ks, ke + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + if (present(rho_ref)) then + rho(is:ie, js:je, ks:ke) = this%density_anomaly_elem(T(is:ie, js:je, ks:ke), & + S(is:ie, js:je, ks:ke), pressure(is:ie, js:je, ks:ke), rho_ref) + else + rho(is:ie, js:je, ks:ke) = this%density_elem(T(is:ie, js:je, ks:ke), & + S(is:ie, js:je, ks:ke), pressure(is:ie, js:je, ks:ke)) + endif + end subroutine a_calculate_density_array_3d + !> In situ specific volume [m3 kg-1] real function a_spec_vol_fn(this, T, S, pressure, spv_ref) class(EOS_base), intent(in) :: this !< This EOS @@ -414,6 +451,35 @@ subroutine a_calculate_density_derivs_2d(this, T, S, pressure, drho_dT, drho_dS, pressure(is:ie, js:je), drho_dt(is:ie, js:je), drho_ds(is:ie, js:je)) end subroutine a_calculate_density_derivs_2d + !> Calculate the derivatives of density with respect to temperature and salinity + !! for 3d array inputs + subroutine a_calculate_density_derivs_3d(this, T, S, pressure, drho_dT, drho_dS, dom) + class(EOS_base), intent(in) :: this + !< This EOS + real, intent(in) :: T(:,:,:) + !< Potential temperature relative to the surface [degC] + real, intent(in) :: S(:,:,:) + !< Salinity [PSU] + real, intent(in) :: pressure(:,:,:) + !< Pressure [Pa] + real, intent(out) :: drho_dT(:,:,:) + !< The partial derivative of density with potential temperature + !! [kg m-3 degC-1] + real, intent(out) :: drho_dS(:,:,:) + !< The partial derivative of density with salinity, in [kg m-3 PSU-1] + integer, intent(in) :: dom(3,2) + !< Index bounds of domain. First index is rank, second is bounds + + integer :: is, ie, js, je, ks, ke + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + ks = dom(3,1) ; ke = dom(3,2) + + call this%calculate_density_derivs_elem(T(is:ie, js:je, ks:ke), S(is:ie, js:je, ks:ke), & + pressure(is:ie, js:je, ks:ke), drho_dt(is:ie, js:je, ks:ke), drho_ds(is:ie, js:je, ks:ke)) + end subroutine a_calculate_density_derivs_3d + !> Calculate the second derivatives of density with respect to temperature, salinity and pressure !! for scalar inputs subroutine a_calculate_density_second_derivs_scalar(this, T, S, pressure, & @@ -471,6 +537,32 @@ subroutine a_calculate_density_second_derivs_array(this, T, S, pressure, & end subroutine a_calculate_density_second_derivs_array + !> Default implementation: calculate second derivatives of density for 2d array inputs. + !! Delegates to the elemental calculate_density_second_derivs_elem via array sections. + subroutine a_calculate_density_second_derivs_2d(this, T, S, pressure, & + drho_ds_ds, drho_ds_dt, drho_dt_dt, drho_ds_dp, drho_dt_dp, dom) + class(EOS_base), intent(in) :: this !< This EOS + real, intent(in) :: T(:,:) !< Potential temperature [degC] + real, intent(in) :: S(:,:) !< Salinity [PSU] + real, intent(in) :: pressure(:,:) !< Pressure [Pa] + real, intent(inout) :: drho_ds_ds(:,:) !< Partial derivative of beta w.r.t. S [kg m-3 PSU-2] + real, intent(inout) :: drho_ds_dt(:,:) !< Partial derivative of beta w.r.t. T [kg m-3 PSU-1 degC-1] + real, intent(inout) :: drho_dt_dt(:,:) !< Partial derivative of alpha w.r.t. T [kg m-3 degC-2] + real, intent(inout) :: drho_ds_dp(:,:) !< Partial derivative of beta w.r.t. p [kg m-3 PSU-1 Pa-1] + real, intent(inout) :: drho_dt_dp(:,:) !< Partial derivative of alpha w.r.t. p [kg m-3 degC-1 Pa-1] + integer, intent(in) :: dom(2,2) !< Index bounds; first index is rank, second is bounds + + integer :: is, ie, js, je + + is = dom(1,1) ; ie = dom(1,2) + js = dom(2,1) ; je = dom(2,2) + + call this%calculate_density_second_derivs_elem( & + T(is:ie, js:je), S(is:ie, js:je), pressure(is:ie, js:je), & + drho_ds_ds(is:ie, js:je), drho_ds_dt(is:ie, js:je), drho_dt_dt(is:ie, js:je), & + drho_ds_dp(is:ie, js:je), drho_dt_dp(is:ie, js:je)) + end subroutine a_calculate_density_second_derivs_2d + !> Calculate the partial derivatives of specific volume with temperature and salinity !! for array inputs subroutine a_calculate_specvol_derivs_array(this, T, S, pressure, dSV_dT, dSV_dS, start, npts) diff --git a/src/framework/MOM_coms.F90 b/src/framework/MOM_coms.F90 index 806c7872f5..a0fac95475 100644 --- a/src/framework/MOM_coms.F90 +++ b/src/framework/MOM_coms.F90 @@ -27,37 +27,56 @@ module MOM_coms public :: EFP_plus, EFP_minus, EFP_to_real, real_to_EFP, EFP_real_diff public :: operator(+), operator(-), assignment(=) public :: query_EFP_overflow_error, reset_EFP_overflow_error -public :: max_count_prec -! This module provides interfaces to the non-domain-oriented communication subroutines. - -integer(kind=int64), parameter :: prec = (2_int64)**46 !< The precision of each integer. -real, parameter :: r_prec=2.0**46 !< A real version of prec [nondim]. -real, parameter :: I_prec=1.0/(2.0**46) !< The inverse of prec [nondim]. -integer, parameter :: max_count_prec=2**(63-46)-1 - !< The number of values that can be added together - !! with the current value of prec before there will - !! be roundoff problems. - -integer, parameter :: ni=6 !< The number of long integers to use to represent - !< a real number. -real, parameter, dimension(ni) :: & - pr = (/ r_prec**2, r_prec, 1.0, 1.0/r_prec, 1.0/r_prec**2, 1.0/r_prec**3 /) - !< An array of the real precision of each of the integers in arbitrary units [a] -real, parameter, dimension(ni) :: & - I_pr = (/ 1.0/r_prec**2, 1.0/r_prec, 1.0, r_prec, r_prec**2, r_prec**3 /) - !< An array of the inverse of the real precision of each of the integers in arbitrary units [a-1] -real, parameter :: max_efp_float = pr(1) * (2.**63 - 1.) - !< The largest float with an EFP representation in arbitrary units [a]. - !! NOTE: Only the first bin can exceed precision, - !! but is bounded by the largest signed integer. - -logical :: overflow_error = .false. !< This becomes true if an overflow is encountered. -logical :: NaN_error = .false. !< This becomes true if a NaN is encountered. -logical :: debug = .false. !< Making this true enables debugging output. +integer, parameter :: accum_width = digits(1_int64) + !< Accumulator width; total available bits for summation (excluding sign bit) +integer, parameter :: prec_width = 46 + !< Precision width; total bits for computed results +integer, parameter :: guard_width = accum_width - prec_width + !< Number of guard bits reserved for carry overflow + +! A sum of N points does N - 1 additions, which at most adds N - 1 carry bits. +! For G guard bits, the maximum value is 2**G - 1. A summation of N values +! therefore requires that N - 1 <= 2**G - 1, or simply N <= 2**G. + +integer, parameter :: max_summands = 2**guard_width + !< Maximum number of summable points that can guarantee no carry overflow. + !! Assumes that guard_bits is less than number of bits in a default integer. + +integer(kind=int64), parameter :: prec = (2_int64)**prec_width + !< EPF upper bound (exclusive). For each EPF bin e(i), 0 <= e(i) < prec. + +real, parameter :: r_prec = 2.**prec_width + !< Real-value of prec [nondim] +real, parameter :: I_prec = 2.**(-prec_width) + !< Inverse real-value of prec [nondim] + +integer, parameter :: efp_digits = 6 + !< The number of base `prec` digits used to represent an EFP value. +real, parameter, dimension(efp_digits) :: & + pr = [r_prec**2, r_prec, 1., r_prec**(-1), r_prec**(-2), r_prec**(-3)] + !< An array of the real precision of each of the integers in arbitrary + !! units [a] +real, parameter, dimension(efp_digits) :: & + I_pr = [r_prec**(-2), r_prec**(-1), 1., r_prec, r_prec**2, r_prec**3] + !< An array of the inverse of the real precision of each of the integers in + !! arbitrary units [a-1] +real, parameter :: max_efp_float = pr(1) * real(huge(1_int64)) + !< The largest float with an EFP representation in arbitrary units [a]. + !! NOTE: Only the first bin can exceed precision, but is bounded by the + !! largest signed integer. !$omp declare target(pr, I_pr) +logical :: overflow_error = .false. + !< This becomes true if an overflow is encountered. +logical :: NaN_error = .false. + !< This becomes true if a NaN is encountered. +logical :: debug = .false. + !< Making this true enables debugging output. + +! This module provides interfaces to the non-domain-oriented communication subroutines. + !> Find an accurate and order-invariant sum of a distributed 2d or 3d field, in some cases after !! undoing the scaling of the input array and restoring that scaling in the returned value interface reproducing_sum @@ -82,7 +101,7 @@ module MOM_coms !! Hallberg, R. & A. Adcroft, 2014: An Order-invariant Real-to-Integer Conversion Sum. !! Parallel Computing, 40(5-6), doi:10.1016/j.parco.2014.04.007. type, public :: EFP_type ; private - integer(kind=int64), dimension(ni) :: v !< The value in this type + integer(kind=int64), dimension(efp_digits) :: v !< The value in this type end type EFP_type !> Add two extended-fixed-point numbers @@ -127,7 +146,7 @@ function reproducing_EFP_sum_2d(array, isr, ier, jsr, jer, overflow_check, err, ! of real numbers to give order-invariant sums that will reproduce ! across PE count. This idea comes from R. Hallberg and A. Adcroft. - integer(kind=int64), dimension(ni) :: ints_sum + integer(kind=int64), dimension(efp_digits) :: ints_sum integer(kind=int64) :: ival, prec_error real :: rs ! The remaining value to add, in arbitrary units [a] real :: max_mag_term ! A running maximum magnitude of the values in arbitrary units [a] @@ -136,11 +155,11 @@ function reproducing_EFP_sum_2d(array, isr, ier, jsr, jer, overflow_check, err, character(len=256) :: mesg integer :: i, j, n, is, ie, js, je, sgn - if (num_PEs() > max_count_prec) call MOM_error(FATAL, & + if (num_PEs() > max_summands) call MOM_error(FATAL, & "reproducing_sum: Too many processors are being used for the value of "//& "prec. Reduce prec to (2^63-1)/num_PEs.") - prec_error = ((2_int64)**62 + ((2_int64)**62 - 1)) / num_PEs() + prec_error = huge(1_int64) / num_PEs() is = 1 ; ie = size(array,1) ; js = 1 ; je = size(array,2) if (present(isr)) then @@ -166,29 +185,16 @@ function reproducing_EFP_sum_2d(array, isr, ier, jsr, jer, overflow_check, err, descale = 1.0 ; if (do_unscale) descale = unscale overflow_error = .false. ; NaN_error = .false. ; max_mag_term = 0.0 + ints_sum(:) = 0 if (over_check) then - if ((je+1-js)*(ie+1-is) < max_count_prec) then - ! Common case: window is small enough that all carrying happens at the tail. - call increment_ints_2d(array, is, ie, js, je, descale, ints_sum, max_mag_term) - call carry_overflow(ints_sum, prec_error) - elseif ((ie+1-is) < max_count_prec) then - do j=js,je - do i=is,ie - call increment_ints_faster(ints_sum, descale*array(i,j), max_mag_term) - enddo - call carry_overflow(ints_sum, prec_error) - enddo - else - do j=js,je ; do i=is,ie - call increment_ints(ints_sum, real_to_ints(descale*array(i,j), prec_error), prec_error) - enddo ; enddo - endif + call increment_block_ints(array, is, ie, js, je, descale, ints_sum, & + max_mag_term, prec_error) else do j=js,je ; do i=is,ie sgn = 1 ; if (array(i,j)<0.0) sgn = -1 rs = abs(descale*array(i,j)) - do n=1,ni + do n=1,efp_digits ival = int(rs*I_pr(n), kind=int64) rs = rs - ival*pr(n) ints_sum(n) = ints_sum(n) + sgn*ival @@ -203,7 +209,7 @@ function reproducing_EFP_sum_2d(array, isr, ier, jsr, jer, overflow_check, err, err = err+2 if (NaN_error) & err = err+4 - if (err > 0) then ; do n=1,ni ; ints_sum(n) = 0 ; enddo ; endif + if (err > 0) then ; do n=1,efp_digits ; ints_sum(n) = 0 ; enddo ; endif else if (NaN_error) then call MOM_error(FATAL, "NaN in input field of reproducing_EFP_sum(_2d).") @@ -217,7 +223,7 @@ function reproducing_EFP_sum_2d(array, isr, ier, jsr, jer, overflow_check, err, endif endif - if (do_sum_across_PEs) call sum_across_PEs(ints_sum, ni) + if (do_sum_across_PEs) call sum_across_PEs(ints_sum, efp_digits) call regularize_ints(ints_sum) @@ -260,7 +266,7 @@ function reproducing_sum_2d(array, isr, ier, jsr, jer, EFP_sum, reproducing, & !! arbitrary units as array [a] or [A ~> a] ! Local variables - integer(kind=int64), dimension(ni) :: ints_sum + integer(kind=int64), dimension(efp_digits) :: ints_sum integer(kind=int64) :: prec_error real :: rsum(1) ! The running sum, in arbitrary units [a] real :: descale ! A local copy of unscale if it is present [a A-1 ~> 1] or 1 @@ -270,11 +276,11 @@ function reproducing_sum_2d(array, isr, ier, jsr, jer, EFP_sum, reproducing, & type(EFP_type) :: EFP_val ! An extended fixed point version of the sum integer :: i, j, is, ie, js, je - if (num_PEs() > max_count_prec) call MOM_error(FATAL, & + if (num_PEs() > max_summands) call MOM_error(FATAL, & "reproducing_sum: Too many processors are being used for the value of "//& "prec. Reduce prec to (2^63-1)/num_PEs.") - prec_error = ((2_int64)**62 + ((2_int64)**62 - 1)) / num_PEs() + prec_error = huge(1_int64) / num_PEs() is = 1 ; ie = size(array,1) ; js = 1 ; je = size(array,2) if (present(isr)) then @@ -334,7 +340,7 @@ function reproducing_sum_2d(array, isr, ier, jsr, jer, EFP_sum, reproducing, & endif if (debug) then - write(mesg,'("2d RS: ", ES24.16, 6 Z17.16)') sum*descale, ints_sum(1:ni) + write(mesg,'("2d RS: ", ES24.16, 6 Z17.16)') sum*descale, ints_sum(1:efp_digits) call MOM_mesg(mesg, 3) endif @@ -376,18 +382,18 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su real :: max_mag_term ! A running maximum magnitude of the val's in arbitrary units [a] real :: descale ! A local copy of unscale if it is present [a A-1 ~> 1] or 1 real :: I_unscale ! The Adcroft reciprocal of unscale [A a-1 ~> 1] - integer(kind=int64), dimension(ni) :: ints_sum - integer(kind=int64), dimension(ni,size(array,3)) :: ints_sums + integer(kind=int64), dimension(efp_digits) :: ints_sum + integer(kind=int64), dimension(efp_digits,size(array,3)) :: ints_sums integer(kind=int64) :: prec_error character(len=256) :: mesg logical :: do_sum_across_PEs, do_unscale integer :: i, j, k, is, ie, js, je, ke, isz, jsz, n - if (num_PEs() > max_count_prec) call MOM_error(FATAL, & + if (num_PEs() > max_summands) call MOM_error(FATAL, & "reproducing_sum: Too many processors are being used for the value of "//& "prec. Reduce prec to (2^63-1)/num_PEs.") - prec_error = ((2_int64)**62 + ((2_int64)**62 - 1)) / num_PEs() + prec_error = huge(1_int64) / num_PEs() max_mag_term = 0.0 is = 1 ; ie = size(array,1) ; js = 1 ; je = size(array,2) ; ke = size(array,3) @@ -420,32 +426,21 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su if (present(EFP_lay_sums)) then ; if (size(EFP_lay_sums) < ke) then call MOM_error(FATAL, "Sums is smaller than the vertical extent of array in reproducing_sum(_3d).") endif ; endif - ints_sums(:,:) = 0 + overflow_error = .false. ; NaN_error = .false. ; max_mag_term = 0.0 - if (jsz*isz < max_count_prec) then - do k=1,ke - call increment_ints_2d(array(:,:,k), is, ie, js, je, descale, ints_sums(:,k), max_mag_term) - call carry_overflow(ints_sums(:,k), prec_error) - enddo - elseif (isz < max_count_prec) then - do k=1,ke ; do j=js,je - do i=is,ie - call increment_ints_faster(ints_sums(:,k), descale*array(i,j,k), max_mag_term) - enddo - call carry_overflow(ints_sums(:,k), prec_error) - enddo ; enddo - else - do k=1,ke ; do j=js,je ; do i=is,ie - call increment_ints(ints_sums(:,k), & - real_to_ints(descale*array(i,j,k), prec_error), prec_error) - enddo ; enddo ; enddo - endif + + ints_sums(:,:) = 0 + do k=1,ke + call increment_block_ints(array(:,:,k), is, ie, js, je, descale, & + ints_sums(:,k), max_mag_term, prec_error) + enddo + if (present(err)) then err = 0 if (abs(max_mag_term) >= prec_error*pr(1)) err = err+1 if (overflow_error) err = err+2 if (NaN_error) err = err+2 - if (err > 0) then ; do k=1,ke ; do n=1,ni ; ints_sums(n,k) = 0 ; enddo ; enddo ; endif + if (err > 0) then ; do k=1,ke ; do n=1,efp_digits ; ints_sums(n,k) = 0 ; enddo ; enddo ; endif else if (NaN_error) call MOM_error(FATAL, "NaN in input field of reproducing_sum(_3d).") if (abs(max_mag_term) >= prec_error*pr(1)) then @@ -455,7 +450,7 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su if (overflow_error) call MOM_error(FATAL, "Overflow in reproducing_sum(_3d).") endif - if (do_sum_across_PEs) call sum_across_PEs(ints_sums(:,1:ke), ni*ke) + if (do_sum_across_PEs) call sum_across_PEs(ints_sums(:,1:ke), efp_digits*ke) sum = 0.0 do k=1,ke @@ -474,38 +469,26 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su endif if (debug) then - do n=1,ni ; ints_sum(n) = 0 ; enddo - do k=1,ke ; do n=1,ni ; ints_sum(n) = ints_sum(n) + ints_sums(n,k) ; enddo ; enddo - write(mesg,'("3D RS: ", ES24.16, 6 Z17.16)') sum, ints_sum(1:ni) + do n=1,efp_digits ; ints_sum(n) = 0 ; enddo + do k=1,ke ; do n=1,efp_digits ; ints_sum(n) = ints_sum(n) + ints_sums(n,k) ; enddo ; enddo + write(mesg,'("3D RS: ", ES24.16, 6 Z17.16)') sum, ints_sum(1:efp_digits) call MOM_mesg(mesg, 3) endif else - ints_sum(:) = 0 overflow_error = .false. ; NaN_error = .false. ; max_mag_term = 0.0 - if (jsz*isz < max_count_prec) then - do k=1,ke - call increment_ints_2d(array(:,:,k), is, ie, js, je, descale, ints_sum, max_mag_term) - call carry_overflow(ints_sum, prec_error) - enddo - elseif (isz < max_count_prec) then - do k=1,ke ; do j=js,je - do i=is,ie - call increment_ints_faster(ints_sum, descale*array(i,j,k), max_mag_term) - enddo - call carry_overflow(ints_sum, prec_error) - enddo ; enddo - else - do k=1,ke ; do j=js,je ; do i=is,ie - call increment_ints(ints_sum, real_to_ints(descale*array(i,j,k), prec_error), & - prec_error) - enddo ; enddo ; enddo - endif + + ints_sum(:) = 0 + do k=1,ke + call increment_block_ints(array(:,:,k), is, ie, js, je, descale, & + ints_sum, max_mag_term, prec_error) + enddo + if (present(err)) then err = 0 if (abs(max_mag_term) >= prec_error*pr(1)) err = err+1 if (overflow_error) err = err+2 if (NaN_error) err = err+2 - if (err > 0) then ; do n=1,ni ; ints_sum(n) = 0 ; enddo ; endif + if (err > 0) then ; do n=1,efp_digits ; ints_sum(n) = 0 ; enddo ; endif else if (NaN_error) call MOM_error(FATAL, "NaN in input field of reproducing_sum(_3d).") if (abs(max_mag_term) >= prec_error*pr(1)) then @@ -515,7 +498,7 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su if (overflow_error) call MOM_error(FATAL, "Overflow in reproducing_sum(_3d).") endif - if (do_sum_across_PEs) call sum_across_PEs(ints_sum, ni) + if (do_sum_across_PEs) call sum_across_PEs(ints_sum, efp_digits) call regularize_ints(ints_sum) sum = ints_to_real(ints_sum) @@ -523,7 +506,7 @@ function reproducing_sum_3d(array, isr, ier, jsr, jer, sums, EFP_sum, EFP_lay_su if (present(EFP_sum)) EFP_sum%v(:) = ints_sum(:) if (debug) then - write(mesg,'("3d RS: ", ES24.16, 6 Z17.16)') sum, ints_sum(1:ni) + write(mesg,'("3d RS: ", ES24.16, 6 Z17.16)') sum, ints_sum(1:efp_digits) call MOM_mesg(mesg, 3) endif endif @@ -548,7 +531,7 @@ function real_to_ints(r, prec_error, overflow) result(ints) !! precision parameter, and is used to detect overflows. logical, optional, intent(inout) :: overflow !< Returns true if the conversion is being !! done on a value that is too large to be represented - integer(kind=int64), dimension(ni) :: ints + integer(kind=int64), dimension(efp_digits) :: ints ! This subroutine converts a real number to an equivalent representation ! using several long integers. @@ -574,7 +557,7 @@ function real_to_ints(r, prec_error, overflow) result(ints) call MOM_error(FATAL,"Overflow in real_to_ints conversion of "//trim(mesg)) endif - do i=1,ni + do i=1,efp_digits ival = int(rs*I_pr(i), kind=int64) rs = rs - ival*pr(i) ints(i) = sgn*ival @@ -585,21 +568,21 @@ end function real_to_ints !> Convert the array of integers that constitute an extended-fixed-point !! representation into a real number function ints_to_real(ints) result(r) - integer(kind=int64), dimension(ni), intent(in) :: ints !< The array of EFP integers + integer(kind=int64), dimension(efp_digits), intent(in) :: ints !< The array of EFP integers real :: r ! The real number that is extracted in arbitrary units [a] ! This subroutine reverses the conversion in real_to_ints. integer :: i r = 0.0 - do i=1,ni ; r = r + pr(i)*ints(i) ; enddo + do i=1,efp_digits ; r = r + pr(i)*ints(i) ; enddo end function ints_to_real !> Increment an array of integers that constitutes an extended-fixed-point !! representation with a another EFP number subroutine increment_ints(int_sum, int2, prec_error) - integer(kind=int64), dimension(ni), intent(inout) :: int_sum !< The array of EFP integers being incremented - integer(kind=int64), dimension(ni), intent(in) :: int2 !< The array of EFP integers being added + integer(kind=int64), dimension(efp_digits), intent(inout) :: int_sum !< The array of EFP integers being incremented + integer(kind=int64), dimension(efp_digits), intent(in) :: int2 !< The array of EFP integers being added integer(kind=int64), optional, intent(in) :: prec_error !< The PE-count dependent precision of the !! integers that is safe from overflows during global !! sums. This will be larger than the compile-time @@ -609,7 +592,7 @@ subroutine increment_ints(int_sum, int2, prec_error) ! representation in real_to_ints. integer :: i - do i=ni,2,-1 + do i=efp_digits,2,-1 int_sum(i) = int_sum(i) + int2(i) ! Carry the local overflow. if (int_sum(i) > prec) then @@ -630,45 +613,12 @@ subroutine increment_ints(int_sum, int2, prec_error) end subroutine increment_ints -!> Increment an EFP number with a real number without doing any carrying of -!! of overflows and using only minimal error checking. -subroutine increment_ints_faster(int_sum, r, max_mag_term) - integer(kind=int64), intent(inout) :: int_sum(ni) - !< The array of EFP integers being incremented - real, intent(in) :: r - !< The real number being added in arbitrary units [a] - real, intent(inout) :: max_mag_term - !< A running maximum magnitude of the r's in arbitrary units [a] - - ! This subroutine increments a number with another, both using the integer - ! representation in real_to_ints, but without doing any carrying of overflow. - ! The per-element decomposition is shared with the GPU kernel via efp_decompose. - integer(kind=int64) :: e(ni) - real :: rmag - integer :: is_nan, is_ovf - - call efp_decompose(r, e, rmag, is_nan, is_ovf) - - if (is_nan /= 0) then ; NaN_error = .true. ; return ; endif - if (rmag > abs(max_mag_term)) max_mag_term = r - - ! Abort if the number has no EFP representation - if (is_ovf /= 0) then ; overflow_error = .true. ; return ; endif - - int_sum(:) = int_sum(:) + e(:) - -end subroutine increment_ints_faster - - -!> Increment an EFP number with a real number over a 2d array without doing any -!! carrying of overflows and using only minimal error checking. -!! modulo: max_mag_term is updated with the magnitude (>=0) rather than the -!! signed last-winner. Only consumer of max_mag_term is abs() in the overflow -!! guard, so values are unchanged; only the sign in one FATAL message can -!! differ. -subroutine increment_ints_2d(array, is, ie, js, je, descale, ints_sum, max_mag_term) +!> Sum the elements of an array in EFP form and append the result to an +!! existing EFP array. +subroutine increment_block_ints(array, is, ie, js, je, descale, ints_sum, & + max_mag_term, prec_error) real, intent(in) :: array(:,:) - !< The field being added, in arbitrary units [a] + !< The field being added, in arbitrary units [A ~> a] integer, intent(in) :: is !< Start i-index of the summed domain integer, intent(in) :: ie @@ -678,54 +628,158 @@ subroutine increment_ints_2d(array, is, ie, js, je, descale, ints_sum, max_mag_t integer, intent(in) :: je !< End j-index of the summed domain real, intent(in) :: descale - !< unscale factor or 1.0 [a A-1 ~> 1] - integer(kind=int64), intent(inout) :: ints_sum(ni) + !< Factor to descale array to physical value [a A-1 ~> 1] + integer(kind=int64), intent(inout) :: ints_sum(efp_digits) !< The array of EFP integers being incremented real, intent(inout) :: max_mag_term !< A running maximum magnitude of the r's, in arbitrary units [a] - - integer :: i, j - integer(kind=int64) :: e(ni) - real :: r, rmag, mmag + integer(kind=int64), intent(in) :: prec_error + !< The maximum resolvable value for a given number of PEs + + integer :: i, j, ib, jb, ibs, ibe, jbs, jbe + ! Loop indices + integer :: b + ! Block counter + integer :: ni, nj + ! Array summation domain size along each axis + integer :: isize_max + ! Largest block size in i. Typically equal to ni + integer :: jsize + ! Number of j-rows per block. + integer :: nblocks, niblocks, njblocks + ! Number of total blocks, and number of blocks in i and j + integer(kind=int64) :: e(efp_digits) + ! The EPF representation of each array element + integer(kind=int64) :: block_sum(efp_digits), array_sum(efp_digits) + ! The cumulant per-block and total array EFP sums + real :: r, rmag + ! Local array element value and its magnitude [a] + real :: max_pos, max_neg, block_max_pos, block_max_neg + ! Largest positive and negative values (whole array and per-block) used to + ! find the largest maximum magnitude of array in a thread-safe manner [a] integer :: inan, iovf, lnan, lovf + ! Thread-safe tracking of NaN and overflow state + integer :: max_sum_count + ! The total number of local sum operations that ensures no carry overflow - mmag = abs(max_mag_term) + max_pos = max(0., max_mag_term) + max_neg = max(0., -max_mag_term) inan = 0 ; iovf = 0 - ! This subroutine increments a number with another, both using the integer - ! representation in real_to_ints, but without doing any carrying of overflow. - do concurrent (j=js:je, i=is:ie) & - DO_LOCALITY(local(r, e, rmag, lnan, lovf)) & - DO_LOCALITY(reduce(+: ints_sum) reduce(max: mmag, inan, iovf)) + ! Reduce the maximum number of summations to account for the cumulant + ! summations of array_sum and ints_sum. + max_sum_count = max_summands - 2 + + ! Get the compute domain size + ni = ie - is + 1 + nj = je - js + 1 + + ! Partition in i so that the widest i-slice fits within max_sum_count. + niblocks = (ni + max_sum_count - 1) / max_sum_count + ! = ⌈ni / max_sum_count⌉ + + ! NOTE: niblocks is typically one, since default max_sum_count is ~130k. + + ! For a balanced i-partition, the number of i-points per block is either + ! ⌊ni / niblocks⌋ or ⌈ni / niblocks⌉. Use the upper bound to find jsize. + + isize_max = (ni + niblocks - 1) / niblocks + ! = ⌈ni / niblocks⌉ + + ! Set jsize so that the widest i-slice times the number of j-rows does not + ! exceed max_sum_count. + jsize = max_sum_count / isize_max + ! = ⌊max_sum_count / isize_max⌋ + + ! Choose enough j-blocks so that no j-block has more than jsize rows. + njblocks = (nj + jsize - 1) / jsize + ! = ⌈nj / jsize⌉ + + nblocks = niblocks * njblocks - r = descale*array(i,j) + ! Abort if the number of blocks also exceeds the carry-bit summation limit. + ! For default settings, this would be over 17 billion points per PE. + if (nblocks > max_sum_count) call MOM_error(FATAL, & + "reproducing sum: Number of blocks exceeds summmation carry limit." & + ) - call efp_decompose(r, e, rmag, lnan, lovf) + array_sum(:) = 0 - inan = max(inan, lnan) - iovf = max(iovf, lovf) + do jb=1,njblocks ; do ib=1,niblocks + ! Use evenly distributed blocks, either ⌊n / nblocks⌋ or ⌈n / nblocks⌉. + jbs = js + ((jb - 1) * nj) / njblocks + jbe = js + (jb * nj) / njblocks - 1 - if (rmag > mmag) mmag = rmag + ibs = is + ((ib - 1) * ni) / niblocks + ibe = is + (ib * ni) / niblocks - 1 - ints_sum(:) = ints_sum(:) + e(:) + block_sum(:) = 0 + block_max_pos = 0. + block_max_neg = 0. + + ! Compute the sum of each block + do concurrent (j=jbs:jbe, i=ibs:ibe) & + DO_LOCALITY(local(r, e, rmag, lnan, lovf)) & + DO_LOCALITY(reduce(+: block_sum)) & + DO_LOCALITY(reduce(max: block_max_pos, block_max_neg, inan, iovf)) + + ! Convert array(i,j) to EFP form + r = descale * array(i,j) + call efp_decompose(r, e, rmag, lnan, lovf) + + ! Verify that the conversion was completed + inan = max(inan, lnan) + iovf = max(iovf, lovf) + + if (r >= 0.) then + if (rmag > block_max_pos) block_max_pos = rmag + else + if (rmag > block_max_neg) block_max_neg = rmag + endif + + ! Add the EFP result (including potential carry bits) + block_sum(:) = block_sum(:) + e(:) + enddo ; enddo + + array_sum(:) = array_sum(:) + block_sum(:) + + ! Redistribute carry bits across bins + ! For the final pass (or single pass) this is handled by ints_sum. + b = (jb - 1) * niblocks + ib + if (b < nblocks) call carry_overflow(array_sum, prec_error) + + ! Update maximum magnitudes + max_pos = max(max_pos, block_max_pos) + max_neg = max(max_neg, block_max_neg) enddo - max_mag_term = mmag + ! Finally, apply the cumulant result + ints_sum(:) = ints_sum(:) + array_sum(:) + + ! Redistribute carry bits to normalize the final result. + call carry_overflow(ints_sum, prec_error) + + ! Extract the maximum value while preserving sign (NOTE: ties to positive) + if (max_pos >= max_neg) then + max_mag_term = max_pos + else + max_mag_term = -max_neg + endif + ! Transfer error/warning signals to module flags if (inan /= 0) NaN_error = .true. if (iovf /= 0) overflow_error = .true. -end subroutine increment_ints_2d +end subroutine increment_block_ints -!> Decompose one real into its 6 signed EFP bin contributions. The function is -!! used for both CPU and GPU kernels. NaNs and overflows are reported by -!! flags, rather than the module-level error logicals, so that the routine is -!! free of side effects. +!> Decompose one real into its 6 signed EFP bin contributions. NaNs and +!! overflows are reported by flags, rather than the module-level error +!! logicals, so that the routine is free of side effects. pure subroutine efp_decompose(r, e, rmag, is_nan, is_ovf) !$omp declare target real, intent(in) :: r !< The real number being decomposed [a] - integer(kind=int64), intent(out) :: e(ni) + integer(kind=int64), intent(out) :: e(efp_digits) !< Signed contribution to EFP bins real, intent(out) :: rmag !< Equals abs(r), or 0 if r is NaN/Inf [a] @@ -759,7 +813,7 @@ pure subroutine efp_decompose(r, e, rmag, is_nan, is_ovf) return endif - do n=1,ni + do n=1,efp_digits ival = int(rs * I_pr(n), kind=int64) rs = rs - ival * pr(n) e(n) = sgn * ival @@ -769,7 +823,7 @@ end subroutine efp_decompose !> This subroutine handles carrying of the overflow. subroutine carry_overflow(int_sum, prec_error) - integer(kind=int64), dimension(ni), intent(inout) :: int_sum !< The array of EFP integers being + integer(kind=int64), dimension(efp_digits), intent(inout) :: int_sum !< The array of EFP integers being !! modified by carries, but without changing value. integer(kind=int64), intent(in) :: prec_error !< The PE-count dependent precision of the !! integers that is safe from overflows during global @@ -779,7 +833,7 @@ subroutine carry_overflow(int_sum, prec_error) ! This subroutine handles carrying of the overflow. integer :: i, num_carry - do i=ni,2,-1 ; if (abs(int_sum(i)) >= prec) then + do i=efp_digits,2,-1 ; if (abs(int_sum(i)) >= prec) then num_carry = int(int_sum(i) * I_prec) int_sum(i) = int_sum(i) - num_carry*prec int_sum(i-1) = int_sum(i-1) + num_carry @@ -793,7 +847,7 @@ end subroutine carry_overflow !> This subroutine carries the overflow, and then makes sure that !! all integers are of the same sign as the overall value. subroutine regularize_ints(int_sum) - integer(kind=int64), dimension(ni), & + integer(kind=int64), dimension(efp_digits), & intent(inout) :: int_sum !< The array of integers being modified to take a !! regular form with all integers of the same sign, !! but without changing value. @@ -803,7 +857,7 @@ subroutine regularize_ints(int_sum) logical :: positive integer :: i, num_carry - do i=ni,2,-1 ; if (abs(int_sum(i)) >= prec) then + do i=efp_digits,2,-1 ; if (abs(int_sum(i)) >= prec) then num_carry = int(int_sum(i) * I_prec) int_sum(i) = int_sum(i) - num_carry*prec int_sum(i-1) = int_sum(i-1) + num_carry @@ -811,7 +865,7 @@ subroutine regularize_ints(int_sum) ! Determine the sign of the final number. positive = .true. - do i=1,ni + do i=1,efp_digits if (abs(int_sum(i)) > 0) then if (int_sum(i) < 0) positive = .false. exit @@ -819,12 +873,12 @@ subroutine regularize_ints(int_sum) enddo if (positive) then - do i=ni,2,-1 ; if (int_sum(i) < 0) then + do i=efp_digits,2,-1 ; if (int_sum(i) < 0) then int_sum(i) = int_sum(i) + prec int_sum(i-1) = int_sum(i-1) - 1 endif ; enddo else - do i=ni,2,-1 ; if (int_sum(i) > 0) then + do i=efp_digits,2,-1 ; if (int_sum(i) > 0) then int_sum(i) = int_sum(i) - prec int_sum(i-1) = int_sum(i-1) + 1 endif ; enddo @@ -862,7 +916,7 @@ function EFP_minus(EFP1, EFP2) !! subtracted from the first extended fixed point number integer :: i - do i=1,ni ; EFP_minus%v(i) = -1*EFP2%v(i) ; enddo + do i=1,efp_digits ; EFP_minus%v(i) = -1*EFP2%v(i) ; enddo call increment_ints(EFP_minus%v(:), EFP1%v(:)) end function EFP_minus @@ -876,7 +930,7 @@ subroutine EFP_assign(EFP1, EFP2) ! variable on the RHS (EFP2) to the components of the variable on the LHS ! (EFP1). - do i=1,ni ; EFP1%v(i) = EFP2%v(i) ; enddo + do i=1,efp_digits ; EFP1%v(i) = EFP2%v(i) ; enddo end subroutine EFP_assign !> Return the real number that an extended-fixed-point number corresponds with @@ -939,29 +993,30 @@ subroutine EFP_list_sum_across_PEs(EFPs, nval, errors) ! This subroutine does a sum across PEs of a list of EFP variables, ! returning the sums in place, with all overflows carried. - integer(kind=int64), dimension(ni,nval) :: ints + integer(kind=int64), dimension(efp_digits,nval) :: ints integer(kind=int64) :: prec_error logical :: error_found character(len=256) :: mesg integer :: i, n - if (num_PEs() > max_count_prec) call MOM_error(FATAL, & + if (num_PEs() > max_summands) call MOM_error(FATAL, & "reproducing_sum: Too many processors are being used for the value of "//& "prec. Reduce prec to (2^63-1)/num_PEs.") - prec_error = ((2_int64)**62 + ((2_int64)**62 - 1)) / num_PEs() + prec_error = huge(1_int64) / num_PEs() + ! overflow_error is an overflow error flag for the whole module. overflow_error = .false. ; error_found = .false. - do i=1,nval ; do n=1,ni ; ints(n,i) = EFPs(i)%v(n) ; enddo ; enddo + do i=1,nval ; do n=1,efp_digits ; ints(n,i) = EFPs(i)%v(n) ; enddo ; enddo - call sum_across_PEs(ints(:,:), ni*nval) + call sum_across_PEs(ints(:,:), efp_digits*nval) if (present(errors)) errors(:) = .false. do i=1,nval overflow_error = .false. call carry_overflow(ints(:,i), prec_error) - do n=1,ni ; EFPs(i)%v(n) = ints(n,i) ; enddo + do n=1,efp_digits ; EFPs(i)%v(n) = ints(n,i) ; enddo if (present(errors)) errors(i) = overflow_error if (overflow_error) then write (mesg,'("EFP_list_sum_across_PEs error at ",i0," val was ",ES12.6, ", prec_error = ",ES12.6)') & @@ -986,29 +1041,30 @@ subroutine EFP_val_sum_across_PEs(EFP, error) ! This subroutine does a sum across PEs of a list of EFP variables, ! returning the sums in place, with all overflows carried. - integer(kind=int64), dimension(ni) :: ints + integer(kind=int64), dimension(efp_digits) :: ints integer(kind=int64) :: prec_error logical :: error_found character(len=256) :: mesg integer :: n - if (num_PEs() > max_count_prec) call MOM_error(FATAL, & + if (num_PEs() > max_summands) call MOM_error(FATAL, & "reproducing_sum: Too many processors are being used for the value of "//& "prec. Reduce prec to (2^63-1)/num_PEs.") - prec_error = ((2_int64)**62 + ((2_int64)**62 - 1)) / num_PEs() + prec_error = huge(1_int64) / num_PEs() + ! overflow_error is an overflow error flag for the whole module. overflow_error = .false. ; error_found = .false. - do n=1,ni ; ints(n) = EFP%v(n) ; enddo + do n=1,efp_digits ; ints(n) = EFP%v(n) ; enddo - call sum_across_PEs(ints(:), ni) + call sum_across_PEs(ints(:), efp_digits) if (present(error)) error = .false. overflow_error = .false. call carry_overflow(ints(:), prec_error) - do n=1,ni ; EFP%v(n) = ints(n) ; enddo + do n=1,efp_digits ; EFP%v(n) = ints(n) ; enddo if (present(error)) error = overflow_error if (overflow_error) then write (mesg,'("EFP_val_sum_across_PEs error val was ",ES12.6, ", prec_error = ",ES12.6)') &