From 93dbbd36e8f59973e288703f148c2fa5214f41c2 Mon Sep 17 00:00:00 2001 From: Edward Yang <94523015+edoyango@users.noreply.github.com> Date: Sat, 4 Jul 2026 04:18:55 +1000 Subject: [PATCH 1/5] Use blocking in k dimension for continuity reconstruction (#165) * block k in ppm_reconstruction * add error for -ve block size * remove nvf inline and replace with intel forceinline Significantly improves performance of blocked zonal/meridional_mass_flux at -O2 * change default block sizes at compile time if __NVCOMPILER_OPENMP_GPU macro is defined (to be replaced at a later time), set default_n?block to 0. Additionally, in general, block size of 0 means the block size is dynamic to the computational domain. --- src/core/MOM_continuity_PPM.F90 | 444 +++++++++++++++++--------------- 1 file changed, 234 insertions(+), 210 deletions(-) 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) From 7c7af5572fd3dc96b8e253a9da075845039ea4d4 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Fri, 3 Jul 2026 14:58:51 -0400 Subject: [PATCH 2/5] EOS: 2D and 3D density implementations of methods (#185) The following have been added to Roquet and Wright equations of state: module procedure calculate_density_3d module procedure calculate_stanley_density_2d module procedure calculate_density_derivs_3d module procedure calculate_density_second_derivs_2d These were extracted from a larger pull request supporting pressure density integrals (#156) --- src/equation_of_state/MOM_EOS.F90 | 241 +++++++++++++++++++ src/equation_of_state/MOM_EOS_Roquet_rho.F90 | 174 +++++++++++-- src/equation_of_state/MOM_EOS_Wright.F90 | 77 ++++++ src/equation_of_state/MOM_EOS_base_type.F90 | 92 +++++++ 4 files changed, 568 insertions(+), 16 deletions(-) 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) From 8593a732aa38c60b384af6cc1ff916c769ceab5e Mon Sep 17 00:00:00 2001 From: Marshall Ward Date: Thu, 2 Jul 2026 17:14:54 -0400 Subject: [PATCH 3/5] MOM_coms: GPU port of block-based repro sum This incorporates several changes related to the block-based reproducing sum implementation, and mirrors a pull request to dev/gfdl. * Reproducing sum now divides domains into "blocks" such that each block does not exceed the maximum number of values before the carry accumulator may overflow. * Several hard-coded parameters are replaced with Fortran intrinsics. * Some parameters are renamed for convenience (esp. ni -> efp_digits) * A max-domain warning was removed from the unit tests, since this now would be automatically divided into blocks. --- .../unit_tests/test_reproducing_sum.F90 | 12 +- src/framework/MOM_coms.F90 | 448 ++++++++++-------- 2 files changed, 256 insertions(+), 204 deletions(-) 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/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)') & From b8c471cfa8ed9520ed06b5c8e39497689118e62a Mon Sep 17 00:00:00 2001 From: Edward Yang <94523015+edoyango@users.noreply.github.com> Date: Wed, 8 Jul 2026 03:18:11 +1000 Subject: [PATCH 4/5] Kblock coradcalc (#167) * Block k in CorAdCalc Add CORAD_NKBLOCK and restructure CorAdCalc around k-blocks so CPU builds keep layer-wise cache behavior while GPU builds can expose k as part of the do concurrent parallel space. Convert per-layer scratch arrays and gradKE to operate on the active k-block, while retaining serial wrappers for OBC, WENO, and KE_UP3 paths that still need dedicated GPU ports. * Rename kstart/kend to k_start/k_end in CorAdCalc * Rename CORAD_NKBLOCK to CORIOLIS_ADV_NKBLOCK * Iterate k-block loops on the block-local index directly * Port the CORIOLIS_EN_DIS uh/vh-ratio loop to do concurrent --------- Co-Authored-By: Claude Sonnet 4.6 Co-Authored-By: OpenAI Codex --- src/core/MOM_CoriolisAdv.F90 | 1081 ++++++++++++++++++---------------- 1 file changed, 585 insertions(+), 496 deletions(-) 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 "//& From 196689e5a77724370169565b0067d05a60d9e2e8 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Wed, 15 Jul 2026 09:51:29 +1000 Subject: [PATCH 5/5] MOM_cpu_clock_infra: emit NVTX ranges for MOM6 cpu clocks Record each clock's name as cpu_clock_id() registers it, then open an NVTX range in cpu_clock_begin() and close it in cpu_clock_end(). Every existing cpu_clock_id() name becomes a named range in an nsys timeline with no call-site changes. The range is opened before, and closed after, the mpp clock so that it encloses it. Begin and end test the same condition, so starts and ends stay balanced for handles that were never named or that fall outside the table. The code is guarded behind -DMOM_USE_NVTX, which is undefined by default: without it the preprocessed source is byte-identical to before and no extra library is linked. Profiling builds need -DMOM_USE_NVTX and -cudalib=nvtx. Applied to both the FMS1 and FMS2 infra shims, which are byte-identical to each other. This also covers the halo and message-passing clocks, which MOM_domain_infra begins and ends through this module directly rather than through the MOM_cpu_clock framework wrapper. Based on Edward Yang's ae67665d3 ("add nvtx markers to clocks", branch benchmark_ALE_nvtx_clocks), which introduced the name-table approach and the range nesting; this adds the build guard so non-NVHPC builds are unaffected. Co-authored-by: Edward Yang Co-Authored-By: Claude Opus 4.8 (1M context) --- config_src/infra/FMS1/MOM_cpu_clock_infra.F90 | 34 +++++++++++++++++++ config_src/infra/FMS2/MOM_cpu_clock_infra.F90 | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) 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