From 9e6723ae992d669ae025c61f7380980425deb492 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sat, 30 May 2026 17:13:49 -0400 Subject: [PATCH 01/45] add j dimension to find_N2 related arrs --- .../vertical/MOM_set_diffusivity.F90 | 256 ++++++++++-------- 1 file changed, 137 insertions(+), 119 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 971e5f6226..16fe41e737 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -239,6 +239,8 @@ module MOM_set_diffusivity integer :: id_clock_kappaShear, id_clock_CVMix_ddiff !>@} +integer, parameter :: njblock = 1 !< Number of j-rows per block in the tiled j-loop + contains subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_int, & @@ -285,10 +287,10 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] ! local variables - real :: N2_bot(SZI_(G)) ! Bottom squared buoyancy frequency [T-2 ~> s-2] - real :: rho_bot(SZI_(G)) ! In situ near-bottom density [T-2 ~> s-2] - real :: h_bot(SZI_(G)) ! Bottom boundary layer thickness [H ~> m or kg m-2] - integer :: k_bot(SZI_(G)) ! Bottom boundary layer thickness top layer index + real :: N2_bot(SZI_(G),njblock) ! Bottom squared buoyancy frequency [T-2 ~> s-2] + real :: rho_bot(SZI_(G),njblock) ! In situ near-bottom density [T-2 ~> s-2] + real :: h_bot(SZI_(G),njblock) ! Bottom boundary layer thickness [H ~> m or kg m-2] + integer :: k_bot(SZI_(G),njblock) ! Bottom boundary layer thickness top layer index type(diffusivity_diags) :: dd ! structure with arrays of available diags @@ -297,9 +299,11 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i T_f, S_f ! Temperature and salinity [C ~> degC] and [S ~> ppt] with properties in massless layers ! filled vertically by diffusion or the properties after full convective adjustment. + real, dimension(SZI_(G),SZK_(GV),njblock) :: & + N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] + Kd_lay_2d !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),SZK_(GV)) :: & - N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] - Kd_lay_2d, & !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] dz, & !< Height change across layers [Z ~> m] maxTKE, & !< Energy required to entrain to h_max [H Z2 T-3 ~> m3 s-3 or W m-2] prof_leak_2d, & !< vertical profile for leakage [Z-1 ~> m-1] @@ -310,16 +314,18 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i TKE_to_Kd !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] - real, dimension(SZI_(G),SZK_(GV)+1) :: & + real, dimension(SZI_(G),SZK_(GV)+1,njblock) :: & N2_int, & !< squared buoyancy frequency associated at interfaces [T-2 ~> s-2] + dRho_int, & !< Locally referenced potential density difference across interfaces [R ~> kg m-3] Kd_int_2d, & !< The interface diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - Kv_bkgnd, & !< The background diffusion related interface viscosities [H Z T-1 ~> m2 s-1 or Pa s] + Kv_bkgnd !< The background diffusion related interface viscosities [H Z T-1 ~> m2 s-1 or Pa s] + + real, dimension(SZI_(G),SZK_(GV)+1) :: & Kd_leak_2d, & !< internal tides leakage diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_quad_2d, & !< internal tides bottom drag diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_itidal_2d, & !< internal tides wave drag diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_Froude_2d, & !< internal tides high Froude diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_slope_2d, & !< internal tides critical slopes diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - dRho_int, & !< Locally referenced potential density difference across interfaces [R ~> kg m-3] KT_extra, & !< Double diffusion diffusivity of temperature [H Z T-1 ~> m2 s-1 or kg m-1 s-1] KS_extra !< Double diffusion diffusivity of salinity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] @@ -333,6 +339,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i logical :: showCallTree ! If true, show the call tree. integer :: i, j, k, is, ie, js, je, nz, isd, ied, jsd, jed + integer :: jstart, jend, jj real :: kappa_dt_fill ! diffusivity times a timestep used to fill massless layers [H Z ~> m2 or kg m-1] @@ -472,32 +479,35 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! be an appropriate place to add a depth-dependent parameterization or another explicit ! parameterization of Kd. - !$OMP parallel do default(shared) private(dRho_int,N2_lay,Kd_lay_2d,Kd_int_2d,Kv_bkgnd,N2_int,dz, & - !$OMP N2_bot,rho_bot,h_bot,k_bot,KT_extra,KS_extra,TKE_to_Kd,maxTKE,dissip,kb) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=js,je + do jstart=js,je,njblock + jend = min(jstart+njblock-1, je) + !$OMP parallel do default(shared) private(dz, & + !$OMP KT_extra,KS_extra,TKE_to_Kd,maxTKE,dissip,kb,jj) & + !$OMP if(.not. CS%use_CVMix_ddiff) + do j=jstart,jend + jj = j - jstart + 1 ! Set up variables related to the stratification. - call find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) + call find_N2(h, tv, T_f, S_f, fluxes, j, njblock, jstart, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) if (associated(dd%N2_3d)) then - do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K) ; enddo ; enddo + do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K,jj) ; enddo ; enddo endif ! Add background mixing - call calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay_2d, Kd_int_2d, Kv_bkgnd, j, G, GV, US, CS%bkgnd_mixing_csp) + call calculate_bkgnd_mixing(h, tv, N2_lay(:,:,jj), Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), Kv_bkgnd(:,:,jj), j, G, GV, US, CS%bkgnd_mixing_csp) ! Update Kv and 3-d diffusivity diagnostics. if (associated(visc%Kv_slow)) then ; do K=1,nz+1 ; do i=is,ie - visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K) + visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K,jj) enddo ; enddo ; endif if (CS%id_Kv_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K) + dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K,jj) enddo ; enddo ; endif if (CS%id_Kd_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K) + dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) enddo ; enddo ; endif if (associated(VBF%Kd_bkgnd)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K) + VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) enddo ; enddo ; endif ! Double-diffusion (old method) @@ -507,13 +517,13 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Kd_extra_T is positive for double diffusive convection. do K=2,nz ; do i=is,ie if (KS_extra(i,K) > KT_extra(i,K)) then ! salt fingering - Kd_lay_2d(i,k-1) = Kd_lay_2d(i,k-1) + 0.5 * KT_extra(i,K) - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * KT_extra(i,K) + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KT_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KT_extra(i,K) Kd_extra_S(i,j,K) = KS_extra(i,K) - KT_extra(i,K) Kd_extra_T(i,j,K) = 0.0 elseif (KT_extra(i,K) > 0.0) then ! double-diffusive convection - Kd_lay_2d(i,k-1) = Kd_lay_2d(i,k-1) + 0.5 * KS_extra(i,K) - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * KS_extra(i,K) + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KS_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KS_extra(i,K) Kd_extra_T(i,j,K) = KT_extra(i,K) - KS_extra(i,K) Kd_extra_S(i,j,K) = 0.0 else ! There is no double diffusion at this interface. @@ -558,7 +568,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Calculate conversion ratios from TKE to layer diffusivities. if (TKE_to_Kd_used) then - call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) + call find_TKE_to_Kd(h, tv, dRho_int(:,:,jj), N2_lay(:,:,jj), j, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) if (associated(dd%maxTKE)) then ; do k=1,nz ; do i=is,ie dd%maxTKE(i,j,k) = maxTKE(i,k) enddo ; enddo ; endif @@ -570,21 +580,21 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Add the input turbulent diffusivity. if (CS%useKappaShear .or. CS%use_CVMix_shear) then do K=2,nz ; do i=is,ie - Kd_int_2d(i,K) = visc%Kd_shear(i,j,K) + 0.5 * (Kd_lay_2d(i,k-1) + Kd_lay_2d(i,k)) + Kd_int_2d(i,K,jj) = visc%Kd_shear(i,j,K) + 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) enddo ; enddo do i=is,ie - Kd_int_2d(i,1) = visc%Kd_shear(i,j,1) ! This isn't actually used. It could be 0. - Kd_int_2d(i,nz+1) = 0.0 + Kd_int_2d(i,1,jj) = visc%Kd_shear(i,j,1) ! This isn't actually used. It could be 0. + Kd_int_2d(i,nz+1,jj) = 0.0 enddo do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) enddo ; enddo else do i=is,ie - Kd_int_2d(i,1) = Kd_lay_2d(i,1) ; Kd_int_2d(i,nz+1) = 0.0 + Kd_int_2d(i,1,jj) = Kd_lay_2d(i,1,jj) ; Kd_int_2d(i,nz+1,jj) = 0.0 enddo do K=2,nz ; do i=is,ie - Kd_int_2d(i,K) = 0.5 * (Kd_lay_2d(i,k-1) + Kd_lay_2d(i,k)) + Kd_int_2d(i,K,jj) = 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) enddo ; enddo endif @@ -594,28 +604,28 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Add the ML_Rad diffusivity. if (CS%ML_radiation) then - call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d, G, GV, US, CS, TKE_to_Kd, Kd_lay_2d) + call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd, Kd_lay_2d(:,:,jj)) endif ! Add the Nikurashin and / or tidal bottom-driven mixing if (CS%use_tidal_mixing) & - call calculate_tidal_mixing(dz, j, N2_bot, rho_bot, N2_lay, N2_int, TKE_to_Kd, & + call calculate_tidal_mixing(dz, j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd, & maxTKE, G, GV, US, CS%tidal_mixing, & - CS%Kd_max, visc%Kv_slow, Kd_lay_2d, Kd_int_2d, VBF) + CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) ! Add diffusivity from internal tides ray tracing if (CS%use_int_tides) then - call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot, k_bot, j, N2_lay, N2_int, TKE_to_Kd, CS%Kd_max, & + call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd, CS%Kd_max, & CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & - Kd_lay_2d, Kd_int_2d, prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & + Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & prof_slope_2d) if (CS%id_kbbl > 0) then ; do i=is,ie - dd%kbbl(i,j) = k_bot(i) + dd%kbbl(i,j) = k_bot(i,jj) enddo ; endif if (CS%id_bbl_thick > 0) then ; do i=is,ie - dd%bbl_thick(i,j) = h_bot(i) + dd%bbl_thick(i,j) = h_bot(i,jj) enddo ; endif if (CS%id_Kd_leak > 0) then ; do K=1,nz+1 ; do i=is,ie dd%Kd_leak(i,j,K) = Kd_leak_2d(i,K) @@ -668,11 +678,11 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! This adds the diffusion sustained by the energy extracted from the flow by the bottom drag. if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0)) then if (CS%use_LOTW_BBL_diffusivity) then - call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bot, Kd_int_2d, & - G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d) + call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), Kd_int_2d(:,:,jj), & + G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj)) else call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, & - maxTKE, kb, rho_bot, G, GV, US, CS, Kd_lay_2d, Kd_int_2d, dd%Kd_BBL) + maxTKE, kb, rho_bot(:,jj), G, GV, US, CS, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), dd%Kd_BBL) endif if (associated(VBF%Kd_BBL)) then ; do K=1,nz+1 ; do i=is,ie VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) @@ -687,24 +697,24 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 3) dissipation corresponding to a (nearly) constant diffusivity. do K=2,nz ; do i=is,ie dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K)), & ! Floor aka Gargett - CS%dissip_N2 * N2_int(i,K)) ! Floor of Kd_min*rho0/F_Ri - Kd_int_2d(i,K) = max(Kd_int_2d(i,K) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K) + Omega2)))) + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_int(i,K,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_int_2d(i,K,jj) = max(Kd_int_2d(i,K,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K,jj) + Omega2)))) enddo ; enddo endif ! Optionally add a uniform diffusivity at the interfaces. if (CS%Kd_add > 0.0) then do K=1,nz+1 ; do i=is,ie - Kd_int_2d(i,K) = Kd_int_2d(i,K) + CS%Kd_add + Kd_int_2d(i,K,jj) = Kd_int_2d(i,K,jj) + CS%Kd_add enddo ; enddo VBF%Kd_add = CS%Kd_add endif ! Copy the 2-d slices into the 3-d array that is exported. do K=1,nz+1 ; do i=is,ie - Kd_int(i,j,K) = Kd_int_2d(i,K) + Kd_int(i,j,K) = Kd_int_2d(i,K,jj) enddo ; enddo if (CS%limit_dissipation) then @@ -715,37 +725,38 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 3) dissipation corresponding to a (nearly) constant diffusivity. do k=2,nz-1 ; do i=is,ie dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k)), & ! Floor aka Gargett - CS%dissip_N2 * N2_lay(i,k)) ! Floor of Kd_min*rho0/F_Ri - Kd_lay_2d(i,k) = max(Kd_lay_2d(i,k) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k) + Omega2)))) + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_lay(i,k,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_lay_2d(i,k,jj) = max(Kd_lay_2d(i,k,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k,jj) + Omega2)))) enddo ; enddo endif if (associated(dd%Kd_Work)) then do k=1,nz ; do i=is,ie - dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k) * N2_lay(i,k) * dz(i,k) ! Watt m-2 = kg s-3 + dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k) ! Watt m-2 = kg s-3 enddo ; enddo endif ! Optionally add a uniform diffusivity to the layers. if ((CS%Kd_add > 0.0) .and. (present(Kd_lay))) then do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + CS%Kd_add + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + CS%Kd_add enddo ; enddo endif if (associated(dd%Kd_Work_added)) then do k=1,nz ; do i=is,ie - dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k) * dz(i,k) ! Watt m-2 = kg s-3 + dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k) ! Watt m-2 = kg s-3 enddo ; enddo endif ! Copy the 2-d slices into the 3-d array that is exported; this was done above for Kd_int. if (present(Kd_lay)) then ; do k=1,nz ; do i=is,ie - Kd_lay(i,j,k) = Kd_lay_2d(i,k) + Kd_lay(i,j,k) = Kd_lay_2d(i,k,jj) enddo ; enddo ; endif - enddo ! j-loop + enddo ! j-loop + enddo ! jstart-loop if (CS%user_change_diff) then call user_change_diff(h, tv, G, GV, US, CS%user_change_diff_CSp, Kd_lay, Kd_int, & @@ -1092,7 +1103,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & end subroutine find_TKE_to_Kd !> Calculate Brunt-Vaisala frequency, N^2. -subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & +subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_int, & N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure @@ -1109,120 +1120,125 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & !! layers filled vertically by diffusion [S ~> ppt]. type(forcing), intent(in) :: fluxes !< A structure of thermodynamic surface fluxes integer, intent(in) :: j !< j-index of row to work on + integer, intent(in) :: nj !< Number of j-rows in this block + integer, intent(in) :: jstart !< Starting j-index of this block type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)+1), & + real, dimension(SZI_(G),SZK_(GV)+1,nj), & intent(out) :: dRho_int !< Change in locally referenced potential density !! across each interface [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)+1), & + real, dimension(SZI_(G),SZK_(GV)+1,nj), & intent(out) :: N2_int !< The squared buoyancy frequency at the interfaces [T-2 ~> s-2]. - real, dimension(SZI_(G),SZK_(GV)), & + real, dimension(SZI_(G),SZK_(GV),nj), & intent(out) :: N2_lay !< The squared buoyancy frequency of the layers [T-2 ~> s-2]. - real, dimension(SZI_(G)), intent(out) :: N2_bot !< The near-bottom squared buoyancy frequency [T-2 ~> s-2]. - real, dimension(SZI_(G)), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. - real, dimension(SZI_(G)), optional, intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. - integer, dimension(SZI_(G)), optional, intent(out) :: k_bot !< Bottom boundary layer top layer index. + real, dimension(SZI_(G),nj), intent(out) :: N2_bot !< The near-bottom squared buoyancy frequency [T-2 ~> s-2]. + real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. + real, dimension(SZI_(G),nj), optional, intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. + integer, dimension(SZI_(G),nj), optional, intent(out) :: k_bot !< Bottom boundary layer top layer index. ! Local variables - real, dimension(SZI_(G),SZK_(GV)+1) :: & - pres, & ! pressure at each interface [R L2 T-2 ~> Pa] + real, dimension(SZI_(G),SZK_(GV)+1,nj) :: & + pres, & ! pressure at each interface [R L2 T-2 ~> Pa] dRho_int_unfilt, & ! unfiltered density differences across interfaces [R ~> kg m-3] dRho_dT, & ! partial derivative of density wrt temp [R C-1 ~> kg m-3 degC-1] dRho_dS ! partial derivative of density wrt saln [R S-1 ~> kg m-3 ppt-1] - real, dimension(SZI_(G),SZK_(GV)) :: & + real, dimension(SZI_(G),SZK_(GV),nj) :: & dz ! Height change across layers [Z ~> m] - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & Temp_int, & ! temperature at each interface [C ~> degC] Salin_int, & ! salinity at each interface [S ~> ppt] drho_bot, & ! A density difference [R ~> kg m-3] h_amp, & ! The topographic roughness amplitude [Z ~> m]. dz_BBL_avg, & ! The distance over which to average to find the near-bottom density [Z ~> m] hb, & ! The thickness of the bottom layer [H ~> m or kg m-2] - z_from_bot ! The height above the bottom [Z ~> m] + z_from_bot, & ! The height above the bottom [Z ~> m] + h_bot_tmp ! temporary h_bot for passing to find_rho_bottom [H ~> m or kg m-2] + integer, dimension(SZI_(G),nj) :: k_bot_tmp ! temporary k_bot for passing to find_rho_bottom real :: dz_int ! Vertical distance associated with an interface [Z ~> m] real :: G_Rho0 ! Gravitational acceleration, perhaps divided by Boussinesq reference density, ! times some unit conversion factors [H T-2 R-1 ~> m4 s-2 kg-1 or m s-2]. real :: H_neglect ! A negligibly small thickness [H ~> m or kg m-2] - logical :: do_i(SZI_(G)), do_any + logical :: do_i(SZI_(G),nj), do_any integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, is, ie, nz + integer :: i, k, is, ie, nz, jj is = G%isc ; ie = G%iec ; nz = GV%ke + jj = j - jstart + 1 G_Rho0 = GV%g_Earth_Z_T2 / GV%H_to_RZ H_neglect = GV%H_subroundoff ! Find the (limited) density jump across each interface. do i=is,ie - dRho_int(i,1) = 0.0 ; dRho_int(i,nz+1) = 0.0 - dRho_int_unfilt(i,1) = 0.0 ; dRho_int_unfilt(i,nz+1) = 0.0 + dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 + dRho_int_unfilt(i,1,jj) = 0.0 ; dRho_int_unfilt(i,nz+1,jj) = 0.0 enddo if (associated(tv%eqn_of_state)) then if (associated(fluxes%p_surf)) then - do i=is,ie ; pres(i,1) = fluxes%p_surf(i,j) ; enddo + do i=is,ie ; pres(i,1,jj) = fluxes%p_surf(i,j) ; enddo else - do i=is,ie ; pres(i,1) = 0.0 ; enddo + do i=is,ie ; pres(i,1,jj) = 0.0 ; enddo endif EOSdom(:) = EOS_domain(G%HI) do K=2,nz do i=is,ie - pres(i,K) = pres(i,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) - Temp_Int(i) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) - Salin_Int(i) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) + pres(i,K,jj) = pres(i,K-1,jj) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) + Temp_Int(i,jj) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) + Salin_Int(i,jj) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) enddo - call calculate_density_derivs(Temp_int, Salin_int, pres(:,K), dRho_dT(:,K), dRho_dS(:,K), & + call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,K,jj), dRho_dT(:,K,jj), dRho_dS(:,K,jj), & tv%eqn_of_state, EOSdom) do i=is,ie - dRho_int(i,K) = max(dRho_dT(i,K)*(T_f(i,j,k) - T_f(i,j,k-1)) + & - dRho_dS(i,K)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) - dRho_int_unfilt(i,K) = max(dRho_dT(i,K)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & - dRho_dS(i,K)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + dRho_int(i,K,jj) = max(dRho_dT(i,K,jj)*(T_f(i,j,k) - T_f(i,j,k-1)) + & + dRho_dS(i,K,jj)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) + dRho_int_unfilt(i,K,jj) = max(dRho_dT(i,K,jj)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & + dRho_dS(i,K,jj)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) enddo enddo else do K=2,nz ; do i=is,ie - dRho_int(i,K) = GV%Rlay(k) - GV%Rlay(k-1) + dRho_int(i,K,jj) = GV%Rlay(k) - GV%Rlay(k-1) enddo ; enddo endif ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) ! Set the buoyancy frequencies. do k=1,nz ; do i=is,ie - N2_lay(i,k) = G_Rho0 * 0.5*(dRho_int(i,K) + dRho_int(i,K+1)) / & - (h(i,j,k) + H_neglect) + N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,K,jj) + dRho_int(i,K+1,jj)) / & + (h(i,j,k) + H_neglect) enddo ; enddo - do i=is,ie ; N2_int(i,1) = 0.0 ; N2_int(i,nz+1) = 0.0 ; enddo + do i=is,ie ; N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 ; enddo do K=2,nz ; do i=is,ie - N2_int(i,K) = G_Rho0 * dRho_int(i,K) / & - (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) + N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & + (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) enddo ; enddo ! Find the bottom boundary layer stratification, and use this in the deepest layers. do i=is,ie - hb(i) = 0.0 ; dRho_bot(i) = 0.0 ; h_amp(i) = 0.0 - z_from_bot(i) = 0.5*dz(i,nz) - do_i(i) = (G%mask2dT(i,j) > 0.0) + hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 ; h_amp(i,jj) = 0.0 + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) enddo - if (CS%use_tidal_mixing) call tidal_mixing_h_amp(h_amp, G, j, CS%tidal_mixing) + if (CS%use_tidal_mixing) call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) do k=nz,2,-1 do_any = .false. - do i=is,ie ; if (do_i(i)) then - dz_int = 0.5*(dz(i,k) + dz(i,k-1)) - z_from_bot(i) = z_from_bot(i) + dz_int ! middle of the layer above + do i=is,ie ; if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - hb(i) = hb(i) + 0.5*(h(i,j,k) + h(i,j,k-1)) - drho_bot(i) = drho_bot(i) + dRho_int(i,K) + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K,jj) - if (z_from_bot(i) > h_amp(i)) then + if (z_from_bot(i,jj) > h_amp(i,jj)) then if (k>2) then ! Always include at least one full layer. - hb(i) = hb(i) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) - drho_bot(i) = drho_bot(i) + dRho_int(i,K-1) + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K-1,jj) endif - do_i(i) = .false. + do_i(i,jj) = .false. else do_any = .true. endif @@ -1231,25 +1247,25 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & enddo do i=is,ie - if (hb(i) > 0.0) then - N2_bot(i) = (G_Rho0 * drho_bot(i)) / hb(i) - else ; N2_bot(i) = 0.0 ; endif - z_from_bot(i) = 0.5*dz(i,nz) - do_i(i) = (G%mask2dT(i,j) > 0.0) + if (hb(i,jj) > 0.0) then + N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) + else ; N2_bot(i,jj) = 0.0 ; endif + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) enddo do k=nz,2,-1 do_any = .false. - do i=is,ie ; if (do_i(i)) then - dz_int = 0.5*(dz(i,k) + dz(i,k-1)) - z_from_bot(i) = z_from_bot(i) + dz_int ! middle of the layer above + do i=is,ie ; if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - N2_int(i,K) = N2_bot(i) - if (k>2) N2_lay(i,k-1) = N2_bot(i) + N2_int(i,K,jj) = N2_bot(i,jj) + if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) - if (z_from_bot(i) > h_amp(i)) then - if (k>2) N2_int(i,K-1) = N2_bot(i) - do_i(i) = .false. + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) N2_int(i,K-1,jj) = N2_bot(i,jj) + do_i(i,jj) = .false. else do_any = .true. endif @@ -1259,13 +1275,15 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & if (associated(tv%eqn_of_state)) then do K=1,nz+1 ; do i=is,ie - dRho_int(i,K) = dRho_int_unfilt(i,K) + dRho_int(i,K,jj) = dRho_int_unfilt(i,K,jj) enddo ; enddo endif ! Average over the larger of the envelope of the topography or a minimal distance. - do i=is,ie ; dz_BBL_avg(i) = max(h_amp(i), CS%dz_BBL_avg_min) ; enddo - call find_rho_bottom(G, GV, US, tv, h, dz, pres, dz_BBL_avg, j, Rho_bot, h_bot, k_bot) + do i=is,ie ; dz_BBL_avg(i,jj) = max(h_amp(i,jj), CS%dz_BBL_avg_min) ; enddo + call find_rho_bottom(G, GV, US, tv, h, dz(:,:,jj), pres(:,:,jj), dz_BBL_avg(:,jj), j, Rho_bot(:,jj), h_bot_tmp(:,jj), k_bot_tmp(:,jj)) + if (present(h_bot)) h_bot(:,jj) = h_bot_tmp(:,jj) + if (present(k_bot)) k_bot(:,jj) = k_bot_tmp(:,jj) end subroutine find_N2 From 98b80e2f9e2b0f8cded16db661b339c60a65ccad Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sat, 30 May 2026 23:41:59 -0400 Subject: [PATCH 02/45] set_diffusivity: block j loops in find_N2 and find_rho_bottom - Block all local arrays in find_N2 with an nj j-dimension; move the find_N2 call before the inner OMP j-loop in set_diffusivity and add an internal j-loop inside find_N2 - Make h_bot/k_bot non-optional in find_N2 - Add find_rho_bottom generic interface in MOM_interface_heights, renaming the existing routine to find_rho_bottom_1d and adding find_rho_bottom_2d that accepts blocked arrays (nj j-dimension); move find_rho_bottom call outside find_N2's j-loop - Replace each outer j-loop in find_N2 and find_rho_bottom_2d with individual do j / do i nested loops placed next to each i-loop Co-Authored-By: Claude Sonnet 4.6 --- src/core/MOM_interface_heights.F90 | 206 ++++++++++++++++- .../vertical/MOM_set_diffusivity.F90 | 217 +++++++++++------- 2 files changed, 332 insertions(+), 91 deletions(-) diff --git a/src/core/MOM_interface_heights.F90 b/src/core/MOM_interface_heights.F90 index c70c5e8f44..861b73a4c8 100644 --- a/src/core/MOM_interface_heights.F90 +++ b/src/core/MOM_interface_heights.F90 @@ -41,6 +41,11 @@ module MOM_interface_heights module procedure thickness_to_dz_3d, thickness_to_dz_jslice end interface thickness_to_dz +!> Computes the density of the near-bottom water in the ocean. +interface find_rho_bottom + module procedure find_rho_bottom_1d, find_rho_bottom_2d +end interface find_rho_bottom + contains !> Calculates the change in height across layers, using the appropriate form for @@ -486,7 +491,7 @@ end subroutine find_col_mass !> Determine the in situ density averaged over a specified distance from the bottom, !! calculating it as the inverse of the mass-weighted average specific volume. -subroutine find_rho_bottom(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h_bot, k_bot) +subroutine find_rho_bottom_1d(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h_bot, k_bot) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -663,8 +668,205 @@ subroutine find_rho_bottom(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h enddo endif -end subroutine find_rho_bottom +end subroutine find_rho_bottom_1d + +!> Computes the density of the near-bottom water for a j-block of rows, using pre-computed +!! blocked arrays with an explicit j-block index dimension. +subroutine find_rho_bottom_2d(G, GV, US, tv, h, dz, pres_int, dz_avg, jstart, jend, nj, Rho_bot, h_bot, k_bot) + type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure + type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + type(thermo_var_ptrs), intent(in) :: tv !< Structure containing pointers to any available + !! thermodynamic fields. + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), & + intent(in) :: h !< Layer thicknesses [H ~> m or kg m-2] + integer, intent(in) :: nj !< Number of j-rows in the block + real, dimension(SZI_(G),SZK_(GV),nj), & + intent(in) :: dz !< Height change across layers [Z ~> m] + real, dimension(SZI_(G),SZK_(GV)+1,nj), & + intent(in) :: pres_int !< Pressure at each interface [R L2 T-2 ~> Pa] + real, dimension(SZI_(G),nj), intent(in) :: dz_avg !< The vertical distance over which to average [Z ~> m] + integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block + real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. + real, dimension(SZI_(G),nj), intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2] + integer, dimension(SZI_(G),nj), intent(out) :: k_bot !< Bottom boundary layer top layer index + + ! Local variables + real :: hb(SZI_(G),nj) ! Running sum of the thickness in the bottom boundary layer [H ~> m or kg m-2] + real :: SpV_h_bot(SZI_(G),nj) ! Running sum of the specific volume times thickness in the bottom + ! boundary layer [H R-1 ~> m4 kg-1 or m] + real :: dz_bbl_rem(SZI_(G),nj) ! Vertical extent of the boundary layer that has yet to be accounted + ! for [Z ~> m] + real :: h_bbl_frac(SZI_(G),nj) ! Thickness of the fractional layer that makes up the top of the + ! boundary layer [H ~> m or kg m-2] + real :: T_bbl(SZI_(G),nj) ! Temperature of the fractional layer that makes up the top of the + ! boundary layer [C ~> degC] + real :: S_bbl(SZI_(G),nj) ! Salinity of the fractional layer that makes up the top of the + ! boundary layer [S ~> ppt] + real :: P_bbl(SZI_(G),nj) ! Pressure the top of the boundary layer [R L2 T-2 ~> Pa] + real :: dp(SZI_(G),nj) ! Pressure change across the fractional layer that makes up the top + ! of the boundary layer [R L2 T-2 ~> Pa] + real :: SpV_bbl(SZI_(G),nj) ! In situ specific volume of the fractional layer that makes up the + ! top of the boundary layer [R-1 ~> m3 kg-1] + real :: frac_in ! The fraction of a layer that is within the bottom boundary layer [nondim] + logical :: do_i(SZI_(G),nj), do_any + logical :: use_EOS + integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state + integer :: i, j, k, is, ie, nz, jj + + is = G%isc ; ie = G%iec ; nz = GV%ke + + use_EOS = associated(tv%T) .and. associated(tv%S) .and. associated(tv%eqn_of_state) + + if (GV%Boussinesq .or. GV%semi_Boussinesq .or. .not.allocated(tv%SpV_avg)) then + ! Obtain bottom boundary layer thickness and index of top layer + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + rho_bot(i,jj) = GV%Rho0 + hb(i,jj) = 0.0 ; h_bot(i,jj) = 0.0 ; k_bot(i,jj) = nz + dz_bbl_rem(i,jj) = G%mask2dT(i,j) * max(0.0, dz_avg(i,jj)) + do_i(i,jj) = .true. + if (G%mask2dT(i,j) <= 0.0) then + h_bbl_frac(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + enddo + enddo + + do j=jstart,jend ; jj = j - jstart + 1 ; do k=nz,1,-1 + do_any = .false. + + do i=is,ie ; if (do_i(i,jj)) then + if (dz(i,k,jj) < dz_bbl_rem(i,jj)) then + ! This layer is fully within the averaging depth. + dz_bbl_rem(i,jj) = dz_bbl_rem(i,jj) - dz(i,k,jj) + hb(i,jj) = hb(i,jj) + h(i,j,k) + k_bot(i,jj) = k + do_any = .true. + else + if (dz(i,k,jj) > 0.0) then + frac_in = dz_bbl_rem(i,jj) / dz(i,k,jj) + if (frac_in >= 0.5) k_bot(i,jj) = k ! update bbl top index if >= 50% of layer + else + frac_in = 0.0 + endif + h_bbl_frac(i,jj) = frac_in * h(i,j,k) + dz_bbl_rem(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + endif ; enddo + enddo + if (.not.do_any) exit + enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + ! The nominal bottom boundary layer is thicker than the water column, but layer 1 is + ! already included in the averages. These values are set so that the call to find + ! the layer-average specific volume will behave sensibly. + h_bbl_frac(i,jj) = 0.0 + endif ; enddo + enddo + + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (hb(i,jj) + h_bbl_frac(i,jj) < GV%H_subroundoff) h_bbl_frac(i,jj) = GV%H_subroundoff + h_bot(i,jj) = hb(i,jj) + h_bbl_frac(i,jj) + enddo + enddo + + else + ! Check that SpV_avg has been set. + if (tv%valid_SpV_halo < 0) call MOM_error(FATAL, & + "find_rho_bottom called in fully non-Boussinesq mode with invalid values of SpV_avg.") + + ! Set the bottom density to the inverse of the in situ specific volume averaged over the + ! specified distance, with care taken to avoid having compressibility lead to an imprint + ! of the layer thicknesses on this density. + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + hb(i,jj) = 0.0 ; SpV_h_bot(i,jj) = 0.0 ; h_bot(i,jj) = 0.0 ; k_bot(i,jj) = nz + dz_bbl_rem(i,jj) = G%mask2dT(i,j) * max(0.0, dz_avg(i,jj)) + do_i(i,jj) = .true. + if (G%mask2dT(i,j) <= 0.0) then + ! Set acceptable values for calling the equation of state over land. + T_bbl(i,jj) = 0.0 ; S_bbl(i,jj) = 0.0 ; dp(i,jj) = 0.0 ; P_bbl(i,jj) = 0.0 + SpV_bbl(i,jj) = 1.0 ! This value is arbitrary, provided it is non-zero. + h_bbl_frac(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + enddo + enddo + + do k=nz,1,-1 + do_any = .false. + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + if (dz(i,k,jj) < dz_bbl_rem(i,jj)) then + ! This layer is fully within the averaging depth. + SpV_h_bot(i,jj) = SpV_h_bot(i,jj) + h(i,j,k) * tv%SpV_avg(i,j,k) + dz_bbl_rem(i,jj) = dz_bbl_rem(i,jj) - dz(i,k,jj) + hb(i,jj) = hb(i,jj) + h(i,j,k) + k_bot(i,jj) = k + do_any = .true. + else + if (dz(i,k,jj) > 0.0) then + frac_in = dz_bbl_rem(i,jj) / dz(i,k,jj) + if (frac_in >= 0.5) k_bot(i,jj) = k ! update bbl top index if >= 50% of layer + else + frac_in = 0.0 + endif + if (use_EOS) then + ! Store the properties of this layer to determine the average + ! specific volume of the portion that is within the BBL. + T_bbl(i,jj) = tv%T(i,j,k) ; S_bbl(i,jj) = tv%S(i,j,k) + dp(i,jj) = frac_in * (GV%g_Earth*GV%H_to_RZ * h(i,j,k)) + P_bbl(i,jj) = pres_int(i,K,jj) + (1.0-frac_in) * (GV%g_Earth*GV%H_to_RZ * h(i,j,k)) + else + SpV_bbl(i,jj) = tv%SpV_avg(i,j,k) + endif + h_bbl_frac(i,jj) = frac_in * h(i,j,k) + dz_bbl_rem(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + endif ; enddo + enddo + if (.not.do_any) exit + enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + ! The nominal bottom boundary layer is thicker than the water column, but layer 1 is + ! already included in the averages. These values are set so that the call to find + ! the layer-average specific volume will behave sensibly. + if (use_EOS) then + T_bbl(i,jj) = tv%T(i,j,1) ; S_bbl(i,jj) = tv%S(i,j,1) + dp(i,jj) = 0.0 + P_bbl(i,jj) = pres_int(i,1,jj) + else + SpV_bbl(i,jj) = tv%SpV_avg(i,j,1) + endif + h_bbl_frac(i,jj) = 0.0 + endif ; enddo + enddo + + if (use_EOS) then + ! Find the average specific volume of the fractional layer atop the BBL. + EOSdom(:) = EOS_domain(G%HI) + do j=jstart,jend ; jj = j - jstart + 1 + call average_specific_vol(T_bbl(:,jj), S_bbl(:,jj), P_bbl(:,jj), dp(:,jj), SpV_bbl(:,jj), tv%eqn_of_state, EOSdom) + enddo + endif + + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (hb(i,jj) + h_bbl_frac(i,jj) < GV%H_subroundoff) h_bbl_frac(i,jj) = GV%H_subroundoff + rho_bot(i,jj) = G%mask2dT(i,j) * (hb(i,jj) + h_bbl_frac(i,jj)) / (SpV_h_bot(i,jj) + h_bbl_frac(i,jj)*SpV_bbl(i,jj)) + h_bot(i,jj) = hb(i,jj) + h_bbl_frac(i,jj) + enddo + enddo + endif +end subroutine find_rho_bottom_2d !> Converts thickness from geometric height units to thickness units, perhaps via an !! inversion of the integral of the density in pressure using variables stored in diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 16fe41e737..46ace68d79 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -481,15 +481,16 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i do jstart=js,je,njblock jend = min(jstart+njblock-1, je) + + ! Set up variables related to the stratification. + call find_N2(h, tv, T_f, S_f, fluxes, njblock, jstart, jend, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) + !$OMP parallel do default(shared) private(dz, & !$OMP KT_extra,KS_extra,TKE_to_Kd,maxTKE,dissip,kb,jj) & !$OMP if(.not. CS%use_CVMix_ddiff) do j=jstart,jend jj = j - jstart + 1 - ! Set up variables related to the stratification. - call find_N2(h, tv, T_f, S_f, fluxes, j, njblock, jstart, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) - if (associated(dd%N2_3d)) then do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K,jj) ; enddo ; enddo endif @@ -1103,7 +1104,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & end subroutine find_TKE_to_Kd !> Calculate Brunt-Vaisala frequency, N^2. -subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_int, & +subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRho_int, & N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure @@ -1119,9 +1120,9 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_i intent(in) :: S_f !< Layer salinities with values in massless !! layers filled vertically by diffusion [S ~> ppt]. type(forcing), intent(in) :: fluxes !< A structure of thermodynamic surface fluxes - integer, intent(in) :: j !< j-index of row to work on integer, intent(in) :: nj !< Number of j-rows in this block integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure real, dimension(SZI_(G),SZK_(GV)+1,nj), & intent(out) :: dRho_int !< Change in locally referenced potential density @@ -1132,8 +1133,8 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_i intent(out) :: N2_lay !< The squared buoyancy frequency of the layers [T-2 ~> s-2]. real, dimension(SZI_(G),nj), intent(out) :: N2_bot !< The near-bottom squared buoyancy frequency [T-2 ~> s-2]. real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. - real, dimension(SZI_(G),nj), optional, intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. - integer, dimension(SZI_(G),nj), optional, intent(out) :: k_bot !< Bottom boundary layer top layer index. + real, dimension(SZI_(G),nj), intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. + integer, dimension(SZI_(G),nj), intent(out) :: k_bot !< Bottom boundary layer top layer index. ! Local variables real, dimension(SZI_(G),SZK_(GV)+1,nj) :: & @@ -1150,9 +1151,7 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_i h_amp, & ! The topographic roughness amplitude [Z ~> m]. dz_BBL_avg, & ! The distance over which to average to find the near-bottom density [Z ~> m] hb, & ! The thickness of the bottom layer [H ~> m or kg m-2] - z_from_bot, & ! The height above the bottom [Z ~> m] - h_bot_tmp ! temporary h_bot for passing to find_rho_bottom [H ~> m or kg m-2] - integer, dimension(SZI_(G),nj) :: k_bot_tmp ! temporary k_bot for passing to find_rho_bottom + z_from_bot ! The height above the bottom [Z ~> m] real :: dz_int ! Vertical distance associated with an interface [Z ~> m] real :: G_Rho0 ! Gravitational acceleration, perhaps divided by Boussinesq reference density, @@ -1161,129 +1160,169 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, nj, jstart, G, GV, US, CS, dRho_i logical :: do_i(SZI_(G),nj), do_any integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, is, ie, nz, jj + integer :: i, j, k, is, ie, nz, jj is = G%isc ; ie = G%iec ; nz = GV%ke - jj = j - jstart + 1 G_Rho0 = GV%g_Earth_Z_T2 / GV%H_to_RZ H_neglect = GV%H_subroundoff ! Find the (limited) density jump across each interface. - do i=is,ie - dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 - dRho_int_unfilt(i,1,jj) = 0.0 ; dRho_int_unfilt(i,nz+1,jj) = 0.0 + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 + dRho_int_unfilt(i,1,jj) = 0.0 ; dRho_int_unfilt(i,nz+1,jj) = 0.0 + enddo enddo if (associated(tv%eqn_of_state)) then if (associated(fluxes%p_surf)) then - do i=is,ie ; pres(i,1,jj) = fluxes%p_surf(i,j) ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; pres(i,1,jj) = fluxes%p_surf(i,j) ; enddo + enddo else - do i=is,ie ; pres(i,1,jj) = 0.0 ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; pres(i,1,jj) = 0.0 ; enddo + enddo endif EOSdom(:) = EOS_domain(G%HI) do K=2,nz - do i=is,ie - pres(i,K,jj) = pres(i,K-1,jj) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) - Temp_Int(i,jj) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) - Salin_Int(i,jj) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) - enddo - call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,K,jj), dRho_dT(:,K,jj), dRho_dS(:,K,jj), & - tv%eqn_of_state, EOSdom) - do i=is,ie - dRho_int(i,K,jj) = max(dRho_dT(i,K,jj)*(T_f(i,j,k) - T_f(i,j,k-1)) + & - dRho_dS(i,K,jj)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) - dRho_int_unfilt(i,K,jj) = max(dRho_dT(i,K,jj)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & - dRho_dS(i,K,jj)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + pres(i,K,jj) = pres(i,K-1,jj) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) + Temp_Int(i,jj) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) + Salin_Int(i,jj) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) + enddo + call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,K,jj), dRho_dT(:,K,jj), dRho_dS(:,K,jj), & + tv%eqn_of_state, EOSdom) + do i=is,ie + dRho_int(i,K,jj) = max(dRho_dT(i,K,jj)*(T_f(i,j,k) - T_f(i,j,k-1)) + & + dRho_dS(i,K,jj)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) + dRho_int_unfilt(i,K,jj) = max(dRho_dT(i,K,jj)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & + dRho_dS(i,K,jj)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + enddo enddo enddo else - do K=2,nz ; do i=is,ie - dRho_int(i,K,jj) = GV%Rlay(k) - GV%Rlay(k-1) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=2,nz + do i=is,ie + dRho_int(i,K,jj) = GV%Rlay(k) - GV%Rlay(k-1) + enddo + enddo + enddo endif ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo ! Set the buoyancy frequencies. - do k=1,nz ; do i=is,ie - N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,K,jj) + dRho_int(i,K+1,jj)) / & - (h(i,j,k) + H_neglect) - enddo ; enddo - do i=is,ie ; N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 ; enddo - do K=2,nz ; do i=is,ie - N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & - (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) - enddo ; enddo + do k=1,nz + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,K,jj) + dRho_int(i,K+1,jj)) / & + (h(i,j,k) + H_neglect) + enddo + enddo + enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 ; enddo + enddo + do K=2,nz + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & + (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) + enddo + enddo + enddo ! Find the bottom boundary layer stratification, and use this in the deepest layers. - do i=is,ie - hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 ; h_amp(i,jj) = 0.0 - z_from_bot(i,jj) = 0.5*dz(i,nz,jj) - do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 ; h_amp(i,jj) = 0.0 + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + enddo enddo - if (CS%use_tidal_mixing) call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) + if (CS%use_tidal_mixing) then + do j=jstart,jend ; jj = j - jstart + 1 + call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) + enddo + endif do k=nz,2,-1 do_any = .false. - do i=is,ie ; if (do_i(i,jj)) then - dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) - z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - - hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) - drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K,jj) - - if (z_from_bot(i,jj) > h_amp(i,jj)) then - if (k>2) then - ! Always include at least one full layer. - hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) - drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K-1,jj) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K,jj) + + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) then + ! Always include at least one full layer. + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K-1,jj) + endif + do_i(i,jj) = .false. + else + do_any = .true. endif - do_i(i,jj) = .false. - else - do_any = .true. - endif - endif ; enddo + endif ; enddo + enddo if (.not.do_any) exit enddo - do i=is,ie - if (hb(i,jj) > 0.0) then - N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) - else ; N2_bot(i,jj) = 0.0 ; endif - z_from_bot(i,jj) = 0.5*dz(i,nz,jj) - do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (hb(i,jj) > 0.0) then + N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) + else ; N2_bot(i,jj) = 0.0 ; endif + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + enddo enddo do k=nz,2,-1 do_any = .false. - do i=is,ie ; if (do_i(i,jj)) then - dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) - z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - N2_int(i,K,jj) = N2_bot(i,jj) - if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) + N2_int(i,K,jj) = N2_bot(i,jj) + if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) - if (z_from_bot(i,jj) > h_amp(i,jj)) then - if (k>2) N2_int(i,K-1,jj) = N2_bot(i,jj) - do_i(i,jj) = .false. - else - do_any = .true. - endif - endif ; enddo + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) N2_int(i,K-1,jj) = N2_bot(i,jj) + do_i(i,jj) = .false. + else + do_any = .true. + endif + endif ; enddo + enddo if (.not.do_any) exit enddo if (associated(tv%eqn_of_state)) then - do K=1,nz+1 ; do i=is,ie - dRho_int(i,K,jj) = dRho_int_unfilt(i,K,jj) - enddo ; enddo + do K=1,nz+1 + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + dRho_int(i,K,jj) = dRho_int_unfilt(i,K,jj) + enddo + enddo + enddo endif ! Average over the larger of the envelope of the topography or a minimal distance. - do i=is,ie ; dz_BBL_avg(i,jj) = max(h_amp(i,jj), CS%dz_BBL_avg_min) ; enddo - call find_rho_bottom(G, GV, US, tv, h, dz(:,:,jj), pres(:,:,jj), dz_BBL_avg(:,jj), j, Rho_bot(:,jj), h_bot_tmp(:,jj), k_bot_tmp(:,jj)) - if (present(h_bot)) h_bot(:,jj) = h_bot_tmp(:,jj) - if (present(k_bot)) k_bot(:,jj) = k_bot_tmp(:,jj) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; dz_BBL_avg(i,jj) = max(h_amp(i,jj), CS%dz_BBL_avg_min) ; enddo + enddo + + call find_rho_bottom(G, GV, US, tv, h, dz, pres, dz_BBL_avg, jstart, jend, nj, Rho_bot, h_bot, k_bot) end subroutine find_N2 From 6c4ed0c29a48abf3992cb0780d7f22d2f71d566a Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sun, 31 May 2026 08:39:00 -0400 Subject: [PATCH 03/45] set_diffusivity: block j loops in calculate_bkgnd_mixing and find_TKE_to_Kd Add njblock j-dimension to local and dummy arrays in calculate_bkgnd_mixing and find_TKE_to_Kd, move j-loops into both routines, and decouple inner loops to individual j-loops in the caller. Also block TKE_to_Kd, maxTKE, and kb arrays in set_diffusivity. Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_bkgnd_mixing.F90 | 240 ++++++----- .../vertical/MOM_set_diffusivity.F90 | 400 ++++++++++-------- 2 files changed, 345 insertions(+), 295 deletions(-) diff --git a/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 b/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 index 5b41cef038..24b4f3aadb 100644 --- a/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 +++ b/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 @@ -309,21 +309,23 @@ subroutine bkgnd_mixing_init(Time, G, GV, US, param_file, diag, CS, physical_OBL end subroutine bkgnd_mixing_init !> Calculates the vertical background diffusivities/viscosities -subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, GV, US, CS) +subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, jstart, jend, nj, G, GV, US, CS) type(ocean_grid_type), intent(in) :: G !< Grid structure. type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h !< Layer thickness [H ~> m or kg m-2]. type(thermo_var_ptrs), intent(in) :: tv !< Thermodynamics structure. - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: N2_lay !< squared buoyancy frequency associated + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: N2_lay !< squared buoyancy frequency associated !! with layers [T-2 ~> s-2] - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: Kd_lay !< The background diapycnal diffusivity of each + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: Kd_lay !< The background diapycnal diffusivity of each !! layer [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(out) :: Kd_int !< The background diapycnal diffusivity of each + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(out) :: Kd_int !< The background diapycnal diffusivity of each !! interface [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(out) :: Kv_bkgnd !< The background vertical viscosity at + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(out) :: Kv_bkgnd !< The background vertical viscosity at !! each interface [H Z T-1 ~> m2 s-1 or Pa s] - integer, intent(in) :: j !< Meridional grid index + integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block + integer, intent(in) :: nj !< Number of j-rows in this block type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type type(bkgnd_mixing_cs), pointer :: CS !< The control structure returned by !! a previous call to bkgnd_mixing_init. @@ -332,9 +334,9 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, real, dimension(SZK_(GV)+1) :: depth_int !< Distance from surface of the interfaces [m] real, dimension(SZK_(GV)+1) :: Kd_col !< Diffusivities at the interfaces [m2 s-1] real, dimension(SZK_(GV)+1) :: Kv_col !< Viscosities at the interfaces [m2 s-1] - real, dimension(SZI_(G)) :: Kd_sfc !< Surface value of the diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G)) :: depth !< Distance from surface of an interface [H ~> m or kg m-2] - real, dimension(SZI_(G),SZK_(GV)) :: dz !< Height change across layers [Z ~> m] + real, dimension(SZI_(G),nj) :: Kd_sfc !< Surface value of the diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),nj) :: depth !< Distance from surface of an interface [H ~> m or kg m-2] + real, dimension(SZI_(G),SZK_(GV),nj) :: dz !< Height change across layers [Z ~> m] real :: depth_c !< depth of the center of a layer [H ~> m or kg m-2] real :: I_Hmix !< inverse of fixed mixed layer thickness [H-1 ~> m-1 or m2 kg-1] real :: I_x30 !< 2/acos(2) = 1/(sin(30 deg) * acosh(1/sin(30 deg))) [nondim] @@ -343,7 +345,7 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, real :: min_sinlat ! The minimum value of the sine of latitude [nondim] real :: bckgrnd_vdc_psin !< PSI diffusivity in northern hemisphere [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real :: bckgrnd_vdc_psis !< PSI diffusivity in southern hemisphere [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - integer :: i, k, is, ie, js, je, nz + integer :: i, j, k, is, ie, js, je, nz, jj is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke @@ -352,108 +354,122 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, min_sinlat = 1.e-10 ! Start with a constant value that may be replaced below. - Kd_lay(:,:) = CS%Kd - Kv_bkgnd(:,:) = 0.0 + do j=jstart,jend ; jj = j - jstart + 1 + Kd_lay(:,:,jj) = CS%Kd + Kv_bkgnd(:,:,jj) = 0.0 + enddo ! Set up the background diffusivity. if (CS%Bryan_Lewis_diffusivity) then - call thickness_to_dz(h, tv, dz, j, G, GV) - - do i=is,ie - depth_int(1) = 0.0 - do k=2,nz+1 - depth_int(k) = depth_int(k-1) + US%Z_to_m*dz(i,k-1) - enddo + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo - call CVMix_init_bkgnd(max_nlev=nz, & - zw = depth_int(:), & !< interface depths relative to the surface in m, must be positive. - bl1 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c1, & - bl2 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c2, & - bl3 = US%m_to_Z*CS%Bryan_Lewis_c3, & - bl4 = US%Z_to_m*CS%Bryan_Lewis_c4, & - prandtl = CS%prandtl_bkgnd) - - Kd_col(:) = 0.0 ; Kv_col(:) = 0.0 ! Is this line necessary? - call CVMix_coeffs_bkgnd(Mdiff_out=Kv_col, Tdiff_out=Kd_col, nlev=nz, max_nlev=nz) - - ! Update Kd and Kv. - do K=1,nz+1 - Kv_bkgnd(i,K) = GV%m2_s_to_HZ_T * Kv_col(K) - Kd_int(i,K) = GV%m2_s_to_HZ_T*Kd_col(K) - enddo - do k=1,nz - Kd_lay(i,k) = Kd_lay(i,k) + 0.5 * GV%m2_s_to_HZ_T * (Kd_col(K) + Kd_col(K+1)) - enddo - enddo ! i loop + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + depth_int(1) = 0.0 + do k=2,nz+1 + depth_int(k) = depth_int(k-1) + US%Z_to_m*dz(i,k-1,jj) + enddo + + call CVMix_init_bkgnd(max_nlev=nz, & + zw = depth_int(:), & !< interface depths relative to the surface in m, must be positive. + bl1 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c1, & + bl2 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c2, & + bl3 = US%m_to_Z*CS%Bryan_Lewis_c3, & + bl4 = US%Z_to_m*CS%Bryan_Lewis_c4, & + prandtl = CS%prandtl_bkgnd) + + Kd_col(:) = 0.0 ; Kv_col(:) = 0.0 ! Is this line necessary? + call CVMix_coeffs_bkgnd(Mdiff_out=Kv_col, Tdiff_out=Kd_col, nlev=nz, max_nlev=nz) + + ! Update Kd and Kv. + do K=1,nz+1 + Kv_bkgnd(i,K,jj) = GV%m2_s_to_HZ_T * Kv_col(K) + Kd_int(i,K,jj) = GV%m2_s_to_HZ_T*Kd_col(K) + enddo + do k=1,nz + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + 0.5 * GV%m2_s_to_HZ_T * (Kd_col(K) + Kd_col(K+1)) + enddo + enddo ! i loop + enddo elseif (CS%horiz_varying_background) then !### Note that there are lots of hard-coded parameters (mostly latitudes and longitudes) here. - do i=is,ie - bckgrnd_vdc_psis = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)+28.9))**2) - bckgrnd_vdc_psin = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)-28.9))**2) - Kd_int(i,1) = (CS%bckgrnd_vdc_eq + bckgrnd_vdc_psin) + bckgrnd_vdc_psis - - if (G%geoLatT(i,j) < -10.0) then - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 - elseif (G%geoLatT(i,j) <= 10.0) then - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 * (G%geoLatT(i,j)/10.0)**2 - else - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 - endif - - ! North Banda Sea - if ( (G%geoLatT(i,j) < -1.0) .and. (G%geoLatT(i,j) > -4.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 103.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 134.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif - - ! Middle Banda Sea - if ( (G%geoLatT(i,j) <= -4.0) .and. (G%geoLatT(i,j) > -7.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 106.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 140.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif - - ! South Banda Sea - if ( (G%geoLatT(i,j) <= -7.0) .and. (G%geoLatT(i,j) > -8.3) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 111.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 142.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + bckgrnd_vdc_psis = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)+28.9))**2) + bckgrnd_vdc_psin = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)-28.9))**2) + Kd_int(i,1,jj) = (CS%bckgrnd_vdc_eq + bckgrnd_vdc_psin) + bckgrnd_vdc_psis + + if (G%geoLatT(i,j) < -10.0) then + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 + elseif (G%geoLatT(i,j) <= 10.0) then + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 * (G%geoLatT(i,j)/10.0)**2 + else + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 + endif + + ! North Banda Sea + if ( (G%geoLatT(i,j) < -1.0) .and. (G%geoLatT(i,j) > -4.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 103.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 134.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + ! Middle Banda Sea + if ( (G%geoLatT(i,j) <= -4.0) .and. (G%geoLatT(i,j) > -7.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 106.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 140.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + + ! South Banda Sea + if ( (G%geoLatT(i,j) <= -7.0) .and. (G%geoLatT(i,j) > -8.3) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 111.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 142.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + + enddo enddo ! Update interior values of Kd and Kv (uniform profile; no interpolation needed) - do K=1,nz+1 ; do i=is,ie - Kd_int(i,K) = Kd_int(i,1) - Kv_bkgnd(i,K) = Kd_int(i,1) * CS%prandtl_bkgnd - enddo ; enddo - do k=1,nz ; do i=is,ie - Kd_lay(i,k) = Kd_int(i,1) - enddo ; enddo + do K=1,nz+1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int(i,K,jj) = Kd_int(i,1,jj) + Kv_bkgnd(i,K,jj) = Kd_int(i,1,jj) * CS%prandtl_bkgnd + enddo ; enddo ; enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_lay(i,k,jj) = Kd_int(i,1,jj) + enddo ; enddo ; enddo else ! Set a potentially spatially varying surface value of diffusivity. if (CS%Henyey_IGW_background) then I_x30 = 2.0 / invcosh(CS%N0_2Omega*2.0) ! This is evaluated at 30 deg. - do i=is,ie - abs_sinlat = abs(sin(G%geoLatT(i,j)*deg_to_rad)) - if (abs(G%geoLatT(i,j))>CS%Henyey_max_lat) abs_sinlat = min_sinlat - Kd_sfc(i) = max(CS%Kd_min, CS%Kd * & - ((abs_sinlat * invcosh(CS%N0_2Omega / max(min_sinlat, abs_sinlat))) * I_x30) ) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + abs_sinlat = abs(sin(G%geoLatT(i,j)*deg_to_rad)) + if (abs(G%geoLatT(i,j))>CS%Henyey_max_lat) abs_sinlat = min_sinlat + Kd_sfc(i,jj) = max(CS%Kd_min, CS%Kd * & + ((abs_sinlat * invcosh(CS%N0_2Omega / max(min_sinlat, abs_sinlat))) * I_x30) ) + enddo enddo elseif (CS%Kd_tanh_lat_fn) then - do i=is,ie - ! The transition latitude and latitude range are hard-scaled here, since - ! this is not really intended for wide-spread use, but rather for - ! comparison with CM2M / CM2.1 settings. - Kd_sfc(i) = max(CS%Kd_min, CS%Kd * (1.0 + & - CS%Kd_tanh_lat_scale * 0.5*tanh((abs(G%geoLatT(i,j)) - 35.0)/5.0) )) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + ! The transition latitude and latitude range are hard-scaled here, since + ! this is not really intended for wide-spread use, but rather for + ! comparison with CM2M / CM2.1 settings. + Kd_sfc(i,jj) = max(CS%Kd_min, CS%Kd * (1.0 + & + CS%Kd_tanh_lat_scale * 0.5*tanh((abs(G%geoLatT(i,j)) - 35.0)/5.0) )) + enddo enddo else ! Use a spatially constant surface value. - do i=is,ie - Kd_sfc(i) = CS%Kd + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + Kd_sfc(i,jj) = CS%Kd + enddo enddo endif @@ -462,33 +478,37 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, ! This is a crude way to put in a diffusive boundary layer without an explicit boundary ! layer turbulence scheme. It should not be used for any realistic ocean models. I_Hmix = 1.0 / (CS%Hmix + GV%H_subroundoff) - do i=is,ie ; depth(i) = 0.0 ; enddo - do k=1,nz ; do i=is,ie - depth_c = depth(i) + 0.5*h(i,j,k) - if (depth_c <= CS%Hmix) then ; Kd_lay(i,k) = CS%Kd_tot_ml - elseif (depth_c >= 2.0*CS%Hmix) then ; Kd_lay(i,k) = Kd_sfc(i) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; depth(i,jj) = 0.0 ; enddo + enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + depth_c = depth(i,jj) + 0.5*h(i,j,k) + if (depth_c <= CS%Hmix) then ; Kd_lay(i,k,jj) = CS%Kd_tot_ml + elseif (depth_c >= 2.0*CS%Hmix) then ; Kd_lay(i,k,jj) = Kd_sfc(i,jj) else - Kd_lay(i,k) = ((Kd_sfc(i) - CS%Kd_tot_ml) * I_Hmix) * depth_c + (2.0*CS%Kd_tot_ml - Kd_sfc(i)) + Kd_lay(i,k,jj) = ((Kd_sfc(i,jj) - CS%Kd_tot_ml) * I_Hmix) * depth_c + (2.0*CS%Kd_tot_ml - Kd_sfc(i,jj)) endif - depth(i) = depth(i) + h(i,j,k) - enddo ; enddo + depth(i,jj) = depth(i,jj) + h(i,j,k) + enddo ; enddo ; enddo else ! There is no vertical structure to the background diffusivity. - do k=1,nz ; do i=is,ie - Kd_lay(i,k) = Kd_sfc(i) - enddo ; enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_lay(i,k,jj) = Kd_sfc(i,jj) + enddo ; enddo ; enddo endif ! Update Kd_int and Kv_bkgnd, based on Kd_lay. These might be just used for diagnostic purposes. - do i=is,ie - Kd_int(i,1) = 0.0 ; Kv_bkgnd(i,1) = 0.0 - Kd_int(i,nz+1) = 0.0 ; Kv_bkgnd(i,nz+1) = 0.0 + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + Kd_int(i,1,jj) = 0.0 ; Kv_bkgnd(i,1,jj) = 0.0 + Kd_int(i,nz+1,jj) = 0.0 ; Kv_bkgnd(i,nz+1,jj) = 0.0 + enddo enddo - do K=2,nz ; do i=is,ie - Kd_int(i,K) = 0.5*(Kd_lay(i,k-1) + Kd_lay(i,k)) - Kv_bkgnd(i,K) = Kd_int(i,K) * CS%prandtl_bkgnd - enddo ; enddo + do K=2,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int(i,K,jj) = 0.5*(Kd_lay(i,k-1,jj) + Kd_lay(i,k,jj)) + Kv_bkgnd(i,K,jj) = Kd_int(i,K,jj) * CS%prandtl_bkgnd + enddo ; enddo ; enddo endif end subroutine calculate_bkgnd_mixing diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 46ace68d79..c68eb28f96 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -301,18 +301,18 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i real, dimension(SZI_(G),SZK_(GV),njblock) :: & N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] - Kd_lay_2d !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + Kd_lay_2d, & !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + maxTKE, & !< Energy required to entrain to h_max [H Z2 T-3 ~> m3 s-3 or W m-2] + TKE_to_Kd !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between + !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] real, dimension(SZI_(G),SZK_(GV)) :: & dz, & !< Height change across layers [Z ~> m] - maxTKE, & !< Energy required to entrain to h_max [H Z2 T-3 ~> m3 s-3 or W m-2] prof_leak_2d, & !< vertical profile for leakage [Z-1 ~> m-1] prof_quad_2d, & !< vertical profile for bottom drag [Z-1 ~> m-1] prof_itidal_2d, & !< vertical profile for wave drag [Z-1 ~> m-1] prof_Froude_2d, & !< vertical profile for Froude drag [Z-1 ~> m-1] - prof_slope_2d, & !< vertical profile for critical slopes [Z-1 ~> m-1] - TKE_to_Kd !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between - !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] + prof_slope_2d !< vertical profile for critical slopes [Z-1 ~> m-1] real, dimension(SZI_(G),SZK_(GV)+1,njblock) :: & N2_int, & !< squared buoyancy frequency associated at interfaces [T-2 ~> s-2] @@ -334,8 +334,8 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i logical :: use_EOS ! If true, compute density from T/S using equation of state. logical :: TKE_to_Kd_used ! If true, TKE_to_Kd and maxTKE need to be calculated. - integer :: kb(SZI_(G)) ! The index of the lightest layer denser than the - ! buffer layer, or -1 without a bulk mixed layer. + integer :: kb(SZI_(G),njblock) ! The index of the lightest layer denser than the + ! buffer layer, or -1 without a bulk mixed layer. logical :: showCallTree ! If true, show the call tree. integer :: i, j, k, is, ie, js, je, nz, isd, ied, jsd, jed @@ -485,99 +485,123 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Set up variables related to the stratification. call find_N2(h, tv, T_f, S_f, fluxes, njblock, jstart, jend, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) - !$OMP parallel do default(shared) private(dz, & - !$OMP KT_extra,KS_extra,TKE_to_Kd,maxTKE,dissip,kb,jj) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=jstart,jend - jj = j - jstart + 1 + ! Add background mixing + call calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay_2d, Kd_int_2d, Kv_bkgnd, jstart, jend, njblock, G, GV, US, CS%bkgnd_mixing_csp) if (associated(dd%N2_3d)) then - do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K,jj) ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K,jj) ; enddo ; enddo + enddo endif - - ! Add background mixing - call calculate_bkgnd_mixing(h, tv, N2_lay(:,:,jj), Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), Kv_bkgnd(:,:,jj), j, G, GV, US, CS%bkgnd_mixing_csp) ! Update Kv and 3-d diffusivity diagnostics. - if (associated(visc%Kv_slow)) then ; do K=1,nz+1 ; do i=is,ie - visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K,jj) - enddo ; enddo ; endif - if (CS%id_Kv_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K,jj) - enddo ; enddo ; endif - if (CS%id_Kd_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) - enddo ; enddo ; endif - if (associated(VBF%Kd_bkgnd)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) - enddo ; enddo ; endif - - ! Double-diffusion (old method) - if (CS%double_diffusion) then - call double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, KT_extra, KS_extra) - ! One of Kd_extra_T and Kd_extra_S is always 0. Kd_extra_S is positive for salt fingering. - ! Kd_extra_T is positive for double diffusive convection. - do K=2,nz ; do i=is,ie - if (KS_extra(i,K) > KT_extra(i,K)) then ! salt fingering - Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KT_extra(i,K) - Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KT_extra(i,K) - Kd_extra_S(i,j,K) = KS_extra(i,K) - KT_extra(i,K) - Kd_extra_T(i,j,K) = 0.0 - elseif (KT_extra(i,K) > 0.0) then ! double-diffusive convection - Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KS_extra(i,K) - Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KS_extra(i,K) - Kd_extra_T(i,j,K) = KT_extra(i,K) - KS_extra(i,K) - Kd_extra_S(i,j,K) = 0.0 - else ! There is no double diffusion at this interface. - Kd_extra_T(i,j,K) = 0.0 - Kd_extra_S(i,j,K) = 0.0 - endif - enddo ; enddo - if (associated(dd%KT_extra)) then ; do K=1,nz+1 ; do i=is,ie - dd%KT_extra(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - - if (associated(dd%KS_extra)) then ; do K=1,nz+1 ; do i=is,ie - dd%KS_extra(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif - - if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif + if (associated(visc%Kv_slow)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K,jj) + enddo ; enddo + enddo + endif + if (CS%id_Kv_bkgnd > 0) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K,jj) + enddo ; enddo + enddo + endif + if (CS%id_Kd_bkgnd > 0) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo + endif + if (associated(VBF%Kd_bkgnd)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo endif - ! Apply double diffusion via CVMix - ! GMM, we need to pass HBL to compute_ddiff_coeffs, but it is not yet available. - if (CS%use_CVMix_ddiff) then - call cpu_clock_begin(id_clock_CVMix_ddiff) - if (associated(dd%drho_rat)) then - call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, & - CS%CVMix_ddiff_csp, dd%drho_rat) - else - call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, CS%CVMix_ddiff_csp) + do j=jstart,jend ; jj = j - jstart + 1 + ! Double-diffusion (old method) + if (CS%double_diffusion) then + call double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, KT_extra, KS_extra) + ! One of Kd_extra_T and Kd_extra_S is always 0. Kd_extra_S is positive for salt fingering. + ! Kd_extra_T is positive for double diffusive convection. + do K=2,nz ; do i=is,ie + if (KS_extra(i,K) > KT_extra(i,K)) then ! salt fingering + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KT_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KT_extra(i,K) + Kd_extra_S(i,j,K) = KS_extra(i,K) - KT_extra(i,K) + Kd_extra_T(i,j,K) = 0.0 + elseif (KT_extra(i,K) > 0.0) then ! double-diffusive convection + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KS_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KS_extra(i,K) + Kd_extra_T(i,j,K) = KT_extra(i,K) - KS_extra(i,K) + Kd_extra_S(i,j,K) = 0.0 + else ! There is no double diffusion at this interface. + Kd_extra_T(i,j,K) = 0.0 + Kd_extra_S(i,j,K) = 0.0 + endif + enddo ; enddo + if (associated(dd%KT_extra)) then ; do K=1,nz+1 ; do i=is,ie + dd%KT_extra(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + + if (associated(dd%KS_extra)) then ; do K=1,nz+1 ; do i=is,ie + dd%KS_extra(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif + + if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif endif - if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif - call cpu_clock_end(id_clock_CVMix_ddiff) - endif + + ! Apply double diffusion via CVMix + ! GMM, we need to pass HBL to compute_ddiff_coeffs, but it is not yet available. + if (CS%use_CVMix_ddiff) then + call cpu_clock_begin(id_clock_CVMix_ddiff) + if (associated(dd%drho_rat)) then + call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, & + CS%CVMix_ddiff_csp, dd%drho_rat) + else + call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, CS%CVMix_ddiff_csp) + endif + if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif + call cpu_clock_end(id_clock_CVMix_ddiff) + endif + enddo ! Calculate conversion ratios from TKE to layer diffusivities. if (TKE_to_Kd_used) then - call find_TKE_to_Kd(h, tv, dRho_int(:,:,jj), N2_lay(:,:,jj), j, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) - if (associated(dd%maxTKE)) then ; do k=1,nz ; do i=is,ie - dd%maxTKE(i,j,k) = maxTKE(i,k) - enddo ; enddo ; endif - if (associated(dd%TKE_to_Kd)) then ; do k=1,nz ; do i=is,ie - dd%TKE_to_Kd(i,j,k) = TKE_to_Kd(i,k) - enddo ; enddo ; endif + call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, njblock, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) + if (associated(dd%maxTKE)) then + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + dd%maxTKE(i,j,k) = maxTKE(i,k,jj) + enddo ; enddo ; enddo + endif + if (associated(dd%TKE_to_Kd)) then + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + dd%TKE_to_Kd(i,j,k) = TKE_to_Kd(i,k,jj) + enddo ; enddo ; enddo + endif endif + !$OMP parallel do default(shared) private(dz, & + !$OMP KT_extra,KS_extra,dissip,jj) & + !$OMP if(.not. CS%use_CVMix_ddiff) + do j=jstart,jend + jj = j - jstart + 1 + ! Add the input turbulent diffusivity. if (CS%useKappaShear .or. CS%use_CVMix_shear) then do K=2,nz ; do i=is,ie @@ -605,19 +629,19 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Add the ML_Rad diffusivity. if (CS%ML_radiation) then - call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd, Kd_lay_2d(:,:,jj)) + call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) endif ! Add the Nikurashin and / or tidal bottom-driven mixing if (CS%use_tidal_mixing) & - call calculate_tidal_mixing(dz, j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd, & - maxTKE, G, GV, US, CS%tidal_mixing, & + call calculate_tidal_mixing(dz, j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & + maxTKE(:,:,jj), G, GV, US, CS%tidal_mixing, & CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) ! Add diffusivity from internal tides ray tracing if (CS%use_int_tides) then - call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd, CS%Kd_max, & + call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), CS%Kd_max, & CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & prof_slope_2d) @@ -682,8 +706,8 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), Kd_int_2d(:,:,jj), & G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj)) else - call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, & - maxTKE, kb, rho_bot(:,jj), G, GV, US, CS, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), dd%Kd_BBL) + call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd(:,:,jj), & + maxTKE(:,:,jj), kb(:,jj), rho_bot(:,jj), G, GV, US, CS, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), dd%Kd_BBL) endif if (associated(VBF%Kd_BBL)) then ; do K=1,nz+1 ; do i=is,ie VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) @@ -880,7 +904,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i end subroutine set_diffusivity !> Convert turbulent kinetic energy to diffusivity -subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & +subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, US, CS, & TKE_to_Kd, maxTKE, kb) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure @@ -889,24 +913,26 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & intent(in) :: h !< Layer thicknesses [H ~> m or kg m-2] type(thermo_var_ptrs), intent(in) :: tv !< Structure containing pointers to any available !! thermodynamic fields. - real, dimension(SZI_(G),SZK_(GV)+1), intent(in) :: dRho_int !< Change in locally referenced potential density + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(in) :: dRho_int !< Change in locally referenced potential density !! across each interface [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: N2_lay !< The squared buoyancy frequency of the + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: N2_lay !< The squared buoyancy frequency of the !! layers [T-2 ~> s-2]. - integer, intent(in) :: j !< j-index of row to work on + integer, intent(in) :: jstart !< Start j-index of the j-block + integer, intent(in) :: jend !< End j-index of the j-block + integer, intent(in) :: nj !< Number of j-rows in the block real, intent(in) :: dt !< Time increment [T ~> s]. type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: TKE_to_Kd !< The conversion rate between the + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: TKE_to_Kd !< The conversion rate between the !! TKE dissipated within a layer and the !! diapycnal diffusivity within that layer, !! usually (~Rho_0 / (G_Earth * dRho_lay)) !! [T2 Z-1 ~> s2 m-1] - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: maxTKE !< The energy required to for a layer to entrain to its + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: maxTKE !< The energy required to for a layer to entrain to its !! maximum realizable thickness [H Z2 T-3 ~> m3 s-3 or W m-2] - integer, dimension(SZI_(G)), intent(out) :: kb !< Index of lightest layer denser than the buffer + integer, dimension(SZI_(G),nj), intent(out) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer. ! Local variables - real, dimension(SZI_(G),SZK_(GV)) :: & + real, dimension(SZI_(G),SZK_(GV),nj) :: & ds_dsp1, & ! coordinate variable (sigma-2) difference across an ! interface divided by the difference across the interface ! below it [nondim] @@ -919,7 +945,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & ! compensating entrainment from above to keep the layer ! density from changing) that will not deplete all of the ! layers above or below a layer within a timestep [H ~> m or kg m-2]. - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & htot, & ! total thickness above or below a layer, or the ! integrated thickness in the BBL [H ~> m or kg m-2]. mFkb, & ! total thickness in the mixed and buffer layers times ds_dsp1 [H ~> m or kg m-2] @@ -939,10 +965,10 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & real :: I_dt ! 1/dt [T-1 ~> s-1] real :: dz_neglect ! A negligibly small height change [Z ~> m] real :: hN2pO2 ! h (N^2 + Omega^2), in [Z T-2 ~> m s-2]. - logical :: do_i(SZI_(G)) + logical :: do_i(SZI_(G),nj) integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, is, ie, nz, i_rem, kmb, kb_min + integer :: i, j, k, is, ie, nz, i_rem, kmb, kb_min, jj is = G%isc ; ie = G%iec ; nz = GV%ke I_dt = 1.0 / dt @@ -956,150 +982,154 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & G_IRho0 = GV%H_to_Z*G_Rho0 endif - ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) + do j=jstart,jend + jj = j - jstart + 1 + ! Find the vertical distances across layers. + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) ! Simple but coordinate-independent estimate of Kd/TKE if (CS%simple_TKE_to_Kd) then do k=1,nz ; do i=is,ie - hN2pO2 = dz(i,k) * (N2_lay(i,k) + Omega2) ! Units of Z T-2. + hN2pO2 = dz(i,k,jj) * (N2_lay(i,k,jj) + Omega2) ! Units of Z T-2. if (hN2pO2 > 0.) then - TKE_to_Kd(i,k) = 1.0 / hN2pO2 ! Units of T2 H-1. - else ; TKE_to_Kd(i,k) = 0. ; endif + TKE_to_Kd(i,k,jj) = 1.0 / hN2pO2 ! Units of T2 H-1. + else ; TKE_to_Kd(i,k,jj) = 0. ; endif ! The maximum TKE conversion we allow is really a statement ! about the upper diffusivity we allow. Kd_max must be set. - maxTKE(i,k) = hN2pO2 * CS%Kd_max ! Units of H Z2 T-3. + maxTKE(i,k,jj) = hN2pO2 * CS%Kd_max ! Units of H Z2 T-3. enddo ; enddo - kb(is:ie) = -1 ! kb should not be used by any code in non-layered mode -AJA - return + kb(is:ie,jj) = -1 ! kb should not be used by any code in non-layered mode -AJA + cycle endif ! Determine kb - the index of the shallowest active interior layer. if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies - do i=is,ie ; p_0(i) = 0.0 ; p_ref(i) = tv%P_Ref ; enddo + do i=is,ie ; p_0(i,jj) = 0.0 ; p_ref(i,jj) = tv%P_Ref ; enddo EOSdom(:) = EOS_domain(G%HI) do k=1,nz - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0, rho_0(:,k), tv%eqn_of_state, EOSdom) + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0(:,jj), rho_0(:,k,jj), tv%eqn_of_state, EOSdom) enddo - call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref, Rcv_kmb, tv%eqn_of_state, EOSdom) + call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref(:,jj), Rcv_kmb(:,jj), tv%eqn_of_state, EOSdom) kb_min = kmb+1 do i=is,ie ! Determine the next denser layer than the buffer layer in the ! coordinate density (sigma-2). - do k=kmb+1,nz-1 ; if (Rcv_kmb(i) <= GV%Rlay(k)) exit ; enddo - kb(i) = k + do k=kmb+1,nz-1 ; if (Rcv_kmb(i,jj) <= GV%Rlay(k)) exit ; enddo + kb(i,jj) = k ! Backtrack, in case there are massive layers above that are stable ! in sigma-0. - do k=kb(i)-1,kmb+1,-1 - if (rho_0(i,kmb) > rho_0(i,k)) exit - if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i) = k + do k=kb(i,jj)-1,kmb+1,-1 + if (rho_0(i,kmb,jj) > rho_0(i,k,jj)) exit + if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i,jj) = k enddo enddo - call set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) + call set_density_ratios(h, tv, kb(:,jj), G, GV, US, CS, j, ds_dsp1(:,:,jj), rho_0(:,:,jj)) else ! not bulkmixedlayer kb_min = 2 ; kmb = 0 - do i=is,ie ; kb(i) = 1 ; enddo - call set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1) + do i=is,ie ; kb(i,jj) = 1 ; enddo + call set_density_ratios(h, tv, kb(:,jj), G, GV, US, CS, j, ds_dsp1(:,:,jj)) endif ! Determine maxEnt - the maximum permitted entrainment from below by each ! interior layer. do k=2,nz-1 ; do i=is,ie - dsp1_ds(i,k) = 1.0 / ds_dsp1(i,k) + dsp1_ds(i,k,jj) = 1.0 / ds_dsp1(i,k,jj) enddo ; enddo - do i=is,ie ; dsp1_ds(i,nz) = 0.0 ; enddo + do i=is,ie ; dsp1_ds(i,nz,jj) = 0.0 ; enddo if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies do i=is,ie - htot(i) = h(i,j,kmb) - mFkb(i) = 0.0 - if (kb(i) < nz) mFkb(i) = ds_dsp1(i,kb(i)) * (h(i,j,kmb) - GV%Angstrom_H) + htot(i,jj) = h(i,j,kmb) + mFkb(i,jj) = 0.0 + if (kb(i,jj) < nz) mFkb(i,jj) = ds_dsp1(i,kb(i,jj),jj) * (h(i,j,kmb) - GV%Angstrom_H) enddo do k=1,kmb-1 ; do i=is,ie - htot(i) = htot(i) + h(i,j,k) - mFkb(i) = mFkb(i) + ds_dsp1(i,k+1)*(h(i,j,k) - GV%Angstrom_H) + htot(i,jj) = htot(i,jj) + h(i,j,k) + mFkb(i,jj) = mFkb(i,jj) + ds_dsp1(i,k+1,jj)*(h(i,j,k) - GV%Angstrom_H) enddo ; enddo else do i=is,i - maxEnt(i,1) = 0.0 ; htot(i) = h(i,j,1) - GV%Angstrom_H + maxEnt(i,1,jj) = 0.0 ; htot(i,jj) = h(i,j,1) - GV%Angstrom_H enddo endif do k=kb_min,nz-1 ; do i=is,ie - if (k == kb(i)) then - maxEnt(i,kb(i)) = mFkb(i) - elseif (k > kb(i)) then + if (k == kb(i,jj)) then + maxEnt(i,kb(i,jj),jj) = mFkb(i,jj) + elseif (k > kb(i,jj)) then if (CS%answer_date < 20190101) then - maxEnt(i,k) = (1.0/dsp1_ds(i,k))*(maxEnt(i,k-1) + htot(i)) + maxEnt(i,k,jj) = (1.0/dsp1_ds(i,k,jj))*(maxEnt(i,k-1,jj) + htot(i,jj)) else - maxEnt(i,k) = ds_dsp1(i,k)*(maxEnt(i,k-1) + htot(i)) + maxEnt(i,k,jj) = ds_dsp1(i,k,jj)*(maxEnt(i,k-1,jj) + htot(i,jj)) endif - htot(i) = htot(i) + (h(i,j,k) - GV%Angstrom_H) + htot(i,jj) = htot(i,jj) + (h(i,j,k) - GV%Angstrom_H) endif enddo ; enddo do i=is,ie - htot(i) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,nz) = 0.0 - do_i(i) = (G%mask2dT(i,j) > 0.0) + htot(i,jj) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,nz,jj) = 0.0 + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) enddo do k=nz-1,kb_min,-1 i_rem = 0 - do i=is,ie ; if (do_i(i)) then - if (k Date: Sun, 31 May 2026 14:25:21 -0400 Subject: [PATCH 04/45] set_diffusivity: block j loops in find_TKE_to_Kd and set_density_ratios Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_set_diffusivity.F90 | 307 +++++++++--------- 1 file changed, 157 insertions(+), 150 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index c68eb28f96..ca93b0b8c6 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -303,11 +303,11 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] Kd_lay_2d, & !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] maxTKE, & !< Energy required to entrain to h_max [H Z2 T-3 ~> m3 s-3 or W m-2] - TKE_to_Kd !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between + TKE_to_Kd, & !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] + dz !< Height change across layers [Z ~> m] real, dimension(SZI_(G),SZK_(GV)) :: & - dz, & !< Height change across layers [Z ~> m] prof_leak_2d, & !< vertical profile for leakage [Z-1 ~> m-1] prof_quad_2d, & !< vertical profile for bottom drag [Z-1 ~> m-1] prof_itidal_2d, & !< vertical profile for wave drag [Z-1 ~> m-1] @@ -596,45 +596,47 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i endif endif - !$OMP parallel do default(shared) private(dz, & - !$OMP KT_extra,KS_extra,dissip,jj) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=jstart,jend - jj = j - jstart + 1 - ! Add the input turbulent diffusivity. if (CS%useKappaShear .or. CS%use_CVMix_shear) then - do K=2,nz ; do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do K=2,nz ; do i=is,ie Kd_int_2d(i,K,jj) = visc%Kd_shear(i,j,K) + 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) - enddo ; enddo - do i=is,ie + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie Kd_int_2d(i,1,jj) = visc%Kd_shear(i,j,1) ! This isn't actually used. It could be 0. Kd_int_2d(i,nz+1,jj) = 0.0 - enddo - do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) + enddo ; enddo ; enddo else - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie Kd_int_2d(i,1,jj) = Kd_lay_2d(i,1,jj) ; Kd_int_2d(i,nz+1,jj) = 0.0 - enddo - do K=2,nz ; do i=is,ie - Kd_int_2d(i,K,jj) = 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do K=2,nz ; do i=is,ie + Kd_int_2d(i,K,jj) = 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) + enddo ; enddo ; enddo endif if (CS%ML_radiation .or. CS%use_tidal_mixing .or. associated(dd%Kd_Work)) then - call thickness_to_dz(h, tv, dz, j, G, GV) + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo endif + !$OMP parallel do default(shared) private( & + !$OMP KT_extra,KS_extra,dissip,jj) & + !$OMP if(.not. CS%use_CVMix_ddiff) + do j=jstart,jend + jj = j - jstart + 1 + ! Add the ML_Rad diffusivity. if (CS%ML_radiation) then - call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) + call add_MLrad_diffusivity(dz(:,:,jj), fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) endif ! Add the Nikurashin and / or tidal bottom-driven mixing if (CS%use_tidal_mixing) & - call calculate_tidal_mixing(dz, j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & + call calculate_tidal_mixing(dz(:,:,jj), j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & maxTKE(:,:,jj), G, GV, US, CS%tidal_mixing, & CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) @@ -759,7 +761,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i if (associated(dd%Kd_Work)) then do k=1,nz ; do i=is,ie - dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k) ! Watt m-2 = kg s-3 + dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 enddo ; enddo endif @@ -772,7 +774,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i if (associated(dd%Kd_Work_added)) then do k=1,nz ; do i=is,ie - dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k) ! Watt m-2 = kg s-3 + dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 enddo ; enddo endif @@ -982,14 +984,13 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, G_IRho0 = GV%H_to_Z*G_Rho0 endif - do j=jstart,jend - jj = j - jstart + 1 - ! Find the vertical distances across layers. + do j=jstart,jend ; jj = j - jstart + 1 call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo ! Simple but coordinate-independent estimate of Kd/TKE if (CS%simple_TKE_to_Kd) then - do k=1,nz ; do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie hN2pO2 = dz(i,k,jj) * (N2_lay(i,k,jj) + Omega2) ! Units of Z T-2. if (hN2pO2 > 0.) then TKE_to_Kd(i,k,jj) = 1.0 / hN2pO2 ! Units of T2 H-1. @@ -997,23 +998,28 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, ! The maximum TKE conversion we allow is really a statement ! about the upper diffusivity we allow. Kd_max must be set. maxTKE(i,k,jj) = hN2pO2 * CS%Kd_max ! Units of H Z2 T-3. - enddo ; enddo - kb(is:ie,jj) = -1 ! kb should not be used by any code in non-layered mode -AJA - cycle + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + kb(is:ie,jj) = -1 ! kb should not be used by any code in non-layered mode -AJA + enddo + return endif ! Determine kb - the index of the shallowest active interior layer. if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies - do i=is,ie ; p_0(i,jj) = 0.0 ; p_ref(i,jj) = tv%P_Ref ; enddo + kb_min = kmb+1 EOSdom(:) = EOS_domain(G%HI) - do k=1,nz + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + p_0(i,jj) = 0.0 ; p_ref(i,jj) = tv%P_Ref + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0(:,jj), rho_0(:,k,jj), tv%eqn_of_state, EOSdom) + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref(:,jj), Rcv_kmb(:,jj), tv%eqn_of_state, EOSdom) enddo - call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref(:,jj), Rcv_kmb(:,jj), tv%eqn_of_state, EOSdom) - - kb_min = kmb+1 - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie ! Determine the next denser layer than the buffer layer in the ! coordinate density (sigma-2). do k=kmb+1,nz-1 ; if (Rcv_kmb(i,jj) <= GV%Rlay(k)) exit ; enddo @@ -1025,39 +1031,42 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, if (rho_0(i,kmb,jj) > rho_0(i,k,jj)) exit if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i,jj) = k enddo - enddo - - call set_density_ratios(h, tv, kb(:,jj), G, GV, US, CS, j, ds_dsp1(:,:,jj), rho_0(:,:,jj)) + enddo ; enddo + call set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp1, rho_0) else ! not bulkmixedlayer kb_min = 2 ; kmb = 0 - do i=is,ie ; kb(i,jj) = 1 ; enddo - call set_density_ratios(h, tv, kb(:,jj), G, GV, US, CS, j, ds_dsp1(:,:,jj)) + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + kb(i,jj) = 1 + enddo ; enddo + call set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp1) endif ! Determine maxEnt - the maximum permitted entrainment from below by each ! interior layer. - do k=2,nz-1 ; do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do k=2,nz-1 ; do i=is,ie dsp1_ds(i,k,jj) = 1.0 / ds_dsp1(i,k,jj) + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + dsp1_ds(i,nz,jj) = 0.0 enddo ; enddo - do i=is,ie ; dsp1_ds(i,nz,jj) = 0.0 ; enddo if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie htot(i,jj) = h(i,j,kmb) mFkb(i,jj) = 0.0 if (kb(i,jj) < nz) mFkb(i,jj) = ds_dsp1(i,kb(i,jj),jj) * (h(i,j,kmb) - GV%Angstrom_H) - enddo - do k=1,kmb-1 ; do i=is,ie + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,kmb-1 ; do i=is,ie htot(i,jj) = htot(i,jj) + h(i,j,k) mFkb(i,jj) = mFkb(i,jj) + ds_dsp1(i,k+1,jj)*(h(i,j,k) - GV%Angstrom_H) - enddo ; enddo + enddo ; enddo ; enddo else - do i=is,i + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie maxEnt(i,1,jj) = 0.0 ; htot(i,jj) = h(i,j,1) - GV%Angstrom_H - enddo + enddo ; enddo endif - do k=kb_min,nz-1 ; do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do k=kb_min,nz-1 ; do i=is,ie if (k == kb(i,jj)) then maxEnt(i,kb(i,jj),jj) = mFkb(i,jj) elseif (k > kb(i,jj)) then @@ -1068,68 +1077,68 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, endif htot(i,jj) = htot(i,jj) + (h(i,j,k) - GV%Angstrom_H) endif - enddo ; enddo + enddo ; enddo ; enddo - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie htot(i,jj) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,nz,jj) = 0.0 do_i(i,jj) = (G%mask2dT(i,j) > 0.0) - enddo - do k=nz-1,kb_min,-1 - i_rem = 0 - do i=is,ie ; if (do_i(i,jj)) then - if (k kg m-3]. ! Local variables real :: g_R0 ! g_R0 is a rescaled version of g/Rho [L2 Z-1 R-1 T-2 ~> m4 kg-1 s-2] real :: eps, tmp ! nondimensional temporary variables [nondim] - real :: a(SZK_(GV)), a_0(SZK_(GV)) ! nondimensional temporary variables [nondim] - real :: p_ref(SZI_(G)) ! an array of tv%P_Ref pressures [R L2 T-2 ~> Pa] - real :: Rcv(SZI_(G),SZK_(GV)) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] + real :: a(SZK_(GV),nj), a_0(SZK_(GV),nj) ! nondimensional temporary variables [nondim] + real :: p_ref(SZI_(G),nj) ! an array of tv%P_Ref pressures [R L2 T-2 ~> Pa] + real :: Rcv(SZI_(G),SZK_(GV),nj) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] real :: I_Drho ! The inverse of the coordinate density difference between ! layers [R-1 ~> m3 kg-1] integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, k3, is, ie, nz, kmb + integer :: i, j, k, k3, is, ie, nz, kmb, jj is = G%isc ; ie = G%iec ; nz = GV%ke - do k=2,nz-1 + do j=jstart,jend ; jj = j - jstart + 1 ; do k=2,nz-1 ; do i=is,ie if (GV%g_prime(k+1) /= 0.0) then if (GV%Boussinesq .or. GV%Semi_Boussinesq) then - do i=is,ie - ds_dsp1(i,k) = GV%g_prime(k) / GV%g_prime(k+1) - enddo + ds_dsp1(i,k,jj) = GV%g_prime(k) / GV%g_prime(k+1) else ! Use a mathematically equivalent form that avoids any dependency on RHO_0. - do i=is,ie - ds_dsp1(i,k) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) - enddo + ds_dsp1(i,k,jj) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) endif else - do i=is,ie - ds_dsp1(i,k) = 1. - enddo + ds_dsp1(i,k,jj) = 1. endif - enddo + enddo ; enddo ; enddo if (CS%bulkmixedlayer) then g_R0 = GV%g_Earth / (GV%Rho0) kmb = GV%nk_rho_varies eps = 0.1 - do i=is,ie ; p_ref(i) = tv%P_Ref ; enddo EOSdom(:) = EOS_domain(G%HI) - do k=1,kmb - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref, Rcv(:,k), tv%eqn_of_state, EOSdom) - enddo - do i=is,ie - if (kb(i) <= nz-1) then + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + p_ref(i,jj) = tv%P_Ref + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,kmb + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref(:,jj), Rcv(:,k,jj), tv%eqn_of_state, EOSdom) + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + if (kb(i,jj) <= nz-1) then ! Set up appropriately limited ratios of the reduced gravities of the ! interfaces above and below the buffer layer and the next denser layer. - k = kb(i) + k = kb(i,jj) if (GV%Boussinesq .or. GV%Semi_Boussinesq) then I_Drho = g_R0 / GV%g_prime(k+1) @@ -2313,44 +2320,44 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) endif ! The indexing convention for a is appropriate for the interfaces. do k3=1,kmb - a(k3+1) = (GV%Rlay(k) - Rcv(i,k3)) * I_Drho + a(k3+1,jj) = (GV%Rlay(k) - Rcv(i,k3,jj)) * I_Drho enddo - if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,k))) then + if ((present(rho_0)) .and. (a(kmb+1,jj) < 2.0*eps*ds_dsp1(i,k,jj))) then ! If the buffer layer nearly matches the density of the layer below in the ! coordinate variable (sigma-2), use the sigma-0-based density ratio if it is ! greater (and stable). - if ((rho_0(i,k) > rho_0(i,kmb)) .and. & - (rho_0(i,k+1) > rho_0(i,k))) then - I_Drho = 1.0 / (rho_0(i,k+1)-rho_0(i,k)) - a_0(kmb+1) = min((rho_0(i,k)-rho_0(i,kmb)) * I_Drho, ds_dsp1(i,k)) - if (a_0(kmb+1) > a(kmb+1)) then + if ((rho_0(i,k,jj) > rho_0(i,kmb,jj)) .and. & + (rho_0(i,k+1,jj) > rho_0(i,k,jj))) then + I_Drho = 1.0 / (rho_0(i,k+1,jj)-rho_0(i,k,jj)) + a_0(kmb+1,jj) = min((rho_0(i,k,jj)-rho_0(i,kmb,jj)) * I_Drho, ds_dsp1(i,k,jj)) + if (a_0(kmb+1,jj) > a(kmb+1,jj)) then do k3=2,kmb - a_0(k3) = a_0(kmb+1) + (rho_0(i,kmb)-rho_0(i,k3-1)) * I_Drho + a_0(k3,jj) = a_0(kmb+1,jj) + (rho_0(i,kmb,jj)-rho_0(i,k3-1,jj)) * I_Drho enddo - if (a(kmb+1) <= eps*ds_dsp1(i,k)) then - do k3=2,kmb+1 ; a(k3) = a_0(k3) ; enddo + if (a(kmb+1,jj) <= eps*ds_dsp1(i,k,jj)) then + do k3=2,kmb+1 ; a(k3,jj) = a_0(k3,jj) ; enddo else -! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,k)) - 1.0)) ) - tmp = a(kmb+1)/(eps*ds_dsp1(i,k)) - 1.0 - do k3=2,kmb+1 ; a(k3) = tmp*a(k3) + (1.0-tmp)*a_0(k3) ; enddo +! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,k,jj)) - 1.0)) ) + tmp = a(kmb+1,jj)/(eps*ds_dsp1(i,k,jj)) - 1.0 + do k3=2,kmb+1 ; a(k3,jj) = tmp*a(k3,jj) + (1.0-tmp)*a_0(k3,jj) ; enddo endif endif endif endif - ds_dsp1(i,k) = MAX(a(kmb+1),1e-5) + ds_dsp1(i,k,jj) = MAX(a(kmb+1,jj),1e-5) do k3=2,kmb -! ds_dsp1(i,k3) = MAX(a(k3),1e-5) +! ds_dsp1(i,k3,jj) = MAX(a(k3,jj),1e-5) ! Deliberately treat convective instabilities of the upper mixed ! and buffer layers with respect to the deepest buffer layer as ! though they don't exist. They will be eliminated by the upcoming ! call to the mixedlayer code anyway. ! The indexing convention is appropriate for the interfaces. - ds_dsp1(i,k3) = MAX(a(k3),ds_dsp1(i,k)) + ds_dsp1(i,k3,jj) = MAX(a(k3,jj),ds_dsp1(i,k,jj)) enddo - endif ! (kb(i) <= nz-1) - enddo ! I-loop. + endif ! (kb(i,jj) <= nz-1) + enddo ; enddo ! i-loop / j-loop endif ! bulkmixedlayer end subroutine set_density_ratios From 6071a19206e415316c6708f0e557137d0b9809b8 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sun, 31 May 2026 14:45:17 -0400 Subject: [PATCH 05/45] set_diffusivity: consolidate thickness_to_dz into single dz array Compute dz once unconditionally in set_diffusivity and pass it as a required argument to find_N2, find_TKE_to_Kd, and add_LOTW_BBL_diffusivity, removing the redundant per-call thickness_to_dz invocations in each routine. set_BBL_TKE retains its own thickness_to_dz since it is called externally. Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_set_diffusivity.F90 | 43 ++++++------------- 1 file changed, 14 insertions(+), 29 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index ca93b0b8c6..3326c2a64f 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -482,8 +482,12 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i do jstart=js,je,njblock jend = min(jstart+njblock-1, je) + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo + ! Set up variables related to the stratification. - call find_N2(h, tv, T_f, S_f, fluxes, njblock, jstart, jend, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) + call find_N2(h, tv, T_f, S_f, fluxes, njblock, jstart, jend, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot, dz) ! Add background mixing call calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay_2d, Kd_int_2d, Kv_bkgnd, jstart, jend, njblock, G, GV, US, CS%bkgnd_mixing_csp) @@ -583,7 +587,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! Calculate conversion ratios from TKE to layer diffusivities. if (TKE_to_Kd_used) then - call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, njblock, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) + call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, njblock, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb, dz) if (associated(dd%maxTKE)) then do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie dd%maxTKE(i,j,k) = maxTKE(i,k,jj) @@ -617,12 +621,6 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i enddo ; enddo ; enddo endif - if (CS%ML_radiation .or. CS%use_tidal_mixing .or. associated(dd%Kd_Work)) then - do j=jstart,jend ; jj = j - jstart + 1 - call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) - enddo - endif - !$OMP parallel do default(shared) private( & !$OMP KT_extra,KS_extra,dissip,jj) & !$OMP if(.not. CS%use_CVMix_ddiff) @@ -706,7 +704,7 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0)) then if (CS%use_LOTW_BBL_diffusivity) then call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), Kd_int_2d(:,:,jj), & - G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj)) + G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj), dz(:,:,jj)) else call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd(:,:,jj), & maxTKE(:,:,jj), kb(:,jj), rho_bot(:,jj), G, GV, US, CS, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), dd%Kd_BBL) @@ -907,7 +905,7 @@ end subroutine set_diffusivity !> Convert turbulent kinetic energy to diffusivity subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, US, CS, & - TKE_to_Kd, maxTKE, kb) + TKE_to_Kd, maxTKE, kb, dz) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -933,6 +931,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, !! maximum realizable thickness [H Z2 T-3 ~> m3 s-3 or W m-2] integer, dimension(SZI_(G),nj), intent(out) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer. + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables real, dimension(SZI_(G),SZK_(GV),nj) :: & ds_dsp1, & ! coordinate variable (sigma-2) difference across an @@ -942,7 +941,6 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, ! across an interface times the difference across the ! interface above it [nondim] rho_0, & ! Layer potential densities relative to surface pressure [R ~> kg m-3] - dz, & ! Height change across layers [Z ~> m] maxEnt ! maxEnt is the maximum value of entrainment from below (with ! compensating entrainment from above to keep the layer ! density from changing) that will not deplete all of the @@ -984,10 +982,6 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, G_IRho0 = GV%H_to_Z*G_Rho0 endif - do j=jstart,jend ; jj = j - jstart + 1 - call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) - enddo - ! Simple but coordinate-independent estimate of Kd/TKE if (CS%simple_TKE_to_Kd) then do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie @@ -1144,7 +1138,7 @@ end subroutine find_TKE_to_Kd !> Calculate Brunt-Vaisala frequency, N^2. subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRho_int, & - N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot) + N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot, dz) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -1174,6 +1168,7 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. real, dimension(SZI_(G),nj), intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. integer, dimension(SZI_(G),nj), intent(out) :: k_bot !< Bottom boundary layer top layer index. + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables real, dimension(SZI_(G),SZK_(GV)+1,nj) :: & @@ -1181,8 +1176,6 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh dRho_int_unfilt, & ! unfiltered density differences across interfaces [R ~> kg m-3] dRho_dT, & ! partial derivative of density wrt temp [R C-1 ~> kg m-3 degC-1] dRho_dS ! partial derivative of density wrt saln [R S-1 ~> kg m-3 ppt-1] - real, dimension(SZI_(G),SZK_(GV),nj) :: & - dz ! Height change across layers [Z ~> m] real, dimension(SZI_(G),nj) :: & Temp_int, & ! temperature at each interface [C ~> degC] Salin_int, & ! salinity at each interface [S ~> ppt] @@ -1250,11 +1243,6 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh enddo endif - ! Find the vertical distances across layers. - do j=jstart,jend ; jj = j - jstart + 1 - call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) - enddo - ! Set the buoyancy frequencies. do k=1,nz do j=jstart,jend ; jj = j - jstart + 1 @@ -1700,7 +1688,7 @@ end subroutine add_drag_diffusivity !! wall turbulent viscosity, up to a BBL height where the energy used for mixing has !! consumed the mechanical TKE input. subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bot, Kd_int, & - G, GV, US, CS, Kd_BBL, Kd_lay) + G, GV, US, CS, Kd_BBL, Kd_lay, dz) type(ocean_grid_type), intent(in) :: G !< Grid structure type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -1725,10 +1713,10 @@ subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bo type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure real, dimension(:,:,:), pointer :: Kd_BBL !< Interface BBL diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real, dimension(SZI_(G),SZK_(GV)), & - optional, intent(inout) :: Kd_lay !< Layer net diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + optional, intent(inout) :: Kd_lay !< Layer net diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),SZK_(GV)), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real :: dz(SZI_(G),SZK_(GV)) ! Height change across layers [Z ~> m] real :: dz_above(SZK_(GV)+1) ! Distance from each interface to the surface [Z ~> m] real :: TKE_column ! net TKE input into the column [H Z2 T-3 ~> m3 s-3 or W m-2] real :: BBL_meanKE_dis ! Sum of tidal and mean kinetic energy dissipation in the bottom boundary layer, which @@ -1767,9 +1755,6 @@ subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bo if (allocated(visc%Ray_u) .and. allocated(visc%Ray_v)) Rayleigh_drag = .true. cdrag_sqrt = sqrt(CS%cdrag) - ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) - do i=G%isc,G%iec ! Developed in single-column mode ! Column-wise parameters. From d22f2ea5f310245f6225a62c7dedea41e98d05cb Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sun, 31 May 2026 14:56:37 -0400 Subject: [PATCH 06/45] set_diffusivity: move ML_radiation, tidal_mixing, int_tides out of OMP loop Give each branch its own j-loop outside the OMP parallel region, with the int_tides diagnostics kept inside that loop alongside the call. Remove stray print statements. Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_set_diffusivity.F90 | 142 +++++++++--------- 1 file changed, 74 insertions(+), 68 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 3326c2a64f..bf7b7103cf 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -621,85 +621,91 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i enddo ; enddo ; enddo endif - !$OMP parallel do default(shared) private( & - !$OMP KT_extra,KS_extra,dissip,jj) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=jstart,jend - jj = j - jstart + 1 - ! Add the ML_Rad diffusivity. if (CS%ML_radiation) then - call add_MLrad_diffusivity(dz(:,:,jj), fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) + do j=jstart,jend ; jj = j - jstart + 1 + call add_MLrad_diffusivity(dz(:,:,jj), fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) + enddo endif ! Add the Nikurashin and / or tidal bottom-driven mixing - if (CS%use_tidal_mixing) & - call calculate_tidal_mixing(dz(:,:,jj), j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & - maxTKE(:,:,jj), G, GV, US, CS%tidal_mixing, & - CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) + if (CS%use_tidal_mixing) then + do j=jstart,jend ; jj = j - jstart + 1 + call calculate_tidal_mixing(dz(:,:,jj), j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & + maxTKE(:,:,jj), G, GV, US, CS%tidal_mixing, & + CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) + enddo + endif ! Add diffusivity from internal tides ray tracing if (CS%use_int_tides) then + do j=jstart,jend ; jj = j - jstart + 1 + call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), CS%Kd_max, & + CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & + Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & + prof_slope_2d) + + if (CS%id_kbbl > 0) then ; do i=is,ie + dd%kbbl(i,j) = k_bot(i,jj) + enddo ; endif + if (CS%id_bbl_thick > 0) then ; do i=is,ie + dd%bbl_thick(i,j) = h_bot(i,jj) + enddo ; endif + if (CS%id_Kd_leak > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_leak(i,j,K) = Kd_leak_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_quad > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_quad(i,j,K) = Kd_quad_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_itidal > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_itidal(i,j,K) = Kd_itidal_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_Froude > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_Froude(i,j,K) = Kd_Froude_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_slope > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_slope(i,j,K) = Kd_slope_2d(i,K) + enddo ; enddo ; endif + if (associated (VBF%Kd_leak)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_leak(i,j,K) = min(Kd_leak_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_quad)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_quad(i,j,K) = min(Kd_quad_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_itidal)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_itidal(i,j,K) = min(Kd_itidal_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_Froude)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_Froude(i,j,K) = min(Kd_Froude_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_slope)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_slope(i,j,K) = min(Kd_slope_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif - call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), CS%Kd_max, & - CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & - Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & - prof_slope_2d) - - if (CS%id_kbbl > 0) then ; do i=is,ie - dd%kbbl(i,j) = k_bot(i,jj) - enddo ; endif - if (CS%id_bbl_thick > 0) then ; do i=is,ie - dd%bbl_thick(i,j) = h_bot(i,jj) - enddo ; endif - if (CS%id_Kd_leak > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_leak(i,j,K) = Kd_leak_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_quad > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_quad(i,j,K) = Kd_quad_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_itidal > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_itidal(i,j,K) = Kd_itidal_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_Froude > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_Froude(i,j,K) = Kd_Froude_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_slope > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_slope(i,j,K) = Kd_slope_2d(i,K) - enddo ; enddo ; endif - if (associated (VBF%Kd_leak)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_leak(i,j,K) = min(Kd_leak_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_quad)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_quad(i,j,K) = min(Kd_quad_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_itidal)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_itidal(i,j,K) = min(Kd_itidal_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_Froude)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_Froude(i,j,K) = min(Kd_Froude_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_slope)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_slope(i,j,K) = min(Kd_slope_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - - if (CS%id_prof_leak > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_leak(i,j,k) = prof_leak_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_quad > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_quad(i,j,k) = prof_quad_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_itidal > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_itidal(i,j,k) = prof_itidal_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_Froude > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_Froude(i,j,k) = prof_Froude_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_slope > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_slope(i,j,k) = prof_slope_2d(i,k) - enddo ; enddo ; endif + if (CS%id_prof_leak > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_leak(i,j,k) = prof_leak_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_quad > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_quad(i,j,k) = prof_quad_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_itidal > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_itidal(i,j,k) = prof_itidal_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_Froude > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_Froude(i,j,k) = prof_Froude_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_slope > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_slope(i,j,k) = prof_slope_2d(i,k) + enddo ; enddo ; endif + enddo endif + !$OMP parallel do default(shared) private( & + !$OMP KT_extra,KS_extra,dissip,jj) & + !$OMP if(.not. CS%use_CVMix_ddiff) + do j=jstart,jend + jj = j - jstart + 1 + ! This adds the diffusion sustained by the energy extracted from the flow by the bottom drag. if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0)) then if (CS%use_LOTW_BBL_diffusivity) then From f93183b901c4c8285192e5b0a9f9b6f876346ca0 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sun, 31 May 2026 15:15:09 -0400 Subject: [PATCH 07/45] set_diffusivity: block j loops in add_drag_diffusivity Add nj j-block dimension to local arrays (htot, rho_htot, gh_sum_top, Rho_top, TKE, I2decay, do_i), promote all dummy array args to blocked, replace scalar j with jstart/jend/nj, and fission the outer j-loop into three separate ji/jki loop nests. Move the call before the OMP region and pass full blocked arrays. Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_set_diffusivity.F90 | 258 +++++++++--------- 1 file changed, 133 insertions(+), 125 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index bf7b7103cf..ab54ee2959 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -700,6 +700,11 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i enddo endif + ! This adds the diffusion sustained by drag TKE. + if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0) .and. .not. CS%use_LOTW_BBL_diffusivity) & + call add_drag_diffusivity(h, u, v, tv, fluxes, visc, jstart, jend, njblock, TKE_to_Kd, & + maxTKE, kb, rho_bot, G, GV, US, CS, Kd_lay_2d, Kd_int_2d, dd%Kd_BBL) + !$OMP parallel do default(shared) private( & !$OMP KT_extra,KS_extra,dissip,jj) & !$OMP if(.not. CS%use_CVMix_ddiff) @@ -711,9 +716,6 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i if (CS%use_LOTW_BBL_diffusivity) then call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), Kd_int_2d(:,:,jj), & G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj), dz(:,:,jj)) - else - call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd(:,:,jj), & - maxTKE(:,:,jj), kb(:,jj), rho_bot(:,jj), G, GV, US, CS, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), dd%Kd_BBL) endif if (associated(VBF%Kd_BBL)) then ; do K=1,nz+1 ; do i=is,ie VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) @@ -1451,7 +1453,7 @@ subroutine double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, Kd_T_dd, Kd_S_dd) end subroutine double_diffusion !> This routine adds diffusion sustained by flow energy extracted by bottom drag. -subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, & +subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, jstart, jend, nj, TKE_to_Kd, maxTKE, & kb, rho_bot, G, GV, US, CS, Kd_lay, Kd_int, Kd_BBL) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure @@ -1467,22 +1469,24 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, type(forcing), intent(in) :: fluxes !< A structure of thermodynamic surface fluxes type(vertvisc_type), intent(in) :: visc !< Structure containing vertical viscosities, bottom !! boundary layer properties and related fields - integer, intent(in) :: j !< j-index of row to work on - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: TKE_to_Kd !< The conversion rate between the TKE + integer, intent(in) :: jstart !< j-index of first row to work on + integer, intent(in) :: jend !< j-index of last row to work on + integer, intent(in) :: nj !< Number of j-rows in a j-block + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: TKE_to_Kd !< The conversion rate between the TKE !! TKE dissipated within a layer and the !! diapycnal diffusivity within that layer, !! usually (~Rho_0 / (G_Earth * dRho_lay)) !! [T2 Z-1 ~> s2 m-1] - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: maxTKE !< The energy required to for a layer to entrain to its + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: maxTKE !< The energy required to for a layer to entrain to its !! maximum-realizable thickness [H Z2 T-3 ~> m3 s-3 or W m-2] - integer, dimension(SZI_(G)), intent(in) :: kb !< Index of lightest layer denser than the buffer + integer, dimension(SZI_(G),nj), intent(in) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer - real, dimension(SZI_(G)), intent(in) :: rho_bot !< In situ density averaged over a near-bottom + real, dimension(SZI_(G),nj), intent(in) :: rho_bot !< In situ density averaged over a near-bottom !! region [R ~> kg m-3] type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)), intent(inout) :: Kd_lay !< The diapycnal diffusivity in layers, + real, dimension(SZI_(G),SZK_(GV),nj), intent(inout) :: Kd_lay !< The diapycnal diffusivity in layers, !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(inout) :: Kd_int !< The diapycnal diffusivity at interfaces, + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(inout) :: Kd_int !< The diapycnal diffusivity at interfaces, !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real, dimension(:,:,:), pointer :: Kd_BBL !< Interface BBL diffusivity !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] @@ -1491,7 +1495,7 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, real, dimension(SZK_(GV)+1) :: & Rint ! coordinate density of an interface [R ~> kg m-3] - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & htot, & ! total thickness above or below a layer, or the ! integrated thickness in the BBL [H ~> m or kg m-2]. rho_htot, & ! running integral with depth of density [R H ~> kg m-2 or kg2 m-5] @@ -1516,10 +1520,10 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, ! defined in visc, on the assumption that this ! extracted energy also drives diapycnal mixing. - logical :: domore, do_i(SZI_(G)) + logical :: domore, do_i(SZI_(G),nj) logical :: do_diag_Kd_BBL - integer :: i, k, is, ie, nz, i_rem, kb_min + integer :: i, j, k, is, ie, nz, i_rem, kb_min, jj is = G%isc ; ie = G%iec ; nz = GV%ke do_diag_Kd_BBL = associated(Kd_BBL) @@ -1539,11 +1543,11 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, ! The turbulence decay scale is 0.5*ustar/f from K&E & MOM_vertvisc.F90 ! Any turbulence that makes it into the mixed layers is assumed ! to be relatively small and is discarded. - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie ustar_h = visc%ustar_BBL(i,j) if (associated(fluxes%ustar_tidal)) then if (allocated(tv%SpV_avg)) then - ustar_h = ustar_h + GV%RZ_to_H*rho_bot(i) * fluxes%ustar_tidal(i,j) + ustar_h = ustar_h + GV%RZ_to_H*rho_bot(i,jj) * fluxes%ustar_tidal(i,j) else ustar_h = ustar_h + GV%Z_to_H * fluxes%ustar_tidal(i,j) endif @@ -1551,142 +1555,146 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, absf = 0.25 * ((abs(G%CoriolisBu(I-1,J-1)) + abs(G%CoriolisBu(I,J))) + & (abs(G%CoriolisBu(I-1,J)) + abs(G%CoriolisBu(I,J-1)))) if ((ustar_h > 0.0) .and. (absf > 0.5*CS%IMax_decay*ustar_h)) then - I2decay(i) = absf / ustar_h + I2decay(i,jj) = absf / ustar_h else ! The maximum decay scale should be something of order 200 m. ! If ustar_h = 0, this is land so this value doesn't matter. - I2decay(i) = 0.5*CS%IMax_decay + I2decay(i,jj) = 0.5*CS%IMax_decay endif if (CS%drag_diff_answer_date <= 20250301) then - TKE(i) = ((CS%BBL_effic * cdrag_sqrt) * exp(-I2decay(i)*h(i,j,nz)) ) * visc%BBL_meanKE_loss_sqrtCd(i,j) + TKE(i,jj) = ((CS%BBL_effic * cdrag_sqrt) * exp(-I2decay(i,jj)*h(i,j,nz)) ) * visc%BBL_meanKE_loss_sqrtCd(i,j) else - TKE(i) = (CS%BBL_effic * exp(-I2decay(i)*h(i,j,nz)) ) * visc%BBL_meanKE_loss(i,j) + TKE(i,jj) = (CS%BBL_effic * exp(-I2decay(i,jj)*h(i,j,nz)) ) * visc%BBL_meanKE_loss(i,j) endif if (associated(fluxes%BBL_tidal_dis)) & - TKE(i) = TKE(i) + fluxes%BBL_tidal_dis(i,j) * GV%RZ_to_H * & - (CS%BBL_effic * exp(-I2decay(i)*h(i,j,nz))) + TKE(i,jj) = TKE(i,jj) + fluxes%BBL_tidal_dis(i,j) * GV%RZ_to_H * & + (CS%BBL_effic * exp(-I2decay(i,jj)*h(i,j,nz))) ! Distribute the work over a BBL of depth 20^2 ustar^2 / g' following ! Killworth & Edwards (1999) and Zilitikevich & Mironov (1996). ! Rho_top is determined by finding the density where ! integral(bottom, Z) (rho(z') - rho(Z)) dz' = rho_0 400 ustar^2 / g - gh_sum_top(i) = R0_g * 400.0 * ustar_h**2 + gh_sum_top(i,jj) = R0_g * 400.0 * ustar_h**2 - do_i(i) = (G%mask2dT(i,j) > 0.0) - htot(i) = h(i,j,nz) - rho_htot(i) = GV%Rlay(nz)*(h(i,j,nz)) - Rho_top(i) = GV%Rlay(1) - if (CS%bulkmixedlayer .and. do_i(i)) Rho_top(i) = GV%Rlay(kb(i)-1) - enddo - - do k=nz-1,2,-1 ; domore = .false. - do i=is,ie ; if (do_i(i)) then - htot(i) = htot(i) + h(i,j,k) - rho_htot(i) = rho_htot(i) + GV%Rlay(k)*(h(i,j,k)) - if (htot(i)*GV%Rlay(k-1) <= (rho_htot(i) - gh_sum_top(i))) then - ! The top of the mixing is in the interface atop the current layer. - Rho_top(i) = (rho_htot(i) - gh_sum_top(i)) / htot(i) - do_i(i) = .false. - elseif (k <= kb(i)) then ; do_i(i) = .false. - else ; domore = .true. ; endif - endif ; enddo - if (.not.domore) exit - enddo ! k-loop - - do i=is,ie ; do_i(i) = (G%mask2dT(i,j) > 0.0) ; enddo - do k=nz-1,kb_min,-1 - i_rem = 0 - do i=is,ie ; if (do_i(i)) then - if (k 0.0) then - if (Rint(K) <= Rho_top(i)) then - TKE_to_layer = TKE(i) - else - dRl = Rint(K+1) - Rint(K) ; dRbot = Rint(K+1) - Rho_top(i) - TKE_to_layer = TKE(i) * dRl * & - (3.0*dRbot*(Rint(K) - Rho_top(i)) + dRl**2) / (dRbot**3) - endif - else ; TKE_to_layer = 0.0 ; endif + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + htot(i,jj) = h(i,j,nz) + rho_htot(i,jj) = GV%Rlay(nz)*(h(i,j,nz)) + Rho_top(i,jj) = GV%Rlay(1) + if (CS%bulkmixedlayer .and. do_i(i,jj)) Rho_top(i,jj) = GV%Rlay(kb(i,jj)-1) + enddo ; enddo - ! TKE_Ray has been initialized to 0 above. - if (Rayleigh_drag) TKE_Ray = 0.5*CS%BBL_effic * G%IareaT(i,j) * & - (((G%areaCu(I-1,j) * visc%Ray_u(I-1,j,k) * u(I-1,j,k)**2) + & - (G%areaCu(I,j) * visc%Ray_u(I,j,k) * u(I,j,k)**2)) + & - ((G%areaCv(i,J-1) * visc%Ray_v(i,J-1,k) * v(i,J-1,k)**2) + & - (G%areaCv(i,J) * visc%Ray_v(i,J,k) * v(i,J,k)**2))) + do j=jstart,jend ; jj = j - jstart + 1 + do k=nz-1,2,-1 ; domore = .false. + do i=is,ie ; if (do_i(i,jj)) then + htot(i,jj) = htot(i,jj) + h(i,j,k) + rho_htot(i,jj) = rho_htot(i,jj) + GV%Rlay(k)*(h(i,j,k)) + if (htot(i,jj)*GV%Rlay(k-1) <= (rho_htot(i,jj) - gh_sum_top(i,jj))) then + ! The top of the mixing is in the interface atop the current layer. + Rho_top(i,jj) = (rho_htot(i,jj) - gh_sum_top(i,jj)) / htot(i,jj) + do_i(i,jj) = .false. + elseif (k <= kb(i,jj)) then ; do_i(i,jj) = .false. + else ; domore = .true. ; endif + endif ; enddo + if (.not.domore) exit + enddo ! k-loop + enddo - if (TKE_to_layer + TKE_Ray > 0.0) then - if (CS%BBL_mixing_as_max) then - if (TKE_to_layer + TKE_Ray > maxTKE(i,k)) & - TKE_to_layer = maxTKE(i,k) - TKE_Ray + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; do_i(i,jj) = (G%mask2dT(i,j) > 0.0) ; enddo + do k=nz-1,kb_min,-1 + i_rem = 0 + do i=is,ie ; if (do_i(i,jj)) then + if (k= 0.0) .and. (delta_Kd > CS%Kd_max)) then - delta_Kd = CS%Kd_max - Kd_lay(i,k) = Kd_lay(i,k) + delta_Kd - else - Kd_lay(i,k) = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k) - endif - Kd_int(i,K) = Kd_int(i,K) + 0.5 * delta_Kd - Kd_int(i,K+1) = Kd_int(i,K+1) + 0.5 * delta_Kd - if (do_diag_Kd_BBL) then - Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd - Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd - endif - endif - else - if (Kd_lay(i,k) >= maxTKE(i,k) * TKE_to_Kd(i,k)) then - TKE_here = 0.0 - TKE(i) = TKE(i) + TKE_Ray - elseif (Kd_lay(i,k) + (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k) > & - maxTKE(i,k) * TKE_to_Kd(i,k)) then - TKE_here = ((TKE_to_layer + TKE_Ray) + Kd_lay(i,k) / TKE_to_Kd(i,k)) - maxTKE(i,k) - TKE(i) = (TKE(i) - TKE_here) + TKE_Ray + ! This is an analytic integral where diffusivity is a quadratic function of + ! rho that goes asymptotically to 0 at Rho_top (vaguely following KPP?). + if (TKE(i,jj) > 0.0) then + if (Rint(K) <= Rho_top(i,jj)) then + TKE_to_layer = TKE(i,jj) else - TKE_here = TKE_to_layer + TKE_Ray - TKE(i) = TKE(i) - TKE_to_layer + dRl = Rint(K+1) - Rint(K) ; dRbot = Rint(K+1) - Rho_top(i,jj) + TKE_to_layer = TKE(i,jj) * dRl * & + (3.0*dRbot*(Rint(K) - Rho_top(i,jj)) + dRl**2) / (dRbot**3) endif - if (TKE(i) < 0.0) TKE(i) = 0.0 ! This should be unnecessary? - - if (TKE_here > 0.0) then - delta_Kd = TKE_here * TKE_to_Kd(i,k) - if (CS%Kd_max >= 0.0) delta_Kd = min(delta_Kd, CS%Kd_max) - Kd_lay(i,k) = Kd_lay(i,k) + delta_Kd - Kd_int(i,K) = Kd_int(i,K) + 0.5 * delta_Kd - Kd_int(i,K+1) = Kd_int(i,K+1) + 0.5 * delta_Kd - if (do_diag_Kd_BBL) then - Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd - Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + else ; TKE_to_layer = 0.0 ; endif + + ! TKE_Ray has been initialized to 0 above. + if (Rayleigh_drag) TKE_Ray = 0.5*CS%BBL_effic * G%IareaT(i,j) * & + (((G%areaCu(I-1,j) * visc%Ray_u(I-1,j,k) * u(I-1,j,k)**2) + & + (G%areaCu(I,j) * visc%Ray_u(I,j,k) * u(I,j,k)**2)) + & + ((G%areaCv(i,J-1) * visc%Ray_v(i,J-1,k) * v(i,J-1,k)**2) + & + (G%areaCv(i,J) * visc%Ray_v(i,J,k) * v(i,J,k)**2))) + + if (TKE_to_layer + TKE_Ray > 0.0) then + if (CS%BBL_mixing_as_max) then + if (TKE_to_layer + TKE_Ray > maxTKE(i,k,jj)) & + TKE_to_layer = maxTKE(i,k,jj) - TKE_Ray + + TKE(i,jj) = TKE(i,jj) - TKE_to_layer + + if (Kd_lay(i,k,jj) < (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj)) then + delta_Kd = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) - Kd_lay(i,k,jj) + if ((CS%Kd_max >= 0.0) .and. (delta_Kd > CS%Kd_max)) then + delta_Kd = CS%Kd_max + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + delta_Kd + else + Kd_lay(i,k,jj) = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) + endif + Kd_int(i,K,jj) = Kd_int(i,K,jj) + 0.5 * delta_Kd + Kd_int(i,K+1,jj) = Kd_int(i,K+1,jj) + 0.5 * delta_Kd + if (do_diag_Kd_BBL) then + Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd + Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + endif + endif + else + if (Kd_lay(i,k,jj) >= maxTKE(i,k,jj) * TKE_to_Kd(i,k,jj)) then + TKE_here = 0.0 + TKE(i,jj) = TKE(i,jj) + TKE_Ray + elseif (Kd_lay(i,k,jj) + (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) > & + maxTKE(i,k,jj) * TKE_to_Kd(i,k,jj)) then + TKE_here = ((TKE_to_layer + TKE_Ray) + Kd_lay(i,k,jj) / TKE_to_Kd(i,k,jj)) - maxTKE(i,k,jj) + TKE(i,jj) = (TKE(i,jj) - TKE_here) + TKE_Ray + else + TKE_here = TKE_to_layer + TKE_Ray + TKE(i,jj) = TKE(i,jj) - TKE_to_layer + endif + if (TKE(i,jj) < 0.0) TKE(i,jj) = 0.0 ! This should be unnecessary? + + if (TKE_here > 0.0) then + delta_Kd = TKE_here * TKE_to_Kd(i,k,jj) + if (CS%Kd_max >= 0.0) delta_Kd = min(delta_Kd, CS%Kd_max) + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + delta_Kd + Kd_int(i,K,jj) = Kd_int(i,K,jj) + 0.5 * delta_Kd + Kd_int(i,K+1,jj) = Kd_int(i,K+1,jj) + 0.5 * delta_Kd + if (do_diag_Kd_BBL) then + Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd + Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + endif endif endif endif - endif - ! This may be risky - in the case that there are exactly zero - ! velocities at 4 neighboring points, but nonzero velocities - ! above the iterations would stop too soon. I don't see how this - ! could happen in practice. RWH - if ((TKE(i)<= 0.0) .and. (TKE_Ray == 0.0)) then - do_i(i) = .false. ; i_rem = i_rem - 1 - endif + ! This may be risky - in the case that there are exactly zero + ! velocities at 4 neighboring points, but nonzero velocities + ! above the iterations would stop too soon. I don't see how this + ! could happen in practice. RWH + if ((TKE(i,jj)<= 0.0) .and. (TKE_Ray == 0.0)) then + do_i(i,jj) = .false. ; i_rem = i_rem - 1 + endif - endif ; enddo - if (i_rem == 0) exit - enddo ! k-loop + endif ; enddo + if (i_rem == 0) exit + enddo ! k-loop + enddo end subroutine add_drag_diffusivity From 9486a85132335ddba9a2cb3d922397dbdb2012bc Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Sun, 31 May 2026 15:37:10 -0400 Subject: [PATCH 08/45] set_diffusivity: move remaining OMP j-loop blocks out into own j-loops Moves add_LOTW_BBL_diffusivity, VBF%Kd_BBL copy, limit_dissipation, Kd_add, Kd_int/Kd_lay copy, Kd_Work, and Kd_Work_added blocks out of the OMP parallel do into individual do j=jstart,jend loops. Removes the now-empty OMP parallel do directive and its enclosing j-loop. Co-Authored-By: Claude Sonnet 4.6 --- .../vertical/MOM_set_diffusivity.F90 | 106 ++++++++++-------- 1 file changed, 61 insertions(+), 45 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index ab54ee2959..3079377a69 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -705,21 +705,20 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i call add_drag_diffusivity(h, u, v, tv, fluxes, visc, jstart, jend, njblock, TKE_to_Kd, & maxTKE, kb, rho_bot, G, GV, US, CS, Kd_lay_2d, Kd_int_2d, dd%Kd_BBL) - !$OMP parallel do default(shared) private( & - !$OMP KT_extra,KS_extra,dissip,jj) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=jstart,jend - jj = j - jstart + 1 - - ! This adds the diffusion sustained by the energy extracted from the flow by the bottom drag. if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0)) then if (CS%use_LOTW_BBL_diffusivity) then - call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), Kd_int_2d(:,:,jj), & - G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj), dz(:,:,jj)) + do j=jstart,jend ; jj = j - jstart + 1 + call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), & + Kd_int_2d(:,:,jj), G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj), dz(:,:,jj)) + enddo + endif + if (associated(VBF%Kd_BBL)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) + enddo ; enddo + enddo endif - if (associated(VBF%Kd_BBL)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) - enddo ; enddo ; endif endif if (CS%limit_dissipation) then @@ -728,27 +727,33 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 1) a global constant, ! 2) a dissipation proportional to N (aka Gargett) and ! 3) dissipation corresponding to a (nearly) constant diffusivity. - do K=2,nz ; do i=is,ie - dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K,jj)), & ! Floor aka Gargett - CS%dissip_N2 * N2_int(i,K,jj)) ! Floor of Kd_min*rho0/F_Ri - Kd_int_2d(i,K,jj) = max(Kd_int_2d(i,K,jj) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K,jj) + Omega2)))) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=2,nz ; do i=is,ie + dissip = max( CS%dissip_min, & ! Const. floor on dissip. + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_int(i,K,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_int_2d(i,K,jj) = max(Kd_int_2d(i,K,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K,jj) + Omega2)))) + enddo ; enddo + enddo endif ! Optionally add a uniform diffusivity at the interfaces. if (CS%Kd_add > 0.0) then - do K=1,nz+1 ; do i=is,ie - Kd_int_2d(i,K,jj) = Kd_int_2d(i,K,jj) + CS%Kd_add - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + Kd_int_2d(i,K,jj) = Kd_int_2d(i,K,jj) + CS%Kd_add + enddo ; enddo + enddo VBF%Kd_add = CS%Kd_add endif ! Copy the 2-d slices into the 3-d array that is exported. - do K=1,nz+1 ; do i=is,ie - Kd_int(i,j,K) = Kd_int_2d(i,K,jj) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + Kd_int(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo if (CS%limit_dissipation) then ! This calculates the layer dissipation ONLY from Kd calculated in this routine @@ -756,39 +761,50 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 1) a global constant, ! 2) a dissipation proportional to N (aka Gargett) and ! 3) dissipation corresponding to a (nearly) constant diffusivity. - do k=2,nz-1 ; do i=is,ie - dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k,jj)), & ! Floor aka Gargett - CS%dissip_N2 * N2_lay(i,k,jj)) ! Floor of Kd_min*rho0/F_Ri - Kd_lay_2d(i,k,jj) = max(Kd_lay_2d(i,k,jj) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k,jj) + Omega2)))) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=2,nz-1 ; do i=is,ie + dissip = max( CS%dissip_min, & ! Const. floor on dissip. + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_lay(i,k,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_lay_2d(i,k,jj) = max(Kd_lay_2d(i,k,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k,jj) + Omega2)))) + enddo ; enddo + enddo endif if (associated(dd%Kd_Work)) then - do k=1,nz ; do i=is,ie - dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 + enddo ; enddo + enddo endif ! Optionally add a uniform diffusivity to the layers. if ((CS%Kd_add > 0.0) .and. (present(Kd_lay))) then - do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + CS%Kd_add - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + CS%Kd_add + enddo ; enddo + enddo endif if (associated(dd%Kd_Work_added)) then - do k=1,nz ; do i=is,ie - dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 + enddo ; enddo + enddo endif ! Copy the 2-d slices into the 3-d array that is exported; this was done above for Kd_int. - if (present(Kd_lay)) then ; do k=1,nz ; do i=is,ie - Kd_lay(i,j,k) = Kd_lay_2d(i,k,jj) - enddo ; enddo ; endif - enddo ! j-loop + if (present(Kd_lay)) then + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + Kd_lay(i,j,k) = Kd_lay_2d(i,k,jj) + enddo ; enddo + enddo + endif enddo ! jstart-loop if (CS%user_change_diff) then From a2438fc092f6bca143d888d70e328452567cd34b Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 1 Jun 2026 07:43:59 -0400 Subject: [PATCH 09/45] ijk arrays in find_n2 --- .../vertical/MOM_set_diffusivity.F90 | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 3079377a69..9fc92e69f5 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -1195,7 +1195,7 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real, dimension(SZI_(G),SZK_(GV)+1,nj) :: & + real, dimension(SZI_(G),nj,SZK_(GV)+1) :: & pres, & ! pressure at each interface [R L2 T-2 ~> Pa] dRho_int_unfilt, & ! unfiltered density differences across interfaces [R ~> kg m-3] dRho_dT, & ! partial derivative of density wrt temp [R C-1 ~> kg m-3 degC-1] @@ -1226,34 +1226,34 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh do j=jstart,jend ; jj = j - jstart + 1 do i=is,ie dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 - dRho_int_unfilt(i,1,jj) = 0.0 ; dRho_int_unfilt(i,nz+1,jj) = 0.0 + dRho_int_unfilt(i,jj,1) = 0.0 ; dRho_int_unfilt(i,jj,nz+1) = 0.0 enddo enddo if (associated(tv%eqn_of_state)) then if (associated(fluxes%p_surf)) then do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; pres(i,1,jj) = fluxes%p_surf(i,j) ; enddo + do i=is,ie ; pres(i,jj,1) = fluxes%p_surf(i,j) ; enddo enddo else do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; pres(i,1,jj) = 0.0 ; enddo + do i=is,ie ; pres(i,jj,1) = 0.0 ; enddo enddo endif EOSdom(:) = EOS_domain(G%HI) do K=2,nz do j=jstart,jend ; jj = j - jstart + 1 do i=is,ie - pres(i,K,jj) = pres(i,K-1,jj) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) + pres(i,jj,K) = pres(i,jj,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) Temp_Int(i,jj) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) Salin_Int(i,jj) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) enddo - call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,K,jj), dRho_dT(:,K,jj), dRho_dS(:,K,jj), & + call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,jj,K), dRho_dT(:,jj,K), dRho_dS(:,jj,K), & tv%eqn_of_state, EOSdom) do i=is,ie - dRho_int(i,K,jj) = max(dRho_dT(i,K,jj)*(T_f(i,j,k) - T_f(i,j,k-1)) + & - dRho_dS(i,K,jj)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) - dRho_int_unfilt(i,K,jj) = max(dRho_dT(i,K,jj)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & - dRho_dS(i,K,jj)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + dRho_int(i,K,jj) = max(dRho_dT(i,jj,K)*(T_f(i,j,k) - T_f(i,j,k-1)) + & + dRho_dS(i,jj,K)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) + dRho_int_unfilt(i,jj,K) = max(dRho_dT(i,jj,K)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & + dRho_dS(i,jj,K)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) enddo enddo enddo @@ -1362,7 +1362,7 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh do K=1,nz+1 do j=jstart,jend ; jj = j - jstart + 1 do i=is,ie - dRho_int(i,K,jj) = dRho_int_unfilt(i,K,jj) + dRho_int(i,K,jj) = dRho_int_unfilt(i,jj,K) enddo enddo enddo From 72c9ed1c856300b08508fa28b7359748d55e7b54 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 1 Jun 2026 07:59:36 -0400 Subject: [PATCH 10/45] kji arrays set_density_ratios --- .../vertical/MOM_set_diffusivity.F90 | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 9fc92e69f5..11c59fa7a3 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -2289,9 +2289,9 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp ! Local variables real :: g_R0 ! g_R0 is a rescaled version of g/Rho [L2 Z-1 R-1 T-2 ~> m4 kg-1 s-2] real :: eps, tmp ! nondimensional temporary variables [nondim] - real :: a(SZK_(GV),nj), a_0(SZK_(GV),nj) ! nondimensional temporary variables [nondim] + real :: a(SZK_(GV)), a_0(SZK_(GV)) ! nondimensional temporary variables [nondim] real :: p_ref(SZI_(G),nj) ! an array of tv%P_Ref pressures [R L2 T-2 ~> Pa] - real :: Rcv(SZI_(G),SZK_(GV),nj) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] + real :: Rcv(SZI_(G),nj,SZK_(GV)) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] real :: I_Drho ! The inverse of the coordinate density difference between ! layers [R-1 ~> m3 kg-1] @@ -2319,8 +2319,8 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie p_ref(i,jj) = tv%P_Ref enddo ; enddo - do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,kmb - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref(:,jj), Rcv(:,k,jj), tv%eqn_of_state, EOSdom) + do k=1,kmb ; do j=jstart,jend ; jj = j - jstart + 1 + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref(:,jj), Rcv(:,jj,k), tv%eqn_of_state, EOSdom) enddo ; enddo do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie if (kb(i,jj) <= nz-1) then @@ -2335,41 +2335,41 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp endif ! The indexing convention for a is appropriate for the interfaces. do k3=1,kmb - a(k3+1,jj) = (GV%Rlay(k) - Rcv(i,k3,jj)) * I_Drho + a(k3+1) = (GV%Rlay(k) - Rcv(i,k3,jj)) * I_Drho enddo - if ((present(rho_0)) .and. (a(kmb+1,jj) < 2.0*eps*ds_dsp1(i,k,jj))) then + if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,k,jj))) then ! If the buffer layer nearly matches the density of the layer below in the ! coordinate variable (sigma-2), use the sigma-0-based density ratio if it is ! greater (and stable). if ((rho_0(i,k,jj) > rho_0(i,kmb,jj)) .and. & (rho_0(i,k+1,jj) > rho_0(i,k,jj))) then I_Drho = 1.0 / (rho_0(i,k+1,jj)-rho_0(i,k,jj)) - a_0(kmb+1,jj) = min((rho_0(i,k,jj)-rho_0(i,kmb,jj)) * I_Drho, ds_dsp1(i,k,jj)) - if (a_0(kmb+1,jj) > a(kmb+1,jj)) then + a_0(kmb+1) = min((rho_0(i,k,jj)-rho_0(i,kmb,jj)) * I_Drho, ds_dsp1(i,k,jj)) + if (a_0(kmb+1) > a(kmb+1)) then do k3=2,kmb - a_0(k3,jj) = a_0(kmb+1,jj) + (rho_0(i,kmb,jj)-rho_0(i,k3-1,jj)) * I_Drho + a_0(k3) = a_0(kmb+1) + (rho_0(i,kmb,jj)-rho_0(i,k3-1,jj)) * I_Drho enddo - if (a(kmb+1,jj) <= eps*ds_dsp1(i,k,jj)) then - do k3=2,kmb+1 ; a(k3,jj) = a_0(k3,jj) ; enddo + if (a(kmb+1) <= eps*ds_dsp1(i,k,jj)) then + do k3=2,kmb+1 ; a(k3) = a_0(k3) ; enddo else ! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,k,jj)) - 1.0)) ) - tmp = a(kmb+1,jj)/(eps*ds_dsp1(i,k,jj)) - 1.0 - do k3=2,kmb+1 ; a(k3,jj) = tmp*a(k3,jj) + (1.0-tmp)*a_0(k3,jj) ; enddo + tmp = a(kmb+1)/(eps*ds_dsp1(i,k,jj)) - 1.0 + do k3=2,kmb+1 ; a(k3) = tmp*a(k3) + (1.0-tmp)*a_0(k3) ; enddo endif endif endif endif - ds_dsp1(i,k,jj) = MAX(a(kmb+1,jj),1e-5) + ds_dsp1(i,k,jj) = MAX(a(kmb+1),1e-5) do k3=2,kmb -! ds_dsp1(i,k3,jj) = MAX(a(k3,jj),1e-5) +! ds_dsp1(i,k3,jj) = MAX(a(k3),1e-5) ! Deliberately treat convective instabilities of the upper mixed ! and buffer layers with respect to the deepest buffer layer as ! though they don't exist. They will be eliminated by the upcoming ! call to the mixedlayer code anyway. ! The indexing convention is appropriate for the interfaces. - ds_dsp1(i,k3,jj) = MAX(a(k3,jj),ds_dsp1(i,k,jj)) + ds_dsp1(i,k3,jj) = MAX(a(k3),ds_dsp1(i,k,jj)) enddo endif ! (kb(i,jj) <= nz-1) enddo ; enddo ! i-loop / j-loop From f679d39916f05992c1fdb14b7ad11caa8b1ab2f8 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 1 Jun 2026 08:16:27 -0400 Subject: [PATCH 11/45] more kji in find_tke_to_kd --- .../vertical/MOM_set_diffusivity.F90 | 80 +++++++++---------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 11c59fa7a3..0f36fb4c4f 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -957,7 +957,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, !! layer, or -1 without a bulk mixed layer. real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real, dimension(SZI_(G),SZK_(GV),nj) :: & + real, dimension(SZI_(G),nj,SZK_(GV)) :: & ds_dsp1, & ! coordinate variable (sigma-2) difference across an ! interface divided by the difference across the interface ! below it [nondim] @@ -1032,7 +1032,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, p_0(i,jj) = 0.0 ; p_ref(i,jj) = tv%P_Ref enddo ; enddo do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0(:,jj), rho_0(:,k,jj), tv%eqn_of_state, EOSdom) + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0(:,jj), rho_0(:,jj,k), tv%eqn_of_state, EOSdom) enddo ; enddo do j=jstart,jend ; jj = j - jstart + 1 call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref(:,jj), Rcv_kmb(:,jj), tv%eqn_of_state, EOSdom) @@ -1046,7 +1046,7 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, ! Backtrack, in case there are massive layers above that are stable ! in sigma-0. do k=kb(i,jj)-1,kmb+1,-1 - if (rho_0(i,kmb,jj) > rho_0(i,k,jj)) exit + if (rho_0(i,jj,kmb) > rho_0(i,jj,k)) exit if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i,jj) = k enddo enddo ; enddo @@ -1061,11 +1061,11 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, ! Determine maxEnt - the maximum permitted entrainment from below by each ! interior layer. - do j=jstart,jend ; jj = j - jstart + 1 ; do k=2,nz-1 ; do i=is,ie - dsp1_ds(i,k,jj) = 1.0 / ds_dsp1(i,k,jj) + do k=2,nz-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + dsp1_ds(i,jj,k) = 1.0 / ds_dsp1(i,jj,k) enddo ; enddo ; enddo do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie - dsp1_ds(i,nz,jj) = 0.0 + dsp1_ds(i,jj,nz) = 0.0 enddo ; enddo if (CS%bulkmixedlayer) then @@ -1073,41 +1073,41 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie htot(i,jj) = h(i,j,kmb) mFkb(i,jj) = 0.0 - if (kb(i,jj) < nz) mFkb(i,jj) = ds_dsp1(i,kb(i,jj),jj) * (h(i,j,kmb) - GV%Angstrom_H) + if (kb(i,jj) < nz) mFkb(i,jj) = ds_dsp1(i,jj,kb(i,jj)) * (h(i,j,kmb) - GV%Angstrom_H) enddo ; enddo - do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,kmb-1 ; do i=is,ie + do k=1,kmb-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie htot(i,jj) = htot(i,jj) + h(i,j,k) - mFkb(i,jj) = mFkb(i,jj) + ds_dsp1(i,k+1,jj)*(h(i,j,k) - GV%Angstrom_H) + mFkb(i,jj) = mFkb(i,jj) + ds_dsp1(i,jj,k+1)*(h(i,j,k) - GV%Angstrom_H) enddo ; enddo ; enddo else do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie - maxEnt(i,1,jj) = 0.0 ; htot(i,jj) = h(i,j,1) - GV%Angstrom_H + maxEnt(i,jj,1) = 0.0 ; htot(i,jj) = h(i,j,1) - GV%Angstrom_H enddo ; enddo endif - do j=jstart,jend ; jj = j - jstart + 1 ; do k=kb_min,nz-1 ; do i=is,ie + do k=kb_min,nz-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie if (k == kb(i,jj)) then - maxEnt(i,kb(i,jj),jj) = mFkb(i,jj) + maxEnt(i,jj,kb(i,jj)) = mFkb(i,jj) elseif (k > kb(i,jj)) then if (CS%answer_date < 20190101) then - maxEnt(i,k,jj) = (1.0/dsp1_ds(i,k,jj))*(maxEnt(i,k-1,jj) + htot(i,jj)) + maxEnt(i,jj,k) = (1.0/dsp1_ds(i,jj,k))*(maxEnt(i,jj,k-1) + htot(i,jj)) else - maxEnt(i,k,jj) = ds_dsp1(i,k,jj)*(maxEnt(i,k-1,jj) + htot(i,jj)) + maxEnt(i,jj,k) = ds_dsp1(i,jj,k)*(maxEnt(i,jj,k-1) + htot(i,jj)) endif htot(i,jj) = htot(i,jj) + (h(i,j,k) - GV%Angstrom_H) endif enddo ; enddo ; enddo do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie - htot(i,jj) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,nz,jj) = 0.0 + htot(i,jj) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,jj,nz) = 0.0 do_i(i,jj) = (G%mask2dT(i,j) > 0.0) enddo ; enddo - do j=jstart,jend ; jj = j - jstart + 1 - do k=nz-1,kb_min,-1 + do k=nz-1,kb_min,-1 + do j=jstart,jend ; jj = j - jstart + 1 i_rem = 0 do i=is,ie ; if (do_i(i,jj)) then if (k kg m-3]. @@ -2302,12 +2302,12 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp do j=jstart,jend ; jj = j - jstart + 1 ; do k=2,nz-1 ; do i=is,ie if (GV%g_prime(k+1) /= 0.0) then if (GV%Boussinesq .or. GV%Semi_Boussinesq) then - ds_dsp1(i,k,jj) = GV%g_prime(k) / GV%g_prime(k+1) + ds_dsp1(i,jj,k) = GV%g_prime(k) / GV%g_prime(k+1) else ! Use a mathematically equivalent form that avoids any dependency on RHO_0. - ds_dsp1(i,k,jj) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) + ds_dsp1(i,jj,k) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) endif else - ds_dsp1(i,k,jj) = 1. + ds_dsp1(i,jj,k) = 1. endif enddo ; enddo ; enddo @@ -2337,39 +2337,39 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp do k3=1,kmb a(k3+1) = (GV%Rlay(k) - Rcv(i,k3,jj)) * I_Drho enddo - if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,k,jj))) then + if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,jj,k))) then ! If the buffer layer nearly matches the density of the layer below in the ! coordinate variable (sigma-2), use the sigma-0-based density ratio if it is ! greater (and stable). - if ((rho_0(i,k,jj) > rho_0(i,kmb,jj)) .and. & - (rho_0(i,k+1,jj) > rho_0(i,k,jj))) then - I_Drho = 1.0 / (rho_0(i,k+1,jj)-rho_0(i,k,jj)) - a_0(kmb+1) = min((rho_0(i,k,jj)-rho_0(i,kmb,jj)) * I_Drho, ds_dsp1(i,k,jj)) + if ((rho_0(i,jj,k) > rho_0(i,jj,kmb)) .and. & + (rho_0(i,jj,k+1) > rho_0(i,jj,k))) then + I_Drho = 1.0 / (rho_0(i,jj,k+1)-rho_0(i,jj,k)) + a_0(kmb+1) = min((rho_0(i,jj,k)-rho_0(i,jj,kmb)) * I_Drho, ds_dsp1(i,jj,k)) if (a_0(kmb+1) > a(kmb+1)) then do k3=2,kmb - a_0(k3) = a_0(kmb+1) + (rho_0(i,kmb,jj)-rho_0(i,k3-1,jj)) * I_Drho + a_0(k3) = a_0(kmb+1) + (rho_0(i,jj,kmb)-rho_0(i,jj,k3-1)) * I_Drho enddo - if (a(kmb+1) <= eps*ds_dsp1(i,k,jj)) then + if (a(kmb+1) <= eps*ds_dsp1(i,jj,k)) then do k3=2,kmb+1 ; a(k3) = a_0(k3) ; enddo else -! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,k,jj)) - 1.0)) ) - tmp = a(kmb+1)/(eps*ds_dsp1(i,k,jj)) - 1.0 +! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,jj,k)) - 1.0)) ) + tmp = a(kmb+1)/(eps*ds_dsp1(i,jj,k)) - 1.0 do k3=2,kmb+1 ; a(k3) = tmp*a(k3) + (1.0-tmp)*a_0(k3) ; enddo endif endif endif endif - ds_dsp1(i,k,jj) = MAX(a(kmb+1),1e-5) + ds_dsp1(i,jj,k) = MAX(a(kmb+1),1e-5) do k3=2,kmb -! ds_dsp1(i,k3,jj) = MAX(a(k3),1e-5) +! ds_dsp1(i,jj,k3) = MAX(a(k3),1e-5) ! Deliberately treat convective instabilities of the upper mixed ! and buffer layers with respect to the deepest buffer layer as ! though they don't exist. They will be eliminated by the upcoming ! call to the mixedlayer code anyway. ! The indexing convention is appropriate for the interfaces. - ds_dsp1(i,k3,jj) = MAX(a(k3),ds_dsp1(i,k,jj)) + ds_dsp1(i,jj,k3) = MAX(a(k3),ds_dsp1(i,jj,k)) enddo endif ! (kb(i,jj) <= nz-1) enddo ; enddo ! i-loop / j-loop From 7b155bc7864cb0bc71da1318d3e858eebc75e1a0 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Wed, 15 Jul 2026 09:51:29 +1000 Subject: [PATCH 12/45] 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 From 471e1ff200d49c954b8cec90729c4d91dcbd58e1 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 12:56:49 +1000 Subject: [PATCH 13/45] kappa_shear: offload h_at_u/h_at_v vertex interpolation to GPU First increment of the MOM_kappa_shear GPU port. Converts the two thickness-to-vertex interpolation loops in Calc_kappa_shear_vertex (both branches of vertex_shear_OBC_bug) from !$OMP parallel do to do concurrent, offloading them to the device. Residency: - h is host-authoritative in the (host-only) diabatic stack and is mutated on host during diabatic (applyBoundaryFluxesInOut), so the globally-resident device copy (CS%h) is refreshed with update to(h) before the kernels; a plain map(to:) does not refresh a present object. - G%mask2dCu/Cv/T are already device-resident (initialize_MOM). - h_at_u/h_at_v are device workspace (map alloc), copied back with update from for the still-host per-column solver below. enter/exit data are mirrored and refcount-neutral (release, not delete). The per-column solver (kappa_shear_column/find_kappa_tke) stays on host for now; later increments move it to the device, at which point the update from(h_at_u,h_at_v) round-trip is removed. Verification: - benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). - nsys --stats=true confirms both do concurrent loops run on the GPU (calc_kappa_shear_vertex_{550,553}_gpu, 12 launches each). - Reviewed; no code defect found. The >=2-rank bitwise run remains the mandatory pre-merge gate and is still pending (single-GPU node). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 1cff115fe5..3f4a301800 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -538,33 +538,37 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Convert layer thicknesses into geometric thickness in height units. call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=1) + ! --- GPU port increment 1: h_at_u/h_at_v interpolation offloaded to device. + ! h is host-authoritative in the (host-only) diabatic stack; refresh the device copy. + ! G%mask2dCu/Cv/T are already device-resident (mapped in initialize_MOM). h_at_u/h_at_v + ! are device workspace, copied back for the (still host) per-column loop below. + !$omp target enter data map(to: h) + !$omp target update to(h) + !$omp target enter data map(alloc: h_at_u, h_at_v) + if (CS%vertex_shear_OBC_bug) then - !$OMP parallel do default(shared) - do k=1,nz - do j=JsB,JeB+1 ; do I=IsB,IeB - h_at_u(I,j,k) = G%mask2dCu(I,j) * (h(i,j,k) + h(i+1,j,k)) * 0.5 - enddo ; enddo - do J=JsB,JeB ; do i=IsB,IeB+1 - h_at_v(i,J,k) = G%mask2dCv(i,J) * (h(i,j,k) + h(i,j+1,k)) * 0.5 - enddo ; enddo + do concurrent (k=1:nz, j=JsB:JeB+1, I=IsB:IeB) + h_at_u(I,j,k) = G%mask2dCu(I,j) * (h(i,j,k) + h(i+1,j,k)) * 0.5 + enddo + do concurrent (k=1:nz, J=JsB:JeB, i=IsB:IeB+1) + h_at_v(i,J,k) = G%mask2dCv(i,J) * (h(i,j,k) + h(i,j+1,k)) * 0.5 enddo else ! Because G%mask2dCu(I,j) is zero if either G%mask2dT(i,j) or G%mask2dT(i+1,j) except at OBC ! faces, the following form give equivalent answers to those above unless OBCs are in use, ! although the former is clearly less complicated and costly. - !$OMP parallel do default(shared) - do k=1,nz - do j=JsB,JeB+1 ; do I=IsB,IeB - h_at_u(I,j,k) = G%mask2dCu(I,j) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j) * h(i+1,j,k)) / & - (G%mask2dT(i,j) + G%mask2dT(i+1,j) + 1.0e-36) - enddo ; enddo - do J=JsB,JeB ; do i=IsB,IeB+1 - h_at_v(i,J,k) = G%mask2dCv(i,J) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) / & - (G%mask2dT(i,j) + G%mask2dT(i,j+1) + 1.0e-36) - enddo ; enddo + do concurrent (k=1:nz, j=JsB:JeB+1, I=IsB:IeB) + h_at_u(I,j,k) = G%mask2dCu(I,j) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j) * h(i+1,j,k)) / & + (G%mask2dT(i,j) + G%mask2dT(i+1,j) + 1.0e-36) + enddo + do concurrent (k=1:nz, J=JsB:JeB, i=IsB:IeB+1) + h_at_v(i,J,k) = G%mask2dCv(i,J) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) / & + (G%mask2dT(i,j) + G%mask2dT(i,j+1) + 1.0e-36) enddo endif + !$omp target update from(h_at_u, h_at_v) + !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & @@ -860,6 +864,10 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) + ! --- GPU port increment 1: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: h_at_u, h_at_v) + !$omp target exit data map(release: h) + end subroutine Calc_kappa_shear_vertex From 304e3a0a17b1a000ecce0eb21ebc7c14d83f929f Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 13:30:54 +1000 Subject: [PATCH 14/45] MOM_EOS: add device-callable density-derivs dispatcher (declare target) Groundwork for whole-column GPU kernels (kappa_shear next) that must compute density derivatives from INSIDE a device region, where the polymorphic EOS%type%... dispatch is not usable. - Expose the existing this-free _loc elementals calculate_density_derivs_elem_{Roquet_rho,buggy_Wright}_loc as public and !$omp declare target. - Add calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_dT, drho_dS): a declare-target subroutine that select-case's on the integer form and calls the matching _loc kernel, returning mks derivatives. Unit rescaling and any `scale` factor remain the caller's responsibility (as in calculate_density_derivs_1d). Forms without a device _loc kernel hit the default branch; a device-using module must FATAL at init on a GPU build before reaching it. This is the reusable, form-general enabler chosen over a Roquet-scoped inline; it is aligned with doc 06 6.4's select-case direction but adds only the device-callable point path (not the full polymorphism removal). No functional change: nothing calls the dispatcher yet. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). Dispatcher is exercised (and Fable-reviewed) in the follow-up that wires kappa_shear to it. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/equation_of_state/MOM_EOS.F90 | 30 ++++++++++++++++++++ src/equation_of_state/MOM_EOS_Roquet_rho.F90 | 4 +++ src/equation_of_state/MOM_EOS_Wright.F90 | 4 +++ 3 files changed, 38 insertions(+) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 754f293b69..fa821f276f 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -9,6 +9,7 @@ module MOM_EOS use MOM_EOS_linear, only : linear_EOS, avg_spec_vol_linear use MOM_EOS_linear, only : int_density_dz_linear, int_spec_vol_dp_linear use MOM_EOS_Wright, only : buggy_Wright_EOS, avg_spec_vol_buggy_Wright +use MOM_EOS_Wright, only : calculate_density_derivs_elem_buggy_Wright_loc use MOM_EOS_Wright, only : int_density_dz_wright, int_spec_vol_dp_wright use MOM_EOS_Wright_full, only : Wright_full_EOS, avg_spec_vol_Wright_full use MOM_EOS_Wright_full, only : int_density_dz_wright_full, int_spec_vol_dp_wright_full @@ -17,6 +18,7 @@ module MOM_EOS use MOM_EOS_Jackett06, only : Jackett06_EOS use MOM_EOS_UNESCO, only : UNESCO_EOS use MOM_EOS_Roquet_rho, only : Roquet_rho_EOS +use MOM_EOS_Roquet_rho, only : calculate_density_derivs_elem_Roquet_rho_loc use MOM_EOS_Roquet_SpV, only : Roquet_SpV_EOS use MOM_EOS_TEOS10, only : TEOS10_EOS use MOM_EOS_TEOS10, only : gsw_sp_from_sr, gsw_pt_from_ct, gsw_sr_from_sp, gsw_ct_from_pt @@ -46,6 +48,7 @@ module MOM_EOS public calculate_density_elem public calculate_density public calculate_density_derivs +public calculate_density_derivs_elem_loc public calculate_density_second_derivs public calculate_spec_vol public calculate_specific_vol_derivs @@ -893,6 +896,33 @@ subroutine calculate_density_derivs_1d(T, S, pressure, drho_dT, drho_dS, EOS, do end subroutine calculate_density_derivs_1d +!> Device-callable dispatcher for density derivatives at a single point, in mks units, +!! selecting the equation-of-state form at runtime by integer id (no polymorphic dispatch) +!! so it can be called from inside a do concurrent / target region by whole-column GPU +!! kernels. Unit rescaling (EOS%*_to_* factors) and any `scale` factor are the caller's +!! responsibility, exactly as in calculate_density_derivs_1d. Forms without a device-callable +!! _loc kernel are not handled here; a device-using module must FATAL at init on a GPU build +!! before reaching this with an unsupported form. +subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_dT, drho_dS) + integer, intent(in) :: form_of_EOS !< The equation of state form (EOS_ROQUET_RHO, EOS_WRIGHT, ...) + real, intent(in) :: T !< Temperature in the EOS kernel's mks units [degC] + real, intent(in) :: S !< Salinity in the EOS kernel's mks units [ppt or g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + real, intent(out) :: drho_dT !< Partial derivative of density wrt temperature [kg m-3 degC-1] + real, intent(out) :: drho_dS !< Partial derivative of density wrt salinity [kg m-3 ppt-1] + !$omp declare target + + select case (form_of_EOS) + case (EOS_ROQUET_RHO) + call calculate_density_derivs_elem_Roquet_rho_loc(T, S, pressure, drho_dT, drho_dS) + case (EOS_WRIGHT) + call calculate_density_derivs_elem_buggy_Wright_loc(T, S, pressure, drho_dT, drho_dS) + case default + drho_dT = 0.0 ; drho_dS = 0.0 + end select + +end subroutine calculate_density_derivs_elem_loc + !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs. subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom) diff --git a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 index 86110d6aa1..8222d71e7d 100644 --- a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 +++ b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 @@ -10,6 +10,10 @@ module MOM_EOS_Roquet_rho implicit none ; private public Roquet_rho_EOS +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels +! can compute density derivatives in-region without polymorphic dispatch. +public calculate_density_derivs_elem_Roquet_rho_loc +!$omp declare target(calculate_density_derivs_elem_Roquet_rho_loc) real, parameter :: Pa2kb = 1.e-8 !< Conversion factor between Pa and kbar [kbar Pa-1] !>@{ Parameters in the Roquet_rho (Roquet density) equation of state diff --git a/src/equation_of_state/MOM_EOS_Wright.F90 b/src/equation_of_state/MOM_EOS_Wright.F90 index c2861c451d..b690c6537b 100644 --- a/src/equation_of_state/MOM_EOS_Wright.F90 +++ b/src/equation_of_state/MOM_EOS_Wright.F90 @@ -12,6 +12,10 @@ module MOM_EOS_Wright implicit none ; private public buggy_Wright_EOS +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels +! can compute density derivatives in-region without polymorphic dispatch. +public calculate_density_derivs_elem_buggy_Wright_loc +!$omp declare target(calculate_density_derivs_elem_buggy_Wright_loc) public int_density_dz_wright, int_spec_vol_dp_wright public avg_spec_vol_buggy_Wright public set_params_buggy_Wright From b4a3110c5d10f04c54b5b2f98f7cd23e96f117d9 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:06:04 +1000 Subject: [PATCH 15/45] kappa_shear: compute density derivs via device-callable EOS dispatcher on GPU builds Wires kappa_shear_column's Boussinesq density-derivative computation to the device-callable path added in 47ac931cb, so the column solver can later run inside a target region (the polymorphic calculate_density_derivs interface is not device-callable). - MOM_EOS: add get_EOS_form_and_scaling(EOS, form, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) so callers can reproduce, host-side, the unit rescaling that calculate_density_derivs_1d applies (EOS_type components are private). - kappa_shear_column: under #ifdef __NVCOMPILER_OPENMP_GPU, replace the calculate_density_derivs array call with an inline K=2..nzc loop calling calculate_density_derivs_elem_loc + the exact scale math of calculate_density_derivs_1d(dom=(/2,nzc/), scale=-g_R0). Non-GPU builds keep the original polymorphic call verbatim for all EOS forms. On GPU builds only ROQUET_RHO and WRIGHT have device _loc kernels; others FATAL. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1) with the GPU path active on host. Reviewed; bitwise equivalence confirmed operation-by-operation, verdict ship. Follow-ups required by the offload step (increment 3), per review: hoist get_EOS_form_and_scaling + the form FATAL out of kappa_shear_column to the driver (both are host-only and cannot run in a device region), passing the form/scaling in; and FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). Untested-by-benchmark but code-identical: the non-unity unit-rescale branch and the WRIGHT dispatcher case. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/equation_of_state/MOM_EOS.F90 | 21 ++++++++++ .../vertical/MOM_kappa_shear.F90 | 40 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index fa821f276f..8b757fde76 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -49,6 +49,7 @@ module MOM_EOS public calculate_density public calculate_density_derivs public calculate_density_derivs_elem_loc +public get_EOS_form_and_scaling public calculate_density_second_derivs public calculate_spec_vol public calculate_specific_vol_derivs @@ -923,6 +924,26 @@ subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_d end subroutine calculate_density_derivs_elem_loc +!> Return the equation-of-state form id and the unit-rescaling factors held in an EOS_type. +!! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion +!! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without +!! needing access to the private components of EOS_type. +subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) + type(EOS_type), intent(in) :: EOS !< Equation of state structure + integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) + real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal density unit R [R m3 kg-1 ~> 1] + real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] + real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] + real, intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1] + + form_of_EOS = EOS%form_of_EOS + kg_m3_to_R = EOS%kg_m3_to_R + C_to_degC = EOS%C_to_degC + S_to_ppt = EOS%S_to_ppt + RL2_T2_to_Pa = EOS%RL2_T2_to_Pa + +end subroutine get_EOS_form_and_scaling + !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs. subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 3f4a301800..9c0e20fbda 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -19,6 +19,8 @@ module MOM_kappa_shear use MOM_verticalGrid, only : verticalGrid_type use MOM_EOS, only : calculate_density_derivs use MOM_EOS, only : calculate_density, calculate_specific_vol_derivs +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT implicit none ; private @@ -978,6 +980,12 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real :: gR0 ! A conversion factor from H to pressure, Rho_0 times g in Boussinesq ! mode, or just g when non-Boussinesq [R L2 T-2 H-1 ~> kg m-2 s-2 or m s-2]. real :: g_R0 ! g_R0 is a rescaled version of g/Rho [Z R-1 T-2 ~> m4 kg-1 s-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + ! Locals for the device-callable EOS derivs path (reproduces calculate_density_derivs_1d). + integer :: eos_form ! The equation-of-state form id. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. + real :: eos_rho_scale, dRdT_scale, dRdS_scale ! Output rescaling factors [various]. +#endif real :: Norm ! A factor that normalizes two weights to 1 [H-2 ~> m-2 or m4 kg-2]. real :: tol_dksrc ! Tolerance for the change in the kappa source within an iteration ! relative to the local source [nondim]. This must be greater than 1. @@ -1141,8 +1149,40 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la Sal_int(K) = 0.5*(Sal(k-1) + Sal(k)) enddo if (GV%Boussinesq .or. GV%semi_Boussinesq) then +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable EOS path: dispatch density derivatives by form id (host-resolved + ! scaling), reproducing calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) + ! bit-for-bit. Needed because the polymorphic calculate_density_derivs interface is + ! not callable from inside the (soon-to-be-offloaded) column device region. + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if ((eos_RL2_T2_to_Pa == 1.0) .and. (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0)) then + do K=2,nzc + call calculate_density_derivs_elem_loc(eos_form, T_int(K), Sal_int(K), pressure(K), & + dbuoy_dT(K), dbuoy_dS(K)) + enddo + else + do K=2,nzc + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_int(K), & + eos_S_to_ppt*Sal_int(K), eos_RL2_T2_to_Pa*pressure(K), dbuoy_dT(K), dbuoy_dS(K)) + enddo + endif + eos_rho_scale = eos_kg_m3_to_R * (-g_R0) + dRdT_scale = eos_rho_scale * eos_C_to_degC + dRdS_scale = eos_rho_scale * eos_S_to_ppt + if ((dRdT_scale /= 1.0) .or. (dRdS_scale /= 1.0)) then + do K=2,nzc + dbuoy_dT(K) = dRdT_scale * dbuoy_dT(K) + dbuoy_dS(K) = dRdS_scale * dbuoy_dS(K) + enddo + endif +#else call calculate_density_derivs(T_int, Sal_int, pressure, dbuoy_dT, dbuoy_dS, & tv%eqn_of_state, (/2,nzc/), scale=-g_R0 ) +#endif else ! These should perhaps be combined into a single call to calculate the thermal expansion ! and haline contraction coefficients? From 0593e85ed2e7f155874dbb75481baecacfa9a58f Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:34:12 +1000 Subject: [PATCH 16/45] kappa_shear: hoist EOS form/scaling resolution to the drivers (offload prep) Prerequisite for offloading the column solver (increment 3): get_EOS_form_and_scaling and MOM_error are host-only and cannot run inside a device region, so move them out of kappa_shear_column up into both drivers (Calculate_kappa_shear and Calc_kappa_shear_vertex). - Resolve eos_form + the four unit-scaling factors once per driver call (guarded by use_temperature), and thread them into kappa_shear_column as new intent(in) args. - Move the "unsupported EOS form on GPU" FATAL to the drivers, and add a FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). - Add the five hoisted scalars to the shared() clause of both drivers' !$OMP parallel do (they are loop-invariant and read-only in the region; default(private) would leave them uninitialized inside the loop -- caught by the bitwise check). No functional change: same values, computed once per call instead of once per column. Verification: benchmark_ALE (vertex driver, ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1). Reviewed; verdict ship, bitwise safety confirmed by inspection for both drivers. Not runtime-exercised by the benchmark (inspection-only): the h-point driver, the non-unity unit-scaling branch. Notes for the offload step: the driver FATAL guards execution not compilation, so declare target on kappa_shear_column will require #ifdef-excluding the non-Boussinesq branch; CS (pointer dummy) and tv (pointer-laden) must be demoted/pre-extracted before declare target. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 70 +++++++++++++++---- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 9c0e20fbda..914a21750e 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -214,11 +214,29 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for ! interpolating back to the original index space [nondim]. integer :: is, ie, js, je, i, j, k, nz, nzc + integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke use_temperature = associated(tv%T) + ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor + ! and MOM_error are not device-callable) so they can be passed into the column solver. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_temperature) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if (.not. (GV%Boussinesq .or. GV%semi_Boussinesq)) call MOM_error(FATAL, & + "kappa_shear GPU build: the non-Boussinesq density-derivs path is not device-callable.") +#endif + endif + k0dt = dt*CS%kappa_0 dz_massless = 0.1*sqrt((US%Z_to_m*GV%m_to_H)*k0dt) @@ -229,6 +247,7 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,tv,G,GV,US, & !$OMP CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do j=js,je @@ -329,7 +348,8 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, & - tv, CS, GV, US) + tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. @@ -520,6 +540,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ real :: H_tiny ! A sub-roundoff thickness to use in the denominator when calculating ! thickness-weighted averages [H ~> m or kg m-2] integer :: IsB, IeB, JsB, JeB, i, j, k, nz, nzc + integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. ! Diagnostics that should be deleted? isB = G%isc-1 ; ieB = G%iecB ; jsB = G%jsc-1 ; jeB = G%jecB ; nz = GV%ke @@ -532,6 +554,23 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ use_temperature = associated(tv%T) + ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor + ! and MOM_error are not device-callable) so they can be passed into the column solver, which + ! will run inside a device region. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_temperature) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if (.not. (GV%Boussinesq .or. GV%semi_Boussinesq)) call MOM_error(FATAL, & + "kappa_shear GPU build: the non-Boussinesq density-derivs path is not device-callable.") +#endif + endif + k0dt = dt*CS%kappa_0 dz_massless = 0.1*sqrt((US%Z_to_m*GV%m_to_H)*k0dt) I_Prandtl = 0.0 ; if (CS%Prandtl_turb > 0.0) I_Prandtl = 1.0 / CS%Prandtl_turb @@ -575,6 +614,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do J=JsB,JeB @@ -708,7 +748,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & - tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US) + tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) ! call cpu_clock_begin(Id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then @@ -876,7 +917,8 @@ end subroutine Calc_kappa_shear_vertex !> This subroutine calculates shear-driven diffusivity and TKE in a single column subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_lay, & u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, tke_avg, N2_init, S2_init, & - N2_mean, S2_mean, tv, CS, GV, US ) + N2_mean, S2_mean, tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ) type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. real, dimension(SZK_(GV)+1), & intent(inout) :: kappa !< The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -918,6 +960,12 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la type(Kappa_shear_CS), pointer :: CS !< The control structure returned by a previous !! call to kappa_shear_init. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + integer, intent(in) :: eos_form !< The equation-of-state form id, resolved + !! host-side for the device-callable EOS derivs path. + real, intent(in) :: eos_kg_m3_to_R !< EOS factor converting kg m-3 to R [R m3 kg-1 ~> 1] + real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] + real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] + real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure to Pa [Pa T2 R-1 L-2 ~> 1] ! Local variables real, dimension(nzc) :: & @@ -981,9 +1029,7 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la ! mode, or just g when non-Boussinesq [R L2 T-2 H-1 ~> kg m-2 s-2 or m s-2]. real :: g_R0 ! g_R0 is a rescaled version of g/Rho [Z R-1 T-2 ~> m4 kg-1 s-2]. #ifdef __NVCOMPILER_OPENMP_GPU - ! Locals for the device-callable EOS derivs path (reproduces calculate_density_derivs_1d). - integer :: eos_form ! The equation-of-state form id. - real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. + ! Locals for the device-callable EOS derivs path (form + scaling are passed in as arguments). real :: eos_rho_scale, dRdT_scale, dRdS_scale ! Output rescaling factors [various]. #endif real :: Norm ! A factor that normalizes two weights to 1 [H-2 ~> m-2 or m4 kg-2]. @@ -1150,15 +1196,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la enddo if (GV%Boussinesq .or. GV%semi_Boussinesq) then #ifdef __NVCOMPILER_OPENMP_GPU - ! Device-callable EOS path: dispatch density derivatives by form id (host-resolved - ! scaling), reproducing calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) - ! bit-for-bit. Needed because the polymorphic calculate_density_derivs interface is - ! not callable from inside the (soon-to-be-offloaded) column device region. - call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & - eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) - if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & - "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & - "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + ! Device-callable EOS path: dispatch density derivatives by form id (form + unit scaling + ! resolved host-side in the driver and passed in), reproducing + ! calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) bit-for-bit. if ((eos_RL2_T2_to_Pa == 1.0) .and. (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0)) then do K=2,nzc call calculate_density_derivs_elem_loc(eos_form, T_int(K), Sal_int(K), pressure(K), & From a3976c9d56c4c8040c2f880b8b7c867c1c02e303 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:55:47 +1000 Subject: [PATCH 17/45] kappa_shear: make the column solver device-callable (declare target) Prepares the per-column solver for offload: the three workers (kappa_shear_column, find_kappa_tke, calculate_projected_state) are now !$omp declare target, so the driver column loop can run inside a target region in the next step. Device code is generated for all three (confirmed by -Minfo). Execution is still on the host (the column loop is still !$OMP parallel do), so this is bitwise-neutral. To make them device-compilable: - Demote CS from `pointer` to `type(Kappa_shear_CS), intent(in)` in kappa_shear_column and find_kappa_tke (the type is all scalars bar a diag pointer that the column math never touches). - Hoist use_temperature to a kappa_shear_column argument (resolved in the drivers) so the device path no longer reads the tv%T pointer via associated(); tv is then untouched on GPU builds. - #ifndef __NVCOMPILER_OPENMP_GPU around the non-Boussinesq branch's polymorphic calculate_specific_vol_derivs/calculate_density calls so they are excluded from device compilation (unreachable on GPU anyway -- the driver FATALs on non-Boussinesq). Notably the 26 early exit/return sites in the two solvers compile for the device without error on NVHPC 26.5 (the 25.11 exit-under-parallel miscompile appears fixed); their runtime correctness on device is gated by the bitwise check in the offload step, not pre-emptive rewrites. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). Next step flips the driver !$OMP parallel do to a target teams loop and adds residency; the solver then runs on the GPU. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 914a21750e..a5f7d916c3 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -131,6 +131,10 @@ module MOM_kappa_shear ! integer :: id_clock_project, id_clock_KQ, id_clock_avg, id_clock_setup +! The per-column solver and its helpers are device-callable so the driver column loop +! can run inside a target region (GPU port increment 3). +!$omp declare target(kappa_shear_column, find_kappa_tke, calculate_projected_state) + contains !> Subroutine for calculating shear-driven diffusivity and TKE in tracer columns @@ -349,7 +353,8 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, & tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature) ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. @@ -749,7 +754,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature) ! call cpu_clock_begin(Id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then @@ -918,7 +924,8 @@ end subroutine Calc_kappa_shear_vertex subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_lay, & u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, tke_avg, N2_init, S2_init, & N2_mean, S2_mean, tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature ) type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. real, dimension(SZK_(GV)+1), & intent(inout) :: kappa !< The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -957,8 +964,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la type(thermo_var_ptrs), intent(in) :: tv !< A structure containing pointers to any !! available thermodynamic fields. Absent fields !! have NULL ptrs. - type(Kappa_shear_CS), pointer :: CS !< The control structure returned by a previous - !! call to kappa_shear_init. + type(Kappa_shear_CS), intent(in) :: CS !< The control structure returned by a previous + !! call to kappa_shear_init. Plain (not pointer) so + !! the routine is device-callable. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type integer, intent(in) :: eos_form !< The equation-of-state form id, resolved !! host-side for the device-callable EOS derivs path. @@ -966,6 +974,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure to Pa [Pa T2 R-1 L-2 ~> 1] + logical, intent(in) :: use_temperature !< If true, temperature and salinity are + !! state variables (resolved host-side to avoid a device + !! read of the tv%T pointer). ! Local variables real, dimension(nzc) :: & @@ -1055,8 +1066,6 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real :: k0dt ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1]. real :: I_lz_rescale_sqr ! The inverse of a rescaling factor for L2_bdry (Lz) squared [nondim]. logical :: valid_dt ! If true, all levels so far exhibit acceptably small changes in k_src. - logical :: use_temperature ! If true, temperature and salinity have been - ! allocated and are being used as state variables. integer :: ks_kappa, ke_kappa ! The k-range with nonzero kappas. integer :: dt_refinements ! The number of 2-fold refinements that will be used ! to estimate the maximum permitted time step. I.e., @@ -1083,7 +1092,6 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la endif tol2 = 2.0*CS%kappa_tol_err dt_refinements = 5 ! Selected so that 1/2^dt_refinements < 1-tol_dksrc_low - use_temperature = .false. ; if (associated(tv%T)) use_temperature = .true. ! Set up Idz as the inverse of layer thicknesses. @@ -1224,6 +1232,7 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la tv%eqn_of_state, (/2,nzc/), scale=-g_R0 ) #endif else +#ifndef __NVCOMPILER_OPENMP_GPU ! These should perhaps be combined into a single call to calculate the thermal expansion ! and haline contraction coefficients? call calculate_specific_vol_derivs(T_int, Sal_int, pressure, dSpV_dT, dSpV_dS, & @@ -1233,6 +1242,11 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la dbuoy_dT(K) = GV%g_Earth_Z_T2 * (rho_int(K) * dSpV_dT(K)) dbuoy_dS(K) = GV%g_Earth_Z_T2 * (rho_int(K) * dSpV_dS(K)) enddo +#else + ! The non-Boussinesq density-derivs path uses the polymorphic EOS interface, which is not + ! device-callable, and is excluded from device compilation. It is unreachable on GPU builds: + ! the driver FATALs on non-Boussinesq before the column solver runs. +#endif endif elseif (GV%Boussinesq .or. GV%semi_Boussinesq) then do K=1,nzc+1 ; dbuoy_dT(K) = -g_R0 ; dbuoy_dS(K) = 0.0 ; enddo @@ -1612,7 +1626,8 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b !! boundaries [H-1 Z-1 ~> m-2 or m kg-1]. real, dimension(nz), intent(in) :: Idz !< The inverse grid spacing of layers [Z-1 ~> m-1]. real, intent(in) :: f2 !< The squared Coriolis parameter [T-2 ~> s-2]. - type(Kappa_shear_CS), pointer :: CS !< A pointer to this module's control structure. + type(Kappa_shear_CS), intent(in) :: CS !< This module's control structure (plain, not pointer, + !! so the routine is device-callable). type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type real, dimension(nz+1), intent(inout) :: K_Q !< The shear-driven diapycnal diffusivity divided by From 0473a5f9a888911509329ca954877f0e17e1805a Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:07:26 +1000 Subject: [PATCH 18/45] kappa_shear: offload the vertex slab interpolation to GPU (3-D slabs) Promote the per-J 2-D vertex slabs (u_2d..rho_2d) in Calc_kappa_shear_vertex to 3-D arrays (u_slab..rho_slab) computed in device do-concurrent passes ahead of the column loop. The loop bodies are verbatim except that I_hwt is inlined as a reciprocal multiply and the use_temperature branch is hoisted so each pass is purely elementwise; answers are bitwise-identical (verified np=1 against repro_ocean.stats on benchmark_ALE, V100/NVHPC 26.5). The key residency fix relative to the earlier reverted attempt: u_in/v_in and T_in/S_in (tv%T, tv%S) are already device-resident (mapped in MOM.F90), so a map(to:) on them does not copy - and the diabatic stack modifies T/S on the host before set_diffusivity, leaving the device copies stale. They now get an explicit "target update to" refresh, which resolves the NaN-diffusivity crash the first attempt hit. dz_3d is a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are now consumed on the device by the slab passes, so increment 1's copy-back is dropped; the slabs are copied back for the (still host) per-column solver loop, which reads them as X_slab(I,J,k). Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 150 ++++++++++-------- 1 file changed, 86 insertions(+), 64 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index a5f7d916c3..8855164d4c 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -494,11 +494,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ h_at_u ! A mask-weighted thickness interpolated to u-points [H ~> m or kg m-2] real, dimension(SZI_(G),SZJB_(G),SZK_(GV)) :: & h_at_v ! A mask-weighted thickness interpolated to v-points [H ~> m or kg m-2] - real, dimension(SZIB_(G),SZK_(GV)) :: & - h_2d, & ! A 2-D version of h interpolated to vertices [H ~> m or kg m-2]. - dz_2d, & ! Vertical distance between interface heights [Z ~> m]. - u_2d, v_2d, & ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1]. - T_2d, S_2d, rho_2d ! 2-D versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. + real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)) :: & + h_slab, & ! A version of h interpolated to vertices [H ~> m or kg m-2]. + dz_slab, & ! Vertical distance between interface heights at vertices [Z ~> m]. + u_slab, v_slab, & ! Versions of u_in and v_in interpolated to vertices [L T-1 ~> m s-1]. + T_slab, S_slab, rho_slab ! Vertex versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. real, dimension(SZIB_(G),SZK_(GV)+1) :: & kappa_2d ! 2-D slice of kappa_vert [H Z T-1 ~> m2 s-1 or Pa s] real, dimension(SZIB_(G),SZK_(GV)+1) :: & @@ -587,7 +587,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! --- GPU port increment 1: h_at_u/h_at_v interpolation offloaded to device. ! h is host-authoritative in the (host-only) diabatic stack; refresh the device copy. ! G%mask2dCu/Cv/T are already device-resident (mapped in initialize_MOM). h_at_u/h_at_v - ! are device workspace, copied back for the (still host) per-column loop below. + ! are device workspace, consumed on the device by the slab interpolation below. !$omp target enter data map(to: h) !$omp target update to(h) !$omp target enter data map(alloc: h_at_u, h_at_v) @@ -613,52 +613,72 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - !$omp target update from(h_at_u, h_at_v) + ! --- GPU port increment 2: the per-J 2-D vertex slabs are promoted to 3-D arrays computed on + ! the device in one pass over J before the column loop. The loop bodies are verbatim from the + ! former per-J loops, except that I_hwt is inlined as a reciprocal multiply (bitwise-identical) + ! and the temperature/salinity branch is hoisted out of the loop so each body is purely + ! elementwise. u_in/v_in (the dycore u,v) and T_in/S_in (tv%T, tv%S) are already device- + ! resident but host-authoritative at this point in the (host-only) diabatic stack, so they + ! need an explicit refresh — a map(to:) on an already-present object does NOT copy. dz_3d is + ! a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are consumed on the + ! device here, so increment 1's copy-back is no longer needed. + !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) + !$omp target update to(u_in, v_in, T_in, S_in) + !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) + + ! Interpolate the various quantities to the corners, using masks. + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + u_slab(I,J,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & + ( (h_at_u(I,j,k) + h_at_u(I,j+1,k)) + H_tiny ) + v_slab(I,J,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & + ( (h_at_v(i,J,k) + h_at_v(i+1,J,k)) + H_tiny ) + + h_slab(I,J,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) + dz_slab(I,J,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * dz_3d(i+1,j,k) + G%mask2dT(i,j+1) * dz_3d(i,j+1,k)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) +! h_slab(I,J,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) +! h_slab(I,J,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & +! ((h(i+1,j,k)**2) + (h(i,j+1,k)**2))) * I_hwt + enddo + if (use_temperature) then + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + T_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & + G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * T_in(i+1,j+1,k))) + & + (G%mask2dT(i+1,j) * (h(i+1,j,k) * T_in(i+1,j,k)) + & + G%mask2dT(i,j+1) * (h(i,j+1,k) * T_in(i,j+1,k))) ) * & + (1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & + GV%H_subroundoff)) + S_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & + G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * S_in(i+1,j+1,k))) + & + (G%mask2dT(i+1,j) * (h(i+1,j,k) * S_in(i+1,j,k)) + & + G%mask2dT(i,j+1) * (h(i,j+1,k) * S_in(i,j+1,k))) ) * & + (1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & + GV%H_subroundoff)) + enddo + else + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + rho_slab(I,J,k) = GV%Rlay(k) + enddo + endif + ! The slabs feed the (still host) per-column loop below; copy them back. + !$omp target update from(u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & + !$OMP u_slab,v_slab,T_slab,S_slab,h_slab,dz_slab,rho_slab, & !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do J=JsB,JeB - ! Interpolate the various quantities to the corners, using masks. - do k=1,nz ; do I=IsB,IeB - u_2d(I,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & - ( (h_at_u(I,j,k) + h_at_u(I,j+1,k)) + H_tiny ) - v_2d(I,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & - ( (h_at_v(i,J,k) + h_at_v(i+1,J,k)) + H_tiny ) - - I_hwt = 1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & - GV%H_subroundoff) - if (use_temperature) then - T_2d(I,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & - G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * T_in(i+1,j+1,k))) + & - (G%mask2dT(i+1,j) * (h(i+1,j,k) * T_in(i+1,j,k)) + & - G%mask2dT(i,j+1) * (h(i,j+1,k) * T_in(i,j+1,k))) ) * I_hwt - S_2d(I,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & - G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * S_in(i+1,j+1,k))) + & - (G%mask2dT(i+1,j) * (h(i+1,j,k) * S_in(i+1,j,k)) + & - G%mask2dT(i,j+1) * (h(i,j+1,k) * S_in(i,j+1,k))) ) * I_hwt - endif - h_2d(I,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) - dz_2d(I,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * dz_3d(i+1,j,k) + G%mask2dT(i,j+1) * dz_3d(i,j+1,k)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) -! h_2d(I,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) -! h_2d(I,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & -! ((h(i+1,j,k)**2) + (h(i,j+1,k)**2))) * I_hwt - enddo ; enddo - if (.not.use_temperature) then ; do k=1,nz ; do I=IsB,IeB - rho_2d(I,k) = GV%Rlay(k) - enddo ; enddo ; endif - !--------------------------------------- ! Work on each column. !--------------------------------------- @@ -675,25 +695,25 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ T0xdz(k) = 0.0 ; S0xdz(k) = 0.0 ! Add a new layer if this one has mass. -! if ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless)) nzc = nzc+1 +! if ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 if ((k>CS%nkml) .and. (h_lay(nzc) > 0.0) .and. & - (h_2d(I,k) > dz_massless)) nzc = nzc+1 + (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 ! Only merge clusters of massless layers. ! if ((h_lay(nzc) > dz_massless) .or. & -! ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless))) nzc = nzc+1 +! ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless))) nzc = nzc+1 kc(k) = nzc - h_lay(nzc) = h_lay(nzc) + h_2d(I,k) - dz_lay(nzc) = dz_lay(nzc) + dz_2d(I,k) - u0xdz(nzc) = u0xdz(nzc) + u_2d(I,k)*h_2d(I,k) - v0xdz(nzc) = v0xdz(nzc) + v_2d(I,k)*h_2d(I,k) + h_lay(nzc) = h_lay(nzc) + h_slab(I,J,k) + dz_lay(nzc) = dz_lay(nzc) + dz_slab(I,J,k) + u0xdz(nzc) = u0xdz(nzc) + u_slab(I,J,k)*h_slab(I,J,k) + v0xdz(nzc) = v0xdz(nzc) + v_slab(I,J,k)*h_slab(I,J,k) if (use_temperature) then - T0xdz(nzc) = T0xdz(nzc) + T_2d(I,k)*h_2d(I,k) - S0xdz(nzc) = S0xdz(nzc) + S_2d(I,k)*h_2d(I,k) + T0xdz(nzc) = T0xdz(nzc) + T_slab(I,J,k)*h_slab(I,J,k) + S0xdz(nzc) = S0xdz(nzc) + S_slab(I,J,k)*h_slab(I,J,k) else - T0xdz(nzc) = T0xdz(nzc) + rho_2d(I,k)*h_2d(I,k) - S0xdz(nzc) = S0xdz(nzc) + rho_2d(I,k)*h_2d(I,k) + T0xdz(nzc) = T0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) + S0xdz(nzc) = S0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) endif enddo kc(nz+1) = nzc+1 @@ -703,28 +723,28 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Now determine kf, the fractional weight of interface kc when ! interpolating between interfaces kc and kc+1. - kf(1) = 0.0 ; dz_in_lay = h_2d(I,1) + kf(1) = 0.0 ; dz_in_lay = h_slab(I,J,1) do k=2,nz if (kc(k) > kc(k-1)) then - kf(k) = 0.0 ; dz_in_lay = h_2d(I,k) + kf(k) = 0.0 ; dz_in_lay = h_slab(I,J,k) else - kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(I,k) + kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_slab(I,J,k) endif enddo kf(nz+1) = 0.0 else do k=1,nz - h_lay(k) = h_2d(I,k) - dz_lay(k) = dz_2d(I,k) - u0xdz(k) = u_2d(I,k)*h_lay(k) ; v0xdz(k) = v_2d(I,k)*h_lay(k) + h_lay(k) = h_slab(I,J,k) + dz_lay(k) = dz_slab(I,J,k) + u0xdz(k) = u_slab(I,J,k)*h_lay(k) ; v0xdz(k) = v_slab(I,J,k)*h_lay(k) enddo if (use_temperature) then do k=1,nz - T0xdz(k) = T_2d(I,k)*h_lay(k) ; S0xdz(k) = S_2d(I,k)*h_lay(k) + T0xdz(k) = T_slab(I,J,k)*h_lay(k) ; S0xdz(k) = S_slab(I,J,k)*h_lay(k) enddo else do k=1,nz - T0xdz(k) = rho_2d(I,k)*h_lay(k) ; S0xdz(k) = rho_2d(I,k)*h_lay(k) + T0xdz(k) = rho_slab(I,J,k)*h_lay(k) ; S0xdz(k) = rho_slab(I,J,k)*h_lay(k) enddo endif nzc = nz @@ -817,7 +837,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Store the 2-d slices back in the 3-d arrays for restarts or interpolation back to tracer points. if (CS%VS_ThicknessMean) then do K=1,nz+1 ; do I=IsB,IeB - h_vert(I,J,k) = h_2d(I,k) + h_vert(I,J,k) = h_slab(I,J,k) enddo ; enddo endif if (CS%VS_viscosity_bug) then @@ -913,7 +933,9 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) - ! --- GPU port increment 1: mirror the enter-data above (balance discipline). + ! --- GPU port increments 1+2: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) + !$omp target exit data map(release: u_in, v_in, T_in, S_in, dz_3d) !$omp target exit data map(release: h_at_u, h_at_v) !$omp target exit data map(release: h) From 114849b1664bb8df15173f720399067200adbec2 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:19:38 +1000 Subject: [PATCH 19/45] kappa_shear: run the vertex column solver on the GPU Convert Calc_kappa_shear_vertex's per-column loop from a host OMP parallel do over J into an "omp target teams loop collapse(2)" over (J,I), calling the declare-target kappa_shear_column on the device. Bitwise-identical (np=1, benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5), which also demonstrates that the solver's 26 early exit/return sites are runtime-correct on 26.5. Restructuring, all bitwise-neutral: - The per-J staging slices kappa_2d/tke_2d become device-resident 3-D staging arrays kappa_3d/tke_3d (map(alloc:)), written per column and consumed by new elementwise do-concurrent write-back passes for kappa_vertex/tke_io/kv_io. - The surface pressure is interpolated to vertices on the host ahead of the region (surface_pres_2d, verbatim expressions), keeping the possibly-null p_surf pointer and the CS%psurf_bug branch out of device code. - The h_vert fill stays on the host (it is only consumed by the host tracer-point averaging when VS_ThicknessMean=true), fed by a guarded update from(h_slab). - All per-column scratch is private at the loop; loop-invariant scalars are firstprivate; CS gets a per-call shallow map(to:) in the driver (all scalars plus a diag pointer that device code never dereferences). kappa_vertex, tke_io and kv_io are mapped to: (not alloc:) so host-set halo/initial values survive the full-array update from. diag_N2/S2_init/mean are mapped after their conditional host zeroing and copied back under one guard matching their consumers (post_data/Bchksum). Known limitation, to be fixed next: the solver's automatic arrays (sized nzc or SZK_(GV)) are device-heap allocated, so runs currently need NVCOMPILER_ACC_CUDA_HEAPSIZE=1073741824 (the default heap is too small: "DEVICE FORTRAN AUTO ALLOCATION FAILED"). Fixed-size locals will remove this. Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 148 +++++++++++------- 1 file changed, 94 insertions(+), 54 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 8855164d4c..6a7738939a 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -499,10 +499,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ dz_slab, & ! Vertical distance between interface heights at vertices [Z ~> m]. u_slab, v_slab, & ! Versions of u_in and v_in interpolated to vertices [L T-1 ~> m s-1]. T_slab, S_slab, rho_slab ! Vertex versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. - real, dimension(SZIB_(G),SZK_(GV)+1) :: & - kappa_2d ! 2-D slice of kappa_vert [H Z T-1 ~> m2 s-1 or Pa s] - real, dimension(SZIB_(G),SZK_(GV)+1) :: & - tke_2d ! 2-D version tke_io [Z2 T-2 ~> m2 s-2]. + real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1) :: & + kappa_3d, & ! Device staging array for the columns' averaged kappa [H Z T-1 ~> m2 s-1 or Pa s] + tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. + real, dimension(SZIB_(G),SZJB_(G)) :: & + surface_pres_2d ! The surface pressure interpolated to vertices [R L2 T-2 ~> Pa]. real, dimension(SZK_(GV)) :: & Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] @@ -527,7 +528,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ real :: dz_in_lay ! The running sum of the thickness in a layer [H ~> m or kg m-2] real :: k0dt ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1] real :: dz_massless ! A layer thickness that is considered massless [H ~> m or kg m-2] - real :: I_hwt ! The inverse of the sum of the adjacent masked thickness weights [H-1 ~> m-1 or m2 kg-1] real :: I_htot ! The inverse of the sum of the thicknesses at adjacent vertices [H-1 ~> m-1 or m2 kg-1] real :: I_Prandtl ! The inverse of the turbulent Prandtl number [nondim]. logical :: use_temperature ! If true, temperature and salinity have been @@ -581,6 +581,28 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ I_Prandtl = 0.0 ; if (CS%Prandtl_turb > 0.0) I_Prandtl = 1.0 / CS%Prandtl_turb H_tiny = 0.5 * GV%H_subroundoff + ! GPU port increment 3: interpolate the surface pressure to the vertices on the host, ahead + ! of the device column region (p_surf is a possibly-unassociated pointer, which is simpler to + ! keep off the device). The expressions are verbatim from the former in-column code, but are + ! evaluated for every vertex instead of only ocean vertices; the extra values are never read. + surface_pres_2d(:,:) = 0.0 + if (associated(p_surf)) then + if (CS%psurf_bug) then + ! This is wrong because it is averaging values from land in some places. + do J=JsB,JeB ; do I=IsB,IeB + surface_pres_2d(I,J) = 0.25 * ((p_surf(i,j) + p_surf(i+1,j+1)) + & + (p_surf(i+1,j) + p_surf(i,j+1))) + enddo ; enddo + else + do J=JsB,JeB ; do I=IsB,IeB + surface_pres_2d(I,J) = ((G%mask2dT(i,j) * p_surf(i,j) + G%mask2dT(i+1,j+1) * p_surf(i+1,j+1)) + & + (G%mask2dT(i+1,j) * p_surf(i+1,j) + G%mask2dT(i,j+1) * p_surf(i,j+1)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) + enddo ; enddo + endif + endif + ! Convert layer thicknesses into geometric thickness in height units. call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=1) @@ -668,20 +690,31 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - ! The slabs feed the (still host) per-column loop below; copy them back. - !$omp target update from(u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) - - !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & - !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & - !$OMP u_slab,v_slab,T_slab,S_slab,h_slab,dz_slab,rho_slab, & - !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & - !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & - !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) - do J=JsB,JeB + ! --- GPU port increment 3: run the per-column solver on the device. The columns iterate as + ! a target teams loop collapsed over (J,I); every piece of per-column scratch is private (the + ! declare-target solver's own locals are automatically private per device thread). CS is all + ! scalars plus a diag pointer that is never dereferenced in the device code, so a per-call + ! shallow map(to:) suffices. kappa_vertex/tke_io/kv_io are mapped to: (not alloc) so that the + ! host-set values - kv_io is intent(inout), kappa_vertex is zeroed on the host - survive the + ! full-array update from below. The diag_* arrays are mapped after their (conditional) host + ! zeroing for the same reason. kappa_3d/tke_3d are device-only staging for what was the per-J + ! kappa_2d/tke_2d, written per column and consumed by the write-back passes below. + !$omp target enter data map(to: CS) + !$omp target enter data map(to: surface_pres_2d) + !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) + !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target enter data map(alloc: kappa_3d, tke_3d) !--------------------------------------- ! Work on each column. !--------------------------------------- + !$omp target teams loop collapse(2) & + !$omp private(nzc, kc, kf, Idz, h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, dz_in_lay, & + !$omp f2, surface_pres, kappa, tke, kappa_avg, tke_avg, N2_init, S2_init, & + !$omp N2_mean, S2_mean, k) & + !$omp firstprivate(nz, dt, k0dt, dz_massless, use_temperature, eos_form, & + !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + do J=JsB,JeB do I=IsB,IeB ; if ((G%mask2dCu(I,j) + G%mask2dCu(I,j+1)) + & (G%mask2dCv(i,J) + G%mask2dCv(i+1,J)) > 0.0) then ! call cpu_clock_begin(Id_clock_setup) @@ -752,19 +785,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ endif f2 = G%Coriolis2Bu(I,J) - surface_pres = 0.0 - if (associated(p_surf)) then - if (CS%psurf_bug) then - ! This is wrong because it is averaging values from land in some places. - surface_pres = 0.25 * ((p_surf(i,j) + p_surf(i+1,j+1)) + & - (p_surf(i+1,j) + p_surf(i,j+1))) - else - surface_pres = ((G%mask2dT(i,j) * p_surf(i,j) + G%mask2dT(i+1,j+1) * p_surf(i+1,j+1)) + & - (G%mask2dT(i+1,j) * p_surf(i+1,j) + G%mask2dT(i,j+1) * p_surf(i,j+1)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) - endif - endif + surface_pres = surface_pres_2d(I,J) ! ---------------------------------------------------- ! Set the initial guess for kappa, here defined at interfaces. @@ -780,11 +801,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_2d(I,K) = kappa_avg(K) + kappa_3d(I,J,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_2d(I,K) = tke(K) + tke_3d(I,J,K) = tke(K) else - tke_2d(I,K) = tke_avg(K) + tke_3d(I,J,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -802,11 +823,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ else do K=1,nz+1 if (kf(K) == 0.0) then - kappa_2d(I,K) = kappa_avg(kc(K)) - tke_2d(I,K) = tke_avg(kc(K)) + kappa_3d(I,J,K) = kappa_avg(kc(K)) + tke_3d(I,J,K) = tke_avg(kc(K)) else - kappa_2d(I,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) - tke_2d(I,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) + kappa_3d(I,J,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) + tke_3d(I,J,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) endif enddo do K=1,nz+1 @@ -830,30 +851,44 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! call cpu_clock_end(Id_clock_setup) else ! Land points, still inside the i-loop. do K=1,nz+1 - kappa_2d(I,K) = 0.0 ; tke_2d(I,K) = 0.0 + kappa_3d(I,J,K) = 0.0 ; tke_3d(I,J,K) = 0.0 enddo endif ; enddo ! i-loop + enddo ! end of J-loop - ! Store the 2-d slices back in the 3-d arrays for restarts or interpolation back to tracer points. - if (CS%VS_ThicknessMean) then + ! Store the columns' results back in the 3-d arrays for restarts or interpolation back to + ! tracer points. h_vert is only used when VS_ThicknessMean is true, and only by the (host) + ! tracer-point averaging below, so its fill stays a host loop fed by a guarded copy-back of + ! the device-computed h_slab. + if (CS%VS_ThicknessMean) then + !$omp target update from(h_slab) + do J=JsB,JeB do K=1,nz+1 ; do I=IsB,IeB h_vert(I,J,k) = h_slab(I,J,k) enddo ; enddo - endif - if (CS%VS_viscosity_bug) then - do K=1,nz+1 ; do I=IsB,IeB - kappa_vertex(I,J,K) = kappa_2d(I,K) - tke_io(I,J,K) = G%mask2dBu(I,J) * tke_2d(I,K) - kv_io(I,J,K) = ( G%mask2dBu(I,J) * kappa_vertex(I,J,K) ) * CS%Prandtl_turb - enddo ; enddo - else - do K=1,nz+1 ; do I=IsB,IeB - kappa_vertex(I,J,K) = kappa_2d(I,K) - tke_io(I,J,K) = tke_2d(I,K) - kv_io(I,J,K) = kappa_vertex(I,J,K) * CS%Prandtl_turb - enddo ; enddo - endif - enddo ! end of J-loop + enddo + endif + if (CS%VS_viscosity_bug) then + do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) + kappa_vertex(I,J,K) = kappa_3d(I,J,K) + tke_io(I,J,K) = G%mask2dBu(I,J) * tke_3d(I,J,K) + kv_io(I,J,K) = ( G%mask2dBu(I,J) * kappa_vertex(I,J,K) ) * CS%Prandtl_turb + enddo + else + do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) + kappa_vertex(I,J,K) = kappa_3d(I,J,K) + tke_io(I,J,K) = tke_3d(I,J,K) + kv_io(I,J,K) = kappa_vertex(I,J,K) * CS%Prandtl_turb + enddo + endif + + ! The vertex-to-tracer-point averaging below, the checksums and post_data are all still on + ! the host. The diag_* transfer is guarded by the same conditions as their consumers. + !$omp target update from(kappa_vertex, tke_io, kv_io) + if ((CS%id_N2_init>0) .or. (CS%id_S2_init>0) .or. (CS%id_N2_mean>0) .or. (CS%id_S2_mean>0) & + .or. CS%debug) then + !$omp target update from(diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + endif ! Set the diffusivities in tracer columns from the values at vertices. @@ -933,7 +968,12 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) - ! --- GPU port increments 1+2: mirror the enter-data above (balance discipline). + ! --- GPU port increments 1-3: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: kappa_3d, tke_3d) + !$omp target exit data map(release: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target exit data map(release: kappa_vertex, tke_io, kv_io) + !$omp target exit data map(release: surface_pres_2d) + !$omp target exit data map(release: CS) !$omp target exit data map(release: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) !$omp target exit data map(release: u_in, v_in, T_in, S_in, dz_3d) !$omp target exit data map(release: h_at_u, h_at_v) From 9d97e79fb47dc3771d7634125a828ca69e1d3634 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:31:50 +1000 Subject: [PATCH 20/45] kappa_shear: fixed-size column arrays in GPU builds (drop device-heap allocs) The vertex driver's per-column private scratch and the column routines' automatic locals are sized by GV%ke/nzc, which is not a compile-time constant, so nvfortran allocated them from the device heap on every column call (NVFORTRAN-W-0155). That exhausted the default device heap ("DEVICE FORTRAN AUTO ALLOCATION FAILED" unless NVCOMPILER_ACC_CUDA_HEAPSIZE was raised) and serialized columns on the device allocator (~2x total model slowdown). In GPU builds (#ifdef __NVCOMPILER_OPENMP_GPU) these arrays now use a compile-time ceiling, the new module parameter GPU_nk_max=128, so each device thread gets stack ("local memory") arrays; kappa_shear_init fatals if GV%ke > GPU_nk_max when the scheme is active. CPU builds keep the exact-size declarations. No code queries these arrays' sizes and all accesses are indexed to nzc/nz, so the enlargement is bitwise-neutral: verified np=1 against repro_ocean.stats at the default device heap, wall time back to the pre-offload baseline (98s -> 45s), and the W-0155 warnings are gone. Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 6a7738939a..c891e41903 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -135,6 +135,13 @@ module MOM_kappa_shear ! can run inside a target region (GPU port increment 3). !$omp declare target(kappa_shear_column, find_kappa_tke, calculate_projected_state) +!> A compile-time ceiling on the number of layers in GPU builds, used to give the +!! device-executed column routines and the driver's per-column private scratch fixed-size +!! (stack) arrays instead of per-call device-heap automatic allocations, which exhaust the +!! default device heap and serialize on the device allocator. Checked against GV%ke in +!! kappa_shear_init. Unused in CPU builds, where the declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 + contains !> Subroutine for calculating shear-driven diffusivity and TKE in tracer columns @@ -504,7 +511,14 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. real, dimension(SZIB_(G),SZJB_(G)) :: & surface_pres_2d ! The surface pressure interpolated to vertices [R L2 T-2 ~> Pa]. + ! In GPU builds the per-column private scratch has compile-time-constant sizes so each + ! device thread gets stack ("local memory") arrays; runtime-sized privates are device-heap + ! allocated per column, which exhausts the default heap and serializes on the allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] dz_lay, & ! The geometric layer thickness in height units [Z ~> m] @@ -512,7 +526,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ v0xdz, & ! The initial meridional velocity times dz [H L T-1 ~> m2 s-1 or kg m-1 s-1] T0xdz, & ! The initial temperature times dz [C H ~> degC m or degC kg m-2] S0xdz ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif kappa, & ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s] tke, & ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2]. kappa_avg, & ! The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -533,10 +551,18 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ logical :: use_temperature ! If true, temperature and salinity have been ! allocated and are being used as state variables. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max+1) :: kc ! The index map between the original +#else integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original +#endif ! interfaces and the interfaces with massless layers ! merged into nearby massive layers. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: kf ! The fractional weight of interface kc+1 for +#else real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for +#endif ! interpolating back to the original index space [nondim]. real :: h_SW, h_SE, h_NW, h_NE ! Thicknesses at adjacent vertices [H ~> m or kg m-2] real :: mks_to_HZ_T ! A factor used to restore dimensional scaling after the geometric mean @@ -1041,7 +1067,14 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la !! read of the tv%T pointer). ! Local variables + ! In GPU builds these locals have compile-time-constant sizes so each device call uses + ! stack ("local memory") arrays; runtime-sized automatics are device-heap allocated per + ! call, which exhausts the default heap and serializes on the device allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(nzc) :: & +#endif u, & ! The zonal velocity after a timestep of mixing [L T-1 ~> m s-1]. v, & ! The meridional velocity after a timestep of mixing [L T-1 ~> m s-1]. Idz, & ! The inverse of the distance between TKE points [Z-1 ~> m-1]. @@ -1050,7 +1083,11 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la u_test, v_test, & ! Temporary velocities [L T-1 ~> m s-1]. T_test, S_test ! Temporary temperatures [C ~> degC] and salinities [S ~> ppt]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nzc+1) :: & +#endif N2, & ! The squared buoyancy frequency at an interface [T-2 ~> s-2]. h_Int, & ! The extent of a finite-volume space surrounding an interface, ! as used in calculating kappa and TKE [H ~> m or kg m-2] @@ -1574,7 +1611,11 @@ subroutine calculate_projected_state(kappa, u0, v0, T0, S0, dt, nz, dz, I_dz_int !! diffusivity. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: c1 ! A tridiagonal variable [nondim] +#else real, dimension(nz+1) :: c1 ! A tridiagonal variable [nondim] +#endif real :: a_a, a_b ! Tridiagonal coupling coefficients [H ~> m or kg m-2] real :: b1, b1nz_0 ! Tridiagonal variables [H-1 ~> m-1 or m2 kg-1] real :: bd1 ! A term in the denominator of b1 [H ~> m or kg m-2] @@ -1707,10 +1748,18 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b ! This subroutine calculates new, consistent estimates of TKE and kappa. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(nz) :: & +#endif aQ, & ! aQ is the coupling between adjacent interfaces in the TKE equations [H T-1 ~> m s-1 or kg m-2 s-1] dQdz ! Half the partial derivative of TKE with depth [Z T-2 ~> m s-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nz+1) :: & +#endif dK, & ! The change in kappa [H Z T-1 ~> m2 s-1 or Pa s]. dQ, & ! The change in TKE [Z2 T-2 ~> m2 s-2]. cQ, cK, & ! cQ and cK are the upward influences in the tridiagonal and @@ -1788,7 +1837,11 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b logical, parameter :: debug_soln = .false. real :: K_err_lin ! The imbalance in the K equation [H T-1 ~> m s-1 or kg m-2 s-1] real :: Q_err_lin ! The imbalance in the Q equation [H Z T-3 ~> m2 s-3 or kg m-1 s-3] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nz+1) :: & +#endif I_Ld2_debug, & ! A separate version of I_Ld2 for debugging [H-1 Z-1 ~> m-2 or m kg-1]. kappa_prev, & ! The value of kappa at the start of the current iteration [H Z T-1 ~> m2 s-1 or Pa s] TKE_prev ! The value of TKE at the start of the current iteration [Z2 T-2 ~> m2 s-2]. @@ -2495,6 +2548,13 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) 's-2', conversion=US%s_to_T**2) endif +#ifdef __NVCOMPILER_OPENMP_GPU + ! The device-executed column routines use fixed-size local arrays in GPU builds. + if (kappa_shear_init .and. (GV%ke > GPU_nk_max)) call MOM_error(FATAL, & + "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max; "//& + "increase GPU_nk_max in MOM_kappa_shear.F90.") +#endif + end function kappa_shear_init !> This function indicates to other modules whether the Jackson et al shear mixing From 79459a67ae4e8b64778d2d166fde3c3736773c19 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:53:14 +1000 Subject: [PATCH 21/45] kappa_shear: review fixes - refresh device tke_io/kv_io, correct T/S residency Review of the column-offload series found one latent halo bug: kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident, mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update, so this routine's map(to: kv_io) copies nothing and the full-array "update from" wrote stale pre-pass_var device halos over the corrected host halos - silently resurrecting the restart non-reproducibility in non-symmetric mode that the pass_var was added to fix. Invisible to a cold-start np=1 run. Fixed with an explicit "update to(tke_io, kv_io)" after the maps (tke_io defensively, so the routine does not depend on it staying unmapped elsewhere). Also corrects the increment-2 residency comment and drops the redundant "update to(T_in, S_in)": the actuals are the caller's convection-filtered T_f/S_f, fresh host locals each call, so their map(to:) already copies (they are not tv%T/tv%S as the comment claimed). This also establishes that the earlier reverted slab attempt's NaN was not T/S staleness; the operative fix was restructuring the interpolation (hoisting the use_temperature branch out of the do concurrent). And the GPU_nk_max FATAL message now notes it covers the shared tracer-point column routines too. Bit-identical at np=1 (benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5). Co-Authored-By: Claude Fable 5 Co-authored-by: Utheri Wagura --- .../vertical/MOM_kappa_shear.F90 | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index c891e41903..7ca8143fba 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -665,13 +665,14 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! the device in one pass over J before the column loop. The loop bodies are verbatim from the ! former per-J loops, except that I_hwt is inlined as a reciprocal multiply (bitwise-identical) ! and the temperature/salinity branch is hoisted out of the loop so each body is purely - ! elementwise. u_in/v_in (the dycore u,v) and T_in/S_in (tv%T, tv%S) are already device- - ! resident but host-authoritative at this point in the (host-only) diabatic stack, so they - ! need an explicit refresh — a map(to:) on an already-present object does NOT copy. dz_3d is - ! a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are consumed on the - ! device here, so increment 1's copy-back is no longer needed. + ! elementwise. u_in/v_in (the dycore u,v) are already device-resident but host- + ! authoritative at this point in the (host-only) diabatic stack, so they need an explicit + ! refresh — a map(to:) on an already-present object does NOT copy. T_in/S_in (the caller's + ! convection-filtered T_f/S_f) and dz_3d are fresh host locals each call, so their map(to:) + ! does copy. h_at_u/h_at_v are consumed on the device here, so increment 1's copy-back is + ! no longer needed. !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) - !$omp target update to(u_in, v_in, T_in, S_in) + !$omp target update to(u_in, v_in) !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) ! Interpolate the various quantities to the corners, using masks. @@ -728,6 +729,13 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ !$omp target enter data map(to: CS) !$omp target enter data map(to: surface_pres_2d) !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) + ! kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident + ! (mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update), so the + ! map(to:) above does not copy it; without this refresh, the full-array update from below + ! would overwrite the corrected host halos with stale device values after a restart in + ! non-symmetric mode. tke_io is not currently mapped elsewhere, but is refreshed too so + ! this routine does not silently depend on that staying true. + !$omp target update to(tke_io, kv_io) !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) !$omp target enter data map(alloc: kappa_3d, tke_3d) @@ -2551,8 +2559,10 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) #ifdef __NVCOMPILER_OPENMP_GPU ! The device-executed column routines use fixed-size local arrays in GPU builds. if (kappa_shear_init .and. (GV%ke > GPU_nk_max)) call MOM_error(FATAL, & - "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max; "//& - "increase GPU_nk_max in MOM_kappa_shear.F90.") + "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max because the "//& + "column routines use fixed-size local arrays on the device (this applies to the "//& + "tracer-point scheme too, which shares those routines); increase GPU_nk_max in "//& + "MOM_kappa_shear.F90 or use a CPU build.") #endif end function kappa_shear_init From c0d2980399e8a60a05eb32786dcfa713fa73d3cb Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Sat, 18 Jul 2026 11:22:48 +1000 Subject: [PATCH 22/45] MOM_EOS: add device-callable in-situ-density dispatcher (declare target) Groundwork for offloading the finite-volume pressure-gradient density integrals (int_density_dz_generic_plm next), which must evaluate the equation of state from INSIDE a device region, where the polymorphic EOS%type%density_elem / density_anomaly_elem dispatch is not usable. - Expose the existing this-free _loc elementals as public and !$omp declare target: density_elem_Roquet_rho_loc and density_anomaly_elem_Roquet_rho_loc (Roquet_rho), and density_elem_buggy_Wright_loc (Wright). buggy_Wright has no anomaly _loc kernel, so its anomaly branch stays unsupported here. - Add calculate_density_elem_loc(form_of_EOS, T, S, pressure, use_rho_ref, rho_ref): a declare-target function that select-case's on the integer form and returns, in mks units, either in-situ density or (when use_rho_ref) the anomaly relative to rho_ref, by forwarding to the matching _loc kernel. It is the density analogue of calculate_density_derivs_elem_loc (47ac931cb) and, like it, leaves unit rescaling (EOS%*_to_* factors) and any `scale` factor to the caller, exactly as calculate_density_elem does. Combinations without a device-callable kernel -- an unsupported form, or the anomaly branch of buggy_Wright -- return 0; a device-using module must FATAL at init on a GPU build before reaching such a combination. Bitwise-safety: this is purely additive. Nothing calls the dispatcher yet, and the exposed elementals are the identical functions the polymorphic host path already wraps (density_anomaly_elem_Roquet_rho => density_anomaly_elem_Roquet_rho_loc, etc.), so no existing arithmetic moves or changes. Scope note (reviewer-directed): no new Wright kernels are added here. benchmark_ALE (ROQUET_RHO) cannot exercise a new Wright anomaly kernel, and untested numerics is scope creep; parity can come later with a config that exercises it. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); DEBUG=True redundancy gate clean (0 'redundant'). "(Ocean pressure force)" clock unchanged at ~20.8 s/24 calls, as expected for an additive change. Co-Authored-By: Claude Opus --- src/equation_of_state/MOM_EOS.F90 | 41 ++++++++++++++++++++ src/equation_of_state/MOM_EOS_Roquet_rho.F90 | 5 +++ src/equation_of_state/MOM_EOS_Wright.F90 | 4 ++ 3 files changed, 50 insertions(+) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 8b757fde76..fcfc5af375 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -10,6 +10,7 @@ module MOM_EOS use MOM_EOS_linear, only : int_density_dz_linear, int_spec_vol_dp_linear use MOM_EOS_Wright, only : buggy_Wright_EOS, avg_spec_vol_buggy_Wright use MOM_EOS_Wright, only : calculate_density_derivs_elem_buggy_Wright_loc +use MOM_EOS_Wright, only : density_elem_buggy_Wright_loc use MOM_EOS_Wright, only : int_density_dz_wright, int_spec_vol_dp_wright use MOM_EOS_Wright_full, only : Wright_full_EOS, avg_spec_vol_Wright_full use MOM_EOS_Wright_full, only : int_density_dz_wright_full, int_spec_vol_dp_wright_full @@ -19,6 +20,7 @@ module MOM_EOS use MOM_EOS_UNESCO, only : UNESCO_EOS use MOM_EOS_Roquet_rho, only : Roquet_rho_EOS use MOM_EOS_Roquet_rho, only : calculate_density_derivs_elem_Roquet_rho_loc +use MOM_EOS_Roquet_rho, only : density_elem_Roquet_rho_loc, density_anomaly_elem_Roquet_rho_loc use MOM_EOS_Roquet_SpV, only : Roquet_SpV_EOS use MOM_EOS_TEOS10, only : TEOS10_EOS use MOM_EOS_TEOS10, only : gsw_sp_from_sr, gsw_pt_from_ct, gsw_sr_from_sp, gsw_ct_from_pt @@ -49,6 +51,7 @@ module MOM_EOS public calculate_density public calculate_density_derivs public calculate_density_derivs_elem_loc +public calculate_density_elem_loc public get_EOS_form_and_scaling public calculate_density_second_derivs public calculate_spec_vol @@ -924,6 +927,44 @@ subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_d end subroutine calculate_density_derivs_elem_loc +!> Device-callable dispatcher for in-situ density (or its anomaly relative to rho_ref) at a single +!! point, in mks units, selecting the equation-of-state form at runtime by integer id (no +!! polymorphic dispatch) so it can be called from inside a do concurrent / target region by GPU +!! kernels (e.g. the finite-volume pressure-gradient density integrals). Unit rescaling +!! (EOS%*_to_* factors) and any `scale` factor are the caller's responsibility, exactly as in +!! calculate_density_elem. If use_rho_ref is true, the density anomaly relative to rho_ref (all in +!! mks units) is returned via the form's device-callable anomaly kernel. Combinations without a +!! device-callable _loc kernel -- an unsupported form, or the anomaly branch of a form that has no +!! anomaly kernel (e.g. buggy_Wright) -- are not handled here and return 0; a device-using module +!! must FATAL at init on a GPU build before reaching this with such a combination. +real function calculate_density_elem_loc(form_of_EOS, T, S, pressure, use_rho_ref, rho_ref) + integer, intent(in) :: form_of_EOS !< The equation of state form (EOS_ROQUET_RHO, EOS_WRIGHT, ...) + real, intent(in) :: T !< Temperature in the EOS kernel's mks units [degC] + real, intent(in) :: S !< Salinity in the EOS kernel's mks units [ppt or g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + logical, intent(in) :: use_rho_ref !< If true, return the density anomaly relative to rho_ref + real, intent(in) :: rho_ref !< A reference density [kg m-3], subtracted when use_rho_ref is true + !$omp declare target + + calculate_density_elem_loc = 0.0 + select case (form_of_EOS) + case (EOS_ROQUET_RHO) + if (use_rho_ref) then + calculate_density_elem_loc = density_anomaly_elem_Roquet_rho_loc(T, S, pressure, rho_ref) + else + calculate_density_elem_loc = density_elem_Roquet_rho_loc(T, S, pressure) + endif + case (EOS_WRIGHT) + ! buggy_Wright has a device-callable in-situ density kernel but no anomaly kernel, so the + ! anomaly branch is unsupported and returns 0 (host FATALs at init on GPU builds). + if (.not. use_rho_ref) & + calculate_density_elem_loc = density_elem_buggy_Wright_loc(T, S, pressure) + case default + calculate_density_elem_loc = 0.0 + end select + +end function calculate_density_elem_loc + !> Return the equation-of-state form id and the unit-rescaling factors held in an EOS_type. !! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion !! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without diff --git a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 index 8222d71e7d..6a4f3cc7a2 100644 --- a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 +++ b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 @@ -14,6 +14,11 @@ module MOM_EOS_Roquet_rho ! can compute density derivatives in-region without polymorphic dispatch. public calculate_density_derivs_elem_Roquet_rho_loc !$omp declare target(calculate_density_derivs_elem_Roquet_rho_loc) +! Exposed as device-callable (declare target) elementals so whole-column GPU kernels can compute +! in-situ density and its anomaly in-region without polymorphic dispatch. +public density_elem_Roquet_rho_loc, density_anomaly_elem_Roquet_rho_loc +!$omp declare target(density_elem_Roquet_rho_loc) +!$omp declare target(density_anomaly_elem_Roquet_rho_loc) real, parameter :: Pa2kb = 1.e-8 !< Conversion factor between Pa and kbar [kbar Pa-1] !>@{ Parameters in the Roquet_rho (Roquet density) equation of state diff --git a/src/equation_of_state/MOM_EOS_Wright.F90 b/src/equation_of_state/MOM_EOS_Wright.F90 index b690c6537b..08c85858d6 100644 --- a/src/equation_of_state/MOM_EOS_Wright.F90 +++ b/src/equation_of_state/MOM_EOS_Wright.F90 @@ -16,6 +16,10 @@ module MOM_EOS_Wright ! can compute density derivatives in-region without polymorphic dispatch. public calculate_density_derivs_elem_buggy_Wright_loc !$omp declare target(calculate_density_derivs_elem_buggy_Wright_loc) +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels can compute +! in-situ density in-region without polymorphic dispatch. (buggy_Wright has no anomaly _loc kernel.) +public density_elem_buggy_Wright_loc +!$omp declare target(density_elem_buggy_Wright_loc) public int_density_dz_wright, int_spec_vol_dp_wright public avg_spec_vol_buggy_Wright public set_params_buggy_Wright From d48ee86eea520c030dca3dabcb3863879dcf3fbc Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Sat, 18 Jul 2026 14:32:36 +1000 Subject: [PATCH 23/45] int_density_dz_generic_plm: element-wise device-callable EOS path (offload prep) Host-side bitwise-safe refactor of the finite-volume PGF density integrals, preparing int_density_dz_generic_plm for offload (increment C next): replace the polymorphic array calculate_density calls -- which are not callable from a device region -- with an element-wise path built on the device-callable dispatcher calculate_density_elem_loc (327f88470). Still host-executed here; directives come in C. - MOM_EOS: get_EOS_form_and_scaling gains an optional R_to_kg_m3 out-arg (backward-compatible; kappa_shear's two callers are unaffected). Needed to reproduce calculate_density_1d's unity fast-path test and non-unity rescale exactly. - int_density_dz_generic_plm: resolve eos_form + the five unit-scaling factors once on the host (the accessor and MOM_error are not device-callable), and precompute eos_unity mirroring calculate_density_1d's no-rescale test. Under #ifdef __NVCOMPILER_OPENMP_GPU, replace the three non-Stanley calculate_density array calls (phase 1 vertical r5; phase 2 x-integral r15; phase 3 y-integral r15) with element loops calling calculate_density_elem_loc, reproducing calculate_density_1d(..., EOSdom_*[, rho_ref]) bit-for-bit: same _loc kernel per point, same rescale math (Ta=C_to_degC*T, ..., rho*=kg_m3_to_R, rho_ref->R_to_kg_m3*rho_ref). The eos_unity branch is hoisted outside the point loop so the loop bodies are branch-free (offload-ready; trap 3). Non-GPU builds keep the original array call verbatim in #else. The Stanley calls, and the int_density_dz_generic_ppm / _pcm siblings, are untouched. - On GPU builds, FATAL host-side (before any device region) when the non-Stanley path hits an unsupported EOS form / anomaly mode: the dispatcher covers only ROQUET_RHO (density and anomaly) and buggy_Wright (density only). Bitwise-safety: (i,n)/(i,n15) loops cover exactly the EOSdom_h5 / EOSdom_q15 / EOSdom_h15 index ranges the array calls work on; per-point density values are order-independent, so the loop order does not matter; the _loc kernels are the identical functions the polymorphic path already wraps. Multiplying by an unscaled factor of exactly 1.0 is bit-exact, so the non-unity branch matches calculate_density_1d's guarded rho_scale multiply. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); the #ifdef GPU element path is the executed arithmetic on the GPU build. DEBUG=True redundancy gate clean (0 'redundant'). "(Ocean pressure force)" clock ~20.4 s/24 calls, unchanged (host execution; directives land in C). Not runtime-exercised by the benchmark (inspection-only, code-identical to the array path): the Stanley branch, the use_rho_ref=.false. (inaccurate-form) branch, and the non-unity unit-scaling branch. Co-Authored-By: Claude Opus --- src/core/MOM_density_integrals.F90 | 87 ++++++++++++++++++++++++++++++ src/equation_of_state/MOM_EOS.F90 | 5 +- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/src/core/MOM_density_integrals.F90 b/src/core/MOM_density_integrals.F90 index bb398d9f00..311658b704 100644 --- a/src/core/MOM_density_integrals.F90 +++ b/src/core/MOM_density_integrals.F90 @@ -10,6 +10,8 @@ module MOM_density_integrals use MOM_EOS, only : analytic_int_density_dz use MOM_EOS, only : analytic_int_specific_vol_dp use MOM_EOS, only : calculate_density +use MOM_EOS, only : calculate_density_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_EOS, only : calculate_spec_vol use MOM_EOS, only : calculate_specific_vol_derivs use MOM_EOS, only : average_specific_vol @@ -532,6 +534,12 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & integer, dimension(2) :: EOSdom_h5 ! The 5-point h-point i-computational domain for the equation of state integer, dimension(2) :: EOSdom_q15 ! The 3x5-point q-point i-computational domain for the equation of state integer, dimension(2) :: EOSdom_h15 ! The 3x5-point h-point i-computational domain for the equation of state + integer :: eos_form ! The equation-of-state form id, resolved host-side for the device-callable density path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, eos_R_to_kg_m3 + ! EOS unit-rescaling factors, resolved host-side. + logical :: eos_unity ! True when all EOS unit-rescaling factors are 1, matching the no-rescale + ! fast path of calculate_density_1d. + integer :: n15 ! A subgrid-point loop index over the 15 q-point slots in a column. integer :: Isq, Ieq, Jsq, Jeq, i, j, m, n, pos Isq = HI%IscB ; Ieq = HI%IecB ; Jsq = HI%JscB ; Jeq = HI%JecB @@ -569,6 +577,31 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & use_varS = associated(tv%varS) endif + ! GPU port increment B: resolve the EOS form + unit scaling once on the host (the accessor is + ! not device-callable) so the non-Stanley density integrals can be evaluated point-by-point via + ! the device-callable dispatcher calculate_density_elem_loc, instead of the polymorphic + ! calculate_density interface, when this is (later) offloaded. eos_unity mirrors + ! calculate_density_1d's no-rescale fast path. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 + eos_RL2_T2_to_Pa = 1.0 ; eos_R_to_kg_m3 = 1.0 + call get_EOS_form_and_scaling(EOS, eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, & + eos_RL2_T2_to_Pa, R_to_kg_m3=eos_R_to_kg_m3) + eos_unity = (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) +#ifdef __NVCOMPILER_OPENMP_GPU + ! On GPU builds the non-Stanley density integrals below are evaluated with the device-callable + ! dispatcher, which only covers ROQUET_RHO (in-situ density and anomaly) and buggy_Wright + ! (in-situ density only). FATAL host-side, before any device region, on an unsupported config. + if (.not. use_stanley_eos) then + if (.not. ((eos_form == EOS_ROQUET_RHO) .or. & + ((eos_form == EOS_WRIGHT) .and. (.not. use_rho_ref)))) call MOM_error(FATAL, & + "int_density_dz_generic_plm GPU build: no device-callable density kernel for this "// & + "EQN_OF_STATE / anomaly mode (only ROQUET_RHO, and buggy_Wright without rho_ref, are "// & + "supported); use a CPU build or add a _loc kernel.") + endif +#endif + T25(:) = 0. TS5(:) = 0. S25(:) = 0. @@ -603,12 +636,32 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & if (use_Stanley_eos) then call calculate_density(T5, S5, p5, T25, TS5, S25, r5, EOS, EOSdom_h5, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d(T5,S5,p5,r5,EOS,EOSdom_h5 + ! [,rho_ref]) bit-for-bit: same _loc kernel per point, same unit rescaling. The eos_unity + ! branch is hoisted outside the point loop so the loop body is branch-free (offload-ready). + ! Iterating (i,n) covers exactly the EOSdom_h5 index range that the array call works on. + if (eos_unity) then + do i=Isq,Ieq+1 ; do n=1,5 + r5(i*5+n) = calculate_density_elem_loc(eos_form, T5(i*5+n), S5(i*5+n), p5(i*5+n), & + use_rho_ref, rho_ref) + enddo ; enddo + else + do i=Isq,Ieq+1 ; do n=1,5 + r5(i*5+n) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T5(i*5+n), eos_S_to_ppt*S5(i*5+n), eos_RL2_T2_to_Pa*p5(i*5+n), & + use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif + if (.not. use_rho_ref) u5(:) = r5(:) - rho_ref +#else if (use_rho_ref) then call calculate_density(T5, S5, p5, r5, EOS, EOSdom_h5, rho_ref=rho_ref) else call calculate_density(T5, S5, p5, r5, EOS, EOSdom_h5) u5(:) = r5(:) - rho_ref endif +#endif endif if (use_rho_ref) then @@ -720,11 +773,28 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & if (use_stanley_eos) then call calculate_density(T15, S15, p15, T215, TS15, S215, r15, EOS, EOSdom_q15, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d(T15,S15,p15,r15,EOS,EOSdom_q15 + ! [,rho_ref]) bit-for-bit. Iterating (I,n15) covers exactly the EOSdom_q15 index range. + if (eos_unity) then + do I=Isq,Ieq ; do n15=1,15 + r15(I*15+n15) = calculate_density_elem_loc(eos_form, T15(I*15+n15), S15(I*15+n15), & + p15(I*15+n15), use_rho_ref, rho_ref) + enddo ; enddo + else + do I=Isq,Ieq ; do n15=1,15 + r15(I*15+n15) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T15(I*15+n15), eos_S_to_ppt*S15(I*15+n15), & + eos_RL2_T2_to_Pa*p15(I*15+n15), use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif +#else if (use_rho_ref) then call calculate_density(T15, S15, p15, r15, EOS, EOSdom_q15, rho_ref=rho_ref) else call calculate_density(T15, S15, p15, r15, EOS, EOSdom_q15) endif +#endif endif do I=Isq,Ieq @@ -834,6 +904,22 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & T215(15*HI%isc+1:), TS15(15*HI%isc+1:), S215(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d over the h15 domain + ! bit-for-bit. Iterating (i,n15) covers exactly the r15(15*HI%isc+1:)/EOSdom_h15 range. + if (eos_unity) then + do i=HI%isc,HI%iec ; do n15=1,15 + r15(i*15+n15) = calculate_density_elem_loc(eos_form, T15(i*15+n15), S15(i*15+n15), & + p15(i*15+n15), use_rho_ref, rho_ref) + enddo ; enddo + else + do i=HI%isc,HI%iec ; do n15=1,15 + r15(i*15+n15) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T15(i*15+n15), eos_S_to_ppt*S15(i*15+n15), & + eos_RL2_T2_to_Pa*p15(i*15+n15), use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif +#else if (use_rho_ref) then call calculate_density(T15(15*HI%isc+1:), S15(15*HI%isc+1:), p15(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15, rho_ref=rho_ref) @@ -841,6 +927,7 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & call calculate_density(T15(15*HI%isc+1:), S15(15*HI%isc+1:), p15(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15) endif +#endif endif do i=HI%isc,HI%iec diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index fcfc5af375..39fa0a86bb 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -969,19 +969,22 @@ end function calculate_density_elem_loc !! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion !! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without !! needing access to the private components of EOS_type. -subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) +subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa, & + R_to_kg_m3) type(EOS_type), intent(in) :: EOS !< Equation of state structure integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal density unit R [R m3 kg-1 ~> 1] real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] real, intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1] + real, optional, intent(out) :: R_to_kg_m3 !< Factor converting the internal density unit R to kg m-3 [kg R-1 m-3 ~> 1] form_of_EOS = EOS%form_of_EOS kg_m3_to_R = EOS%kg_m3_to_R C_to_degC = EOS%C_to_degC S_to_ppt = EOS%S_to_ppt RL2_T2_to_Pa = EOS%RL2_T2_to_Pa + if (present(R_to_kg_m3)) R_to_kg_m3 = EOS%R_to_kg_m3 end subroutine get_EOS_form_and_scaling From f56c33fda70998db8cf9fc2379727a00eb621717 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Sat, 18 Jul 2026 16:30:43 +1000 Subject: [PATCH 24/45] int_density_dz_generic_plm: offload phase 1 (vertical integrals) to GPU First quadrature-phase offload of the finite-volume PGF density integrals. The phase-1 vertical-integral loop (-> dpa, intz_dpa) now runs in a !$omp target teams loop collapse(2) over (j,i) when the configuration is device-supported; the host loop is unchanged for every other case. - New host-resolved gate offload_phase1 = (.not.use_stanley_eos) .and. eos_unity .and. use_rho_ref, set only on GPU builds (#ifdef). When false (CPU build, Stanley, non-unity units, or inaccurate-form) the original host loop runs verbatim -- no code duplicated (single else block). - Device kernel: per-(i,j) size-5 quadrature locals T5l/S5l/p5l/r5l (compile-time sized, so no device auto-allocation) reproduce the host point-for-point, call the declare-target dispatcher calculate_density_elem_loc, and write dpa(i,j)/intz_dpa(i,j) via Boole's rule. dpa/intz_dpa are independent across (i,j). - Residency: e, dpa, intz_dpa are already device-resident (caller maps); the region carries no map() clauses and relies on presence. z0pres (a host-computed local) is mapped to/released around the region. The layer edge reconstructions T_t/T_b/S_t/S_b are mapped to() once per call by the caller (PressureForce_FV_Bouss) before the k-loop and released after (guarded to the plm/ppm path); they are read-only and identical for all k. - After the region, target update from(dpa) (and intz_dpa when present) refreshes the host: phases 2 & 3 (still host) read dpa, and the caller's kept "update to(dpa,...)" would otherwise overwrite the device values with a stale host copy. This intermediate copy-back is removed once phases 2/3 also run on device (increment D) and the caller round-trip is deleted (E). Bitwise-safety: the size-5 locals hold the identical per-point values as the old i*5+n slices; rho_anom/dpa/intz_dpa are computed independently per (i,j) from those points, so collapsing the (j,i) loops reorders no FP op; the _loc kernel and Boole coefficients are unchanged. The device path is entered only for the unscaled, accurate, non-Stanley configuration; all others keep the verbatim host arithmetic. Sole caller: PressureForce_FV_Bouss (audited: only call site of int_density_dz_generic_plm), where all no-clause arrays are device-present, so the compiler's implicit map(tofrom:) are runtime no-ops. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/incC_verify.log. Execution proof: NVCOMPILER_ACC_NOTIFY=1 shows the F1L639 kernel launching (tmp_local_artifacts/notify_incC.log). Compile proof: -Minfo reports "Generating NVIDIA GPU code / Loop parallelized across teams collapse(2)" at the teams loop. DEBUG=True redundancy gate clean, 0 'redundant' (tmp_local_artifacts/debug_incC.log). "(Ocean pressure force)" clock 19.49 s/24 (was 20.40), consistent with phase 1 on device; the ~2-3 s target awaits D+E (round-trip removal). Not benchmark-exercised (inspection-only): Stanley, inaccurate-form, and non-unity-scaling paths (host, unchanged). Co-Authored-By: Claude Opus --- src/core/MOM_PressureForce_FV.F90 | 5 ++++ src/core/MOM_density_integrals.F90 | 45 ++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/core/MOM_PressureForce_FV.F90 b/src/core/MOM_PressureForce_FV.F90 index 1dcb872597..11414856f8 100644 --- a/src/core/MOM_PressureForce_FV.F90 +++ b/src/core/MOM_PressureForce_FV.F90 @@ -1304,6 +1304,10 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, !$omp target enter data map(to: tv_tmp, tv_tmp%T, tv_tmp%S) & !$omp if(.not.(use_ALE .and. CS%Recon_Scheme > 0)) + ! The layer edge T/S reconstructions are read on the device by the (offloaded) generic_plm + ! density integrals; map them once here (they are the same for every layer k). + !$omp target enter data map(to: T_t, T_b, S_t, S_b) if(use_ALE .and. CS%Recon_Scheme > 0) + ! The following routine computes the integrals that are needed to ! calculate the pressure gradient force. Linear profiles for T and S are ! assumed when regridding is activated. Otherwise, the previous version @@ -1346,6 +1350,7 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, G%HI, MassWt_u(:,:,k), MassWt_v(:,:,k), & MassWghtInterpVanOnly=CS%MassWghtInterpVanOnly, h_nv=CS%h_nonvanished) enddo + !$omp target exit data map(release: T_t, T_b, S_t, S_b) if(use_ALE .and. CS%Recon_Scheme > 0) !$omp target exit data map(release: tv_tmp, tv_tmp%T, tv_tmp%S) & !$omp if(.not.(use_ALE .and. CS%Recon_Scheme > 0)) else diff --git a/src/core/MOM_density_integrals.F90 b/src/core/MOM_density_integrals.F90 index 311658b704..ccf04fbd5f 100644 --- a/src/core/MOM_density_integrals.F90 +++ b/src/core/MOM_density_integrals.F90 @@ -539,7 +539,12 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & ! EOS unit-rescaling factors, resolved host-side. logical :: eos_unity ! True when all EOS unit-rescaling factors are 1, matching the no-rescale ! fast path of calculate_density_1d. + logical :: do_intz ! present(intz_dpa), hoisted host-side for use inside the device region. + logical :: offload_phase1 ! True on GPU builds when phase 1 (vertical integrals) runs on device. integer :: n15 ! A subgrid-point loop index over the 15 q-point slots in a column. + real :: T5l(5), S5l(5), p5l(5), r5l(5) ! Per-(i,j) size-5 quadrature locals for the device phase-1 + ! kernel (compile-time sized to avoid device auto-allocation). + real :: dz_l ! A per-(i,j) layer thickness for the device phase-1 kernel [Z ~> m]. integer :: Isq, Ieq, Jsq, Jeq, i, j, m, n, pos Isq = HI%IscB ; Ieq = HI%IecB ; Jsq = HI%JscB ; Jeq = HI%JecB @@ -589,7 +594,12 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & eos_RL2_T2_to_Pa, R_to_kg_m3=eos_R_to_kg_m3) eos_unity = (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) + do_intz = present(intz_dpa) ! Hoisted host-side so the device phase-1 region needs no present(). + ! Phase 1 runs on device only on GPU builds and only for the device-supported configuration + ! (non-Stanley, unscaled EOS units, accurate anomaly form); everything else keeps the host path. + offload_phase1 = .false. #ifdef __NVCOMPILER_OPENMP_GPU + offload_phase1 = (.not. use_stanley_eos) .and. eos_unity .and. use_rho_ref ! On GPU builds the non-Stanley density integrals below are evaluated with the device-callable ! dispatcher, which only covers ROQUET_RHO (in-situ density and anomaly) and buggy_Wright ! (in-situ density only). FATAL host-side, before any device region, on an unsupported config. @@ -620,6 +630,40 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & EOSdom_h15(1) = 1 ; EOSdom_h15(2) = 15*(HI%iec-HI%isc+1) ! 1. Compute vertical integrals + if (offload_phase1) then + ! Offload phase 1 (vertical integrals -> dpa, intz_dpa). Per-(i,j) size-5 quadrature locals + ! reproduce the host loop below point-for-point; dpa/intz_dpa are independent across (i,j). + ! Residency: e, dpa, intz_dpa are already device-resident (mapped by the caller); S_t/S_b/T_t/T_b + ! are mapped by the caller before the k-loop; z0pres (a host-computed local) is mapped here. + !$omp target enter data map(to: z0pres) + !$omp target teams loop collapse(2) & + !$omp private(T5l, S5l, p5l, r5l, rho_anom, dz_l, n) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, do_intz, k, eos_form, wt_t, wt_b) + do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 + dz_l = e(i,j,K) - e(i,j,K+1) + do n=1,5 + p5l(n) = -GxRho*((e(i,j,K) - z0pres(i,j)) - 0.25*real(n-1)*dz_l) + S5l(n) = wt_t(n) * S_t(i,j,k) + wt_b(n) * S_b(i,j,k) + T5l(n) = wt_t(n) * T_t(i,j,k) + wt_b(n) * T_b(i,j,k) + enddo + do n=1,5 + r5l(n) = calculate_density_elem_loc(eos_form, T5l(n), S5l(n), p5l(n), use_rho_ref, rho_ref) + enddo + ! Boole's rule for the pressure anomaly change (use_rho_ref accurate form). + rho_anom = C1_90*(7.0*(r5l(1)+r5l(5)) + 32.0*(r5l(2)+r5l(4)) + 12.0*r5l(3)) + dpa(i,j) = G_e*dz_l*rho_anom + if (do_intz) intz_dpa(i,j) = 0.5*G_e*dz_l**2 * & + (rho_anom - C1_90*(16.0*(r5l(4)-r5l(2)) + 7.0*(r5l(5)-r5l(1))) ) + enddo ; enddo + ! Pull the device-computed results back to the host: phases 2 & 3 below (still host) read dpa, + ! and the caller's kept "update to(dpa,...)" would otherwise overwrite the device values with a + ! stale host copy. This intermediate copy-back is removed once phases 2/3 also run on device. + !$omp target update from(dpa) + if (do_intz) then + !$omp target update from(intz_dpa) + endif + !$omp target exit data map(release: z0pres) + else do j=Jsq,Jeq+1 do i = Isq,Ieq+1 dz(i) = e(i,j,K) - e(i,j,K+1) @@ -691,6 +735,7 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & enddo endif enddo ! end loops on j + endif ! offload_phase1 ! 2. Compute horizontal integrals in the x direction if (present(intx_dpa)) then ; do j=HI%jsc,HI%jec From 88d34bc8f00d6bee5aa06327ce7769ea0c06cf17 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Sat, 18 Jul 2026 19:46:14 +1000 Subject: [PATCH 25/45] int_density_dz_generic_plm: offload phases 2 & 3 (horizontal x/y integrals) to GPU Second and third quadrature-phase offload of the finite-volume PGF density integrals. The x-integral (-> intx_dpa) and y-integral (-> inty_dpa) loops now each run in a !$omp target teams loop collapse(2) when the configuration is device-supported; the host loops are unchanged for every other case. - New host-resolved gates offload_phase2/offload_phase3 = offload_phase1's condition, kept as separate flags so phase 3 can be forced off for x-only failure bisection. - Device kernels: per-cell size-15 q-point locals T15l/S15l/p15l/r15l and dz_x_l/dz_y_l (compile-time sized) reproduce the host loops point-for-point, call the declare-target dispatcher calculate_density_elem_loc for all 15 points, and write intx_dpa(I,j)/inty_dpa(i,J) via Boole's rule. Both the hWght>0 mass-weighting branch and its else branch are carried verbatim from the host path (they compute the same set of private corner scalars, so there is no divergent conditional array write). intx_dpa/inty_dpa are independent across cells. - Residency: dpa (written on the device by phase 1) is now read by phases 2/3 on the device with no host round-trip. z0pres is mapped once to bracket all three phases (its enter-data moved above phase 1, release moved after phase 3). e/dpa/intx_dpa/inty_dpa are already device-resident; T_t/T_b/S_t/S_b are mapped by the caller (unchanged from the phase-1 increment). - Phase 1's interim "update from(dpa,intz_dpa)" is replaced by a single consolidated "update from(dpa, intz_dpa, intx_dpa, inty_dpa)" after phase 3, so the caller's kept "update to(dpa,...)" still sees current host values. That consolidated copy-back and the caller round-trip are removed in increment E. Bitwise-safety: the size-15 locals hold the identical per-point values as the old i*15+(m-2)*5 slices; the pressure recurrence over n and the Boole quadrature are copied verbatim and run serially within each cell's thread; intx_dpa/inty_dpa are computed independently per cell, so collapsing the (j,I)/(J,i) loops reorders no FP op. The device path is entered only for the unscaled, accurate, non-Stanley configuration; all others keep the verbatim host arithmetic (host else-paths are byte-identical to the prior increment, verified by diff). Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/incD_verify.log. Compile proof: -Minfo generates GPU kernels F1L750 (x) and F1L946 (y), each "Loop parallelized across teams collapse(2)". DEBUG=True redundancy gate clean, 0 'redundant' (tmp_local_artifacts/debug_incD.log). "(Ocean pressure force)" clock 9.66 s/24 (was 19.49), halved by moving the horizontal integrals on-device; the remaining host round-trip is removed in increment E. Not benchmark-exercised (inspection-only): Stanley, inaccurate-form, and non-unity-scaling host paths. Co-Authored-By: Claude Opus --- src/core/MOM_density_integrals.F90 | 188 ++++++++++++++++++++++++++--- 1 file changed, 173 insertions(+), 15 deletions(-) diff --git a/src/core/MOM_density_integrals.F90 b/src/core/MOM_density_integrals.F90 index ccf04fbd5f..e44d0e4521 100644 --- a/src/core/MOM_density_integrals.F90 +++ b/src/core/MOM_density_integrals.F90 @@ -541,10 +541,14 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & ! fast path of calculate_density_1d. logical :: do_intz ! present(intz_dpa), hoisted host-side for use inside the device region. logical :: offload_phase1 ! True on GPU builds when phase 1 (vertical integrals) runs on device. + logical :: offload_phase2, offload_phase3 ! True on GPU builds when phases 2/3 (x/y integrals) run on device. integer :: n15 ! A subgrid-point loop index over the 15 q-point slots in a column. real :: T5l(5), S5l(5), p5l(5), r5l(5) ! Per-(i,j) size-5 quadrature locals for the device phase-1 ! kernel (compile-time sized to avoid device auto-allocation). real :: dz_l ! A per-(i,j) layer thickness for the device phase-1 kernel [Z ~> m]. + real :: T15l(15), S15l(15), p15l(15), r15l(15) ! Per-(I,j)/(J,i) size-15 q-point locals for the + ! device phase-2/3 kernels (compile-time sized). + real :: dz_x_l(5), dz_y_l(5) ! Per-cell subgrid layer thicknesses for the device x/y kernels [Z ~> m]. integer :: Isq, Ieq, Jsq, Jeq, i, j, m, n, pos Isq = HI%IscB ; Ieq = HI%IecB ; Jsq = HI%JscB ; Jeq = HI%JecB @@ -597,9 +601,11 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & do_intz = present(intz_dpa) ! Hoisted host-side so the device phase-1 region needs no present(). ! Phase 1 runs on device only on GPU builds and only for the device-supported configuration ! (non-Stanley, unscaled EOS units, accurate anomaly form); everything else keeps the host path. - offload_phase1 = .false. + offload_phase1 = .false. ; offload_phase2 = .false. ; offload_phase3 = .false. #ifdef __NVCOMPILER_OPENMP_GPU offload_phase1 = (.not. use_stanley_eos) .and. eos_unity .and. use_rho_ref + offload_phase2 = offload_phase1 ! (settable independently for x-only failure bisection) + offload_phase3 = offload_phase1 ! On GPU builds the non-Stanley density integrals below are evaluated with the device-callable ! dispatcher, which only covers ROQUET_RHO (in-situ density and anomaly) and buggy_Wright ! (in-situ density only). FATAL host-side, before any device region, on an unsupported config. @@ -629,13 +635,17 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & EOSdom_q15(1) = 1 ; EOSdom_q15(2) = 15*(Ieq-Isq+1) EOSdom_h15(1) = 1 ; EOSdom_h15(2) = 15*(HI%iec-HI%isc+1) + ! z0pres is read by all offloaded phases (1-3); map it once here and release it after phase 3. + if (offload_phase1) then + !$omp target enter data map(to: z0pres) + endif + ! 1. Compute vertical integrals if (offload_phase1) then ! Offload phase 1 (vertical integrals -> dpa, intz_dpa). Per-(i,j) size-5 quadrature locals ! reproduce the host loop below point-for-point; dpa/intz_dpa are independent across (i,j). ! Residency: e, dpa, intz_dpa are already device-resident (mapped by the caller); S_t/S_b/T_t/T_b - ! are mapped by the caller before the k-loop; z0pres (a host-computed local) is mapped here. - !$omp target enter data map(to: z0pres) + ! are mapped by the caller before the k-loop; z0pres is mapped just above (bracketing phases 1-3). !$omp target teams loop collapse(2) & !$omp private(T5l, S5l, p5l, r5l, rho_anom, dz_l, n) & !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, do_intz, k, eos_form, wt_t, wt_b) @@ -655,14 +665,8 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & if (do_intz) intz_dpa(i,j) = 0.5*G_e*dz_l**2 * & (rho_anom - C1_90*(16.0*(r5l(4)-r5l(2)) + 7.0*(r5l(5)-r5l(1))) ) enddo ; enddo - ! Pull the device-computed results back to the host: phases 2 & 3 below (still host) read dpa, - ! and the caller's kept "update to(dpa,...)" would otherwise overwrite the device values with a - ! stale host copy. This intermediate copy-back is removed once phases 2/3 also run on device. - !$omp target update from(dpa) - if (do_intz) then - !$omp target update from(intz_dpa) - endif - !$omp target exit data map(release: z0pres) + ! dpa/intz_dpa stay on the device; phases 2 & 3 (also offloaded) read dpa there, and a single + ! consolidated copy-back to the host happens after phase 3. else do j=Jsq,Jeq+1 do i = Isq,Ieq+1 @@ -738,7 +742,73 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & endif ! offload_phase1 ! 2. Compute horizontal integrals in the x direction - if (present(intx_dpa)) then ; do j=HI%jsc,HI%jec + if (present(intx_dpa)) then + if (offload_phase2) then + ! Offload phase 2 (x-integrals -> intx_dpa). Per-(I,j) size-15 q-point locals reproduce the host + ! loop below point-for-point; intx_dpa is independent across (I,j). dpa was written on the device + ! by phase 1 and is read here without a host round-trip. Both hWght branches are carried verbatim. + !$omp target teams loop collapse(2) & + !$omp private(T15l, S15l, p15l, r15l, dz_x_l, intz, hWght, hWghtTop, hL, hR, iDenom, & + !$omp Ttl, Ttr, Tbl, Tbr, Stl, Str, Sbl, Sbr, w_left, w_right, m, n, n15, pos) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, k, eos_form, massWeightToggle, & + !$omp TopWeightToggle, massWeightNVonlyToggle, h_nonvanished, dz_subroundoff, wt_t, wt_b) + do j=HI%jsc,HI%jec ; do I=Isq,Ieq + hWght = massWeightToggle * & + max(0., -bathyT(i,j)-e(i+1,j,K), -bathyT(i+1,j)-e(i,j,K)) + hWghtTop = TopWeightToggle * & + max(0., e(i+1,j,K+1)-e(i,j,1), e(i,j,K+1)-e(i+1,j,1)) + hWght = max(hWght, hWghtTop) + if (((e(i,j,K) - e(i,j,K+1)) > h_nonvanished) .and. ((e(i+1,j,K) - e(i+1,j,K+1)) > h_nonvanished)) then + hWght = massWeightNVonlyToggle * hWght + endif + if (hWght > 0.) then + hL = (e(i,j,K) - e(i,j,K+1)) + dz_subroundoff + hR = (e(i+1,j,K) - e(i+1,j,K+1)) + dz_subroundoff + hWght = hWght * ( (hL-hR)/(hL+hR) )**2 + iDenom = 1./( hWght*(hR + hL) + hL*hR ) + Ttl = ( (hWght*hR)*T_t(i+1,j,k) + (hWght*hL + hR*hL)*T_t(i,j,k) ) * iDenom + Ttr = ( (hWght*hL)*T_t(i,j,k) + (hWght*hR + hR*hL)*T_t(i+1,j,k) ) * iDenom + Tbl = ( (hWght*hR)*T_b(i+1,j,k) + (hWght*hL + hR*hL)*T_b(i,j,k) ) * iDenom + Tbr = ( (hWght*hL)*T_b(i,j,k) + (hWght*hR + hR*hL)*T_b(i+1,j,k) ) * iDenom + Stl = ( (hWght*hR)*S_t(i+1,j,k) + (hWght*hL + hR*hL)*S_t(i,j,k) ) * iDenom + Str = ( (hWght*hL)*S_t(i,j,k) + (hWght*hR + hR*hL)*S_t(i+1,j,k) ) * iDenom + Sbl = ( (hWght*hR)*S_b(i+1,j,k) + (hWght*hL + hR*hL)*S_b(i,j,k) ) * iDenom + Sbr = ( (hWght*hL)*S_b(i,j,k) + (hWght*hR + hR*hL)*S_b(i+1,j,k) ) * iDenom + else + Ttl = T_t(i,j,k) ; Tbl = T_b(i,j,k) ; Ttr = T_t(i+1,j,k) ; Tbr = T_b(i+1,j,k) + Stl = S_t(i,j,k) ; Sbl = S_b(i,j,k) ; Str = S_t(i+1,j,k) ; Sbr = S_b(i+1,j,k) + endif + do m=2,4 + w_left = wt_t(m) ; w_right = wt_b(m) + dz_x_l(m) = (w_left*(e(i,j,K) - e(i,j,K+1))) + (w_right*(e(i+1,j,K) - e(i+1,j,K+1))) + pos = (m-2)*5 + T15l(pos+1) = (w_left*Ttl) + (w_right*Ttr) + T15l(pos+5) = (w_left*Tbl) + (w_right*Tbr) + S15l(pos+1) = (w_left*Stl) + (w_right*Str) + S15l(pos+5) = (w_left*Sbl) + (w_right*Sbr) + p15l(pos+1) = -GxRho * ((w_left*(e(i,j,K)-z0pres(i,j))) + (w_right*(e(i+1,j,K)-z0pres(i+1,j)))) + do n=2,5 + p15l(pos+n) = p15l(pos+n-1) + GxRho*0.25*dz_x_l(m) + enddo + do n=2,4 + S15l(pos+n) = wt_t(n) * S15l(pos+1) + wt_b(n) * S15l(pos+5) + T15l(pos+n) = wt_t(n) * T15l(pos+1) + wt_b(n) * T15l(pos+5) + enddo + enddo + do n15=1,15 + r15l(n15) = calculate_density_elem_loc(eos_form, T15l(n15), S15l(n15), p15l(n15), use_rho_ref, rho_ref) + enddo + intz(1) = dpa(i,j) ; intz(5) = dpa(i+1,j) + do m=2,4 + pos = (m-2)*5 + intz(m) = (G_e*dz_x_l(m)*( C1_90*(7.0*(r15l(pos+1)+r15l(pos+5)) + 32.0*(r15l(pos+2)+r15l(pos+4)) + & + 12.0*r15l(pos+3)) )) + enddo + intx_dpa(I,j) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & + 12.0*intz(3)) + enddo ; enddo + else + do j=HI%jsc,HI%jec do I=Isq,Ieq ! Corner values of T and S ! hWght is the distance measure by which the cell is violation of @@ -863,10 +933,79 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & intx_dpa(I,j) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & 12.0*intz(3)) enddo - enddo ; endif + enddo + endif ! offload_phase2 + endif ! present(intx_dpa) ! 3. Compute horizontal integrals in the y direction - if (present(inty_dpa)) then ; do J=Jsq,Jeq + if (present(inty_dpa)) then + if (offload_phase3) then + ! Offload phase 3 (y-integrals -> inty_dpa). Per-(J,i) size-15 q-point locals reproduce the host + ! loop below point-for-point; inty_dpa is independent across (J,i). dpa (device, from phase 1) is + ! read here without a host round-trip. Both hWght branches are carried verbatim. + !$omp target teams loop collapse(2) & + !$omp private(T15l, S15l, p15l, r15l, dz_y_l, intz, hWght, hWghtTop, hL, hR, iDenom, & + !$omp Ttl, Ttr, Tbl, Tbr, Stl, Str, Sbl, Sbr, w_left, w_right, m, n, n15, pos) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, k, eos_form, massWeightToggle, & + !$omp TopWeightToggle, massWeightNVonlyToggle, h_nonvanished, dz_subroundoff, wt_t, wt_b) + do J=Jsq,Jeq ; do i=HI%isc,HI%iec + hWght = massWeightToggle * & + max(0., -bathyT(i,j)-e(i,j+1,K), -bathyT(i,j+1)-e(i,j,K)) + hWghtTop = TopWeightToggle * & + max(0., e(i,j+1,K+1)-e(i,j,1), e(i,j,K+1)-e(i,j+1,1)) + hWght = max(hWght, hWghtTop) + if (((e(i,j,K) - e(i,j,K+1)) > h_nonvanished) .and. ((e(i,j+1,K) - e(i,j+1,K+1)) > h_nonvanished)) then + hWght = massWeightNVonlyToggle * hWght + endif + if (hWght > 0.) then + hL = (e(i,j,K) - e(i,j,K+1)) + dz_subroundoff + hR = (e(i,j+1,K) - e(i,j+1,K+1)) + dz_subroundoff + hWght = hWght * ( (hL-hR)/(hL+hR) )**2 + iDenom = 1./( hWght*(hR + hL) + hL*hR ) + Ttl = ( (hWght*hR)*T_t(i,j+1,k) + (hWght*hL + hR*hL)*T_t(i,j,k) ) * iDenom + Ttr = ( (hWght*hL)*T_t(i,j,k) + (hWght*hR + hR*hL)*T_t(i,j+1,k) ) * iDenom + Tbl = ( (hWght*hR)*T_b(i,j+1,k) + (hWght*hL + hR*hL)*T_b(i,j,k) ) * iDenom + Tbr = ( (hWght*hL)*T_b(i,j,k) + (hWght*hR + hR*hL)*T_b(i,j+1,k) ) * iDenom + Stl = ( (hWght*hR)*S_t(i,j+1,k) + (hWght*hL + hR*hL)*S_t(i,j,k) ) * iDenom + Str = ( (hWght*hL)*S_t(i,j,k) + (hWght*hR + hR*hL)*S_t(i,j+1,k) ) * iDenom + Sbl = ( (hWght*hR)*S_b(i,j+1,k) + (hWght*hL + hR*hL)*S_b(i,j,k) ) * iDenom + Sbr = ( (hWght*hL)*S_b(i,j,k) + (hWght*hR + hR*hL)*S_b(i,j+1,k) ) * iDenom + else + Ttl = T_t(i,j,k) ; Tbl = T_b(i,j,k) ; Ttr = T_t(i,j+1,k) ; Tbr = T_b(i,j+1,k) + Stl = S_t(i,j,k) ; Sbl = S_b(i,j,k) ; Str = S_t(i,j+1,k) ; Sbr = S_b(i,j+1,k) + endif + do m=2,4 + w_left = wt_t(m) ; w_right = wt_b(m) + dz_y_l(m) = (w_left*(e(i,j,K) - e(i,j,K+1))) + (w_right*(e(i,j+1,K) - e(i,j+1,K+1))) + pos = (m-2)*5 + T15l(pos+1) = (w_left*Ttl) + (w_right*Ttr) + T15l(pos+5) = (w_left*Tbl) + (w_right*Tbr) + S15l(pos+1) = (w_left*Stl) + (w_right*Str) + S15l(pos+5) = (w_left*Sbl) + (w_right*Sbr) + p15l(pos+1) = -GxRho * ((w_left*(e(i,j,K)-z0pres(i,j))) + (w_right*(e(i,j+1,K)-z0pres(i,j+1)))) + do n=2,5 + p15l(pos+n) = p15l(pos+n-1) + GxRho*0.25*dz_y_l(m) + enddo + do n=2,4 + S15l(pos+n) = wt_t(n) * S15l(pos+1) + wt_b(n) * S15l(pos+5) + T15l(pos+n) = wt_t(n) * T15l(pos+1) + wt_b(n) * T15l(pos+5) + enddo + enddo + do n15=1,15 + r15l(n15) = calculate_density_elem_loc(eos_form, T15l(n15), S15l(n15), p15l(n15), use_rho_ref, rho_ref) + enddo + intz(1) = dpa(i,j) ; intz(5) = dpa(i,j+1) + do m=2,4 + pos = (m-2)*5 + intz(m) = (G_e*dz_y_l(m)*( C1_90*(7.0*(r15l(pos+1)+r15l(pos+5)) + & + 32.0*(r15l(pos+2)+r15l(pos+4)) + & + 12.0*r15l(pos+3)) )) + enddo + inty_dpa(i,J) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & + 12.0*intz(3)) + enddo ; enddo + else + do J=Jsq,Jeq do i=HI%isc,HI%iec ! Corner values of T and S ! hWght is the distance measure by which the cell is violation of @@ -998,7 +1137,26 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & inty_dpa(i,J) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & 12.0*intz(3)) enddo - enddo ; endif + enddo + endif ! offload_phase3 + endif ! present(inty_dpa) + + ! Consolidated device->host copy-back: the device now holds dpa/intz_dpa (phase 1) and + ! intx_dpa/inty_dpa (phases 2/3). Refresh the host so the caller's kept "update to(dpa,...)" + ! does not overwrite them with stale host values. Removed in increment E with the round-trip. + if (offload_phase1) then + !$omp target update from(dpa) + if (do_intz) then + !$omp target update from(intz_dpa) + endif + !$omp target exit data map(release: z0pres) + endif + if (offload_phase2 .and. present(intx_dpa)) then + !$omp target update from(intx_dpa) + endif + if (offload_phase3 .and. present(inty_dpa)) then + !$omp target update from(inty_dpa) + endif end subroutine int_density_dz_generic_plm From 0d216aa0e83613dc0239b1f6bc6b67d50f2babb1 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 07:25:05 +1000 Subject: [PATCH 26/45] PressureForce_FV_Bouss: consume the PLM density integrals on-device (increment E, caller side) Removes the caller half of the int_density_dz_generic_plm device round-trip: the four integral outputs (dpa, intz_dpa, intx_dpa, inty_dpa) are now consumed entirely in the caller's own device regions (interface-pressure accumulation and the PFu/PFv loops), so the host no longer pushes them to the device before use. - New host-resolved gate idz_on_device, set only on GPU builds (#ifdef __NVCOMPILER_OPENMP_GPU), that mirrors int_density_dz_generic_plm's internal offload_phase1 condition EXACTLY on the generic_plm call path: use_EOS .and. use_ALE .and. Recon in {1,3} .and. .not.use_stanley_pgf .and. .not.use_inaccurate_pgf_rho_anom .and. the four EOS unit-rescale factors == 1 (queried via get_EOS_form_and_scaling). When true, the routine leaves its outputs device-resident; when false (CPU build, Recon scheme 2, Stanley, inaccurate-form, non-unity units) the host round-trip is kept verbatim. - Drop the per-call "defensive" update to(dpa,intx_dpa,inty_dpa,intz_dpa) when idz_on_device: the generic_plm kernels already wrote them on the device, so the copy would clobber them with a stale host copy. Kept (guarded) for the host path. - Narrow the full update from(e): e is built on the device, and under idz_on_device the blanket ALE pull is skipped. The pull is retained for its genuine host consumers -- host-run int_density, the correction/reset-intxpa reference-interface searches, the MassWt diagnostics, and (this is the stale-host fix) the SAL/tides "new answers" blocks that read e(:,:,1) and read-modify e(:,:,K) on the host. - Add a guarded update from(intx_dpa,inty_dpa) at routine exit for the SAL/tides horizontal-gradient diagnostics (the only host readers of those two arrays), guarded to idz_on_device .and. .not.bq_sal_tides .and. (calculate_SAL .or. tides). Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/incE_fix_verify.log. The SAL/tides and correction/reset guard terms are false in benchmark_ALE, so this is inspection-only for those configs; "(Ocean pressure force)" clock ~2.3 s/24 (was ~20). Note the SAL/tides-new-answers GPU write-back (host e/pa mods pushed back to the device) is a separate, pre-existing gap -- idz_on_device is the fully-GPU, non-SAL/tides PGF path. Remaining gates: np>=2 multi-GPU verify (1 V100 here) and a plain-CPU compile of the #else branches. Co-Authored-By: Claude Opus --- src/core/MOM_PressureForce_FV.F90 | 48 +++++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/src/core/MOM_PressureForce_FV.F90 b/src/core/MOM_PressureForce_FV.F90 index 11414856f8..ed136a802c 100644 --- a/src/core/MOM_PressureForce_FV.F90 +++ b/src/core/MOM_PressureForce_FV.F90 @@ -19,6 +19,7 @@ module MOM_PressureForce_FV use MOM_variables, only : thermo_var_ptrs, accel_diag_ptrs use MOM_verticalGrid, only : verticalGrid_type use MOM_EOS, only : calculate_density, calculate_spec_vol, EOS_domain +use MOM_EOS, only : get_EOS_form_and_scaling use MOM_density_integrals, only : int_density_dz, int_specific_vol_dp use MOM_density_integrals, only : int_density_dz_generic_plm, int_density_dz_generic_ppm use MOM_density_integrals, only : int_spec_vol_dp_generic_plm @@ -1101,6 +1102,12 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, integer, dimension(2) :: EOSdom_u ! The i-computational domain for the equation of state at u-velocity points integer, dimension(2) :: EOSdom_v ! The i-computational domain for the equation of state at v-velocity points integer :: EOSdom2d(2,2) ! The 2D compute domain for the equation of state + logical :: idz_on_device ! If true, int_density_dz_generic_plm leaves its outputs (dpa, intz_dpa, + ! intx_dpa, inty_dpa) device-resident, so the caller neither pushes e to the + ! host for it nor pulls the integrals back (GPU builds, supported config). + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, eos_R_to_kg_m3 + ! EOS unit-rescaling factors, used only to test eos_unity for idz_on_device. + integer :: eos_form_pf ! EOS form id (from get_EOS_form_and_scaling; not otherwise used here). integer :: is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, nkmb integer :: i, j, k, m @@ -1295,10 +1302,34 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, ! Calculate 4 integrals through the layer that are required in the ! subsequent calculation. + + ! Determine whether int_density_dz_generic_plm will leave its outputs (dpa, intz_dpa, intx_dpa, + ! inty_dpa) on the device: GPU build, ALE PLM reconstruction (scheme 1 or 3), non-Stanley, + ! accurate anomaly form, and unscaled EOS units. This mirrors that routine's internal offload + ! gate so the caller can drop the host round-trip for exactly that case. It is k-independent. + idz_on_device = .false. +#ifdef __NVCOMPILER_OPENMP_GPU + if (use_EOS .and. use_ALE .and. (CS%Recon_Scheme == 1 .or. CS%Recon_Scheme == 3) .and. & + (.not. CS%use_stanley_pgf) .and. (.not. CS%use_inaccurate_pgf_rho_anom)) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form_pf, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa, R_to_kg_m3=eos_R_to_kg_m3) + idz_on_device = (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) + endif +#endif + !$omp target enter data map(alloc: dpa, intx_dpa, inty_dpa, intz_dpa) if (use_EOS) then - !$omp target update from(e) if( (use_ALE .and. CS%Recon_Scheme > 0) .or. & + ! e is pulled to the host only for host-side consumers: a host-run int_density (when not + ! idz_on_device -- e.g. Recon scheme 2, Stanley, non-unity units), the MassWt diagnostics, the + ! correction/reset-intxpa reference-interface searches, or the SAL/tides "new answers" blocks + ! (which read e(:,:,1) and read-modify e(:,:,K) on the host). When idz_on_device would otherwise + ! suppress the pull, those SAL/tides host reads would see stale e, so keep the pull for them. + !$omp target update from(e) if( (use_ALE .and. CS%Recon_Scheme > 0 .and. (.not. idz_on_device)) .or. & + !$omp CS%correction_intxpa .or. CS%reset_intxpa_integral .or. & + !$omp (CS%calculate_SAL .and. CS%tides_answer_date>20250131) .or. & + !$omp (CS%tides .and. CS%tides_answer_date>20250131) .or. & !$omp (CS%id_MassWt_u > 0) .or. (CS%id_MassWt_v > 0)) ! transfer tv_tmp%* only if int_density_dz is called !$omp target enter data map(to: tv_tmp, tv_tmp%T, tv_tmp%S) & @@ -1331,8 +1362,13 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, MassWghtInterp=CS%MassWghtInterp, Z_0p=Z_0p, & MassWghtInterpVanOnly=CS%MassWghtInterpVanOnly, h_nv=dz_nonvanished) endif - ! defensive update - not sure if it works - !$omp target update to(dpa, intx_dpa, inty_dpa, intz_dpa) + ! Push the host-computed integrals to the device only when int_density ran on the host + ! (Recon scheme 2 / Stanley / non-unity units). When idz_on_device the generic_plm kernels + ! already wrote them on the device, so this copy is both unnecessary and would clobber them + ! with a stale host copy. + if (.not. idz_on_device) then + !$omp target update to(dpa, intx_dpa, inty_dpa, intz_dpa) + endif else call int_density_dz(tv_tmp%T(:,:,k), tv_tmp%S(:,:,k), e(:,:,K), e(:,:,K+1), & rho_ref, rho0_int_density, GV%g_Earth, G%HI, tv%eqn_of_state, US, dpa(:,:,k), & @@ -1954,6 +1990,12 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, !$omp target exit data map(delete: Z_0p) if (use_EOS) + ! When idz_on_device, intx_dpa/inty_dpa were written on the device and are not otherwise copied + ! back. The SAL/tides gradient diagnostics below read them on the host, so pull them back before + ! the device mapping is deleted -- guarded to exactly that host-consuming configuration. + !$omp target update from(intx_dpa, inty_dpa) & + !$omp if(idz_on_device .and. (.not. CS%bq_sal_tides) .and. (CS%calculate_SAL .or. CS%tides)) + !$omp target exit data & !$omp map(delete: pa, dpa) & !$omp map(delete: intx_pa, inty_pa, intx_dpa, inty_dpa, intz_dpa) From c59bbb18d37999519d040be5c3d74637d69a559d Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 07:25:06 +1000 Subject: [PATCH 27/45] int_density_dz_generic_plm: leave integral outputs device-resident (increment E) Removes the routine half of the device round-trip. Phases 1-3 already compute dpa/intz_dpa/intx_dpa/inty_dpa on the device; this drops the consolidated device->host copy-back at the end of the routine. The sole caller (PressureForce_FV_Bouss) now consumes all four outputs in its own device regions (prior commit), so no host refresh is needed. Only the per-call z0pres scratch release is kept. Also rewrites the offload_phase2/offload_phase3 comment: they must track offload_phase1 now that the round-trip is gone. Decoupling a phase would leave its host-computed integral stranded on the host while the caller (which mirrors this gate as idz_on_device and no longer copies back) reads a stale device copy. The "settable independently for x-only failure bisection" affordance no longer holds. Bitwise-safety: no arithmetic changes; only OpenMP data-movement directives are removed. On CPU builds (offload_phase1/2/3 all false) the removed blocks were already skipped, so host builds are unaffected. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq, PLM) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/incE_fix_verify.log. Co-Authored-By: Claude Opus --- src/core/MOM_density_integrals.F90 | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/src/core/MOM_density_integrals.F90 b/src/core/MOM_density_integrals.F90 index e44d0e4521..c039d44cfb 100644 --- a/src/core/MOM_density_integrals.F90 +++ b/src/core/MOM_density_integrals.F90 @@ -604,7 +604,11 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & offload_phase1 = .false. ; offload_phase2 = .false. ; offload_phase3 = .false. #ifdef __NVCOMPILER_OPENMP_GPU offload_phase1 = (.not. use_stanley_eos) .and. eos_unity .and. use_rho_ref - offload_phase2 = offload_phase1 ! (settable independently for x-only failure bisection) + ! These must track offload_phase1: the caller (MOM_PressureForce_FV) mirrors this gate to decide + ! idz_on_device and no longer copies the outputs back to the host, so decoupling a phase would + ! leave its host-computed integral (intx/inty_dpa) stranded on the host while the caller reads a + ! stale device copy. Do not set them independently now that the device->host round-trip is gone. + offload_phase2 = offload_phase1 offload_phase3 = offload_phase1 ! On GPU builds the non-Stanley density integrals below are evaluated with the device-callable ! dispatcher, which only covers ROQUET_RHO (in-situ density and anomaly) and buggy_Wright @@ -1141,22 +1145,12 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & endif ! offload_phase3 endif ! present(inty_dpa) - ! Consolidated device->host copy-back: the device now holds dpa/intz_dpa (phase 1) and - ! intx_dpa/inty_dpa (phases 2/3). Refresh the host so the caller's kept "update to(dpa,...)" - ! does not overwrite them with stale host values. Removed in increment E with the round-trip. + ! The outputs dpa/intz_dpa/intx_dpa/inty_dpa are left device-resident: the caller consumes them + ! entirely in its own device regions (interface-pressure accumulation and the PFu/PFv loops), so + ! there is no device->host copy-back here. Only the per-call z0pres scratch mapping is released. if (offload_phase1) then - !$omp target update from(dpa) - if (do_intz) then - !$omp target update from(intz_dpa) - endif !$omp target exit data map(release: z0pres) endif - if (offload_phase2 .and. present(intx_dpa)) then - !$omp target update from(intx_dpa) - endif - if (offload_phase3 .and. present(inty_dpa)) then - !$omp target update from(inty_dpa) - endif end subroutine int_density_dz_generic_plm From 5b5cc1d1cd327c18b618846ec41bc6c390f90240 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 08:09:54 +1000 Subject: [PATCH 28/45] kappa_shear: run the h-point column solver (Calculate_kappa_shear) on the GPU Mirrors the Calc_kappa_shear_vertex offload for the tracer-point driver used when VERTEX_SHEAR=False (the plain benchmark case). All the shared machinery was already in place -- the three column workers are declare target, the density-derivs element dispatcher covers WRIGHT and ROQUET_RHO, eos_form/scaling are hoisted host-side in this driver too, and GPU_nk_max sizes the per-column scratch -- so this is only the final loop-flip. - Hoist thickness_to_dz out of the j-loop to a whole-domain dz_3d (halo_size=0), and precompute surface_pres_2d from the (possibly-null) p_surf pointer on the host, ahead of the device region. - Flip the column loop from !$OMP parallel do (host) to !$omp target teams loop collapse(2) over (j,i). The per-J 2-D slabs (h_2d/u_2d/...) are gone: the device loop reads the mapped 3-D inputs (h, u_in, v_in, tv%T, tv%S, dz_3d) directly -- the h-point driver has no vertex interpolation, so no slab staging is needed. Results stage in device-only kappa_3d/tke_3d (was the per-J kappa_2d/tke_2d); a masked do concurrent then writes kappa_io/tke_io/kv_io. - Per-column private scratch (Idz/h_lay/.../kappa/tke/.../kc/kf) is GPU_nk_max-sized under #ifdef __NVCOMPILER_OPENMP_GPU (compile-time constant -> device stack, not the serialized device-heap path); SZK_(GV) on CPU builds. Bitwise-safe (no size()/bounds intrinsics on these in the routine). - Residency: h host-authoritative -> map(to:)+update to; u_in/v_in (set_diffusivity's u_h/v_h) and dz_3d are fresh host locals -> map(to:); CS shallow map(to:); outputs map(to:)+update to (preserve host halos) then update from; diag_* mapped after their conditional host zeroing, copy-back guarded by the same predicate. tv%T/tv%S are read directly on the device here (this path passes raw tv, no full_convection filtering, unlike the vertex driver's fresh T_f/S_f): the T/S data is persistently mapped via MOM's CS%tv but on a different descriptor, so this routine's tv dummy is explicitly mapped map(to: tv, tv%T, tv%S) + update to (the MOM_set_viscosity idiom) so the device can resolve tv%T/tv%S -- update to alone does not attach the descriptor. Verification: benchmark (VERTEX_SHEAR=False, EQN_OF_STATE=WRIGHT, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/kshpoint_verify2.log. Execution proof: NVCOMPILER_ACC_NOTIFY shows the device kernel calculate_kappa_shear line 328 (target teams loop) launching grid=64800 (360x180 columns) once per call, plus the write-back do concurrent at line 469. "(Ocean kappa_shear)" ~0.78 s/6 hits (the benchmark kappa_shear workload is small; the kernel is teams-only block=(1,1,1), the same occupancy headroom the vertex driver carries). Untested-but-code-identical: the non-unity unit-rescale branch and the non-temperature (rho) branch. Deferred pre-merge gate: >=2-GPU / np>=2 verify (this node has 1 V100). Co-Authored-By: Claude Opus Co-authored-by: Utheri Wagura --- .../vertical/MOM_kappa_shear.F90 | 197 ++++++++++++------ 1 file changed, 129 insertions(+), 68 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 7ca8143fba..b5306d6b88 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -183,15 +183,24 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & diag_S2_init, & ! Diagnostic of S2 as provided to this routine [T-2 ~> s-2] diag_N2_mean, & ! Diagnostic of N2 averaged over the timestep applied [T-2 ~> s-2] diag_S2_mean ! Diagnostic of S2 averaged over the timestep applied [T-2 ~> s-2] - real, dimension(SZI_(G),SZK_(GV)) :: & - h_2d, & ! A 2-D version of h [H ~> m or kg m-2]. - dz_2d, & ! Vertical distance between interface heights [Z ~> m]. - u_2d, v_2d, & ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1]. - T_2d, S_2d, rho_2d ! 2-D versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)+1) :: & - kappa_2d, & ! 2-D version of kappa_io [H Z T-1 ~> m2 s-1 or Pa s] - tke_2d ! 2-D version tke_io [Z2 T-2 ~> m2 s-2]. + ! GPU port: the per-J 2-D slabs (h_2d/dz_2d/u_2d/v_2d/T_2d/S_2d/rho_2d) and the per-J kappa_2d/ + ! tke_2d are gone; the device column loop reads the mapped 3-D inputs (h, u_in, v_in, tv%T, tv%S, + ! dz_3d) directly and stages its results in kappa_3d/tke_3d (was the per-J kappa_2d/tke_2d). + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: & + dz_3d ! Vertical distance between interface heights [Z ~> m]. + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: & + kappa_3d, & ! Device staging array for the columns' averaged kappa [H Z T-1 ~> m2 s-1 or Pa s] + tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. + real, dimension(SZI_(G),SZJ_(G)) :: & + surface_pres_2d ! The surface pressure at tracer points [R L2 T-2 ~> Pa]. + ! In GPU builds the per-column private scratch has compile-time-constant sizes so each device + ! thread gets stack ("local memory") arrays; runtime-sized privates are device-heap allocated + ! per column, which exhausts the default heap and serializes on the allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] dz_lay, & ! The geometric layer thickness in height units [Z ~> m] @@ -200,7 +209,11 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & T0xdz, & ! The initial temperature times thickness [C H ~> degC m or degC kg m-2] or if ! temperature is not a state variable, the density times thickness [R H ~> kg m-2 or kg2 m-5] S0xdz ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif kappa, & ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s] tke, & ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2]. kappa_avg, & ! The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -219,10 +232,18 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & logical :: use_temperature ! If true, temperature and salinity have been ! allocated and are being used as state variables. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max+1) :: kc ! The index map between the original +#else integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original +#endif ! interfaces and the interfaces with massless layers ! merged into nearby massive layers. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: kf ! The fractional weight of interface kc+1 for +#else real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for +#endif ! interpolating back to the original index space [nondim]. integer :: is, ie, js, je, i, j, k, nz, nzc integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. @@ -256,30 +277,62 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & if (CS%id_N2_mean>0) diag_N2_mean(:,:,:) = 0.0 if (CS%id_S2_mean>0) diag_S2_mean(:,:,:) = 0.0 - !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,tv,G,GV,US, & - !$OMP CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io, & - !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & - !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) - do j=js,je - - ! Convert layer thicknesses into geometric thickness in height units. - call thickness_to_dz(h, tv, dz_2d, j, G, GV) - - do k=1,nz ; do i=is,ie - h_2d(i,k) = h(i,j,k) - u_2d(i,k) = u_in(i,j,k) ; v_2d(i,k) = v_in(i,j,k) + ! GPU port: interpolate p_surf to tracer points on the host, ahead of the device column region + ! (p_surf is a possibly-unassociated pointer, simpler kept off the device). Verbatim from the + ! former in-column expression, evaluated for every tracer point; land values are never read. + surface_pres_2d(:,:) = 0.0 + if (associated(p_surf)) then + do j=js,je ; do i=is,ie + surface_pres_2d(i,j) = p_surf(i,j) enddo ; enddo - if (use_temperature) then ; do k=1,nz ; do i=is,ie - T_2d(i,k) = tv%T(i,j,k) ; S_2d(i,k) = tv%S(i,j,k) - enddo ; enddo ; else ; do k=1,nz ; do i=is,ie - rho_2d(i,k) = GV%Rlay(k) ! Could be tv%Rho(i,j,k) ? - enddo ; enddo ; endif + endif + + ! Convert layer thicknesses into geometric thickness in height units, over the whole compute + ! domain (was a per-J call inside the loop; hoisted so dz_3d can be mapped once). + call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=0) + + ! --- GPU port: run the per-column solver on the device. The columns iterate as a target teams + ! loop collapsed over (j,i); every piece of per-column scratch is private (the declare-target + ! solver's own locals are automatically private per device thread). h is host-authoritative in + ! the (host) diabatic stack -> refresh. tv%T/tv%S are persistently mapped (MOM.F90) but this + ! non-full_convection path reads the diabatic-mutated fields directly, so refresh them too + ! (a map(to:) on an already-present object copies nothing). u_in/v_in (set_diffusivity's + ! u_h/v_h) and dz_3d are fresh host locals. CS is all scalars plus a diag pointer that is never + ! dereferenced in device code, so a per-call shallow map(to:) suffices. kappa_io/tke_io/kv_io + ! are mapped to: (not alloc) and refreshed so the host values (halos, and kappa_io's previous + ! contents outside the compute domain) survive the full-array update from below. diag_* are + ! mapped after their (conditional) host zeroing. kappa_3d/tke_3d are device-only staging for + ! what was the per-J kappa_2d/tke_2d. + !$omp target enter data map(to: h) + !$omp target update to(h) + ! tv%T/tv%S are read on the device in the column setup below. The T/S *data* is persistently + ! mapped via MOM's CS%tv (MOM.F90), but that attaches it to a different descriptor; this routine's + ! tv dummy must be mapped here so the device can resolve tv%T/tv%S (refcount bump + attach, no + ! copy), then update to refreshes the diabatic-mutated values. (set_viscosity uses this idiom.) + if (use_temperature) then + !$omp target enter data map(to: tv, tv%T, tv%S) + !$omp target update to(tv%T, tv%S) + endif + !$omp target enter data map(to: u_in, v_in, dz_3d) + !$omp target update to(u_in, v_in) + !$omp target enter data map(to: CS) + !$omp target enter data map(to: surface_pres_2d) + !$omp target enter data map(to: kappa_io, tke_io, kv_io) + !$omp target update to(kappa_io, tke_io, kv_io) + !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target enter data map(alloc: kappa_3d, tke_3d) !--------------------------------------- ! Work on each column. !--------------------------------------- + !$omp target teams loop collapse(2) & + !$omp private(nzc, kc, kf, Idz, h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, dz_in_lay, & + !$omp f2, surface_pres, kappa, tke, kappa_avg, tke_avg, N2_init, S2_init, & + !$omp N2_mean, S2_mean, k) & + !$omp firstprivate(nz, dt, k0dt, dz_massless, use_temperature, eos_form, & + !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + do j=js,je do i=is,ie ; if (G%mask2dT(i,j) > 0.0) then - ! call cpu_clock_begin(id_clock_setup) ! Store a transposed version of the initial arrays. ! Any elimination of massless layers would occur here. @@ -291,25 +344,20 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & T0xdz(k) = 0.0 ; S0xdz(k) = 0.0 ! Add a new layer if this one has mass. -! if ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless)) nzc = nzc+1 if ((k>CS%nkml) .and. (h_lay(nzc) > 0.0) .and. & - (h_2d(i,k) > dz_massless)) nzc = nzc+1 - - ! Only merge clusters of massless layers. -! if ((h_lay(nzc) > dz_massless) .or. & -! ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless))) nzc = nzc+1 + (h(i,j,k) > dz_massless)) nzc = nzc+1 kc(k) = nzc - h_lay(nzc) = h_lay(nzc) + h_2d(i,k) - dz_lay(nzc) = dz_lay(nzc) + dz_2d(i,k) - u0xdz(nzc) = u0xdz(nzc) + u_2d(i,k)*h_2d(i,k) - v0xdz(nzc) = v0xdz(nzc) + v_2d(i,k)*h_2d(i,k) + h_lay(nzc) = h_lay(nzc) + h(i,j,k) + dz_lay(nzc) = dz_lay(nzc) + dz_3d(i,j,k) + u0xdz(nzc) = u0xdz(nzc) + u_in(i,j,k)*h(i,j,k) + v0xdz(nzc) = v0xdz(nzc) + v_in(i,j,k)*h(i,j,k) if (use_temperature) then - T0xdz(nzc) = T0xdz(nzc) + T_2d(i,k)*h_2d(i,k) - S0xdz(nzc) = S0xdz(nzc) + S_2d(i,k)*h_2d(i,k) + T0xdz(nzc) = T0xdz(nzc) + tv%T(i,j,k)*h(i,j,k) + S0xdz(nzc) = S0xdz(nzc) + tv%S(i,j,k)*h(i,j,k) else - T0xdz(nzc) = T0xdz(nzc) + rho_2d(i,k)*h_2d(i,k) - S0xdz(nzc) = S0xdz(nzc) + rho_2d(i,k)*h_2d(i,k) + T0xdz(nzc) = T0xdz(nzc) + GV%Rlay(k)*h(i,j,k) + S0xdz(nzc) = S0xdz(nzc) + GV%Rlay(k)*h(i,j,k) endif enddo kc(nz+1) = nzc+1 @@ -319,28 +367,28 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & ! Now determine kf, the fractional weight of interface kc when ! interpolating between interfaces kc and kc+1. - kf(1) = 0.0 ; dz_in_lay = h_2d(i,1) + kf(1) = 0.0 ; dz_in_lay = h(i,j,1) do k=2,nz if (kc(k) > kc(k-1)) then - kf(k) = 0.0 ; dz_in_lay = h_2d(i,k) + kf(k) = 0.0 ; dz_in_lay = h(i,j,k) else - kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(i,k) + kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h(i,j,k) endif enddo kf(nz+1) = 0.0 else do k=1,nz - h_lay(k) = h_2d(i,k) - dz_lay(k) = dz_2d(i,k) - u0xdz(k) = u_2d(i,k)*h_lay(k) ; v0xdz(k) = v_2d(i,k)*h_lay(k) + h_lay(k) = h(i,j,k) + dz_lay(k) = dz_3d(i,j,k) + u0xdz(k) = u_in(i,j,k)*h_lay(k) ; v0xdz(k) = v_in(i,j,k)*h_lay(k) enddo if (use_temperature) then do k=1,nz - T0xdz(k) = T_2d(i,k)*h_lay(k) ; S0xdz(k) = S_2d(i,k)*h_lay(k) + T0xdz(k) = tv%T(i,j,k)*h_lay(k) ; S0xdz(k) = tv%S(i,j,k)*h_lay(k) enddo else do k=1,nz - T0xdz(k) = rho_2d(i,k)*h_lay(k) ; S0xdz(k) = rho_2d(i,k)*h_lay(k) + T0xdz(k) = GV%Rlay(k)*h_lay(k) ; S0xdz(k) = GV%Rlay(k)*h_lay(k) enddo endif nzc = nz @@ -348,12 +396,9 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & endif f2 = 0.25 * ((G%Coriolis2Bu(I,J) + G%Coriolis2Bu(I-1,J-1)) + & (G%Coriolis2Bu(I,J-1) + G%Coriolis2Bu(I-1,J))) - surface_pres = 0.0 ; if (associated(p_surf)) surface_pres = p_surf(i,j) - - ! ---------------------------------------------------- I_Ld2_1d, dz_Int_1d + surface_pres = surface_pres_2d(i,j) ! Set the initial guess for kappa, here defined at interfaces. - ! ---------------------------------------------------- do K=1,nzc+1 ; kappa(K) = CS%kappa_seed ; enddo call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & @@ -363,15 +408,14 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & use_temperature) - ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_2d(i,K) = kappa_avg(K) + kappa_3d(i,j,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_2d(i,K) = tke(K) + tke_3d(i,j,K) = tke(K) else - tke_2d(i,K) = tke_avg(K) + tke_3d(i,j,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -389,11 +433,11 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & else do K=1,nz+1 if (kf(K) == 0.0) then - kappa_2d(i,K) = kappa_avg(kc(K)) - tke_2d(i,K) = tke_avg(kc(K)) + kappa_3d(i,j,K) = kappa_avg(kc(K)) + tke_3d(i,j,K) = tke_avg(kc(K)) else - kappa_2d(i,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) - tke_2d(i,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) + kappa_3d(i,j,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) + tke_3d(i,j,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) endif enddo do K=1,nz+1 @@ -414,21 +458,38 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & endif enddo endif - ! call cpu_clock_end(id_clock_setup) else ! Land points, still inside the i-loop. do K=1,nz+1 - kappa_2d(i,K) = 0.0 ; tke_2d(i,K) = 0.0 + kappa_3d(i,j,K) = 0.0 ; tke_3d(i,j,K) = 0.0 enddo endif ; enddo ! i-loop + enddo ! end of j-loop - do K=1,nz+1 ; do i=is,ie - kappa_io(i,j,K) = G%mask2dT(i,j) * kappa_2d(i,K) - tke_io(i,j,K) = G%mask2dT(i,j) * tke_2d(i,K) - kv_io(i,j,K) = ( G%mask2dT(i,j) * kappa_2d(i,K) ) * CS%Prandtl_turb + ! Store the columns' results back in the output arrays (masked), on the device. + do concurrent (K=1:nz+1, j=js:je, i=is:ie) + kappa_io(i,j,K) = G%mask2dT(i,j) * kappa_3d(i,j,K) + tke_io(i,j,K) = G%mask2dT(i,j) * tke_3d(i,j,K) + kv_io(i,j,K) = ( G%mask2dT(i,j) * kappa_3d(i,j,K) ) * CS%Prandtl_turb + enddo - enddo ; enddo + ! The checksums and post_data below are still on the host; copy back the outputs (and, guarded + ! by the same conditions as their consumers, the diag_* arrays), then mirror the enter data. + !$omp target update from(kappa_io, tke_io, kv_io) + if ((CS%id_N2_init>0) .or. (CS%id_S2_init>0) .or. (CS%id_N2_mean>0) .or. (CS%id_S2_mean>0) & + .or. CS%debug) then + !$omp target update from(diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + endif - enddo ! end of j-loop + !$omp target exit data map(release: kappa_3d, tke_3d) + !$omp target exit data map(release: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target exit data map(release: kappa_io, tke_io, kv_io) + !$omp target exit data map(release: surface_pres_2d) + !$omp target exit data map(release: CS) + !$omp target exit data map(release: u_in, v_in, dz_3d) + if (use_temperature) then + !$omp target exit data map(release: tv, tv%T, tv%S) + endif + !$omp target exit data map(release: h) if (CS%debug) then call hchksum(diag_N2_init, "kappa_shear N2_init", G%HI, unscale=US%s_to_T**2) From 4a26c7b358bd4419f5ac01df43b99e85ff65d848 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 09:54:43 +1000 Subject: [PATCH 29/45] tracer_hor_diff: fixed-size epipycnal private arrays in GPU builds (fix W-0155) The two !$omp target teams loop regions in tracer_epipycnal_ML_diff had per-column private arrays declared dimension(SZK_(GV)) / dimension(SZK_(GV)*2), i.e. sized by the runtime GV%ke. NVFORTRAN cannot place a runtime-sized private array in local (stack) memory (NVFORTRAN-W-0155) and spills it to slow device global memory, with a "results may be incorrect" warning. Fix mirrors the kappa_shear GPU_nk_max pattern (KNOWLEDGE row 21): under #ifdef __NVCOMPILER_OPENMP_GPU the per-column scratch (h_demand_L/R, h_used_L/R, h_supply_frac_L/R real, and kbs_Lp/Rp integer + left_set/right_set logical, the last four dimension GPU_nk_max*2) is given the compile-time-constant bound GPU_nk_max=128; CPU builds keep the exact SZK_(GV) sizes. A GPU-only FATAL guard in tracer_epipycnal_ML_diff catches GV%ke > GPU_nk_max instead of silently corrupting. Bitwise-safe: no size()/ubound()/lbound() on any of these arrays (loops are bound by num_srt(i,j) <= nz), so enlarging the declared extent changes no arithmetic. Pre-existing issue (not from the pressure-force/kappa stack); surfaced by a clean rebuild. Verification: benchmark_ALE (which exercises "Ocean epipycnal diffuse tracer") ocean.stats bit-identical to the dev/gpu reference (np=1); the ten W-0155 lines are gone from the build log. Deferred: >=2-GPU verify (1 V100 here). Co-Authored-By: Claude Opus --- src/tracer/MOM_tracer_hor_diff.F90 | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/tracer/MOM_tracer_hor_diff.F90 b/src/tracer/MOM_tracer_hor_diff.F90 index 93bcd8549f..ac46b4d621 100644 --- a/src/tracer/MOM_tracer_hor_diff.F90 +++ b/src/tracer/MOM_tracer_hor_diff.F90 @@ -39,6 +39,13 @@ module MOM_tracer_hor_diff public tracer_hordiff, tracer_hor_diff_init, tracer_hor_diff_end +!> On GPU builds this makes the per-column private scratch in tracer_epipycnal_ML_diff's device +!! regions fixed-size (stack) arrays instead of runtime-sized privates, which nvfortran cannot +!! place in local memory (NVFORTRAN-W-0155) and instead spills to slow device global memory. +!! Checked against GV%ke in tracer_epipycnal_ML_diff. Unused in CPU builds, where the +!! declarations keep their exact SZK_(GV) sizes. +integer, parameter :: GPU_nk_max = 128 + !> The control structure for along-layer and epineutral tracer diffusion type, public :: tracer_hor_diff_CS ; private real :: KhTr !< The along-isopycnal tracer diffusivity [L2 T-1 ~> m2 s-1]. @@ -800,7 +807,11 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & integer, dimension(SZI_(G),SZK_(GV), SZJ_(G)) :: & k0_srt ! The original k-index that each layer of the sorted column corresponds to. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif h_demand_L, & ! The thickness in the left column that is demanded to match the thickness ! in the counterpart [H ~> m or kg m-2]. h_demand_R, & ! The thickness in the right column that is demanded to match the thickness @@ -843,10 +854,18 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & ! The total number of pairings is usually much less than twice the number of layers, but ! the memory in these 1-d columns of pairings can be allocated generously for safety. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max*2) :: & +#else integer, dimension(SZK_(GV)*2) :: & +#endif kbs_Lp, & ! The sorted indices of the Left and Right columns for kbs_Rp ! each pairing. +#ifdef __NVCOMPILER_OPENMP_GPU + logical, dimension(GPU_nk_max*2) :: & +#else logical, dimension(SZK_(GV)*2) :: & +#endif left_set, & ! If true, the left or right point determines the density of right_set ! of the trio. If densities are exactly equal, both are true. @@ -863,6 +882,11 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed IsdB = G%IsdB ; IedB = G%IedB ; JsdB = G%JsdB ; JedB = G%JedB +#ifdef __NVCOMPILER_OPENMP_GPU + if (nz > GPU_nk_max) call MOM_error(FATAL, & + "tracer_epipycnal_ML_diff: GPU builds require GV%ke <= GPU_nk_max because the device column "//& + "regions use fixed-size private scratch arrays; increase GPU_nk_max in MOM_tracer_hor_diff.F90.") +#endif Idt = 1.0 / dt nkmb = GV%nk_rho_varies From 812dad7d3e2ec1a2aafc3627d11c7ac77c2a3169 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 09:54:43 +1000 Subject: [PATCH 30/45] MOM_regridding/coord_zlike: extract device-callable _loc kernels for the ZSTAR regrid Verbatim refactor (no OpenMP directives, no behaviour change) that prepares the ALE ZSTAR regrid for GPU offload. The per-column regrid kernels dereference the regridding control structure on entry, and regridding_CS holds a nested zlike_CS (a pointer) whose coordinateResolution is allocatable -- exactly the nested-derived-type-with-allocatable shape that is fragile to map onto the device. Rather than map it, the scalar/array parameters are hoisted so the kernels take them explicitly. - coord_zlike: new build_zstar_column_loc(nk, min_thick, coordinateResolution, ...) holds the former build_zstar_column body verbatim; build_zstar_column becomes a thin wrapper passing CS%nk/CS%min_thickness/CS%coordinateResolution. _loc is public. - MOM_regridding: new filtered_grid_motion_loc(nk_tgt, filt_shallow, filt_deep, old_grid_wt, ...) and adjust_interface_motion_loc(nk_tgt, min_thick, ...) hold the former bodies verbatim (CS%nk -> nk_tgt, CS%depth_of_time_filter_* -> filt_*, CS%old_grid_weight -> old_grid_wt, CS%min_thickness -> min_thick); the CS-taking filtered_grid_motion / adjust_interface_motion become thin wrappers. No call site changes anywhere: every existing caller (build_{zstar,sigma,rho,...}_grid, MOM_diag_remap) keeps using the CS-taking wrappers unchanged. This is the first of two commits; the device directives + residency come next (GPU-gated). Bitwise-safety: the _loc bodies are the originals with CS%x textually replaced by the matching argument; the wrappers pass exactly those CS fields, so the arithmetic and the values are identical. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1) -- CPU-path gate (the _loc kernels still run on the host; no directives yet). log tmp_local_artifacts/ale_refactor_verify3.log. Co-Authored-By: Claude Opus --- src/ALE/MOM_regridding.F90 | 75 ++++++++++++++++++++++++++++---------- src/ALE/coord_zlike.F90 | 55 ++++++++++++++++++++++------ 2 files changed, 99 insertions(+), 31 deletions(-) diff --git a/src/ALE/MOM_regridding.F90 b/src/ALE/MOM_regridding.F90 index 59ec8d4d7d..7da0b16472 100644 --- a/src/ALE/MOM_regridding.F90 +++ b/src/ALE/MOM_regridding.F90 @@ -1465,6 +1465,28 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) real, dimension(CS%nk+1), intent(in) :: z_new !< New grid position before filtering [H ~> m or kg m-2] real, dimension(CS%nk+1), intent(inout) :: dz_g !< Change in interface positions including !! the effects of filtering [H ~> m or kg m-2] + + call filtered_grid_motion_loc( CS%nk, CS%depth_of_time_filter_shallow, CS%depth_of_time_filter_deep, & + CS%old_grid_weight, nk, z_old, z_new, dz_g ) + +end subroutine filtered_grid_motion + +!> Element-wise, control-structure-free form of filtered_grid_motion: it takes the (host-resolved) +!! target level count and filter parameters explicitly so it can be made device-callable, avoiding a +!! dereference of regridding_CS inside a device region. The arithmetic is verbatim. +subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_wt, nk, z_old, z_new, dz_g ) + integer, intent(in) :: nk_tgt !< Number of cells in the target grid (was CS%nk) + real, intent(in) :: filt_shallow !< Depth at which the shallow filter + !! timescale applies [H ~> m or kg m-2] (was CS%depth_of_time_filter_shallow) + real, intent(in) :: filt_deep !< Depth at which the deep filter timescale + !! applies [H ~> m or kg m-2] (was CS%depth_of_time_filter_deep) + real, intent(in) :: old_grid_wt !< Weight given to the old grid when + !! filtering [nondim] (was CS%old_grid_weight) + integer, intent(in) :: nk !< Number of cells in source grid + real, dimension(nk+1), intent(in) :: z_old !< Old grid position [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(in) :: z_new !< New grid position before filtering [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(inout) :: dz_g !< Change in interface positions including + !! the effects of filtering [H ~> m or kg m-2] ! Local variables real :: sgn ! The sign convention for downward [nondim]. real :: dz_tgt ! The target grid movement of the unfiltered grid [H ~> m or kg m-2] @@ -1491,19 +1513,19 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) logical :: debug = .false. integer :: k - if ((z_old(nk+1) - z_old(1)) * (z_new(CS%nk+1) - z_new(1)) < 0.0) then + if ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) < 0.0) then call MOM_error(FATAL, "filtered_grid_motion: z_old and z_new use different sign conventions.") - elseif ((z_old(nk+1) - z_old(1)) * (z_new(CS%nk+1) - z_new(1)) == 0.0) then + elseif ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) == 0.0) then ! This is a massless column, so do nothing and return. - do k=1,CS%nk+1 ; dz_g(k) = 0.0 ; enddo ; return - elseif ((z_old(nk+1) - z_old(1)) + (z_new(CS%nk+1) - z_new(1)) > 0.0) then + do k=1,nk_tgt+1 ; dz_g(k) = 0.0 ; enddo ; return + elseif ((z_old(nk+1) - z_old(1)) + (z_new(nk_tgt+1) - z_new(1)) > 0.0) then sgn = 1.0 else sgn = -1.0 endif if (debug) then - do k=2,CS%nk+1 + do k=2,nk_tgt+1 if (sgn*(z_new(k)-z_new(k-1)) < -5e-16*(abs(z_new(k))+abs(z_new(k-1))) ) & call MOM_error(FATAL, "filtered_grid_motion: z_new is tangled.") enddo @@ -1514,9 +1536,9 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) ! ddz_g_s(:) = 0.0 ; ddz_g_d(:) = 0.0 endif - zs = CS%depth_of_time_filter_shallow - zd = CS%depth_of_time_filter_deep - wtd = 1.0 - CS%old_grid_weight + zs = filt_shallow + zd = filt_deep + wtd = 1.0 - old_grid_wt Iwtd = 1.0 / wtd dzwt = (zd - zs) @@ -1526,7 +1548,7 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) dz_g(1) = 0.0 z_old_k = z_old(1) - do k = 2,CS%nk+1 + do k = 2,nk_tgt+1 if (k<=nk+1) z_old_k = z_old(k) ! This allows for virtual z_old interface at bottom of the model ! zr1 is positive and increases with depth, and dz_tgt is positive downward. dz_tgt = sgn*(z_new(k) - z_old_k) @@ -1595,17 +1617,17 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) if (debug) then z_old_k = z_old(1) - do k=1,CS%nk+1 + do k=1,nk_tgt+1 if (k<=nk+1) z_old_k = z_old(k) ! This allows for virtual z_old interface at bottom of the model z_act(k) = z_old_k + dz_g(k) enddo - do k=2,CS%nk+1 + do k=2,nk_tgt+1 if (sgn*((z_act(k))-z_act(k-1)) < -1e-15*(abs(z_act(k))+abs(z_act(k-1))) ) & call MOM_error(FATAL, "filtered_grid_motion: z_output is tangled.") enddo endif -end subroutine filtered_grid_motion +end subroutine filtered_grid_motion_loc !> Builds a z*-coordinate grid with partial steps (Adcroft and Campin, 2004). !! z* is defined as @@ -2164,6 +2186,21 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) real, dimension(nk), intent(in) :: h_old !< Layer thicknesses on the old grid [H ~> m or kg m-2] real, dimension(CS%nk+1), intent(inout) :: dz_int !< Interface movements, adjusted to keep the thicknesses !! thicker than their minimum value [H ~> m or kg m-2] + + call adjust_interface_motion_loc( CS%nk, CS%min_thickness, nk, h_old, dz_int ) + +end subroutine adjust_interface_motion + +!> Element-wise, control-structure-free form of adjust_interface_motion: it takes the (host-resolved) +!! target level count and minimum thickness explicitly so it can be made device-callable. Verbatim. +subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) + integer, intent(in) :: nk_tgt !< Number of cells in the target grid (was CS%nk) + real, intent(in) :: min_thick !< Minimum allowed thickness [H ~> m or kg m-2] + !! (was CS%min_thickness) + integer, intent(in) :: nk !< Number of layers in h_old + real, dimension(nk), intent(in) :: h_old !< Layer thicknesses on the old grid [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(inout) :: dz_int !< Interface movements, adjusted to keep the thicknesses + !! thicker than their minimum value [H ~> m or kg m-2] ! Local variables real :: h_new ! A layer thickness on the new grid [H ~> m or kg m-2] real :: eps ! A tiny relative thickness [nondim] @@ -2175,7 +2212,7 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) eps = 1. ; eps = epsilon(eps) h_total = 0. ; h_err = 0. - do k = 1, min(CS%nk,nk) + do k = 1, min(nk_tgt,nk) h_total = h_total + h_old(k) h_err = h_err + max( h_old(k), abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) @@ -2187,8 +2224,8 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) 'implied h<0 is larger than roundoff!') endif enddo - if (CS%nk>nk) then - do k = nk+1, CS%nk + if (nk_tgt>nk) then + do k = nk+1, nk_tgt h_err = h_err + max( abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = ( dz_int(k) - dz_int(k+1) ) if (h_new < -3.0*h_err) then @@ -2200,10 +2237,10 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) endif enddo endif - do k = min(CS%nk,nk),2,-1 + do k = min(nk_tgt,nk),2,-1 h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) - if (h_new Builds a z* coordinate with a minimum thickness +!> Builds a z* coordinate with a minimum thickness (thin wrapper preserving the CS-taking API; +!! the body lives in the device-callable build_zstar_column_loc, GPU port increment 1). subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & z_rigid_top, eta_orig, zScale) type(zlike_CS), intent(in) :: CS !< Coordinate control structure @@ -78,6 +79,36 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & real, optional, intent(in) :: zScale !< Scaling factor from the target coordinate resolution !! in Z to desired units for zInterface, perhaps Z_to_H, !! often [nondim] or [H Z-1 ~> 1 or kg m-3] + + call build_zstar_column_loc(CS%nk, CS%min_thickness, CS%coordinateResolution, depth, & + total_thickness, zInterface, z_rigid_top=z_rigid_top, & + eta_orig=eta_orig, zScale=zScale) + +end subroutine build_zstar_column + +!> Builds a z* coordinate with a minimum thickness. This is the element-wise, control-structure-free +!! form of build_zstar_column: it takes the (host-resolved) scalar and array parameters explicitly so +!! it can be made device-callable, avoiding a dereference of the nested (pointer + allocatable) zlike_CS +!! inside a device region. The arithmetic is verbatim from the former build_zstar_column body. +subroutine build_zstar_column_loc(nk, min_thick, coordinateResolution, depth, total_thickness, & + zInterface, z_rigid_top, eta_orig, zScale) + integer, intent(in) :: nk !< Number of levels to be generated + real, intent(in) :: min_thick !< Minimum thickness allowed for layers, in the + !! same units as depth [Z ~> m] or [H ~> m or kg m-2] + real, dimension(nk), intent(in) :: coordinateResolution !< Target coordinate resolution [Z ~> m] + real, intent(in) :: depth !< Depth of ocean bottom (positive downward in the + !! output units), units may be [Z ~> m] or [H ~> m or kg m-2] + real, intent(in) :: total_thickness !< Column thickness (positive definite in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, dimension(nk+1), intent(inout) :: zInterface !< Absolute positions of interfaces (in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, optional, intent(in) :: z_rigid_top !< The height of a rigid top (positive upward in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, optional, intent(in) :: eta_orig !< The actual original height of the top (in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, optional, intent(in) :: zScale !< Scaling factor from the target coordinate resolution + !! in Z to desired units for zInterface, perhaps Z_to_H, + !! often [nondim] or [H Z-1 ~> 1 or kg m-3] ! Local variables real :: eta ! Free surface height [Z ~> m] or [H ~> m or kg m-2] real :: stretching ! A stretching factor for the coordinate [nondim] @@ -88,7 +119,7 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & z_scale = 1.0 ; if (present(zScale)) z_scale = zScale new_zstar_def = .false. - min_thickness = min( CS%min_thickness, total_thickness/real(CS%nk) ) + min_thickness = min( min_thick, total_thickness/real(nk) ) z0_top = 0. if (present(z_rigid_top)) then z0_top = z_rigid_top @@ -112,37 +143,37 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & ! z_star is the notional z* coordinate in absence of upper/lower topography z_star = 0. ! z*=0 at the free-surface zInterface(1) = eta ! The actual position of the top of the column - do k = 2,CS%nk - z_star = z_star - CS%coordinateResolution(k-1)*z_scale + do k = 2,nk + z_star = z_star - coordinateResolution(k-1)*z_scale ! This ensures that z is below a rigid upper surface (ice shelf bottom) zInterface(k) = min( eta + stretching * ( z_star - z0_top ), z0_top ) ! This ensures that the layer in inflated zInterface(k) = min( zInterface(k), zInterface(k-1) - min_thickness ) ! This ensures that z is above or at the topography - zInterface(k) = max( zInterface(k), -depth + real(CS%nk+1-k) * min_thickness ) + zInterface(k) = max( zInterface(k), -depth + real(nk+1-k) * min_thickness ) enddo - zInterface(CS%nk+1) = -depth + zInterface(nk+1) = -depth else ! Integrate down from the top for a notional new grid, ignoring topography ! The starting position is offset by z0_top which, if z0_top<0, will place ! interfaces above the rigid boundary. zInterface(1) = eta - do k = 1,CS%nk - dh = stretching * CS%coordinateResolution(k)*z_scale ! Notional grid spacing + do k = 1,nk + dh = stretching * coordinateResolution(k)*z_scale ! Notional grid spacing zInterface(k+1) = zInterface(k) - dh enddo ! Integrating up from the bottom adjusting interface position to accommodate ! inflating layers without disturbing the interface above - zInterface(CS%nk+1) = -depth - do k = CS%nk,1,-1 + zInterface(nk+1) = -depth + do k = nk,1,-1 if ( zInterface(k) < (zInterface(k+1) + min_thickness) ) then zInterface(k) = zInterface(k+1) + min_thickness endif enddo endif -end subroutine build_zstar_column +end subroutine build_zstar_column_loc end module coord_zlike From 06d3ef777fb2c100a83490d252c9560e35fcf65c Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 10:15:16 +1000 Subject: [PATCH 31/45] MOM_regridding/coord_zlike: run the ZSTAR regrid column loop on the GPU Second of two commits (the first extracted the device-callable _loc kernels). This offloads build_zstar_grid's per-column loop to the device for the ZSTAR coordinate (the mode benchmark_ALE uses), building on the _loc refactor. - build_zstar_column_loc / filtered_grid_motion_loc / adjust_interface_motion_loc are now !$omp declare target. The "should never happen" grid sign/tangling/negative-thickness MOM_error + write + stop asserts, and the debug-only blocks (debug=.false.) and their z_act scratch, are #ifndef __NVCOMPILER_OPENMP_GPU (kept on the host, compiled out of device builds). filtered_grid_motion_loc's massless-column early return is expressed as an if-guard (if (.not. massless)) since device loops cannot return early. - build_zstar_column_loc's zScale is now a required positional arg (z_scale) ahead of the two optional shelf args, so the device call sites need no keyword arguments (nvfortran rejects keyword-form args inside a target region). - build_zstar_grid: the !$OMP parallel do over (j,i) becomes a !$omp target teams loop collapse(2); the land-mask `cycle` is an if-guard; the regridding_CS scalars + coordinateResolution are hoisted into plain locals on the host (the private component is fetched via the new get_zlike_coord_res accessor) so the device loop never dereferences the nested (pointer + allocatable) control structure; per-column zNew/zOld are fixed-size (GPU_nk_max=128, module parameter) under #ifdef with a FATAL guard when CS%nk/GV%ke exceed it. - Residency: h is host-authoritative in the (host) ALE stack -> update to(h); nom_depth_H and the hoisted coordRes -> map(to:); dzInterface is device-written and copied back (update from) for the subsequent host remap that reads it (ALE_remap_set_h_vel_via_dz, remap_OBC). calc_h_new_by_dz and the rest of ALE stay on the host for now. Verification: benchmark_ALE (ZSTAR, PLM, ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/ale_port_verify3.log. Execution proof: NVCOMPILER_ACC_NOTIFY shows the device kernel build_zstar_grid line 1728 launching grid=2432 per call (teams-only block=(1,1,1); the _loc kernels are inlined). Untested-but-code-identical: the ice-shelf (frac_shelf_h) branch and the RHO/SIGMA/HYCOM coordinate builders (host, unchanged). Deferred pre-merge gate: >=2-GPU / np>=2 verify (this node has 1 V100). Co-Authored-By: Claude Opus --- src/ALE/MOM_regridding.F90 | 112 +++++++++++++++++++++++++++++-------- src/ALE/coord_zlike.F90 | 30 ++++++---- 2 files changed, 110 insertions(+), 32 deletions(-) diff --git a/src/ALE/MOM_regridding.F90 b/src/ALE/MOM_regridding.F90 index 7da0b16472..e489bc69c0 100644 --- a/src/ALE/MOM_regridding.F90 +++ b/src/ALE/MOM_regridding.F90 @@ -32,6 +32,7 @@ module MOM_regridding use coord_zlike, only : zlike_CS use coord_zlike, only : init_coord_zlike, set_zlike_params, build_zstar_column, end_coord_zlike +use coord_zlike, only : build_zstar_column_loc, get_zlike_coord_res use coord_sigma, only : sigma_CS use coord_sigma, only : init_coord_sigma, set_sigma_params, build_sigma_column, end_coord_sigma use coord_rho, only : init_coord_rho, rho_CS, set_rho_params, build_rho_column, end_coord_rho @@ -48,6 +49,13 @@ module MOM_regridding #include +!> On GPU builds this fixes the per-column private scratch in the regrid column loops +!! (build_zstar_grid's zNew/zOld) to a compile-time-constant size, so each device thread gets +!! stack (local memory) arrays instead of runtime-sized device-heap allocations, which nvfortran +!! cannot place in local memory (NVFORTRAN-W-0155). Checked against CS%nk / GV%ke in +!! build_zstar_grid. Unused in CPU builds, where the declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 + ! A note on unit descriptions in comments: MOM6 uses units that can be rescaled for dimensional ! consistency testing. These are noted in comments with units like Z, H, L, and T, along with ! their mks counterparts with notation like "a velocity [Z T-1 ~> m s-1]". If the units @@ -1508,22 +1516,35 @@ subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_w real :: Int_zd ! A depth integral of the weights in [H ~> m or kg m-2] real :: dInt_zs_zd ! The depth integral of the weights between the deep and shallow depths in [H ~> m or kg m-2] ! For debugging: +#ifndef __NVCOMPILER_OPENMP_GPU real, dimension(nk+1) :: z_act ! The final grid positions after the filtered movement [H ~> m or kg m-2] ! real, dimension(nk+1) :: ddz_g_s, ddz_g_d logical :: debug = .false. +#endif + logical :: massless ! True for a massless column, in which case no grid motion is applied. integer :: k + !$omp declare target + ! GPU port: the massless-column early return is expressed as an if-guard (device loops cannot + ! return early), and the "should never happen" sign/tangling MOM_error asserts are compiled out + ! of device builds (kept on the host). + massless = .false. ; sgn = 1.0 if ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) < 0.0) then +#ifndef __NVCOMPILER_OPENMP_GPU call MOM_error(FATAL, "filtered_grid_motion: z_old and z_new use different sign conventions.") +#endif elseif ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) == 0.0) then - ! This is a massless column, so do nothing and return. - do k=1,nk_tgt+1 ; dz_g(k) = 0.0 ; enddo ; return + ! This is a massless column, so do nothing (the filtering block below is skipped). + massless = .true. ; do k=1,nk_tgt+1 ; dz_g(k) = 0.0 ; enddo elseif ((z_old(nk+1) - z_old(1)) + (z_new(nk_tgt+1) - z_new(1)) > 0.0) then sgn = 1.0 else sgn = -1.0 endif + if (.not. massless) then + +#ifndef __NVCOMPILER_OPENMP_GPU if (debug) then do k=2,nk_tgt+1 if (sgn*(z_new(k)-z_new(k-1)) < -5e-16*(abs(z_new(k))+abs(z_new(k-1))) ) & @@ -1535,6 +1556,7 @@ subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_w enddo ! ddz_g_s(:) = 0.0 ; ddz_g_d(:) = 0.0 endif +#endif zs = filt_shallow zd = filt_deep @@ -1615,6 +1637,7 @@ subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_w enddo !dz_g(CS%nk+1) = 0.0 +#ifndef __NVCOMPILER_OPENMP_GPU if (debug) then z_old_k = z_old(1) do k=1,nk_tgt+1 @@ -1626,6 +1649,9 @@ subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_w call MOM_error(FATAL, "filtered_grid_motion: z_output is tangled.") enddo endif +#endif + + endif ! .not. massless end subroutine filtered_grid_motion_loc @@ -1656,8 +1682,21 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ #ifdef __DO_SAFETY_CHECKS__ real :: dh ! The larger of the total column thickness or bathymetric depth [H ~> m or kg m-2] #endif + ! GPU port: per-column private scratch gets compile-time-constant sizes on GPU builds so each + ! device thread uses stack (local memory) arrays rather than runtime-sized device-heap allocations. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: zOld ! Previous coordinate interface heights [H ~> m or kg m-2] + real, dimension(GPU_nk_max+1) :: zNew ! New coordinate interface heights [H ~> m or kg m-2] +#else real, dimension(SZK_(GV)+1) :: zOld ! Previous coordinate interface heights [H ~> m or kg m-2] real, dimension(CS%nk+1) :: zNew ! New coordinate interface heights [H ~> m or kg m-2] +#endif + ! GPU port: the regridding_CS scalars + coordinateResolution are hoisted into plain locals so the + ! device column loop never dereferences the nested (pointer + allocatable) control structure. + real, dimension(CS%nk) :: coordRes ! A local copy of CS%zlike_CS%coordinateResolution [Z ~> m] + real :: min_thick, filt_shallow, filt_deep, old_grid_wt, z_scale_l ! Hoisted CS scalars + integer :: nk_tgt ! A local copy of CS%nk (target level count) + logical :: do_filter, use_adjust ! Hoisted CS control flags integer :: i, j, k, nz logical :: ice_shelf @@ -1665,20 +1704,35 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ minThickness = CS%min_thickness ice_shelf = present(frac_shelf_h) - !$OMP parallel do default(none) shared(G,GV,dzInterface,CS,nz,h,frac_shelf_h, & - !$OMP ice_shelf,minThickness,zScale,nom_depth_H) & - !$OMP private(nominalDepth,totalThickness, & -#ifdef __DO_SAFETY_CHECKS__ - !$OMP dh, & + ! Hoist the control-structure reads onto the host; the device column loop reads only plain locals. + nk_tgt = CS%nk + min_thick = CS%min_thickness + filt_shallow = CS%depth_of_time_filter_shallow + filt_deep = CS%depth_of_time_filter_deep + old_grid_wt = CS%old_grid_weight + do_filter = CS%use_depth_based_time_filter .or. CS%old_grid_weight > 0. + use_adjust = CS%use_adjust_interface_motion + z_scale_l = 1.0 ; if (present(zScale)) z_scale_l = zScale + call get_zlike_coord_res(CS%zlike_CS, coordRes) +#ifdef __NVCOMPILER_OPENMP_GPU + if (nk_tgt > GPU_nk_max .or. nz > GPU_nk_max) call MOM_error(FATAL, & + "build_zstar_grid: GPU builds require CS%nk and GV%ke <= GPU_nk_max; increase GPU_nk_max "//& + "in MOM_regridding.F90.") #endif - !$OMP zNew,zOld) + + ! GPU port: run the per-column regrid on the device. h is host-authoritative in the (host) ALE + ! stack -> refresh; nom_depth_H/coordRes are fresh host locals; dzInterface is device-written and + ! copied back for the (host) remap that follows. + !$omp target update to(h) + !$omp target enter data map(to: nom_depth_H, coordRes) map(alloc: dzInterface) + !$omp target teams loop collapse(2) & + !$omp private(nominalDepth, totalThickness, zNew, zOld, k) & + !$omp firstprivate(nz, nk_tgt, min_thick, filt_shallow, filt_deep, old_grid_wt, & + !$omp do_filter, use_adjust, ice_shelf, z_scale_l) do j = G%jsc-1,G%jec+1 do i = G%isc-1,G%iec+1 - if (G%mask2dT(i,j)==0.) then - dzInterface(i,j,:) = 0. - cycle - endif + if (G%mask2dT(i,j)/=0.) then ! Local depth (positive downward) nominalDepth = nom_depth_H(i,j) @@ -1699,21 +1753,21 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ if (ice_shelf) then if (frac_shelf_h(i,j) > 0.) then ! under ice shelf - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, zNew, & - z_rigid_top=totalThickness-nominalDepth, & - eta_orig=zOld(1), zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, zNew, & + z_scale_l, totalThickness-nominalDepth, zOld(1)) else - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, & - zNew, zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, & + zNew, z_scale_l) endif else - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, & - zNew, zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, & + zNew, z_scale_l) endif ! Calculate the final change in grid position after blending new and old grids - if (CS%use_depth_based_time_filter .or. CS%old_grid_weight>0.) & - call filtered_grid_motion(CS, nz, zOld, zNew, dzInterface(i,j,:)) + if (do_filter) & + call filtered_grid_motion_loc(nk_tgt, filt_shallow, filt_deep, old_grid_wt, nz, zOld, zNew, & + dzInterface(i,j,:)) #ifdef __DO_SAFETY_CHECKS__ dh = max(nominalDepth,totalThickness) @@ -1738,11 +1792,18 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ endif #endif - if (CS%use_adjust_interface_motion) call adjust_interface_motion( CS, nz, h(i,j,:), dzInterface(i,j,:) ) + if (use_adjust) call adjust_interface_motion_loc( nk_tgt, min_thick, nz, h(i,j,:), dzInterface(i,j,:) ) + + else + dzInterface(i,j,:) = 0. + endif ! G%mask2dT(i,j) /= 0. enddo enddo + !$omp target update from(dzInterface) + !$omp target exit data map(release: nom_depth_H, coordRes, dzInterface) + end subroutine build_zstar_grid !------------------------------------------------------------------------------ @@ -2208,6 +2269,7 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) real :: h_err ! An error tolerance that use used to flag unacceptably large negative layer thicknesses ! that can not be explained by roundoff errors [H ~> m or kg m-2] integer :: k + !$omp declare target eps = 1. ; eps = epsilon(eps) @@ -2216,6 +2278,7 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) h_total = h_total + h_old(k) h_err = h_err + max( h_old(k), abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) +#ifndef __NVCOMPILER_OPENMP_GPU if (h_new < -3.0*h_err) then write(0,*) 'h<0 at k=',k,'h_old=',h_old(k), & 'wup=',dz_int(k),'wdn=',dz_int(k+1),'dw_dz=',dz_int(k) - dz_int(k+1), & @@ -2223,11 +2286,13 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) call MOM_error( FATAL, 'MOM_regridding: adjust_interface_motion() - '//& 'implied h<0 is larger than roundoff!') endif +#endif enddo if (nk_tgt>nk) then do k = nk+1, nk_tgt h_err = h_err + max( abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = ( dz_int(k) - dz_int(k+1) ) +#ifndef __NVCOMPILER_OPENMP_GPU if (h_new < -3.0*h_err) then write(0,*) 'h<0 at k=',k,'h_old was empty',& 'wup=',dz_int(k),'wdn=',dz_int(k+1),'dw_dz=',dz_int(k) - dz_int(k+1), & @@ -2235,6 +2300,7 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) call MOM_error( FATAL, 'MOM_regridding: adjust_interface_motion() - '//& 'implied h<0 is larger than roundoff!') endif +#endif enddo endif do k = min(nk_tgt,nk),2,-1 @@ -2245,6 +2311,7 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) if (h_new<0.) & dz_int(k) = ( 1. - eps ) * ( dz_int(k+1) - h_old(k) ) ! Backup in case min_thickness==0 h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) +#ifndef __NVCOMPILER_OPENMP_GPU if (h_new<0.) then write(0,*) 'h<0 at k=',k,'h_old=',h_old(k), & 'wup=',dz_int(k),'wdn=',dz_int(k+1),'dw_dz=',dz_int(k) - dz_int(k+1), & @@ -2253,6 +2320,7 @@ subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) call MOM_error( FATAL, 'MOM_regridding: adjust_interface_motion() - '//& 'Repeated adjustment for roundoff h<0 failed!') endif +#endif enddo !if (dz_int(1)/=0.) stop 'MOM_regridding: adjust_interface_motion() surface moved' diff --git a/src/ALE/coord_zlike.F90 b/src/ALE/coord_zlike.F90 index 31e8bb0bcb..28b544151d 100644 --- a/src/ALE/coord_zlike.F90 +++ b/src/ALE/coord_zlike.F90 @@ -23,7 +23,8 @@ module coord_zlike real, allocatable, dimension(:) :: coordinateResolution end type zlike_CS -public init_coord_zlike, set_zlike_params, build_zstar_column, build_zstar_column_loc, end_coord_zlike +public init_coord_zlike, set_zlike_params, build_zstar_column, build_zstar_column_loc +public get_zlike_coord_res, end_coord_zlike contains @@ -79,19 +80,29 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & real, optional, intent(in) :: zScale !< Scaling factor from the target coordinate resolution !! in Z to desired units for zInterface, perhaps Z_to_H, !! often [nondim] or [H Z-1 ~> 1 or kg m-3] + real :: z_scale_w ! zScale with its default resolved on the host + z_scale_w = 1.0 ; if (present(zScale)) z_scale_w = zScale call build_zstar_column_loc(CS%nk, CS%min_thickness, CS%coordinateResolution, depth, & - total_thickness, zInterface, z_rigid_top=z_rigid_top, & - eta_orig=eta_orig, zScale=zScale) + total_thickness, zInterface, z_scale_w, z_rigid_top, eta_orig) end subroutine build_zstar_column +!> Copy the target coordinate resolution out of a zlike_CS (a private component), for callers that +!! need it as a plain array -- e.g. to hoist it onto a device region without dereferencing the CS. +subroutine get_zlike_coord_res(CS, coordinateResolution) + type(zlike_CS), intent(in) :: CS !< Coordinate control structure + real, dimension(:), intent(out) :: coordinateResolution !< Target coordinate resolution [Z ~> m] + integer :: k + do k=1,CS%nk ; coordinateResolution(k) = CS%coordinateResolution(k) ; enddo +end subroutine get_zlike_coord_res + !> Builds a z* coordinate with a minimum thickness. This is the element-wise, control-structure-free !! form of build_zstar_column: it takes the (host-resolved) scalar and array parameters explicitly so !! it can be made device-callable, avoiding a dereference of the nested (pointer + allocatable) zlike_CS !! inside a device region. The arithmetic is verbatim from the former build_zstar_column body. subroutine build_zstar_column_loc(nk, min_thick, coordinateResolution, depth, total_thickness, & - zInterface, z_rigid_top, eta_orig, zScale) + zInterface, z_scale, z_rigid_top, eta_orig) integer, intent(in) :: nk !< Number of levels to be generated real, intent(in) :: min_thick !< Minimum thickness allowed for layers, in the !! same units as depth [Z ~> m] or [H ~> m or kg m-2] @@ -102,21 +113,20 @@ subroutine build_zstar_column_loc(nk, min_thick, coordinateResolution, depth, to !! units as depth) [Z ~> m] or [H ~> m or kg m-2] real, dimension(nk+1), intent(inout) :: zInterface !< Absolute positions of interfaces (in the same !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, intent(in) :: z_scale !< Scaling factor from the target coordinate resolution + !! in Z to desired units for zInterface (was optional zScale; + !! required here so device call sites need no keyword args) real, optional, intent(in) :: z_rigid_top !< The height of a rigid top (positive upward in the same !! units as depth) [Z ~> m] or [H ~> m or kg m-2] real, optional, intent(in) :: eta_orig !< The actual original height of the top (in the same !! units as depth) [Z ~> m] or [H ~> m or kg m-2] - real, optional, intent(in) :: zScale !< Scaling factor from the target coordinate resolution - !! in Z to desired units for zInterface, perhaps Z_to_H, - !! often [nondim] or [H Z-1 ~> 1 or kg m-3] ! Local variables real :: eta ! Free surface height [Z ~> m] or [H ~> m or kg m-2] real :: stretching ! A stretching factor for the coordinate [nondim] - real :: dh, min_thickness, z0_top, z_star, z_scale ! Thicknesses or heights [Z ~> m] or [H ~> m or kg m-2] + real :: dh, min_thickness, z0_top, z_star ! Thicknesses or heights [Z ~> m] or [H ~> m or kg m-2] integer :: k logical :: new_zstar_def - - z_scale = 1.0 ; if (present(zScale)) z_scale = zScale + !$omp declare target new_zstar_def = .false. min_thickness = min( min_thick, total_thickness/real(nk) ) From b2706c273cfe98b0e8e9e6bcb5614229ef7e05ae Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 10:59:33 +1000 Subject: [PATCH 32/45] set_diffusivity: run find_N2 on the GPU (benchmark_ALE active path) First device increment of the set_diffusivity port, on Ed's j-blocked groundwork. find_N2 is the dominant, self-contained, EOS-bearing piece of the active benchmark_ALE path (find_TKE_to_Kd and set_density_ratios are never called there and stay host). - EOS: the per-interface calculate_density_derivs (ROQUET_RHO) is evaluated on the device through the declare-target calculate_density_derivs_elem_loc dispatcher. eos_form + unit scaling are resolved once on the host (get_EOS_form_and_scaling, with a GPU-only FATAL if the form is not ROQUET_RHO/WRIGHT); the scaling is always applied, so the unscaled (unity) case reproduces the former array call bit-for-bit. This is the kappa_shear increment-2b pattern. - Loop forms: the pressure recurrence (serial down/up-K) + the two bottom-boundary-layer do k=nz,2,-1 loops become per-column !$omp target teams loop collapse(2) over (j,i) with a serial k inside; the former cross-column do_any/exit early-out is dropped -- each device thread runs its own k loop and terminates via its do_i flag, bit-for-bit equivalent. The elementwise N2_lay/N2_int, dRho_int init/copy and dz_BBL_avg loops become do concurrent. - Residency: the full-domain EOS inputs (T_f/S_f and the tv descriptor) are mapped ONCE by the caller (set_diffusivity) around the whole jstart loop, and tv%T/tv%S are refreshed there (update to) since they are the diabatic-mutated persistently-mapped fields; h is already persistently device-resident. find_N2 itself maps only the per-block scratch (dz/pres/h_amp + the 3-D block arrays); doing the full-domain maps per block would re-copy them on every j-row (njblock=1) and dominate the runtime. dRho_int/N2_int/N2_lay/N2_bot (find_N2 outputs) and pres/dz_BBL_avg (read by the host find_rho_bottom that follows) are update from before the per-block region closes. find_rho_bottom stays on the host. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1); logs tmp_local_artifacts/setdiff_findN2_verify.log (+ ..._perf_verify.log for the residency fix). Execution proof: NVCOMPILER_ACC_NOTIFY shows the device kernels find_n2 lines ~1284 (pres+EOS), ~1333 and ~1368 (bottom-BL) launching per block. Timing: "(Ocean set_diffusivity)" 3.4 s/24 (host baseline 3.2 s; find_N2-on-GPU is ~break-even at this size because the bkgnd/assembly/add_LOTW_BBL tail is still host, so find_N2 pays a host<->device boundary each call and the njblock=1 kernels are teams-only/low-occupancy; the win materialises once the surrounding path is on-device). Remaining set_diffusivity tail left on host by design. Deferred: >=2-GPU verify (1 V100 here). Co-Authored-By: Claude Opus --- .../vertical/MOM_set_diffusivity.F90 | 268 ++++++++++-------- 1 file changed, 155 insertions(+), 113 deletions(-) diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 0f36fb4c4f..84a48cee3d 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -18,6 +18,8 @@ module MOM_set_diffusivity use MOM_diagnose_kdwork, only : vbf_CS use MOM_debugging, only : hchksum, uvchksum, Bchksum, hchksum_pair use MOM_EOS, only : calculate_density, calculate_density_derivs, EOS_domain +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_error_handler, only : MOM_error, is_root_pe, FATAL, WARNING, NOTE use MOM_error_handler, only : callTree_showQuery use MOM_error_handler, only : callTree_enter, callTree_leave, callTree_waypoint @@ -479,6 +481,16 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! be an appropriate place to add a depth-dependent parameterization or another explicit ! parameterization of Kd. + ! GPU port: map the full-domain EOS inputs (T_f/S_f and the tv descriptor) once for the whole + ! j-loop and refresh the diabatic-mutated tv%T/tv%S here. find_N2 offloads per block (njblock=1), + ! so mapping/refreshing these inside it would re-copy full-domain arrays on every j-row and dominate + ! the runtime. h is already persistently device-resident. + !$omp target enter data map(to: T_f, S_f) + if (associated(tv%eqn_of_state)) then + !$omp target enter data map(to: tv, tv%T, tv%S) + !$omp target update to(tv%T, tv%S) + endif + do jstart=js,je,njblock jend = min(jstart+njblock-1, je) @@ -807,6 +819,11 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i endif enddo ! jstart-loop + !$omp target exit data map(release: T_f, S_f) + if (associated(tv%eqn_of_state)) then + !$omp target exit data map(release: tv, tv%T, tv%S) + endif + if (CS%user_change_diff) then call user_change_diff(h, tv, G, GV, US, CS%user_change_diff_CSp, Kd_lay, Kd_int, & T_f, S_f, dd%Kd_user) @@ -1217,162 +1234,187 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRh logical :: do_i(SZI_(G),nj), do_any integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state integer :: i, j, k, is, ie, nz, jj + ! GPU port: EOS form + unit scaling resolved once on the host so the density derivatives can be + ! evaluated on the device through the declare-target dispatcher; per-interface scalar scratch. + integer :: eos_form + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa + real :: T_int_s, S_int_s, dRdT_s, dRdS_s + real :: dz_BBL_min_l ! A host copy of CS%dz_BBL_avg_min, for use in a device region [Z ~> m] + logical :: have_p_surf is = G%isc ; ie = G%iec ; nz = GV%ke G_Rho0 = GV%g_Earth_Z_T2 / GV%H_to_RZ H_neglect = GV%H_subroundoff - ! Find the (limited) density jump across each interface. + ! GPU port: resolve the EOS form + unit scaling on the host (the accessor is not device-callable). + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (associated(tv%eqn_of_state)) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "find_N2 GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") +#endif + endif + have_p_surf = associated(fluxes%p_surf) + dz_BBL_min_l = CS%dz_BBL_avg_min + + ! GPU port: the surface pressure (K=1, from the possibly-null p_surf pointer) and the topographic + ! roughness h_amp (set by the host tidal_mixing_h_amp when active) are computed on the host and + ! mapped in; the per-column pressure recurrence + density derivatives + buoyancy frequencies then + ! run on the device. One device data region brackets all of find_N2's per-column loops. do j=jstart,jend ; jj = j - jstart + 1 do i=is,ie - dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 - dRho_int_unfilt(i,jj,1) = 0.0 ; dRho_int_unfilt(i,jj,nz+1) = 0.0 + if (have_p_surf) then ; pres(i,jj,1) = fluxes%p_surf(i,j) ; else ; pres(i,jj,1) = 0.0 ; endif + h_amp(i,jj) = 0.0 enddo enddo + if (CS%use_tidal_mixing) then + do j=jstart,jend ; jj = j - jstart + 1 + call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) + enddo + endif + + ! h, T_f, S_f and tv/tv%T/tv%S are mapped/refreshed once by the caller (set_diffusivity) around the + ! whole jstart loop -- doing it here would re-copy full-domain arrays on every block (njblock=1 => a + ! per-j-row copy), which dominates the runtime. Only the per-block scratch is mapped here. + !$omp target enter data map(to: dz, pres, h_amp) & + !$omp map(alloc: dRho_int, dRho_int_unfilt, N2_int, N2_lay, N2_bot, hb, drho_bot, z_from_bot, do_i, dz_BBL_avg) + + ! Find the (limited) density jump across each interface. + do concurrent (jj=1:nj, i=is:ie) + dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 + dRho_int_unfilt(i,jj,1) = 0.0 ; dRho_int_unfilt(i,jj,nz+1) = 0.0 + enddo if (associated(tv%eqn_of_state)) then - if (associated(fluxes%p_surf)) then - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; pres(i,jj,1) = fluxes%p_surf(i,j) ; enddo - enddo - else - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; pres(i,jj,1) = 0.0 ; enddo - enddo - endif - EOSdom(:) = EOS_domain(G%HI) - do K=2,nz - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - pres(i,jj,K) = pres(i,jj,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) - Temp_Int(i,jj) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) - Salin_Int(i,jj) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) - enddo - call calculate_density_derivs(Temp_int(:,jj), Salin_int(:,jj), pres(:,jj,K), dRho_dT(:,jj,K), dRho_dS(:,jj,K), & - tv%eqn_of_state, EOSdom) - do i=is,ie - dRho_int(i,K,jj) = max(dRho_dT(i,jj,K)*(T_f(i,j,k) - T_f(i,j,k-1)) + & - dRho_dS(i,jj,K)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) - dRho_int_unfilt(i,jj,K) = max(dRho_dT(i,jj,K)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & - dRho_dS(i,jj,K)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + ! Per-column pressure recurrence (serial in K) + density derivatives via the device-callable + ! EOS dispatcher. The unit-scaling is always applied; for the unscaled (unity) case the factors + ! are exactly 1.0, so this reproduces the former array calculate_density_derivs call bit-for-bit. + !$omp target teams loop collapse(2) private(jj, K, T_int_s, S_int_s, dRdT_s, dRdS_s) + do j=jstart,jend + do i=is,ie + jj = j - jstart + 1 + do K=2,nz + pres(i,jj,K) = pres(i,jj,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,K-1) + T_int_s = 0.5 * (T_f(i,j,K) + T_f(i,j,K-1)) + S_int_s = 0.5 * (S_f(i,j,K) + S_f(i,j,K-1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_int_s, & + eos_S_to_ppt*S_int_s, eos_RL2_T2_to_Pa*pres(i,jj,K), dRdT_s, dRdS_s) + dRdT_s = (eos_kg_m3_to_R*eos_C_to_degC) * dRdT_s + dRdS_s = (eos_kg_m3_to_R*eos_S_to_ppt) * dRdS_s + dRho_int(i,K,jj) = max(dRdT_s*(T_f(i,j,K) - T_f(i,j,K-1)) + & + dRdS_s*(S_f(i,j,K) - S_f(i,j,K-1)), 0.0) + dRho_int_unfilt(i,jj,K) = max(dRdT_s*(tv%T(i,j,K) - tv%T(i,j,K-1)) + & + dRdS_s*(tv%S(i,j,K) - tv%S(i,j,K-1)), 0.0) enddo enddo enddo else - do j=jstart,jend ; jj = j - jstart + 1 - do K=2,nz - do i=is,ie - dRho_int(i,K,jj) = GV%Rlay(k) - GV%Rlay(k-1) - enddo - enddo + do concurrent (jj=1:nj, K=2:nz, i=is:ie) + dRho_int(i,K,jj) = GV%Rlay(K) - GV%Rlay(K-1) enddo endif ! Set the buoyancy frequencies. - do k=1,nz - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,K,jj) + dRho_int(i,K+1,jj)) / & - (h(i,j,k) + H_neglect) - enddo - enddo + do concurrent (k=1:nz, jj=1:nj, i=is:ie) + N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,k,jj) + dRho_int(i,k+1,jj)) / & + (h(i,jstart+jj-1,k) + H_neglect) enddo - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 ; enddo + do concurrent (jj=1:nj, i=is:ie) + N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 enddo - do K=2,nz - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & - (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) - enddo - enddo + do concurrent (K=2:nz, jj=1:nj, i=is:ie) + N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & + (0.5*(h(i,jstart+jj-1,K-1) + h(i,jstart+jj-1,K) + H_neglect)) enddo ! Find the bottom boundary layer stratification, and use this in the deepest layers. - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 ; h_amp(i,jj) = 0.0 - z_from_bot(i,jj) = 0.5*dz(i,nz,jj) - do_i(i,jj) = (G%mask2dT(i,j) > 0.0) - enddo + ! (h_amp was set on the host above, including the tidal_mixing_h_amp path.) + do concurrent (jj=1:nj, i=is:ie) + hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,jstart+jj-1) > 0.0) enddo - if (CS%use_tidal_mixing) then - do j=jstart,jend ; jj = j - jstart + 1 - call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) - enddo - endif - do k=nz,2,-1 - do_any = .false. - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; if (do_i(i,jj)) then - dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) - z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - - hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) - drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K,jj) - - if (z_from_bot(i,jj) > h_amp(i,jj)) then - if (k>2) then - ! Always include at least one full layer. - hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) - drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,K-1,jj) + ! Accumulate the near-bottom thickness/density difference per column (serial down-K); the former + ! cross-column do_any/exit early-out is dropped -- each device thread runs its own K loop and stops + ! via its do_i flag, which is bit-for-bit equivalent to the original. + !$omp target teams loop collapse(2) private(jj, k, dz_int) + do j=jstart,jend + do i=is,ie + jj = j - jstart + 1 + do k=nz,2,-1 + if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,k,jj) + + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) then + ! Always include at least one full layer. + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,k-1,jj) + endif + do_i(i,jj) = .false. endif - do_i(i,jj) = .false. - else - do_any = .true. endif - endif ; enddo + enddo enddo - if (.not.do_any) exit enddo - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - if (hb(i,jj) > 0.0) then - N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) - else ; N2_bot(i,jj) = 0.0 ; endif - z_from_bot(i,jj) = 0.5*dz(i,nz,jj) - do_i(i,jj) = (G%mask2dT(i,j) > 0.0) - enddo + do concurrent (jj=1:nj, i=is:ie) + if (hb(i,jj) > 0.0) then + N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) + else ; N2_bot(i,jj) = 0.0 ; endif + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,jstart+jj-1) > 0.0) enddo - do k=nz,2,-1 - do_any = .false. - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; if (do_i(i,jj)) then - dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) - z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above - - N2_int(i,K,jj) = N2_bot(i,jj) - if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) - - if (z_from_bot(i,jj) > h_amp(i,jj)) then - if (k>2) N2_int(i,K-1,jj) = N2_bot(i,jj) - do_i(i,jj) = .false. - else - do_any = .true. + ! Overwrite the near-bottom interface/layer N2 with the bottom-BL value (serial down-K per column; + ! cross-column do_any/exit dropped, per-column do_i flag preserves bit-for-bit behaviour). + !$omp target teams loop collapse(2) private(jj, k, dz_int) + do j=jstart,jend + do i=is,ie + jj = j - jstart + 1 + do k=nz,2,-1 + if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + + N2_int(i,k,jj) = N2_bot(i,jj) + if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) + + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) N2_int(i,k-1,jj) = N2_bot(i,jj) + do_i(i,jj) = .false. + endif endif - endif ; enddo + enddo enddo - if (.not.do_any) exit enddo if (associated(tv%eqn_of_state)) then - do K=1,nz+1 - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie - dRho_int(i,K,jj) = dRho_int_unfilt(i,jj,K) - enddo - enddo + do concurrent (K=1:nz+1, jj=1:nj, i=is:ie) + dRho_int(i,K,jj) = dRho_int_unfilt(i,jj,K) enddo endif ! Average over the larger of the envelope of the topography or a minimal distance. - do j=jstart,jend ; jj = j - jstart + 1 - do i=is,ie ; dz_BBL_avg(i,jj) = max(h_amp(i,jj), CS%dz_BBL_avg_min) ; enddo + do concurrent (jj=1:nj, i=is:ie) + dz_BBL_avg(i,jj) = max(h_amp(i,jj), dz_BBL_min_l) enddo + ! GPU port: copy the device results back for the host find_rho_bottom and the (host) set_diffusivity + ! consumers (dRho_int/N2_int/N2_lay/N2_bot are find_N2's outputs; pres/dz_BBL_avg feed find_rho_bottom), + ! then close the device data region. + !$omp target update from(dRho_int, N2_int, N2_lay, N2_bot, pres, dz_BBL_avg) + !$omp target exit data map(release: dz, pres, h_amp, hb, drho_bot, z_from_bot, do_i, & + !$omp dRho_int, dRho_int_unfilt, N2_int, N2_lay, N2_bot, dz_BBL_avg) + call find_rho_bottom(G, GV, US, tv, h, dz, pres, dz_BBL_avg, jstart, jend, nj, Rho_bot, h_bot, k_bot) end subroutine find_N2 From b8bc685c00714b6973ed73bbc7765d6f15845bf9 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 12:34:47 +1000 Subject: [PATCH 33/45] ePBL: make the column-solver helpers + Langmuir kernels device-callable (increment 1) First increment of the CPU-preserving ePBL GPU port (the ePBL half of the mislabeled "(Ocean vert remap)" region). Bitwise-neutral enablers: !$omp declare target the pure per-column kernels so ePBL_column can call them once it (and the driver loop) are offloaded in later increments. The same source serves the host and device paths. - MOM_energetic_PBL.F90: module-scope declare target on the pure scalar helpers exp_decay_TKE_adjust, find_PE_chg, find_PE_chg_orig, find_Kd_from_PE_chg, find_mstar, mstar_Langmuir (no host-sinks, no runtime-sized array locals -> nothing to gate). - MOM_wave_interface.F90: declare target on get_Langmuir_Number, get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5 (the LF17 statistical-wave Langmuir path that benchmark_ALE uses: USE_LA_LI2016=True, USE_WAVES=False). get_Langmuir_Number is #ifdef-split: on GPU builds only the LF17 branch (call get_StokesSL_LiFoxKemper) is compiled -- the wave-model branches (TESTPROF/ SURFBANDS/DHH85, their US_H/VS_H(SZK) + allocatable StkBand scratch, the Get_SL_Average_* calls, the Null-method + misalignment MOM_errors) are host-only. This reproduces the host LF17 result bit-for-bit (Dpt_LASL is unused on the LF17 path; LA_Misalignment defaults false). ust_2_u10_coare3p5's vonKar<0 config MOM_error is #ifndef-gated off device. TODO before merge: a GPU build currently compiles the LF17-only get_Langmuir_Number for all callers; add a host-side FATAL guard for GPU builds with a non-LF17 WaveMethod or LA_Misalignment (benign for benchmark_ALE, which is LF17/no-misalignment). Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/epbl_inc1_verify.log. This device-compiles the full ePBL+wave kernel set (the key de-risking) while being bitwise-neutral (device versions unused until ePBL_column + the driver loop are offloaded). Remaining: increment 2 = declare target ePBL_column + fixed-size its SZK locals; increment 3 = flip the energetic_PBL driver loop to a target region + residency. Deferred: >=2-GPU verify. Co-Authored-By: Claude Opus --- .../vertical/MOM_energetic_PBL.F90 | 6 +++++ src/user/MOM_wave_interface.F90 | 22 +++++++++++++++++-- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index 6930007bd1..3b4d422e92 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -321,6 +321,12 @@ module MOM_energetic_PBL integer :: BBL_its !< The number of iterations used to find a self-consistent bottom boundary layer depth end type ePBL_column_diags +!> GPU port: make the per-column ePBL helper kernels device-callable so ePBL_column (which runs on the +!! device once the driver loop is offloaded) can call them. They are pure scalar computations, so the +!! same source serves the host and device paths. +!$omp declare target(exp_decay_TKE_adjust, find_PE_chg, find_PE_chg_orig, find_Kd_from_PE_chg) +!$omp declare target(find_mstar, mstar_Langmuir) + contains !> This subroutine determines the diffusivities from the integrated energetics diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index 23f4c8cb7d..acc61550d7 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -276,6 +276,11 @@ module MOM_wave_interface character*(4), parameter :: LF17_STRING = "LF17" !< LF17 wave method string character*(7), parameter :: EFACTOR_STRING = "EFACTOR" !< EFACTOR (based on vr12-ma) wave method string +!> GPU port: the Li & Fox-Kemper statistical-wave Langmuir-number kernels are device-callable so the +!! ePBL column solver can evaluate the Langmuir number on the device; on GPU builds only the LF17 +!! (USE_LA_LI2016) statistical-wave path is compiled (the wave-model branches are host-only). +!$omp declare target(get_Langmuir_Number, get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) + contains !> Initializes parameters related to MOM_wave_interface @@ -1216,16 +1221,26 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & !Local Variables + real :: LA_STK ! Surface-layer averaged Stokes drift magnitude [L T-1 ~> m s-1] +#ifndef __NVCOMPILER_OPENMP_GPU real :: Top, Bottom, MidPoint ! Positions within each layer [Z ~> m] real :: Dpt_LASL ! Averaging depth for Stokes drift [Z ~> m] real :: ShearDirection ! Shear angular direction from atan2 [radians] real :: WaveDirection ! Wave angular direction from atan2 [radians] - real :: LA_STKx, LA_STKy, LA_STK ! Stokes velocities in [L T-1 ~> m s-1] + real :: LA_STKx, LA_STKy ! Stokes velocities in [L T-1 ~> m s-1] logical :: ContinueLoop, USE_MA real, dimension(SZK_(GV)) :: US_H, VS_H ! Profiles of Stokes velocities [L T-1 ~> m s-1] real, allocatable :: StkBand_X(:), StkBand_Y(:) ! Stokes drifts by band [L T-1 ~> m s-1] integer :: k, BB - +#endif + +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device path: only the LF17 statistical-wave method (USE_LA_LI2016) is supported, and misalignment + ! is not applied; other WaveMethods / LA_Misalignment on a GPU build are rejected by a host-side + ! guard. This reproduces the host LF17 branch (get_StokesSL_LiFoxKemper sets LA directly; Dpt_LASL + ! is not used on that path). + call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, Waves, LA_STK, LA) +#else ! Compute averaging depth for Stokes drift (negative) Dpt_LASL = -1.0*max(Waves%LA_FracHBL*HBL, Waves%LA_HBL_min) @@ -1304,6 +1319,7 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & WaveDirection = atan2(LA_STKy, LA_STKx) LA = LA / sqrt(max(1.e-8, cos( WaveDirection - ShearDirection))) endif +#endif end subroutine get_Langmuir_Number @@ -2084,8 +2100,10 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) ! Note in Edson et al. 2013, eq. 13 m is given as 0.017. However, ! m=0.0017 reproduces the curve in their figure 6. +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%vonKar < 0.0) call MOM_error(FATAL, & "ust_2_u10_coare3p5 called with a negative value of Waves%vonKar") +#endif z0sm = 0.11 * CS%nu_air / USTair ! Compute z0smooth from ustar guess u10a = 1000.0*US%m_s_to_L_T ! An insanely large upper bound for u10. From 1f9063b4032663918ebcec74395ad090dbd987ea Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 12:50:15 +1000 Subject: [PATCH 34/45] ePBL: make ePBL_column device-callable with fixed-size scratch (increment 2) Second increment of the CPU-preserving ePBL GPU port. ePBL_column (the ~1061-line iterative single-column solver: outer MLD iteration + down-K TKE budget with an inner Newton solve) is now !$omp declare target, so it device-compiles with its declare-target helper kernels (find_mstar, find_PE_chg*, get_Langmuir_Number, ...) from increment 1. Still bitwise-neutral: ePBL_column is host-called until the driver loop is offloaded (increment 3); on GPU builds its host version uses the #ifdef paths below, which are bit-identical for benchmark_ALE. - GPU_nk_max=128 module parameter; ePBL_column's per-column private automatic arrays (pres_Z/hb_hs, the dT_to_dColHt.../hp_a.../Sh_b block, MixLen_shape/h_dz_int/Kddt_h, and mech_TKE_k/conv_PErel_k/ nstar_k/dT_expect/dS_expect/num_itts) are given compile-time-constant sizes under #ifdef __NVCOMPILER_OPENMP_GPU (KNOWLEDGE row 21), keeping SZK_(GV) on CPU builds. The dummy args and the dimension(20) iteration arrays are already fine. (ePBL_BBL_column's analogous locals are fixed-sized too via the shared decls; it is dead in benchmark_ALE, not declare target.) - The equation-discovery (eqdisc) mixing-length/velocity calls (kappa_eqdisc/get_eqdisc_v0/v0h) are #ifndef-gated off device (dead in benchmark_ALE: EPBL_EQD_* absent -> .false.); the transLay mixing-length + v0_scale=0 paths run, bit-identical. - FATAL guards in energetic_PBL_init on GPU builds: GV%ke > GPU_nk_max, and any eqdisc* enabled. TODO before merge (carried from increment 1): host-side guard for GPU builds with a non-LF17 WaveMethod / LA_Misalignment (get_Langmuir_Number compiles the LF17-only body on GPU). Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1); log tmp_local_artifacts/epbl_inc2_verify.log. ePBL_column device-compiles cleanly (the key de-risking). Remaining: increment 3 = flip the energetic_PBL driver loop to a target teams loop + one driver-level data region (residency) so the solver runs on the device. Deferred: >=2-GPU verify. Co-Authored-By: Claude Opus --- .../vertical/MOM_energetic_PBL.F90 | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index 3b4d422e92..6aec0a56d1 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -326,6 +326,13 @@ module MOM_energetic_PBL !! same source serves the host and device paths. !$omp declare target(exp_decay_TKE_adjust, find_PE_chg, find_PE_chg_orig, find_Kd_from_PE_chg) !$omp declare target(find_mstar, mstar_Langmuir) +!$omp declare target(ePBL_column) + +!> GPU port: fixes the per-column private scratch in ePBL_column to a compile-time-constant size on +!! GPU builds, so each device thread gets stack (local memory) arrays rather than runtime-sized +!! device-heap automatics (NVFORTRAN-W-0155 / KNOWLEDGE row 21). Checked against GV%ke in +!! energetic_PBL_init. Unused in CPU builds, where the local declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 contains @@ -972,7 +979,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, ! mixing. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif pres_Z, & ! Interface pressures with a rescaling factor to convert interface height ! movements into changes in column potential energy [R Z2 T-2 ~> kg m-1 s-2]. hb_hs ! The distance from the bottom over the thickness of the @@ -989,7 +1000,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real :: Idecay_len_TKE ! The inverse of a turbulence decay length scale [H-1 ~> m-1 or m2 kg-1]. real :: dz_sum ! The total thickness of the water column [Z ~> m]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif dT_to_dColHt, & ! Partial derivative of the total column height with the temperature changes ! within a layer [Z C-1 ~> m degC-1]. dS_to_dColHt, & ! Partial derivative of the total column height with the salinity changes @@ -1026,7 +1041,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, ! mixing effects with other yet lower layers [C H ~> degC m or degC kg m-2]. Sh_b ! An effective salinity times a thickness in the layer below, including implicit ! mixing effects with other yet lower layers [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif MixLen_shape, & ! A nondimensional shape factor for the mixing length that ! gives it an appropriate asymptotic value at the bottom of ! the boundary layer [nondim]. @@ -1159,6 +1178,10 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real, dimension(20) :: Kddt_h_itt ! The value of Kddt_h_guess after each iteration [H ~> m or kg m-2] real, dimension(20) :: dPEa_dKd_itt ! The value of dPEc_dKd after each iteration [R Z3 T-2 H-1 ~> J m-3 or J kg-1] real, dimension(20) :: MKE_src_itt ! The value of MKE_src after each iteration [R Z3 T-2 ~> J m-2] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: mech_TKE_k, conv_PErel_k, nstar_k, dT_expect, dS_expect + integer, dimension(GPU_nk_max) :: num_itts +#else real, dimension(SZK_(GV)) :: mech_TKE_k ! The mechanically generated turbulent kinetic energy ! available for mixing over a time step for each layer [R Z3 T-2 ~> J m-2]. real, dimension(SZK_(GV)) :: conv_PErel_k ! The potential energy that has been convectively released @@ -1168,6 +1191,7 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real, dimension(SZK_(GV)) :: dT_expect ! Expected temperature changes [C ~> degC] real, dimension(SZK_(GV)) :: dS_expect ! Expected salinity changes [S ~> ppt] integer, dimension(SZK_(GV)) :: num_itts +#endif integer :: k, nz, itt, max_itt @@ -1321,7 +1345,9 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, dz_rsum = 0.0 MixLen_shape(1) = 1.0 if (CS%eqdisc) then ! update Kd as per Machine Learning equation discovery +#ifndef __NVCOMPILER_OPENMP_GPU call kappa_eqdisc(MixLen_shape, CS, GV, h, absf, B_flux, u_star, MLD_guess) +#endif else do K=2,nz+1 dz_rsum = dz_rsum + dz(k-1) @@ -1337,11 +1363,13 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, endif v0_ML_turb_vel_scale = 0.0 ! a variable that gets passed on to get_eqdisc_v0 & get_eqdisc_v0h +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%eqdisc_v0) then call get_eqdisc_v0(CS,absf,B_flux,u_star,v0_ML_turb_vel_scale) elseif (CS%eqdisc_v0h) then call get_eqdisc_v0h(CS,B_flux,u_star,MLD_guess,v0_ML_turb_vel_scale) endif +#endif Kd(1) = 0.0 ; Kddt_h(1) = 0.0 hp_a(1) = h(1) @@ -2027,7 +2055,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & ! energy that is supplied as an argument to this routine. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif pres_Z, & ! Interface pressures with a rescaling factor to convert interface height ! movements into changes in column potential energy [R Z2 T-2 ~> kg m-1 s-2]. dztop_dztot ! The distance from the surface divided by the thickness of the @@ -2045,7 +2077,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & real :: Idecay_len_TKE ! The inverse of a turbulence decay length scale [H-1 ~> m-1 or m2 kg-1]. real :: dz_sum ! The total thickness of the water column [Z ~> m]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif dT_to_dColHt, & ! Partial derivative of the total column height with the temperature changes ! within a layer [Z C-1 ~> m degC-1]. dS_to_dColHt, & ! Partial derivative of the total column height with the salinity changes @@ -2095,7 +2131,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & ! mixing effects with other yet lower layers [C H ~> degC m or degC kg m-2]. Sh_b ! An effective salinity times a thickness in the layer below, including implicit ! mixing effects with other yet lower layers [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif MixLen_shape, & ! A nondimensional shape factor for the mixing length that ! gives it an appropriate asymptotic value at the bottom of ! the boundary layer [nondim]. @@ -4437,6 +4477,15 @@ subroutine energetic_PBL_init(Time, G, GV, US, param_file, diag, CS) call safe_alloc_alloc(CS%ML_depth, isd, ied, jsd, jed) call safe_alloc_alloc(CS%BBL_depth, isd, ied, jsd, jed) +#ifdef __NVCOMPILER_OPENMP_GPU + ! GPU port guards: ePBL_column runs on the device with fixed-size (GPU_nk_max) per-column scratch, + ! and the equation-discovery (eqdisc) mixing-length/velocity paths are not device-callable. + if (GV%ke > GPU_nk_max) call MOM_error(FATAL, & + "energetic_PBL GPU build: GV%ke exceeds GPU_nk_max; increase GPU_nk_max in MOM_energetic_PBL.F90.") + if (CS%eqdisc .or. CS%eqdisc_v0 .or. CS%eqdisc_v0h) call MOM_error(FATAL, & + "energetic_PBL GPU build: the EPBL_EQD_DIFFUSIVITY (equation-discovery) paths are not device-callable.") +#endif + end subroutine energetic_PBL_init !> Clean up and deallocate memory associated with the energetic_PBL module. From fe44661f075543aa6dee760518bc5cca5efe1dea Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 15:39:12 +1000 Subject: [PATCH 35/45] ePBL: offload the energetic_PBL driver loop to GPU (increment 3) Flip the energetic_PBL driver (do j / do i over ocean columns calling the now-declare-target ePBL_column) to a single `!$omp target teams loop collapse(2)` with one driver-level data region. Columns are read directly from the 3-d inputs (2-d staging slabs and the per-j Kd_2d staging removed, bit-identical); thickness_to_dz is hoisted whole-domain (dz_3d); the ~14 per-column scratch arrays are teams-loop private + fixed-size under GPU_nk_max. Residency: map(to:) the inputs + CS + tv/tv%T/tv%S (with update-to for the map-to-present trap) + Waves; map(from:) Kd_int / ML_depth / BBL_depth. Forcing (fluxes%ustar) and the MLD first guess are hoisted to plain host arrays so the device never touches a fluxes% pointer member or a CS allocatable. Dead-in-benchmark_ALE branches (BBL mixing, options_diff, stochastics, per-column diagnostics, report_avg_its) are #ifndef'd off the device path and backed by runtime FATAL guards so an unsupported config fails loudly rather than silently wrong. CPU path is preserved verbatim under #else. Device-callable Langmuir chain (MOM_wave_interface): the LF17 path dereferenced the Waves *pointer* on-device, which nvfortran will not attach under -gpu=mem:separate. Replace it with a pointer-free `wave_LF17_params` bundle (allocatable-free, maps cleanly) populated host-side by set_wave_LF17_params and threaded through get_Langmuir_Number_LF17 -> get_StokesSL_LiFoxKemper -> ust_2_u10_coare3p5. get_Langmuir_Number_LF17 uses the in-routine `!$omp declare target` block form (the module name-list form silently failed to emit its device symbol). get_Langmuir_Number is restored to its full CPU form (no longer declare-target). STATUS: ePBL executes correctly on the device (no illegal-address; T/S/mass bit-identical to displayed precision). NOT YET bitwise: the energy diagnostic differs at ~1e-13, the documented host-vs-device x**(1/n) transcendental floor (KNOWLEDGE row 5) from the 4 `(...)**C1_3` cube-roots in the ePBL_column vstar budget and the Langmuir `LA**(-1.33)` power. Closing it needs cuberoot() for the **C1_3 sites + an explicit-form rewrite of the arbitrary power, applied to the one CPU+GPU source, plus a regenerated golden (the fix changes the CPU answer). Deferred by decision. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_energetic_PBL.F90 | 223 ++++++++++++++---- src/user/MOM_wave_interface.F90 | 79 +++++-- 2 files changed, 243 insertions(+), 59 deletions(-) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index 6aec0a56d1..a319c196ae 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -22,7 +22,8 @@ module MOM_energetic_PBL use MOM_unit_scaling, only : unit_scale_type use MOM_variables, only : thermo_var_ptrs, vertvisc_type use MOM_verticalGrid, only : verticalGrid_type -use MOM_wave_interface, only : wave_parameters_CS, Get_Langmuir_Number +use MOM_wave_interface, only : wave_parameters_CS, Get_Langmuir_Number, Get_Langmuir_Number_LF17 +use MOM_wave_interface, only : wave_LF17_params, set_wave_LF17_params use MOM_stochastics, only : stochastic_CS implicit none ; private @@ -277,6 +278,10 @@ module MOM_energetic_PBL ! The next options are used when passively diagnosing sensitivities from parameter choices integer :: id_opt_diff_Kd_ePBL = -1, id_opt_maxdiff_Kd_ePBL = -1, id_opt_diff_hML_depth = -1 !>@} + type(wave_LF17_params) :: wave_lf17 !< A device-mappable bundle of the loop-invariant LF17 wave + !! scalars, populated from the Waves CS each call so the on-device + !! Langmuir path reads these (allocatable-free, maps cleanly) instead + !! of dereferencing the Waves pointer. end type energetic_PBL_CS !>@{ Enumeration values for mstar_scheme @@ -419,7 +424,16 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, v_2d ! A 2-d slice of the meridional velocity [L T-1 ~> m s-1]. real, dimension(SZI_(G),SZK_(GV)+1) :: & Kd_2d ! A 2-d version of the diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: & + dz_3d ! The vertical distance across layers over the whole domain [Z ~> m], computed + ! once per call to replace the former per-j thickness_to_dz slice. +#ifdef __NVCOMPILER_OPENMP_GPU + ! On GPU these per-column column scratch arrays are teams-loop private() and so must have a + ! compile-time-constant size (NVFORTRAN-W-0155); GPU_nk_max is guarded by a FATAL on GV%ke. + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif h, & ! The layer thickness [H ~> m or kg m-2]. dz, & ! The vertical distance across layers [Z ~> m]. T0, & ! The initial layer temperatures [C ~> degC]. @@ -429,13 +443,14 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, TKE_forcing, & ! Forcing of the TKE in the layer coming from TKE_forced [R Z3 T-2 ~> J m-2]. u, & ! The zonal velocity [L T-1 ~> m s-1]. v ! The meridional velocity [L T-1 ~> m s-1]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif Kd, & ! The diapycnal diffusivity due to ePBL [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. mixvel, & ! A turbulent mixing velocity [Z T-1 ~> m s-1]. mixlen, & ! A turbulent mixing length [Z ~> m]. - mixvel_BBL, & ! A bottom boundary layer turbulent mixing velocity [Z T-1 ~> m s-1]. - mixlen_BBL, & ! A bottom boundary layer turbulent mixing length [Z ~> m]. - Kd_BBL, & ! The bottom boundary layer diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. SpV_dt, & ! Specific volume interpolated to interfaces divided by dt or 1.0 / (dt * Rho0), ! in [R-1 T-1 ~> m3 kg-1 s-1], used to convert local TKE into a turbulence velocity cubed. SpV_dt_cf ! Specific volume interpolated to interfaces divided by dt or 1.0 / (dt * Rho0) @@ -443,6 +458,12 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! [m3 Z-3 R-1 T2 s-3 ~> m3 kg-1 s-1] or without the conversion factors for ! answer dates of 20240101 and later in [R-1 T-1 ~> m3 kg-1 s-1], used to ! convert local TKE into a turbulence velocity cubed. + ! These bottom-boundary-layer column arrays are only used on the host (BBL_mixing / options_diff + ! paths, both FATAL-guarded off the GPU port), so they keep the runtime size. + real, dimension(SZK_(GV)+1) :: & + mixvel_BBL, & ! A bottom boundary layer turbulent mixing velocity [Z T-1 ~> m s-1]. + mixlen_BBL, & ! A bottom boundary layer turbulent mixing length [Z ~> m]. + Kd_BBL ! The bottom boundary layer diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. real :: h_neglect ! A thickness that is so small it is usually lost ! in roundoff and can be neglected [H ~> m or kg m-2]. @@ -520,6 +541,16 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! can be modified to test for sensitivities logical :: BBL_mixing ! If true, there is bottom boundary layer mixing. integer :: i, j, k, is, ie, js, je, nz +#ifdef __NVCOMPILER_OPENMP_GPU + ! Plain-array hoists for the GPU port: pointer-member forcing and the CS%ML_depth first guess are + ! copied to these on the host (where associated() is evaluated) so the device loop touches neither + ! fluxes% pointer members nor the CS allocatable components. + real, dimension(SZI_(G),SZJ_(G)) :: & + ustar_2d, & ! Surface friction velocity [Z T-1 ~> m s-1] + ustar_gustless_2d, & ! Gustless surface friction velocity [Z T-1 ~> m s-1] + ML_depth_2d, & ! Mixed layer depth guess (in) / result (out) [H ~> m or kg m-2] + BBL_depth_2d ! Bottom boundary layer depth result [H ~> m or kg m-2] +#endif is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke @@ -600,50 +631,114 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, if (CS%id_opt_diff_hML_depth > 0) diff_hML_depth(:,:) = 0.0 endif - !!OMP parallel do default(private) shared(js,je,nz,is,ie,h_3d,u_3d,v_3d,tv,dt,I_dt,BBL_mixing, & - !!OMP CS,G,GV,US,fluxes,TKE_forced,dSV_dT,dSV_dS,Kd_int) - do j=js,je - ! Copy the thicknesses and other fields to 2-d arrays. - do k=1,nz ; do i=is,ie - h_2d(i,k) = h_3d(i,j,k) ; u_2d(i,k) = u_3d(i,j,k) ; v_2d(i,k) = v_3d(i,j,k) - T_2d(i,k) = tv%T(i,j,k) ; S_2d(i,k) = tv%S(i,j,k) - TKE_forced_2d(i,k) = TKE_forced(i,j,k) - dSV_dT_2d(i,k) = dSV_dT(i,j,k) ; dSV_dS_2d(i,k) = dSV_dS(i,j,k) - enddo ; enddo - call thickness_to_dz(h_3d, tv, dz_2d, j, G, GV) + ! Compute the vertical layer extent over the whole domain once, replacing the per-j + ! thickness_to_dz slice that used to sit inside the j-loop (identical elementwise result). + call thickness_to_dz(h_3d, tv, dz_3d, G, GV, US) - ! Set the inverse density used to translating local TKE into a turbulence velocity - SpV_dt(:) = 0.0 - if ((dt > 0.0) .and. GV%Boussinesq .or. .not.allocated(tv%SpV_avg)) then - if (CS%answer_date < 20240101) then - do K=1,nz+1 - SpV_dt(K) = 1.0 / (dt*GV%Rho0) - enddo - else - do K=1,nz+1 - SpV_dt(K) = I_rho0dt - enddo - endif - endif - - ! Determine the initial mech_TKE and conv_PErel, including the energy required - ! to mix surface heating through the topmost cell, the energy released by mixing - ! surface cooling & brine rejection down through the topmost cell, and - ! homogenizing the shortwave heating within that cell. This sets the energy - ! and ustar and wstar available to drive mixing at the first interior - ! interface. +#ifdef __NVCOMPILER_OPENMP_GPU + ! ---- GPU port: FATAL-guard the features this port does not implement on device ---- + if (.not.GV%Boussinesq) call MOM_error(FATAL, & + "energetic_PBL: the GPU port only supports Boussinesq mode.") + if (BBL_mixing) call MOM_error(FATAL, & + "energetic_PBL: bottom boundary layer mixing is not supported in the GPU port.") + if (CS%options_diff > 0) call MOM_error(FATAL, & + "energetic_PBL: options_diff>0 (sensitivity diagnostics) is not supported in the GPU port.") + if (stoch_CS%pert_epbl) call MOM_error(FATAL, & + "energetic_PBL: stochastic ePBL perturbations are not supported in the GPU port.") + if (CS%TKE_diagnostics) call MOM_error(FATAL, & + "energetic_PBL: the ePBL TKE-budget diagnostics are not supported in the GPU port.") + if (report_avg_its) call MOM_error(FATAL, & + "energetic_PBL: report_avg_its is not supported in the GPU port.") + if (CS%debug) call MOM_error(FATAL, & + "energetic_PBL: CS%debug is not supported in the GPU port.") + if ((CS%id_Mixing_Length>0) .or. (CS%id_Velocity_Scale>0) .or. (CS%id_ustar_ePBL>0) .or. & + (CS%id_Kd_ePBL_col_by_col>0) .or. (CS%id_mstar_sfc>0) .or. (CS%id_mstar_bbl>0) .or. & + (CS%id_mstar_LT>0) .or. (CS%id_LA>0) .or. (CS%id_LA_mod>0)) call MOM_error(FATAL, & + "energetic_PBL: the per-column ePBL diagnostics are not supported in the GPU port.") + + ! Hoist pointer-member forcing and the CS%ML_depth first guess onto plain arrays (associated() + ! tests happen here on the host); the device loop then reads only these plain arrays and never + ! dereferences a fluxes% pointer member or a CS allocatable component. + if (.not.(associated(fluxes%ustar) .and. (GV%Boussinesq .or. .not.associated(fluxes%tau_mag)))) & + call MOM_error(FATAL, "energetic_PBL: the GPU port requires the fluxes%ustar forcing path.") + if (associated(fluxes%ustar_shelf) .and. associated(fluxes%frac_shelf_h)) call MOM_error(FATAL, & + "energetic_PBL: ice-shelf ustar blending is not supported in the GPU port.") + do j=js,je ; do i=is,ie + ustar_2d(i,j) = fluxes%ustar(i,j) ; ustar_gustless_2d(i,j) = fluxes%ustar_gustless(i,j) + ML_depth_2d(i,j) = CS%ML_depth(i,j) + enddo ; enddo + ! Bundle the loop-invariant LF17 wave scalars into CS%wave_lf17 (allocatable-free, so it rides the + ! CS map onto the device) — the on-device Langmuir path reads these instead of the Waves pointer. + if (associated(Waves)) call set_wave_LF17_params(Waves, CS%wave_lf17) + + ! One driver-level data region (find_N2 lesson: map full-domain arrays ONCE, never per-column). + !$omp target enter data map(to: h_3d, u_3d, v_3d, dz_3d, TKE_forced, dSV_dT, dSV_dS, buoy_flux, & + !$omp ustar_2d, ustar_gustless_2d, ML_depth_2d, CS, tv, tv%T, tv%S) & + !$omp map(alloc: BBL_depth_2d, Kd_int) + if (associated(Waves)) then + !$omp target enter data map(to: Waves) + endif + ! tv%T/tv%S are persistently resident via CS%tv and were mutated on the host earlier in diabatic, + ! so refresh them here (the map-to-present-no-copy trap); h_3d may likewise be persistently mapped. + !$omp target update to(tv%T, tv%S, h_3d) + + !$omp target teams loop collapse(2) & + !$omp private(eCD, u_star, u_star_mean, mech_TKE, absf, B_flux, MLD_io, BBLD_io, MLD_in, & + !$omp h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, TKE_forcing, Kd, mixvel, mixlen, & + !$omp SpV_dt, SpV_dt_cf) & + !$omp firstprivate(is, ie, nz, dt, I_dt, I_rho, I_rho0dt) +#endif + do j=js,je + ! For each ocean column determine the initial mech_TKE and conv_PErel, including the energy + ! required to mix surface heating through the topmost cell, the energy released by mixing surface + ! cooling & brine rejection down through the topmost cell, and homogenizing the shortwave heating + ! within that cell. This sets the energy and ustar and wstar available to drive mixing at the + ! first interior interface. do i=is,ie ; if (G%mask2dT(i,j) > 0.0) then - ! Copy the thicknesses and other fields to 1-d arrays. + ! Copy the thicknesses and other fields to 1-d column arrays, read directly from the 3-d + ! inputs (the former per-j 2-d staging slabs are no longer needed). do k=1,nz - h(k) = h_2d(i,k) + GV%H_subroundoff ; dz(k) = dz_2d(i,k) + GV%dZ_subroundoff - u(k) = u_2d(i,k) ; v(k) = v_2d(i,k) - T0(k) = T_2d(i,k) ; S0(k) = S_2d(i,k) ; TKE_forcing(k) = TKE_forced_2d(i,k) - dSV_dT_1d(k) = dSV_dT_2d(i,k) ; dSV_dS_1d(k) = dSV_dS_2d(i,k) + h(k) = h_3d(i,j,k) + GV%H_subroundoff ; dz(k) = dz_3d(i,j,k) + GV%dZ_subroundoff + u(k) = u_3d(i,j,k) ; v(k) = v_3d(i,j,k) + T0(k) = tv%T(i,j,k) ; S0(k) = tv%S(i,j,k) ; TKE_forcing(k) = TKE_forced(i,j,k) + dSV_dT_1d(k) = dSV_dT(i,j,k) ; dSV_dS_1d(k) = dSV_dS(i,j,k) enddo do K=1,nz+1 ; Kd(K) = 0.0 ; enddo + ! Set the inverse density used to translate local TKE into a turbulence velocity. +#ifdef __NVCOMPILER_OPENMP_GPU + do K=1,nz+1 ; SpV_dt(K) = 0.0 ; enddo + if (dt > 0.0) then + if (CS%answer_date < 20240101) then + do K=1,nz+1 ; SpV_dt(K) = 1.0 / (dt*GV%Rho0) ; enddo + else + do K=1,nz+1 ; SpV_dt(K) = I_rho0dt ; enddo + endif + endif +#else + SpV_dt(:) = 0.0 + if ((dt > 0.0) .and. GV%Boussinesq .or. .not.allocated(tv%SpV_avg)) then + if (CS%answer_date < 20240101) then + do K=1,nz+1 + SpV_dt(K) = 1.0 / (dt*GV%Rho0) + enddo + else + do K=1,nz+1 + SpV_dt(K) = I_rho0dt + enddo + endif + endif +#endif + ! Make local copies of surface forcing and process them. +#ifdef __NVCOMPILER_OPENMP_GPU + ! GPU port: Boussinesq fluxes%ustar path only (guarded above); read the hoisted plain arrays. + u_star = ustar_2d(i,j) + u_star_Mean = ustar_gustless_2d(i,j) + mech_TKE = dt * GV%Rho0 * u_star**3 + B_flux = buoy_flux(i,j) +#else if (associated(fluxes%ustar) .and. (GV%Boussinesq .or. .not.associated(fluxes%tau_mag))) then u_star = fluxes%ustar(i,j) u_star_Mean = fluxes%ustar_gustless(i,j) @@ -674,6 +769,7 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, u_star = (1.0 - fluxes%frac_shelf_h(i,j)) * u_star + & fluxes%frac_shelf_h(i,j) * fluxes%ustar_shelf(i,j) endif +#endif if (u_star < CS%ustar_min) u_star = CS%ustar_min if (CS%omega_frac >= 1.0) then absf = 2.0*CS%omega @@ -686,7 +782,11 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! Perhaps provide a first guess for MLD based on a stored previous value. MLD_io = -1.0 +#ifdef __NVCOMPILER_OPENMP_GPU + if (CS%MLD_iteration_guess .and. (ML_depth_2d(i,j) > 0.0)) MLD_io = ML_depth_2d(i,j) +#else if (CS%MLD_iteration_guess .and. (CS%ML_depth(i,j) > 0.0)) MLD_io = CS%ML_depth(i,j) +#endif BBLD_io = 0.0 ! Store the initial guesses at the boundary layer depths for testing sensitivities. @@ -697,19 +797,24 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, else do K=1,nz+1 ; SpV_dt_cf(K) = SpV_dt(K) ; enddo endif +#ifndef __NVCOMPILER_OPENMP_GPU if (stoch_CS%pert_epbl) then ! stochastics are active call ePBL_column(h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, SpV_dt_cf, TKE_forcing, B_flux, absf, & u_star, u_star_mean, mech_TKE, dt, MLD_io, Kd, mixvel, mixlen, GV, & US, CS, eCD, Waves, G, i, j, & TKE_gen_stoch=stoch_CS%epbl1_wts(i,j), TKE_diss_stoch=stoch_CS%epbl2_wts(i,j)) else +#endif call ePBL_column(h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, SpV_dt_cf, TKE_forcing, B_flux, absf, & u_star, u_star_mean, mech_TKE, dt, MLD_io, Kd, mixvel, mixlen, GV, & US, CS, eCD, Waves, G, i, j) +#ifndef __NVCOMPILER_OPENMP_GPU endif if (CS%id_Kd_ePBL_col_by_col > 0) & call post_data_3d_by_column(CS%id_Kd_ePBL_col_by_col, Kd, CS%diag, i, j) +#endif +#ifndef __NVCOMPILER_OPENMP_GPU ! Add the diffusivity due to bottom boundary layer mixing, if there is energy to drive this mixing. if (BBL_mixing) then if (CS%MLD_iteration_guess .and. (CS%BBL_depth(i,j) > 0.0)) BBLD_io = CS%BBL_depth(i,j) @@ -747,14 +852,21 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, if ((CS%id_BBL_decay_scale > 0) .and. (CS%TKE_decay * absf > 0)) & diag_BBL_decay_scale(i,j) = u_star_BBL / (CS%TKE_decay * absf) endif +#endif - ! Copy the diffusivities to a 2-d array. + ! Copy the diffusivities to the output interface diffusivity array. do K=1,nz+1 - Kd_2d(i,K) = Kd(K) + Kd_int(i,j,K) = Kd(K) enddo +#ifdef __NVCOMPILER_OPENMP_GPU + ML_depth_2d(i,j) = MLD_io + BBL_depth_2d(i,j) = BBLD_io +#else CS%ML_depth(i,j) = MLD_io CS%BBL_depth(i,j) = BBLD_io +#endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%TKE_diagnostics) then diag_TKE_MKE(i,j) = diag_TKE_MKE(i,j) + eCD%dTKE_MKE diag_TKE_conv(i,j) = diag_TKE_conv(i,j) + eCD%dTKE_conv @@ -833,17 +945,34 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, endif if (CS%id_opt_diff_hML_depth > 0) diff_hML_depth(i,j) = BLD_1 - BLD_2 endif +#endif else ! End of the ocean-point part of the i-loop ! For masked points, Kd_int must still be set (to 0) because it has intent out. - do K=1,nz+1 ; Kd_2d(i,K) = 0. ; enddo + do K=1,nz+1 ; Kd_int(i,j,K) = 0. ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + ML_depth_2d(i,j) = 0.0 + BBL_depth_2d(i,j) = 0.0 +#else CS%ML_depth(i,j) = 0.0 CS%BBL_depth(i,j) = 0.0 +#endif endif ; enddo ! Close of i-loop - Note the unusual loop order, with k-loops inside i-loops. - do K=1,nz+1 ; do i=is,ie ; Kd_int(i,j,K) = Kd_2d(i,K) ; enddo ; enddo - enddo ! j-loop + +#ifdef __NVCOMPILER_OPENMP_GPU + ! Copy the results back to the host and tear down the driver-level data region. + !$omp target exit data map(from: Kd_int, ML_depth_2d, BBL_depth_2d) & + !$omp map(release: h_3d, u_3d, v_3d, dz_3d, TKE_forced, dSV_dT, dSV_dS, buoy_flux, & + !$omp ustar_2d, ustar_gustless_2d, CS, tv%T, tv%S, tv) + if (associated(Waves)) then + !$omp target exit data map(release: Waves) + endif + do j=js,je ; do i=is,ie + CS%ML_depth(i,j) = ML_depth_2d(i,j) ; CS%BBL_depth(i,j) = BBL_depth_2d(i,j) + enddo ; enddo +#endif if (CS%id_Kd_ePBL_col_by_col > 0) call post_data_3d_final(CS%id_Kd_ePBL_col_by_col, CS%diag) if (CS%debug .and. BBL_mixing) then @@ -1281,8 +1410,12 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, !/ Here we get mstar, which is the ratio of convective TKE driven mixing to UStar**3 if (CS%Use_LT) then +#ifdef __NVCOMPILER_OPENMP_GPU + call get_Langmuir_Number_LF17(LA, GV, US, abs(MLD_guess), u_star_mean, CS%wave_lf17) +#else call get_Langmuir_Number(LA, G, GV, US, abs(MLD_guess), u_star_mean, i, j, dz, Waves, & U_H=u, V_H=v) +#endif call find_mstar(CS, US, B_flux, u_star, MLD_guess, absf, .false., & mstar_total, Langmuir_Number=La, Convect_Langmuir_Number=LAmod,& mstar_LT=mstar_LT) diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index acc61550d7..fa377a338d 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -37,6 +37,8 @@ module MOM_wave_interface ! called in step_mom. public get_Langmuir_Number ! Public interface to compute Langmuir number called from ! ePBL or KPP routines. +public get_Langmuir_Number_LF17 ! Device-callable LF17-only Langmuir-number kernel (GPU offload). +public set_wave_LF17_params ! Populate a device-mappable wave_LF17_params bundle from a Waves CS. public Stokes_PGF ! Public interface to compute Stokes-shear induced pressure gradient force anomaly public StokesMixing ! NOT READY - Public interface to add down-Stokes gradient ! momentum mixing (e.g. the approach of Harcourt 2013/2015) @@ -260,6 +262,26 @@ module MOM_wave_interface end type wave_parameters_CS +!> A small, allocatable-free bundle of the loop-invariant wave_parameters_CS scalars needed by the +!! LF17 statistical-wave Langmuir path. Unlike the wave_parameters_CS pointer (whose allocatable +!! components do not attach under nvfortran `-gpu=mem:separate`), a plain scalar derived type maps to +!! the device cleanly, so the on-device LF17 chain reads these instead of dereferencing Waves. +type, public :: wave_LF17_params + real :: LA_FracHBL !< Fraction of OSBL for averaging Langmuir number [nondim] + real :: rho_air !< A typical density of air at sea level [R ~> kg m-3] + real :: nu_air !< The viscosity of air [Z2 T-1 ~> m2 s-1] + real :: rho_ocn !< A typical surface density of seawater [R ~> kg m-3] + real :: SWH_from_u10sq !< Factor converting the square of the 10 m wind speed to + !! significant wave height [Z T2 L-2 ~> s2 m-1] + real :: vonKar !< The von Karman coefficient as used in the wave code [nondim] + real :: Charnock_slope_U10 !< Slope of the Charnock coefficient/U10 relationship [T L-1 ~> s m-1] + real :: Charnock_min !< Minimum value of the Charnock coefficient [nondim] + real :: Charnock_intercept !< Intercept of the Charnock coefficient/U10 relationship [nondim] + real :: I_g_Earth !< The inverse of the gravitational acceleration [T2 Z L-2 ~> s2 m-1] + integer :: answer_date !< The vintage of the order of arithmetic and expressions in the + !! wave calculations. +end type wave_LF17_params + ! Switches needed in import_stokes_drift !>@{ Enumeration values for the wave method integer, parameter :: TESTPROF = 0, SURFBANDS = 1, DHH85 = 2, LF17 = 3, EFACTOR = 4, NULL_WaveMethod = -99 @@ -279,7 +301,7 @@ module MOM_wave_interface !> GPU port: the Li & Fox-Kemper statistical-wave Langmuir-number kernels are device-callable so the !! ePBL column solver can evaluate the Langmuir number on the device; on GPU builds only the LF17 !! (USE_LA_LI2016) statistical-wave path is compiled (the wave-model branches are host-only). -!$omp declare target(get_Langmuir_Number, get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) +!$omp declare target(get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) contains @@ -1222,7 +1244,7 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & !Local Variables real :: LA_STK ! Surface-layer averaged Stokes drift magnitude [L T-1 ~> m s-1] -#ifndef __NVCOMPILER_OPENMP_GPU + type(wave_LF17_params) :: p_lf17 ! Device-mappable bundle of LF17 wave scalars (for the LF17 branch) real :: Top, Bottom, MidPoint ! Positions within each layer [Z ~> m] real :: Dpt_LASL ! Averaging depth for Stokes drift [Z ~> m] real :: ShearDirection ! Shear angular direction from atan2 [radians] @@ -1232,15 +1254,7 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & real, dimension(SZK_(GV)) :: US_H, VS_H ! Profiles of Stokes velocities [L T-1 ~> m s-1] real, allocatable :: StkBand_X(:), StkBand_Y(:) ! Stokes drifts by band [L T-1 ~> m s-1] integer :: k, BB -#endif -#ifdef __NVCOMPILER_OPENMP_GPU - ! Device path: only the LF17 statistical-wave method (USE_LA_LI2016) is supported, and misalignment - ! is not applied; other WaveMethods / LA_Misalignment on a GPU build are rejected by a host-side - ! guard. This reproduces the host LF17 branch (get_StokesSL_LiFoxKemper sets LA directly; Dpt_LASL - ! is not used on that path). - call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, Waves, LA_STK, LA) -#else ! Compute averaging depth for Stokes drift (negative) Dpt_LASL = -1.0*max(Waves%LA_FracHBL*HBL, Waves%LA_HBL_min) @@ -1301,7 +1315,8 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & call Get_SL_Average_Prof( GV, Dpt_LASL, dz, VS_H, LA_STKy) LA_STK = sqrt((LA_STKX**2) + (LA_STKY**2)) elseif (Waves%WaveMethod==LF17) then - call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, Waves, LA_STK, LA) + call set_wave_LF17_params(Waves, p_lf17) + call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, p_lf17, LA_STK, LA) elseif (Waves%WaveMethod==Null_WaveMethod) then call MOM_error(FATAL, "Get_Langmuir_number called without defining a WaveMethod. "//& "Suggest to make sure USE_LT is set/overridden to False or choose "//& @@ -1319,10 +1334,46 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & WaveDirection = atan2(LA_STKy, LA_STKx) LA = LA / sqrt(max(1.e-8, cos( WaveDirection - ShearDirection))) endif -#endif end subroutine get_Langmuir_Number +!> Device-callable Langmuir-number kernel for the LF17 statistical-wave method (USE_LA_LI2016). +!! This is the GPU-offload entry point (e.g. from ePBL_column): it has no optional arguments — which +!! nvfortran cannot emit `declare target` device code for — and covers only the LF17 path, which is +!! the only WaveMethod supported on the GPU port. It is bit-for-bit the LF17 branch of +!! get_Langmuir_Number (get_StokesSL_LiFoxKemper sets LA directly; no misalignment is applied). +subroutine get_Langmuir_Number_LF17(LA, GV, US, HBL, ustar, p) + real, intent(out) :: LA !< Langmuir number [nondim] + type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + real, intent(in) :: HBL !< (Positive) thickness of boundary layer [Z ~> m] + real, intent(in) :: ustar !< Friction velocity [Z T-1 ~> m s-1] + type(wave_LF17_params), intent(in) :: p !< Device-mappable bundle of LF17 wave scalars. +!$omp declare target + + real :: LA_STK ! Surface-layer averaged Stokes drift magnitude [L T-1 ~> m s-1] + + call get_StokesSL_LiFoxKemper(ustar, HBL*p%LA_FracHBL, GV, US, p, LA_STK, LA) +end subroutine get_Langmuir_Number_LF17 + +!> Populate a device-mappable wave_LF17_params bundle from a Waves control structure (host-side). +subroutine set_wave_LF17_params(Waves, p) + type(Wave_parameters_CS), pointer, intent(in) :: Waves !< Surface wave control structure. + type(wave_LF17_params), intent(out) :: p !< The populated scalar bundle. + + p%LA_FracHBL = Waves%LA_FracHBL + p%rho_air = Waves%rho_air + p%nu_air = Waves%nu_air + p%rho_ocn = Waves%rho_ocn + p%SWH_from_u10sq = Waves%SWH_from_u10sq + p%vonKar = Waves%vonKar + p%Charnock_slope_U10 = Waves%Charnock_slope_U10 + p%Charnock_min = Waves%Charnock_min + p%Charnock_intercept = Waves%Charnock_intercept + p%I_g_Earth = Waves%I_g_Earth + p%answer_date = Waves%answer_date +end subroutine set_wave_LF17_params + !> function to return the wave method string set in the param file function get_wave_method(CS) character(:), allocatable :: get_wave_method @@ -1369,7 +1420,7 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) real, intent(in) :: hbl !< boundary layer depth [Z ~> m]. type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type - type(wave_parameters_CS), pointer :: CS !< Wave parameter Control structure + type(wave_LF17_params), intent(in) :: CS !< Device-mappable bundle of LF17 wave scalars real, intent(out) :: UStokes_SL !< Surface layer averaged Stokes drift [L T-1 ~> m s-1] real, intent(out) :: LA !< Langmuir number [nondim] ! Local variables @@ -2080,7 +2131,7 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) real, intent(out) :: U10 !< 10-m neutral wind speed [L T-1 ~> m s-1] type(verticalGrid_type), intent(in) :: GV !< vertical grid type type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type - type(wave_parameters_CS), pointer :: CS !< Wave parameter Control structure + type(wave_LF17_params), intent(in) :: CS !< Device-mappable bundle of LF17 wave scalars ! Local variables real :: z0sm, z0, z0rough ! Roughness lengths [Z ~> m] From 393b8256a6ceb17268489cbfb12a78d93cdb4c50 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 16:19:28 +1000 Subject: [PATCH 36/45] thickness_diffuse: route density derivatives through the device EOS dispatcher (increment 1) Offload prep for the thickness_diffuse_full GPU port. Swap the three live calculate_density_derivs generic array calls (the u-point slopes :933, the v-point slopes :1246, and the surface u block :1545) for the element-wise device-callable dispatcher calculate_density_derivs_elem_loc, folded into the existing fill loops so the evaluated range matches the array call exactly. get_EOS_form_and_scaling resolves the form + unit factors once (host); the scaling is always applied (unity-exact for the unscaled case), reproducing the former array calls bit-for-bit. Adds the same GPU-build FATAL guard as find_N2 for EOS forms without a device kernel. This mirrors the find_N2/kappa EOS-dispatcher pattern; the calculate_density _second_derivs (Stanley) path is left untouched (dead in benchmark_ALE). No device directives execute yet (elem_loc runs on the host here) -- this is prep so the u/v slope loops can be offloaded next without a polymorphic EOS call on the device. The GPU port of this file follows the structure of Ed Yang's edoyango/port/thickness_diffuse branch (the (i,j,k)-blocking and directive placement) as a reference; his branch is not merged (it is stale and overlaps our density/EOS work) but informs the approach. BITWISE: neutral -- benchmark_ALE np=1 ocean.stats is byte-identical to the pre-increment (ePBL-device) baseline (tmp_local_artifacts/td_baseline_ocean .stats). The only residual vs the canonical golden repro_ocean.stats is the separately-deferred ePBL x**(1/n) transcendental floor (~1e-13 in energy), which this commit does not touch. Co-authored-by: Edward Yang --- .../lateral/MOM_thickness_diffuse.F90 | 47 +++++++++++++++---- 1 file changed, 39 insertions(+), 8 deletions(-) diff --git a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 index 15cb9bc9ea..0393cd2c5c 100644 --- a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 +++ b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 @@ -13,6 +13,8 @@ module MOM_thickness_diffuse use MOM_error_handler, only : MOM_error, FATAL, WARNING, is_root_pe use MOM_EOS, only : calculate_density, calculate_density_derivs, EOS_domain use MOM_EOS, only : calculate_density_second_derivs +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_file_parser, only : get_param, log_version, param_file_type use MOM_grid, only : ocean_grid_type use MOM_io, only : MOM_read_data, slasher @@ -784,6 +786,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV integer :: nk_linear ! The number of layers over which the streamfunction goes to 0. real :: G_rho0 ! g/Rho0 [L2 R-1 Z-1 T-2 ~> m4 kg-1 s-2]. real :: Rho_avg ! The in situ density averaged to an interface [R ~> kg m-3] + integer :: eos_form ! The equation-of-state form code, for the device-callable EOS dispatcher. + real :: eos_kg_m3_to_R ! Factor converting the EOS kernel's density to model units [R m3 kg-1 ~> 1] + real :: eos_C_to_degC ! Factor converting model temperature to the EOS kernel's degC [degC C-1 ~> 1] + real :: eos_S_to_ppt ! Factor converting model salinity to the EOS kernel's ppt [ppt S-1 ~> 1] + real :: eos_RL2_T2_to_Pa ! Factor converting model pressure to Pa [Pa T2 R-1 L-2 ~> 1] real :: N2_floor ! A floor for N2 to avoid degeneracy in the elliptic solver ! times unit conversion factors [L2 Z-2 T-2 ~> s-2] real :: N2_unlim ! An unlimited estimate of the buoyancy frequency @@ -898,12 +905,29 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV EOSdom_v(:) = EOS_domain(G%HI) EOSdom_h1(:) = EOS_domain(G%HI, halo=1) + ! Resolve the EOS form and its unit-scaling once (host v-table lookup), so the density derivatives + ! can be evaluated through the device-callable elementwise dispatcher instead of the polymorphic + ! array interface. The scaling is always applied; for the unscaled case the factors are exactly + ! 1.0, reproducing the former array calculate_density_derivs calls bit-for-bit. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_EOS) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") +#endif + endif + !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,G,GV,US,pres,T,S, & !$OMP nk_linear,IsdB,tv,h,h_neglect,e,dz,dz_neglect,dz_neglect2, & !$OMP h_neglect2,hn_2,I_slope_max2,int_slope_u,KH_u,uhtot, & !$OMP h_frac,h_avail_rsum,uhD,h_avail,Work_u,CS,slope_x,cg1, & !$OMP diag_sfn_x,diag_sfn_unlim_x,N2_floor,EOSdom_u,EOSdom_h1, & - !$OMP use_stanley,present_slope_x,G_rho0,Slope_x_PE,hN2_x_PE) & + !$OMP use_stanley,present_slope_x,G_rho0,Slope_x_PE,hN2_x_PE, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & !$OMP private(drdiA,drdiB,drdkL,drdkR,pres_u,T_u,S_u,G_scale, & !$OMP drho_dT_u,drho_dS_u,hg2A,hg2B,hg2L,hg2R,haA, & !$OMP drho_dT_dT_h,scrap,pres_h,T_h,S_h,N2_unlim, & @@ -929,9 +953,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV pres_u(I) = 0.5*(pres(i,j,K) + pres(i+1,j,K)) T_u(I) = 0.25*((T(i,j,k) + T(i+1,j,k)) + (T(i,j,k-1) + T(i+1,j,k-1))) S_u(I) = 0.25*((S(i,j,k) + S(i+1,j,k)) + (S(i,j,k-1) + S(i+1,j,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_u(I), eos_S_to_ppt*S_u(I), & + eos_RL2_T2_to_Pa*pres_u(I), drho_dT_u(I), drho_dS_u(I)) + drho_dT_u(I) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_u(I) + drho_dS_u(I) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_u(I) enddo - call calculate_density_derivs(T_u, S_u, pres_u, drho_dT_u, drho_dS_u, & - tv%eqn_of_state, EOSdom_u) endif if (use_stanley) then do i=is-1,ie+1 @@ -1218,7 +1244,8 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV !$OMP h_neglect2,int_slope_v,KH_v,vhtot,h_frac,h_avail_rsum, & !$OMP I_slope_max2,vhD,h_avail,Work_v,CS,slope_y,cg1,hn_2,& !$OMP diag_sfn_y,diag_sfn_unlim_y,N2_floor,EOSdom_v,use_stanley,& - !$OMP present_slope_y,G_rho0,Slope_y_PE,hN2_y_PE) & + !$OMP present_slope_y,G_rho0,Slope_y_PE,hN2_y_PE, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & !$OMP private(drdjA,drdjB,drdkL,drdkR,pres_v,T_v,S_v,S_h,S_hr, & !$OMP drho_dT_v,drho_dS_v,hg2A,hg2B,hg2L,hg2R,haA,G_scale, & !$OMP drho_dT_dT_h,drho_dT_dT_hr,scrap,pres_h,T_h,T_hr, & @@ -1242,9 +1269,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV pres_v(i) = 0.5*(pres(i,j,K) + pres(i,j+1,K)) T_v(i) = 0.25*((T(i,j,k) + T(i,j+1,k)) + (T(i,j,k-1) + T(i,j+1,k-1))) S_v(i) = 0.25*((S(i,j,k) + S(i,j+1,k)) + (S(i,j,k-1) + S(i,j+1,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_v(i), eos_S_to_ppt*S_v(i), & + eos_RL2_T2_to_Pa*pres_v(i), drho_dT_v(i), drho_dS_v(i)) + drho_dT_v(i) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_v(i) + drho_dS_v(i) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_v(i) enddo - call calculate_density_derivs(T_v, S_v, pres_v, drho_dT_v, drho_dS_v, & - tv%eqn_of_state, EOSdom_v) endif if (use_stanley) then do i=is,ie @@ -1541,9 +1570,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV pres_u(I) = 0.5*(pres(i,j,1) + pres(i+1,j,1)) T_u(I) = 0.5*(T(i,j,1) + T(i+1,j,1)) S_u(I) = 0.5*(S(i,j,1) + S(i+1,j,1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_u(I), eos_S_to_ppt*S_u(I), & + eos_RL2_T2_to_Pa*pres_u(I), drho_dT_u(I), drho_dS_u(I)) + drho_dT_u(I) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_u(I) + drho_dS_u(I) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_u(I) enddo - call calculate_density_derivs(T_u, S_u, pres_u, drho_dT_u, drho_dS_u, & - tv%eqn_of_state, EOSdom_u ) endif do I=is-1,ie uhD(I,j,1) = -uhtot(I,j) From 77eb4422cde084261250cc04a7eec049be9cc270 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 16:56:29 +1000 Subject: [PATCH 37/45] thickness_diffuse: promote slope scratch to 3-D (increment 2a, offload prep) Promote the ten per-column slope-scratch arrays of thickness_diffuse_full from 2-D (i,K) to 3-D (i,j,K) so the u/v slope+flux loops can later be fully collapsed onto the device with the scratch device-resident across the slope->flux passes: drdi_u/drdj_v, drdkDe_u/_v, dzN2_u/_v, Sfn_unlim_u/_v, slope2_Ratio_u/_v (v-arrays carry the SZJB_ symmetric-memory j-extent). Declarations + every index site updated (~50 sites across the u-block and v-block). c2_dz_u/_v are left 2-D (FGNV path, dead in benchmark_ALE, host-only). No directives yet: this is a pure storage transform, computation unchanged. The promoted arrays remain in the CPU !$OMP parallel do private() clauses transitionally (per-thread copies, still bit-for-bit; the do-j loop writes disjoint j-slices) -- the target-teams-loop offload (increment 2b) replaces those clauses. BITWISE: neutral -- benchmark_ALE np=1 ocean.stats byte-identical to the pre-increment ePBL-device baseline (tmp_local_artifacts/td_baseline_ocean.stats, diff empty). Only residual vs the canonical golden is the separately-deferred ePBL transcendental floor. The GPU offload of this file follows the structure of Ed Yang's edoyango/port/thickness_diffuse (3-D scratch promotion + full collapse) as a reference. Co-authored-by: Edward Yang --- .../lateral/MOM_thickness_diffuse.F90 | 126 +++++++++--------- 1 file changed, 63 insertions(+), 63 deletions(-) diff --git a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 index 0393cd2c5c..f03e7a6b78 100644 --- a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 +++ b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 @@ -735,11 +735,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV real :: drdjA, drdjB ! Along layer meridional potential density gradients in the layers above (A) ! and below (B) the interface times the grid spacing [R ~> kg m-3]. real :: drdkL, drdkR ! Vertical density differences across an interface [R ~> kg m-3]. - real :: drdi_u(SZIB_(G),SZK_(GV)) ! Copy of drdi at u-points [R ~> kg m-3]. - real :: drdj_v(SZI_(G),SZK_(GV)) ! Copy of drdj at v-points [R ~> kg m-3]. - real :: drdkDe_u(SZIB_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at u-points + real :: drdi_u(SZIB_(G),SZJ_(G),SZK_(GV)) ! Copy of drdi at u-points [R ~> kg m-3]. + real :: drdj_v(SZI_(G),SZJB_(G),SZK_(GV)) ! Copy of drdj at v-points [R ~> kg m-3]. + real :: drdkDe_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at u-points ! [Z R ~> kg m-2]. - real :: drdkDe_v(SZI_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at v-points + real :: drdkDe_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at v-points ! [Z R ~> kg m-2]. real :: hg2A, hg2B, hg2L, hg2R ! Squares of geometric mean thicknesses [H2 ~> m2 or kg2 m-4]. real :: haA, haB, haL, haR ! Arithmetic mean thicknesses [H ~> m or kg m-2]. @@ -753,16 +753,16 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV real :: dz_harm ! Harmonic mean layer vertical extent [Z ~> m]. real :: c2_dz_u(SZIB_(G),SZK_(GV)+1) ! Wave speed squared divided by dz at u-points [L2 Z-1 T-2 ~> m s-2] real :: c2_dz_v(SZI_(G),SZK_(GV)+1) ! Wave speed squared divided by dz at v-points [L2 Z-1 T-2 ~> m s-2] - real :: dzN2_u(SZIB_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above u-points times + real :: dzN2_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above u-points times ! rescaling factors from vertical to horizontal distances [L2 Z-1 T-2 ~> m s-2] - real :: dzN2_v(SZI_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above v-points times + real :: dzN2_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above v-points times ! rescaling factors from vertical to horizontal distances [L2 Z-1 T-2 ~> m s-2] real :: Sfn_est ! A preliminary estimate (before limiting) of the overturning ! streamfunction [H L2 T-1 ~> m3 s-1 or kg s-1]. - real :: Sfn_unlim_u(SZIB_(G),SZK_(GV)+1) ! Volume streamfunction for u-points [Z L2 T-1 ~> m3 s-1] - real :: Sfn_unlim_v(SZI_(G),SZK_(GV)+1) ! Volume streamfunction for v-points [Z L2 T-1 ~> m3 s-1] - real :: slope2_Ratio_u(SZIB_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] - real :: slope2_Ratio_v(SZI_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] + real :: Sfn_unlim_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Volume streamfunction for u-points [Z L2 T-1 ~> m3 s-1] + real :: Sfn_unlim_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Volume streamfunction for v-points [Z L2 T-1 ~> m3 s-1] + real :: slope2_Ratio_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] + real :: slope2_Ratio_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] real :: Sfn_in_h ! The overturning streamfunction [H L2 T-1 ~> m3 s-1 or kg s-1] (note that ! the units are different from other Sfn vars). real :: Sfn_safe ! The streamfunction that goes linearly back to 0 at the surface @@ -937,7 +937,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV !$OMP Sfn_unlim_u,Rho_avg,drdi_u,drdkDe_u,c2_dz_u, & !$OMP Sfn_safe,Sfn_est,Sfn_in_h,calc_derivatives) do j=js,je - do I=is-1,ie ; dzN2_u(I,1) = 0. ; dzN2_u(I,nz+1) = 0. ; enddo + do I=is-1,ie ; dzN2_u(I,j,1) = 0. ; dzN2_u(I,j,nz+1) = 0. ; enddo do K=nz,2,-1 if (find_work .and. .not.(use_EOS)) then drdiA = 0.0 ; drdiB = 0.0 @@ -986,9 +986,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drho_dS_u(I) * (S(i,j,k)-S(i,j,k-1))) drdkR = (drho_dT_u(I) * (T(i+1,j,k)-T(i+1,j,k-1)) + & drho_dS_u(I) * (S(i+1,j,k)-S(i+1,j,k-1))) - drdkDe_u(I,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) + drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode - drdkDe_u(I,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) + drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) endif if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in @@ -998,7 +998,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drdiB = drdiB + 0.5 * ((drho_dT_dT_h(i+1) * tv%varT(i+1,j,k)) - & (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif - if (find_work) drdi_u(I,k) = drdiB + if (find_work) drdi_u(I,j,k) = drdiB if (k > nk_linear) then if (use_EOS) then @@ -1041,14 +1041,14 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzaA = 0.5*(dz(i,j,k-1) + dz(i+1,j,k-1)) + dz_neglect dzaB = 0.5*(dz(i,j,k) + dz(i+1,j,k)) + dz_neglect ! dzN2_u is used with the FGNV streamfunction formulation - dzN2_u(I,K) = (0.5 * ( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) + dzN2_u(I,j,K) = (0.5 * ( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) if (find_work .and. CS%GM_src_alt) & hN2_x_PE(I,j,k) = (0.5 * ( hg2A / haA + hg2B / haB )) * max(N2_unlim, N2_floor) endif if (present_slope_x) then Slope = slope_x(I,j,k) - slope2_Ratio_u(I,K) = Slope**2 * I_slope_max2 + slope2_Ratio_u(I,j,K) = Slope**2 * I_slope_max2 else ! Use the harmonic mean thicknesses to weight the horizontal gradients. ! These unnormalized weights have been rearranged to minimize divisions. @@ -1062,10 +1062,10 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV mag_grad2 = (US%Z_to_L*drdx)**2 + drdz**2 if (mag_grad2 > 0.0) then Slope = drdx / sqrt(mag_grad2) - slope2_Ratio_u(I,K) = Slope**2 * I_slope_max2 + slope2_Ratio_u(I,j,K) = Slope**2 * I_slope_max2 else ! Just in case mag_grad2 = 0 ever. Slope = 0.0 - slope2_Ratio_u(I,K) = 1.0e20 ! Force the use of the safe streamfunction. + slope2_Ratio_u(I,j,K) = 1.0e20 ! Force the use of the safe streamfunction. endif endif @@ -1073,7 +1073,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! that ignore density gradients along layers. Slope = (1.0 - int_slope_u(I,j,K)) * Slope + & int_slope_u(I,j,K) * ((e(i+1,j,K)-e(i,j,K)) * G%IdxCu(I,j)) - slope2_Ratio_u(I,K) = (1.0 - int_slope_u(I,j,K)) * slope2_Ratio_u(I,K) + slope2_Ratio_u(I,j,K) = (1.0 - int_slope_u(I,j,K)) * slope2_Ratio_u(I,j,K) if (CS%MEKE_src_slope_bug) then Slope_x_PE(I,j,k) = MIN(Slope, CS%slope_max) @@ -1085,24 +1085,24 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope ! Estimate the streamfunction at each interface [H L2 T-1 ~> m3 s-1 or kg s-1]. - Sfn_unlim_u(I,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope + Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope ! Avoid moving dense water upslope from below the level of ! the bottom on the receiving side. - if (Sfn_unlim_u(I,K) > 0.0) then ! The flow below this interface is positive. + if (Sfn_unlim_u(I,j,K) > 0.0) then ! The flow below this interface is positive. if (e(i,j,K) < e(i+1,j,nz+1)) then - Sfn_unlim_u(I,K) = 0.0 ! This is not uhtot, because it may compensate for + Sfn_unlim_u(I,j,K) = 0.0 ! This is not uhtot, because it may compensate for ! deeper flow in very unusual cases. elseif (e(i+1,j,nz+1) > e(i,j,K+1)) then ! Scale the transport with the fraction of the donor layer above ! the bottom on the receiving side. - Sfn_unlim_u(I,K) = Sfn_unlim_u(I,K) * ((e(i,j,K) - e(i+1,j,nz+1)) / & + Sfn_unlim_u(I,j,K) = Sfn_unlim_u(I,j,K) * ((e(i,j,K) - e(i+1,j,nz+1)) / & ((e(i,j,K) - e(i,j,K+1)) + dz_neglect)) endif else - if (e(i+1,j,K) < e(i,j,nz+1)) then ; Sfn_unlim_u(I,K) = 0.0 + if (e(i+1,j,K) < e(i,j,nz+1)) then ; Sfn_unlim_u(I,j,K) = 0.0 elseif (e(i,j,nz+1) > e(i+1,j,K+1)) then - Sfn_unlim_u(I,K) = Sfn_unlim_u(I,K) * ((e(i+1,j,K) - e(i,j,nz+1)) / & + Sfn_unlim_u(I,j,K) = Sfn_unlim_u(I,j,K) * ((e(i+1,j,K) - e(i,j,nz+1)) / & ((e(i+1,j,K) - e(i+1,j,K+1)) + dz_neglect)) endif endif @@ -1114,14 +1114,14 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV Slope = (e(i+1,j,K)-e(i,j,K)) * G%IdxCu_OBCmask(I,j) endif if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope - Sfn_unlim_u(I,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope - dzN2_u(I,K) = GV%g_prime(K) + Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope + dzN2_u(I,j,K) = GV%g_prime(K) endif ! if (use_EOS) else ! if (k > nk_linear) - dzN2_u(I,K) = N2_floor * dz_neglect - Sfn_unlim_u(I,K) = 0. + dzN2_u(I,j,K) = N2_floor * dz_neglect + Sfn_unlim_u(I,j,K) = 0. endif ! if (k > nk_linear) - if (CS%id_sfn_unlim_x>0) diag_sfn_unlim_x(I,j,K) = Sfn_unlim_u(I,K) + if (CS%id_sfn_unlim_x>0) diag_sfn_unlim_x(I,j,K) = Sfn_unlim_u(I,j,K) enddo ! i-loop enddo ! k-loop @@ -1136,12 +1136,12 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV do I=is-1,ie if (G%OBCmaskCu(I,j)>0.) then do K=2,nz - Sfn_unlim_u(I,K) = (1. + CS%FGNV_scale) * Sfn_unlim_u(I,K) + Sfn_unlim_u(I,j,K) = (1. + CS%FGNV_scale) * Sfn_unlim_u(I,j,K) enddo - call streamfn_solver(nz, c2_dz_u(I,:), dzN2_u(I,:), Sfn_unlim_u(I,:)) + call streamfn_solver(nz, c2_dz_u(I,:), dzN2_u(I,j,:), Sfn_unlim_u(I,j,:)) else do K=2,nz - Sfn_unlim_u(I,K) = 0. + Sfn_unlim_u(I,j,K) = 0. enddo endif enddo @@ -1171,9 +1171,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif ! Determine the actual streamfunction at each interface. - Sfn_est = (Z_to_H*Sfn_unlim_u(I,K) + slope2_Ratio_u(I,K)*Sfn_safe) / (1.0 + slope2_Ratio_u(I,K)) + Sfn_est = (Z_to_H*Sfn_unlim_u(I,j,K) + slope2_Ratio_u(I,j,K)*Sfn_safe) / (1.0 + slope2_Ratio_u(I,j,K)) else ! When use_EOS is false, the layers are constant density. - Sfn_est = Z_to_H*Sfn_unlim_u(I,K) + Sfn_est = Z_to_H*Sfn_unlim_u(I,j,K) endif ! Make sure that there is enough mass above to allow the streamfunction @@ -1228,8 +1228,8 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif Work_u(I,j) = Work_u(I,j) + G_scale * & - ( uhtot(I,j) * drdkDe_u(I,K) - & - (uhD(I,j,k) * drdi_u(I,k)) * 0.25 * & + ( uhtot(I,j) * drdkDe_u(I,j,K) - & + (uhD(I,j,k) * drdi_u(I,j,k)) * 0.25 * & ((e(i,j,K) + e(i,j,K+1)) + (e(i+1,j,K) + e(i+1,j,K+1))) ) endif @@ -1308,9 +1308,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drho_dS_v(i) * (S(i,j,k)-S(i,j,k-1))) drdkR = (drho_dT_v(i) * (T(i,j+1,k)-T(i,j+1,k-1)) + & drho_dS_v(i) * (S(i,j+1,k)-S(i,j+1,k-1))) - drdkDe_v(i,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) + drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode - drdkDe_v(i,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) + drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) endif if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in @@ -1321,7 +1321,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif - if (find_work) drdj_v(i,k) = drdjB + if (find_work) drdj_v(i,J,k) = drdjB if (k > nk_linear) then if (use_EOS) then @@ -1366,13 +1366,13 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzaB = 0.5*(dz(i,j,k) + dz(i,j+1,k)) + dz_neglect ! dzN2_v is used with the FGNV streamfunction formulation - dzN2_v(i,K) = (0.5*( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) + dzN2_v(i,J,K) = (0.5*( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) if (find_work .and. CS%GM_src_alt) & hN2_y_PE(i,J,k) = (0.5*( hg2A / haA + hg2B / haB )) * max(N2_unlim, N2_floor) endif if (present_slope_y) then Slope = slope_y(i,J,k) - slope2_Ratio_v(i,K) = Slope**2 * I_slope_max2 + slope2_Ratio_v(i,J,K) = Slope**2 * I_slope_max2 else ! Use the harmonic mean thicknesses to weight the horizontal gradients. ! These unnormalized weights have been rearranged to minimize divisions. @@ -1386,10 +1386,10 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV mag_grad2 = (US%Z_to_L*drdy)**2 + drdz**2 if (mag_grad2 > 0.0) then Slope = drdy / sqrt(mag_grad2) - slope2_Ratio_v(i,K) = Slope**2 * I_slope_max2 + slope2_Ratio_v(i,J,K) = Slope**2 * I_slope_max2 else ! Just in case mag_grad2 = 0 ever. Slope = 0.0 - slope2_Ratio_v(i,K) = 1.0e20 ! Force the use of the safe streamfunction. + slope2_Ratio_v(i,J,K) = 1.0e20 ! Force the use of the safe streamfunction. endif endif @@ -1397,7 +1397,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! that ignore density gradients along layers. Slope = (1.0 - int_slope_v(i,J,K)) * Slope + & int_slope_v(i,J,K) * ((e(i,j+1,K)-e(i,j,K)) * G%IdyCv(i,J)) - slope2_Ratio_v(i,K) = (1.0 - int_slope_v(i,J,K)) * slope2_Ratio_v(i,K) + slope2_Ratio_v(i,J,K) = (1.0 - int_slope_v(i,J,K)) * slope2_Ratio_v(i,J,K) if (CS%MEKE_src_slope_bug) then Slope_y_PE(i,J,k) = MIN(Slope, CS%slope_max) @@ -1408,24 +1408,24 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope - Sfn_unlim_v(i,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) + Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) ! Avoid moving dense water upslope from below the level of ! the bottom on the receiving side. - if (Sfn_unlim_v(i,K) > 0.0) then ! The flow below this interface is positive. + if (Sfn_unlim_v(i,J,K) > 0.0) then ! The flow below this interface is positive. if (e(i,j,K) < e(i,j+1,nz+1)) then - Sfn_unlim_v(i,K) = 0.0 ! This is not vhtot, because it may compensate for + Sfn_unlim_v(i,J,K) = 0.0 ! This is not vhtot, because it may compensate for ! deeper flow in very unusual cases. elseif (e(i,j+1,nz+1) > e(i,j,K+1)) then ! Scale the transport with the fraction of the donor layer above ! the bottom on the receiving side. - Sfn_unlim_v(i,K) = Sfn_unlim_v(i,K) * ((e(i,j,K) - e(i,j+1,nz+1)) / & + Sfn_unlim_v(i,J,K) = Sfn_unlim_v(i,J,K) * ((e(i,j,K) - e(i,j+1,nz+1)) / & ((e(i,j,K) - e(i,j,K+1)) + dz_neglect)) endif else - if (e(i,j+1,K) < e(i,j,nz+1)) then ; Sfn_unlim_v(i,K) = 0.0 + if (e(i,j+1,K) < e(i,j,nz+1)) then ; Sfn_unlim_v(i,J,K) = 0.0 elseif (e(i,j,nz+1) > e(i,j+1,K+1)) then - Sfn_unlim_v(i,K) = Sfn_unlim_v(i,K) * ((e(i,j+1,K) - e(i,j,nz+1)) / & + Sfn_unlim_v(i,J,K) = Sfn_unlim_v(i,J,K) * ((e(i,j+1,K) - e(i,j,nz+1)) / & ((e(i,j+1,K) - e(i,j+1,K+1)) + dz_neglect)) endif endif @@ -1437,14 +1437,14 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV Slope = (e(i,j+1,K)-e(i,j,K)) * G%IdyCv_OBCmask(i,J) endif if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope - Sfn_unlim_v(i,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) - dzN2_v(i,K) = GV%g_prime(K) + Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) + dzN2_v(i,J,K) = GV%g_prime(K) endif ! if (use_EOS) else ! if (k > nk_linear) - dzN2_v(i,K) = N2_floor * dz_neglect - Sfn_unlim_v(i,K) = 0. + dzN2_v(i,J,K) = N2_floor * dz_neglect + Sfn_unlim_v(i,J,K) = 0. endif ! if (k > nk_linear) - if (CS%id_sfn_unlim_y>0) diag_sfn_unlim_y(i,J,K) = Sfn_unlim_v(i,K) + if (CS%id_sfn_unlim_y>0) diag_sfn_unlim_y(i,J,K) = Sfn_unlim_v(i,J,K) enddo ! i-loop enddo ! k-loop @@ -1459,12 +1459,12 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV do i=is,ie if (G%OBCmaskCv(i,J)>0.) then do K=2,nz - Sfn_unlim_v(i,K) = (1. + CS%FGNV_scale) * Sfn_unlim_v(i,K) + Sfn_unlim_v(i,J,K) = (1. + CS%FGNV_scale) * Sfn_unlim_v(i,J,K) enddo - call streamfn_solver(nz, c2_dz_v(i,:), dzN2_v(i,:), Sfn_unlim_v(i,:)) + call streamfn_solver(nz, c2_dz_v(i,:), dzN2_v(i,J,:), Sfn_unlim_v(i,J,:)) else do K=2,nz - Sfn_unlim_v(i,K) = 0. + Sfn_unlim_v(i,J,K) = 0. enddo endif enddo @@ -1493,9 +1493,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif ! Find the actual streamfunction at each interface. - Sfn_est = (Z_to_H*Sfn_unlim_v(i,K) + slope2_Ratio_v(i,K)*Sfn_safe) / (1.0 + slope2_Ratio_v(i,K)) + Sfn_est = (Z_to_H*Sfn_unlim_v(i,J,K) + slope2_Ratio_v(i,J,K)*Sfn_safe) / (1.0 + slope2_Ratio_v(i,J,K)) else ! When use_EOS is false, the layers are constant density. - Sfn_est = Z_to_H*Sfn_unlim_v(i,K) + Sfn_est = Z_to_H*Sfn_unlim_v(i,J,K) endif ! Make sure that there is enough mass above to allow the streamfunction @@ -1548,8 +1548,8 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif Work_v(i,J) = Work_v(i,J) + G_scale * & - ( vhtot(i,J) * drdkDe_v(i,K) - & - (vhD(i,J,k) * drdj_v(i,k)) * 0.25 * & + ( vhtot(i,J) * drdkDe_v(i,J,K) - & + (vhD(i,J,k) * drdj_v(i,J,k)) * 0.25 * & ((e(i,j,K) + e(i,j,K+1)) + (e(i,j+1,K) + e(i,j+1,K+1))) ) endif From e10ee720e3db346c879ce5a306dbcc14c9d7cc69 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 18:27:07 +1000 Subject: [PATCH 38/45] thickness_diffuse: offload thickness_diffuse_full to GPU (increment 2b) Offload the whole thickness_diffuse_full compute (setup + u/v slope + streamfunction + flux + MEKE work) to the device, one target-data region, building on the 3-D scratch of increment 2a and the elem_loc EOS of increment 1. - CPU-neutral refactor first (verified bitwise-neutral before any directive): fuse the per-K EOS-fill do-I loop into the slope do-I loop (scalarizing the per-I EOS temporaries) so the slope nest is a tight do j / do K / do I; split the outer do-j into slope and flux j-loops; swap the flux loops to do I / serial do K (per-column uhtot/vhtot recurrence order preserved); move the now-3-D scratch out of the OMP private() clauses to shared. - Loop forms (all `#ifdef __NVCOMPILER_OPENMP_GPU`, original !$OMP parallel do kept under #else so CPU threading is retained): slope loops -> `target teams loop collapse(3)` over (j,K,I) with per-iteration scalars private; the uhtot/vhtot flux recurrence + the h_avail_rsum setup running sum -> collapse(2) + serial do K (KNOWLEDGE 4.1 branch 5); the remaining elementwise/init/assembly loops -> collapse(2). 16 device kernels generate, no W-0155. - Dead-in-benchmark_ALE paths #ifndef'd off the device with host FATAL guards: use_stanley (second derivs), FGNV streamfn + streamfn_solver + c2_dz, skeb_use_gm/STOCH, non-Boussinesq tv%SpV_avg, and the slope/sfn diagnostics. The layer-1 v-point EOS array call is converted to the elem_loc dispatcher. - Residency: one target enter/exit data region; map(to:) h (+update to), e, dz, T, S, Kh_u/v, int_slope_u/v, CS, slope_x/y (present-guarded) and the host-zeroed Slope_/hN2 PE arrays; map(alloc:) pres, dz-scratch, h_avail*, uhtot/vhtot, Work_u/v, uhD/vhD and the 3-D slope scratch; map(from:) uhD/vhD. All exit maps use `release` (never `delete`). - Pointer/allocatable hoist (the map-to-present / Waves-pointer-trap class): MEKE%GM_src, CS%GMwork and tv%p_surf hoisted to plain local arrays (copied in, written/read on device, copied back) so no allocatable component of a pointer/derived type is dereferenced on the device. Two fixes found via the crash -> compute-sanitizer -> fix loop: (1) tv%p_surf IS associated in benchmark_ALE, so hoist it rather than FATAL; (2) an illegal address that surfaced downstream in set_viscous_BBL (pbv attach) was traced (sanitizer + a revert-to-2a control build) to this region -- fixed by making the whole region pure OpenMP target (no `do concurrent` inside the target-data region) and by using map(release:) instead of map(delete:) on exit (KNOWLEDGE 8: a per-call delete can zero an outer/persistent mapping refcount). Note tv/tv%T/tv%S are deliberately NOT mapped here: on device the density inputs are the host-filled local arrays T/S from vert_fill_TS (mapped to: fresh), so tv%T/tv%S are never read on device. BITWISE: neutral -- benchmark_ALE np=1 ocean.stats byte-identical to the pre-increment ePBL-device baseline (tmp_local_artifacts/td_baseline_ocean.stats, diff empty), verified on a clean from-scratch rebuild. Device execution confirmed (16 kernels; full 24-step run, no CUDA fault; compute-sanitizer memcheck clean). Only residual vs the canonical golden is the separately- deferred ePBL x**(1/n) transcendental floor. Structure follows Ed Yang's edoyango/port/thickness_diffuse (3-D scratch + full collapse); our EOS goes through the elem_loc dispatcher rather than his EOS-3d stack. Co-authored-by: Edward Yang --- .../lateral/MOM_thickness_diffuse.F90 | 593 ++++++++++++------ 1 file changed, 411 insertions(+), 182 deletions(-) diff --git a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 index f03e7a6b78..0cc0382581 100644 --- a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 +++ b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 @@ -708,6 +708,12 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drho_dT_dT_hr ! The second derivative of density with temperature at h (+1) points [R C-2 ~> kg m-3 degC-2] real :: uhtot(SZIB_(G),SZJ_(G)) ! The vertical sum of uhD [H L2 T-1 ~> m3 s-1 or kg s-1]. real :: vhtot(SZI_(G),SZJB_(G)) ! The vertical sum of vhD [H L2 T-1 ~> m3 s-1 or kg s-1]. + real :: pres_us, T_us, S_us ! Scalar (per-iteration) pressure, temperature and salinity on the + ! interface at a u-point, for the fused device-callable EOS path. + real :: dT_us, dS_us ! Scalar density derivatives with T and S at a u-point. + real :: pres_vs, T_vs, S_vs ! Scalar (per-iteration) pressure, temperature and salinity on the + ! interface at a v-point, for the fused device-callable EOS path. + real :: dT_vs, dS_vs ! Scalar density derivatives with T and S at a v-point. real, dimension(SZIB_(G)) :: & T_u, & ! Temperature on the interface at the u-point [C ~> degC]. S_u, & ! Salinity on the interface at the u-point [S ~> ppt]. @@ -808,6 +814,17 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! applying limiters [H L2 T-1 ~> m3 s-1 or kg s-1] real, allocatable :: skeb_gm_work(:,:) ! Temp array to hold GM work for SKEB real, allocatable :: skeb_ebt_norm2(:,:) ! Used to normalize EBT for SKEB + real :: GM_src_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of MEKE%GM_src [R Z L2 T-3 ~> W m-2]. + ! MEKE%GM_src is an allocatable component of a pointer/derived + ! type, so it is copied to/from this plain array around the + ! device region rather than deep-mapped (the pointer trap). + real :: GMwork_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of CS%GMwork [R Z L2 T-3 ~> W m-2]. + real :: p_surf_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of the surface pressure tv%p_surf + ! [R L2 T-2 ~> Pa] (0 where tv%p_surf is not associated), so the + ! device reads a plain array rather than a pointer component of tv. + logical :: have_GM_src ! True if MEKE%GM_src is allocated (evaluated on the host). + logical :: have_GMwork ! True if CS%GMwork is allocated (evaluated on the host). + logical :: have_p_surf ! True if tv%p_surf is associated (evaluated on the host). logical :: present_slope_x, present_slope_y, calc_derivatives integer, dimension(2) :: EOSdom_u ! The shifted I-computational domain to use for equation of @@ -864,37 +881,98 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (CS%use_FGNV_streamfn .and. .not. associated(cg1)) call MOM_error(FATAL, & "cg1 must be associated when using FGNV streamfunction.") - !$OMP parallel default(shared) + ! --------------------------------------------------------------------------------------------- + ! GPU offload of thickness_diffuse_full. One target data region spans the whole compute below. + ! Notes on residency: tv%T/tv%S are NOT mapped — the density inputs on device are the host-filled + ! local arrays T and S (from vert_fill_TS), copied in fresh with map(to:). MEKE%GM_src and + ! CS%GMwork are allocatable components of a pointer/derived type, so they are hoisted to the plain + ! local arrays GM_src_loc/GMwork_loc (mapped, written on device, copied back) rather than + ! deep-mapped (the pointer-in-derived-type trap). The persistent, diabatic-mutated h is refreshed + ! with update to(h) (map-to-present is a no-copy). + ! --------------------------------------------------------------------------------------------- +#ifdef __NVCOMPILER_OPENMP_GPU + if (use_stanley) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: USE_STANLEY_GM is not supported on device.") + if (CS%use_FGNV_streamfn) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the FGNV streamfunction is not supported on device.") + if (skeb_use_gm) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: SKEB GM work is not supported on device.") + if (allocated(tv%SpV_avg)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the non-Boussinesq SpV_avg path is not supported on device.") + if ((CS%id_slope_x>0) .or. (CS%id_slope_y>0) .or. (CS%id_sfn_x>0) .or. (CS%id_sfn_y>0) .or. & + (CS%id_sfn_unlim_x>0) .or. (CS%id_sfn_unlim_y>0)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the streamfunction/slope diagnostics are not supported on device.") +#endif + + have_p_surf = associated(tv%p_surf) + have_GM_src = allocated(MEKE%GM_src) + have_GMwork = allocated(CS%GMwork) + + ! Hoist MEKE%GM_src / CS%GMwork to plain local arrays (copied back after the region). + GM_src_loc(:,:) = 0.0 ; GMwork_loc(:,:) = 0.0 + if (have_GM_src) then ; do j=js,je ; do i=is,ie ; GM_src_loc(i,j) = MEKE%GM_src(i,j) ; enddo ; enddo ; endif + if (have_GMwork) then ; do j=js,je ; do i=is,ie ; GMwork_loc(i,j) = CS%GMwork(i,j) ; enddo ; enddo ; endif + + ! Hoist the surface pressure onto a plain array so the device never dereferences the tv%p_surf pointer. + p_surf_loc(:,:) = 0.0 + if (have_p_surf) then + do j=js-1,je+1 ; do i=is-1,ie+1 ; p_surf_loc(i,j) = tv%p_surf(i,j) ; enddo ; enddo + endif + +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target enter data map(to: h, e, dz, T, S, Kh_u, Kh_v, int_slope_u, int_slope_v, CS, & + !$omp Slope_x_PE, Slope_y_PE, hN2_x_PE, hN2_y_PE, GM_src_loc, GMwork_loc, p_surf_loc) & + !$omp map(alloc: pres, h_avail, h_frac, h_avail_rsum, uhtot, vhtot, Work_u, Work_v, & + !$omp uhD, vhD, drdi_u, drdj_v, drdkDe_u, drdkDe_v, dzN2_u, dzN2_v, & + !$omp Sfn_unlim_u, Sfn_unlim_v, slope2_Ratio_u, slope2_Ratio_v) + if (present_slope_x) then + !$omp target enter data map(to: slope_x) + endif + if (present_slope_y) then + !$omp target enter data map(to: slope_y) + endif + ! h is persistently mapped and mutated on the host by diabatic, so refresh it (map-to-present + ! does not copy). + !$omp target update to(h) +#endif + ! Find the maximum and minimum permitted streamfunction. - !$OMP do +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js-1,je+1 ; do i=is-1,ie+1 h_avail_rsum(i,j,1) = 0.0 - pres(i,j,1) = 0.0 - if (associated(tv%p_surf)) then ; pres(i,j,1) = tv%p_surf(i,j) ; endif - + pres(i,j,1) = p_surf_loc(i,j) h_avail(i,j,1) = max(I4dt*G%areaT(i,j)*(h(i,j,1)-GV%Angstrom_H),0.0) h_avail_rsum(i,j,2) = h_avail(i,j,1) h_frac(i,j,1) = 1.0 pres(i,j,2) = pres(i,j,1) + (GV%g_Earth*GV%H_to_RZ) * h(i,j,1) enddo ; enddo - do j=js-1,je+1 - do k=2,nz ; do i=is-1,ie+1 + ! Running sum up the column: a downward vertical recurrence, so collapse(2) over (j,i) + serial K. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif + do j=js-1,je+1 ; do i=is-1,ie+1 + do k=2,nz h_avail(i,j,k) = max(I4dt*G%areaT(i,j)*(h(i,j,k)-GV%Angstrom_H),0.0) h_avail_rsum(i,j,k+1) = h_avail_rsum(i,j,k) + h_avail(i,j,k) h_frac(i,j,k) = 0.0 ; if (h_avail(i,j,k) > 0.0) & h_frac(i,j,k) = h_avail(i,j,k) / h_avail_rsum(i,j,k+1) pres(i,j,K+1) = pres(i,j,K) + (GV%g_Earth*GV%H_to_RZ) * h(i,j,k) - enddo ; enddo - enddo - !$OMP do + enddo + enddo ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js,je ; do I=is-1,ie uhtot(I,j) = 0.0 ; Work_u(I,j) = 0.0 enddo ; enddo - !$OMP do +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do J=js-1,je ; do i=is,ie vhtot(i,J) = 0.0 ; Work_v(i,J) = 0.0 enddo ; enddo - !$OMP end parallel if (CS%id_sfn_x > 0) then ; diag_sfn_x(:,:,1) = 0.0 ; diag_sfn_x(:,:,nz+1) = 0.0 ; endif if (CS%id_sfn_y > 0) then ; diag_sfn_y(:,:,1) = 0.0 ; diag_sfn_y(:,:,nz+1) = 0.0 ; endif @@ -921,44 +999,43 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV #endif endif + ! Zero the FGNV vertical-extent-times-N2 boundary interface values. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif + do j=js,je ; do I=is-1,ie + dzN2_u(I,j,1) = 0. ; dzN2_u(I,j,nz+1) = 0. + enddo ; enddo + + ! Calculate the zonal density gradients, slopes and unlimited streamfunction. The 3-D promoted + ! scratch (drdi_u, drdkDe_u, dzN2_u, slope2_Ratio_u, Sfn_unlim_u, Slope_x_PE, hN2_x_PE) is shared + ! so it survives into the separate flux j-loop below; each j slice is written disjointly. This + ! loop has no vertical recurrence, so on the GPU it collapses over (j,K,I). +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(3) & + !$omp private(drdiA,drdiB,drdkL,drdkR,pres_us,T_us,S_us,dT_us,dS_us, & + !$omp hg2A,hg2B,hg2L,hg2R,haA, & + !$omp N2_unlim,haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$omp dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$omp drdx,mag_grad2,Slope,calc_derivatives) +#else !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,G,GV,US,pres,T,S, & !$OMP nk_linear,IsdB,tv,h,h_neglect,e,dz,dz_neglect,dz_neglect2, & - !$OMP h_neglect2,hn_2,I_slope_max2,int_slope_u,KH_u,uhtot, & - !$OMP h_frac,h_avail_rsum,uhD,h_avail,Work_u,CS,slope_x,cg1, & - !$OMP diag_sfn_x,diag_sfn_unlim_x,N2_floor,EOSdom_u,EOSdom_h1, & + !$OMP h_neglect2,hn_2,I_slope_max2,int_slope_u,KH_u, & + !$OMP h_frac,h_avail_rsum,uhD,h_avail,CS,slope_x,cg1, & + !$OMP diag_sfn_unlim_x,N2_floor,EOSdom_u,EOSdom_h1, & !$OMP use_stanley,present_slope_x,G_rho0,Slope_x_PE,hN2_x_PE, & + !$OMP drdi_u,drdkDe_u,dzN2_u,slope2_Ratio_u,Sfn_unlim_u, & !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & - !$OMP private(drdiA,drdiB,drdkL,drdkR,pres_u,T_u,S_u,G_scale, & - !$OMP drho_dT_u,drho_dS_u,hg2A,hg2B,hg2L,hg2R,haA, & - !$OMP drho_dT_dT_h,scrap,pres_h,T_h,S_h,N2_unlim, & - !$OMP haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & - !$OMP dzg2A,dzg2B,dzaA,dzaB,dz_harm,Z_to_H, & - !$OMP drdx,mag_grad2,Slope,slope2_Ratio_u,dzN2_u, & - !$OMP Sfn_unlim_u,Rho_avg,drdi_u,drdkDe_u,c2_dz_u, & - !$OMP Sfn_safe,Sfn_est,Sfn_in_h,calc_derivatives) + !$OMP private(drdiA,drdiB,drdkL,drdkR,pres_us,T_us,S_us,dT_us,dS_us, & + !$OMP hg2A,hg2B,hg2L,hg2R,haA,drho_dT_dT_h,scrap,pres_h,T_h,S_h, & + !$OMP N2_unlim,haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$OMP dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$OMP drdx,mag_grad2,Slope,calc_derivatives) +#endif do j=js,je - do I=is-1,ie ; dzN2_u(I,j,1) = 0. ; dzN2_u(I,j,nz+1) = 0. ; enddo do K=nz,2,-1 - if (find_work .and. .not.(use_EOS)) then - drdiA = 0.0 ; drdiB = 0.0 - drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL - endif - - calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & - (find_work .or. .not. present_slope_x .or. CS%use_FGNV_streamfn .or. use_stanley) - - ! Calculate the zonal fluxes and gradients. - if (calc_derivatives) then - do I=is-1,ie - pres_u(I) = 0.5*(pres(i,j,K) + pres(i+1,j,K)) - T_u(I) = 0.25*((T(i,j,k) + T(i+1,j,k)) + (T(i,j,k-1) + T(i+1,j,k-1))) - S_u(I) = 0.25*((S(i,j,k) + S(i+1,j,k)) + (S(i,j,k-1) + S(i+1,j,k-1))) - call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_u(I), eos_S_to_ppt*S_u(I), & - eos_RL2_T2_to_Pa*pres_u(I), drho_dT_u(I), drho_dS_u(I)) - drho_dT_u(I) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_u(I) - drho_dS_u(I) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_u(I) - enddo - endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then do i=is-1,ie+1 pres_h(i) = pres(i,j,K) @@ -972,24 +1049,45 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV scrap, scrap, drho_dT_dT_h, scrap, scrap, & tv%eqn_of_state, EOSdom_h1) endif +#endif do I=is-1,ie + calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & + (find_work .or. .not. present_slope_x .or. CS%use_FGNV_streamfn .or. use_stanley) + + if (find_work .and. .not.(use_EOS)) then + drdiA = 0.0 ; drdiB = 0.0 + drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL + endif + + ! Fill the interface T, S and pressure at the u-point and evaluate the density derivatives + ! through the device-callable elementwise EOS dispatcher (fused into this loop from the + ! former separate do-I EOS-fill loop, which read the same I in the same iteration). if (calc_derivatives) then + pres_us = 0.5*(pres(i,j,K) + pres(i+1,j,K)) + T_us = 0.25*((T(i,j,k) + T(i+1,j,k)) + (T(i,j,k-1) + T(i+1,j,k-1))) + S_us = 0.25*((S(i,j,k) + S(i+1,j,k)) + (S(i,j,k-1) + S(i+1,j,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_us, eos_S_to_ppt*S_us, & + eos_RL2_T2_to_Pa*pres_us, dT_us, dS_us) + dT_us = (eos_kg_m3_to_R*eos_C_to_degC) * dT_us + dS_us = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_us + ! Estimate the horizontal density gradients along layers. - drdiA = drho_dT_u(I) * (T(i+1,j,k-1)-T(i,j,k-1)) + & - drho_dS_u(I) * (S(i+1,j,k-1)-S(i,j,k-1)) - drdiB = drho_dT_u(I) * (T(i+1,j,k)-T(i,j,k)) + & - drho_dS_u(I) * (S(i+1,j,k)-S(i,j,k)) + drdiA = dT_us * (T(i+1,j,k-1)-T(i,j,k-1)) + & + dS_us * (S(i+1,j,k-1)-S(i,j,k-1)) + drdiB = dT_us * (T(i+1,j,k)-T(i,j,k)) + & + dS_us * (S(i+1,j,k)-S(i,j,k)) ! Estimate the vertical density gradients times the grid spacing. - drdkL = (drho_dT_u(I) * (T(i,j,k)-T(i,j,k-1)) + & - drho_dS_u(I) * (S(i,j,k)-S(i,j,k-1))) - drdkR = (drho_dT_u(I) * (T(i+1,j,k)-T(i+1,j,k-1)) + & - drho_dS_u(I) * (S(i+1,j,k)-S(i+1,j,k-1))) + drdkL = (dT_us * (T(i,j,k)-T(i,j,k-1)) + & + dS_us * (S(i,j,k)-S(i,j,k-1))) + drdkR = (dT_us * (T(i+1,j,k)-T(i+1,j,k-1)) + & + dS_us * (S(i+1,j,k)-S(i+1,j,k-1))) drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in ! the EOS rectifying SGS temperature anomalies @@ -998,6 +1096,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drdiB = drdiB + 0.5 * ((drho_dT_dT_h(i+1) * tv%varT(i+1,j,k)) - & (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif +#endif if (find_work) drdi_u(I,j,k) = drdiB if (k > nk_linear) then @@ -1082,7 +1181,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (Slope > CS%slope_max) Slope_x_PE(I,j,k) = CS%slope_max if (Slope < -CS%slope_max) Slope_x_PE(I,j,k) = -CS%slope_max endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope +#endif ! Estimate the streamfunction at each interface [H L2 T-1 ~> m3 s-1 or kg s-1]. Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope @@ -1113,7 +1214,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Slope = (e(i+1,j,K)-e(i,j,K)) * G%IdxCu_OBCmask(I,j) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope +#endif Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope dzN2_u(I,j,K) = GV%g_prime(K) endif ! if (use_EOS) @@ -1121,10 +1224,28 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzN2_u(I,j,K) = N2_floor * dz_neglect Sfn_unlim_u(I,j,K) = 0. endif ! if (k > nk_linear) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_unlim_x>0) diag_sfn_unlim_x(I,j,K) = Sfn_unlim_u(I,j,K) +#endif enddo ! i-loop enddo ! k-loop + enddo ! end of j-loop + ! Vertically accumulate the zonal transports (limited by the available mass) and, if requested, the + ! energy conversion. uhtot has a downward vertical recurrence, so this loop is collapse(2) over + ! (j,I) with a serial K inside on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) & + !$omp private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale) +#else + !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,GV,G,US,tv,h,e,hn_2, & + !$OMP nk_linear,h_frac,h_avail_rsum,h_avail,uhtot,uhD,Work_u, & + !$OMP slope2_Ratio_u,Sfn_unlim_u,drdkDe_u,drdi_u,dzN2_u,cg1, & + !$OMP dz,dz_neglect,CS,diag_sfn_x) & + !$OMP private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale,dz_harm,c2_dz_u) +#endif + do j=js,je +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%use_FGNV_streamfn) then do k=1,nz ; do I=is-1,ie ; if (G%OBCmaskCu(I,j)>0.) then dz_harm = max( dz_neglect, & @@ -1146,10 +1267,14 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif enddo endif +#endif - do K=nz,2,-1 - do I=is-1,ie + do I=is-1,ie + do K=nz,2,-1 +#ifdef __NVCOMPILER_OPENMP_GPU + Z_to_H = GV%Z_to_H +#else if (allocated(tv%SpV_avg) .and. (find_work .or. (k > nk_linear)) ) then Rho_avg = ( ((h(i,j,k) + h(i,j,k-1)) + (h(i+1,j,k) + h(i+1,j,k-1))) + 4.0*hn_2 ) / & ( (((h(i,j,k)+hn_2) * tv%SpV_avg(i,j,k)) + ((h(i,j,k-1)+hn_2) * tv%SpV_avg(i,j,k-1))) + & @@ -1159,6 +1284,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Z_to_H = GV%Z_to_H endif +#endif if (k > nk_linear) then if (use_EOS) then @@ -1185,7 +1311,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV uhD(I,j,k) = max(min((Sfn_in_H - uhtot(I,j)), h_avail(i,j,k)), & -h_avail(i+1,j,k)) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_x>0) diag_sfn_x(I,j,K) = diag_sfn_x(I,j,K+1) + uhD(I,j,k) +#endif ! sfn_x(I,j,K) = max(min(Sfn_in_h, uhtot(I,j)+h_avail(i,j,k)), & ! uhtot(I,j)-h_avail(i+1,j,K)) ! sfn_slope_x(I,j,K) = max(uhtot(I,j)-h_avail(i+1,j,k), & @@ -1203,7 +1331,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV uhD(I,j,k) = -uhtot(I,j) * h_frac(i+1,j,k) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_x>0) diag_sfn_x(I,j,K) = diag_sfn_x(I,j,K+1) + uhD(I,j,k) +#endif ! if (sfn_slope_x(I,j,K+1) <= 0.0) then ! sfn_slope_x(I,j,K) = sfn_slope_x(I,j,K+1) * (1.0 - h_frac(i,j,k)) ! else @@ -1221,11 +1351,15 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! A second order centered estimate is used for the density transferred ! between water columns. +#ifdef __NVCOMPILER_OPENMP_GPU + G_scale = GV%g_Earth * GV%H_to_Z +#else if (allocated(tv%SpV_avg)) then G_scale = GV%H_to_RZ * GV%g_Earth / Rho_avg else G_scale = GV%g_Earth * GV%H_to_Z endif +#endif Work_u(I,j) = Work_u(I,j) + G_scale * & ( uhtot(I,j) * drdkDe_u(I,j,K) - & @@ -1233,48 +1367,41 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ((e(i,j,K) + e(i,j,K+1)) + (e(i+1,j,K) + e(i+1,j,K+1))) ) endif - enddo - enddo ! end of k-loop + enddo ! end of k-loop + enddo ! end of i-loop enddo ! end of j-loop ! Calculate the meridional fluxes and gradients. + ! As for the u-points: 3-D scratch (drdj_v, drdkDe_v, dzN2_v, slope2_Ratio_v, Sfn_unlim_v, + ! Slope_y_PE, hN2_y_PE) is shared so it survives into the flux j-loop; no vertical recurrence, so + ! this loop collapses over (J,K,i) on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(3) & + !$omp private(drdjA,drdjB,drdkL,drdkR,pres_vs,T_vs,S_vs,dT_vs,dS_vs, & + !$omp hg2A,hg2B,hg2L,hg2R,haA, & + !$omp haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$omp dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$omp drdy,mag_grad2,Slope,N2_unlim,calc_derivatives) +#else !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,G,GV,US,pres,T,S,dz, & !$OMP nk_linear,IsdB,tv,h,h_neglect,e,dz_neglect,dz_neglect2, & - !$OMP h_neglect2,int_slope_v,KH_v,vhtot,h_frac,h_avail_rsum, & - !$OMP I_slope_max2,vhD,h_avail,Work_v,CS,slope_y,cg1,hn_2,& - !$OMP diag_sfn_y,diag_sfn_unlim_y,N2_floor,EOSdom_v,use_stanley,& + !$OMP h_neglect2,int_slope_v,KH_v, & + !$OMP I_slope_max2,CS,slope_y, & + !$OMP diag_sfn_unlim_y,N2_floor,EOSdom_v,use_stanley,& !$OMP present_slope_y,G_rho0,Slope_y_PE,hN2_y_PE, & + !$OMP drdj_v,drdkDe_v,dzN2_v,slope2_Ratio_v,Sfn_unlim_v, & !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & - !$OMP private(drdjA,drdjB,drdkL,drdkR,pres_v,T_v,S_v,S_h,S_hr, & - !$OMP drho_dT_v,drho_dS_v,hg2A,hg2B,hg2L,hg2R,haA,G_scale, & + !$OMP private(drdjA,drdjB,drdkL,drdkR,pres_vs,T_vs,S_vs,dT_vs,dS_vs, & + !$OMP S_h,S_hr,hg2A,hg2B,hg2L,hg2R,haA, & !$OMP drho_dT_dT_h,drho_dT_dT_hr,scrap,pres_h,T_h,T_hr, & !$OMP haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz,pres_hr, & - !$OMP dzg2A,dzg2B,dzaA,dzaB,dz_harm,Z_to_H, & - !$OMP drdy,mag_grad2,Slope,slope2_Ratio_v,dzN2_v,N2_unlim, & - !$OMP Sfn_unlim_v,Rho_avg,drdj_v,drdkDe_v,c2_dz_v, & - !$OMP Sfn_safe,Sfn_est,Sfn_in_h,calc_derivatives) + !$OMP dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$OMP drdy,mag_grad2,Slope,N2_unlim,calc_derivatives) +#endif do J=js-1,je do K=nz,2,-1 - if (find_work .and. .not.(use_EOS)) then - drdjA = 0.0 ; drdjB = 0.0 - drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL - endif - - calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & - (find_work .or. .not. present_slope_y .or. CS%use_FGNV_streamfn .or. use_stanley) - - if (calc_derivatives) then - do i=is,ie - pres_v(i) = 0.5*(pres(i,j,K) + pres(i,j+1,K)) - T_v(i) = 0.25*((T(i,j,k) + T(i,j+1,k)) + (T(i,j,k-1) + T(i,j+1,k-1))) - S_v(i) = 0.25*((S(i,j,k) + S(i,j+1,k)) + (S(i,j,k-1) + S(i,j+1,k-1))) - call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_v(i), eos_S_to_ppt*S_v(i), & - eos_RL2_T2_to_Pa*pres_v(i), drho_dT_v(i), drho_dS_v(i)) - drho_dT_v(i) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_v(i) - drho_dS_v(i) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_v(i) - enddo - endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then do i=is,ie pres_h(i) = pres(i,j,K) @@ -1295,23 +1422,43 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV scrap, scrap, drho_dT_dT_hr, scrap, scrap, & tv%eqn_of_state, EOSdom_v) endif +#endif do i=is,ie + calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & + (find_work .or. .not. present_slope_y .or. CS%use_FGNV_streamfn .or. use_stanley) + + if (find_work .and. .not.(use_EOS)) then + drdjA = 0.0 ; drdjB = 0.0 + drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL + endif + + ! Fill the interface T, S and pressure at the v-point and evaluate the density derivatives + ! through the device-callable elementwise EOS dispatcher (fused from the former EOS-fill loop). if (calc_derivatives) then + pres_vs = 0.5*(pres(i,j,K) + pres(i,j+1,K)) + T_vs = 0.25*((T(i,j,k) + T(i,j+1,k)) + (T(i,j,k-1) + T(i,j+1,k-1))) + S_vs = 0.25*((S(i,j,k) + S(i,j+1,k)) + (S(i,j,k-1) + S(i,j+1,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_vs, eos_S_to_ppt*S_vs, & + eos_RL2_T2_to_Pa*pres_vs, dT_vs, dS_vs) + dT_vs = (eos_kg_m3_to_R*eos_C_to_degC) * dT_vs + dS_vs = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_vs + ! Estimate the horizontal density gradients along layers. - drdjA = drho_dT_v(i) * (T(i,j+1,k-1)-T(i,j,k-1)) + & - drho_dS_v(i) * (S(i,j+1,k-1)-S(i,j,k-1)) - drdjB = drho_dT_v(i) * (T(i,j+1,k)-T(i,j,k)) + & - drho_dS_v(i) * (S(i,j+1,k)-S(i,j,k)) + drdjA = dT_vs * (T(i,j+1,k-1)-T(i,j,k-1)) + & + dS_vs * (S(i,j+1,k-1)-S(i,j,k-1)) + drdjB = dT_vs * (T(i,j+1,k)-T(i,j,k)) + & + dS_vs * (S(i,j+1,k)-S(i,j,k)) ! Estimate the vertical density gradients times the grid spacing. - drdkL = (drho_dT_v(i) * (T(i,j,k)-T(i,j,k-1)) + & - drho_dS_v(i) * (S(i,j,k)-S(i,j,k-1))) - drdkR = (drho_dT_v(i) * (T(i,j+1,k)-T(i,j+1,k-1)) + & - drho_dS_v(i) * (S(i,j+1,k)-S(i,j+1,k-1))) + drdkL = (dT_vs * (T(i,j,k)-T(i,j,k-1)) + & + dS_vs * (S(i,j,k)-S(i,j,k-1))) + drdkR = (dT_vs * (T(i,j+1,k)-T(i,j+1,k-1)) + & + dS_vs * (S(i,j+1,k)-S(i,j+1,k-1))) drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in ! the EOS rectifying SGS temperature anomalies @@ -1320,6 +1467,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drdjB = drdjB + 0.5 * ((drho_dT_dT_hr(i) * tv%varT(i,j+1,k)) - & (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif +#endif if (find_work) drdj_v(i,J,k) = drdjB @@ -1406,7 +1554,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (Slope > CS%slope_max) Slope_y_PE(i,J,k) = CS%slope_max if (Slope < -CS%slope_max) Slope_y_PE(i,J,k) = -CS%slope_max endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope +#endif Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) @@ -1436,7 +1586,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Slope = (e(i,j+1,K)-e(i,j,K)) * G%IdyCv_OBCmask(i,J) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope +#endif Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) dzN2_v(i,J,K) = GV%g_prime(K) endif ! if (use_EOS) @@ -1444,10 +1596,27 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzN2_v(i,J,K) = N2_floor * dz_neglect Sfn_unlim_v(i,J,K) = 0. endif ! if (k > nk_linear) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_unlim_y>0) diag_sfn_unlim_y(i,J,K) = Sfn_unlim_v(i,J,K) +#endif enddo ! i-loop enddo ! k-loop + enddo ! end of j-loop + ! Vertically accumulate the meridional transports (limited by available mass) and, if requested, + ! the energy conversion. vhtot recurs downward, so collapse(2) over (J,i) + serial K on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) & + !$omp private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale) +#else + !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,GV,G,US,tv,h,e,hn_2, & + !$OMP nk_linear,h_frac,h_avail_rsum,h_avail,vhtot,vhD,Work_v, & + !$OMP slope2_Ratio_v,Sfn_unlim_v,drdkDe_v,drdj_v,dzN2_v,cg1, & + !$OMP dz,dz_neglect,CS,diag_sfn_y) & + !$OMP private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale,dz_harm,c2_dz_v) +#endif + do J=js-1,je +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%use_FGNV_streamfn) then do k=1,nz ; do i=is,ie ; if (G%OBCmaskCv(i,J)>0.) then dz_harm = max( dz_neglect, & @@ -1469,9 +1638,13 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif enddo endif +#endif - do K=nz,2,-1 - do i=is,ie + do i=is,ie + do K=nz,2,-1 +#ifdef __NVCOMPILER_OPENMP_GPU + Z_to_H = GV%Z_to_H +#else if (allocated(tv%SpV_avg) .and. (find_work .or. (k > nk_linear)) ) then Rho_avg = ( ((h(i,j,k) + h(i,j,k-1)) + (h(i,j+1,k) + h(i,j+1,k-1))) + 4.0*hn_2 ) / & ( (((h(i,j,k)+hn_2) * tv%SpV_avg(i,j,k)) + ((h(i,j,k-1)+hn_2) * tv%SpV_avg(i,j,k-1))) + & @@ -1481,6 +1654,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Z_to_H = GV%Z_to_H endif +#endif if (k > nk_linear) then if (use_EOS) then @@ -1506,7 +1680,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! neighboring grid cells. vhD(i,J,k) = max(min((Sfn_in_H - vhtot(i,J)), h_avail(i,j,k)), -h_avail(i,j+1,k)) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_y>0) diag_sfn_y(i,J,K) = diag_sfn_y(i,J,K+1) + vhD(i,J,k) +#endif ! sfn_y(i,J,K) = max(min(Sfn_in_h, vhtot(i,J)+h_avail(i,j,k)), & ! vhtot(i,J)-h_avail(i,j+1,k)) ! sfn_slope_y(i,J,K) = max(vhtot(i,J)-h_avail(i,j+1,k), & @@ -1524,7 +1700,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV vhD(i,J,k) = -vhtot(i,J) * h_frac(i,j+1,k) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_y>0) diag_sfn_y(i,J,K) = diag_sfn_y(i,J,K+1) + vhD(i,J,k) +#endif ! if (sfn_slope_y(i,J,K+1) <= 0.0) then ! sfn_slope_y(i,J,K) = sfn_slope_y(i,J,K+1) * (1.0 - h_frac(i,j,k)) ! else @@ -1541,11 +1719,15 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! A second order centered estimate is used for the density transferred ! between water columns. +#ifdef __NVCOMPILER_OPENMP_GPU + G_scale = GV%g_Earth * GV%H_to_Z +#else if (allocated(tv%SpV_avg)) then G_scale = GV%H_to_RZ * GV%g_Earth / Rho_avg else G_scale = GV%g_Earth * GV%H_to_Z endif +#endif Work_v(i,J) = Work_v(i,J) + G_scale * & ( vhtot(i,J) * drdkDe_v(i,J,K) - & @@ -1553,104 +1735,118 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ((e(i,j,K) + e(i,j,K+1)) + (e(i,j+1,K) + e(i,j+1,K+1))) ) endif - enddo - enddo ! end of k-loop + enddo ! end of k-loop + enddo ! end of i-loop enddo ! end of j-loop ! In layer 1, enforce the boundary conditions that Sfn(z=0) = 0.0 if (.not.find_work .or. .not.(use_EOS)) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js,je ; do I=is-1,ie ; uhD(I,j,1) = -uhtot(I,j) ; enddo ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do J=js-1,je ; do i=is,ie ; vhD(i,J,1) = -vhtot(i,J) ; enddo ; enddo else - EOSdom_u(1) = (is-1) - (G%IsdB-1) ; EOSdom_u(2) = ie - (G%IsdB-1) - !$OMP parallel do default(shared) private(pres_u,T_u,S_u,drho_dT_u,drho_dS_u,drdiB,G_scale) - do j=js,je - if (use_EOS) then - do I=is-1,ie - pres_u(I) = 0.5*(pres(i,j,1) + pres(i+1,j,1)) - T_u(I) = 0.5*(T(i,j,1) + T(i+1,j,1)) - S_u(I) = 0.5*(S(i,j,1) + S(i+1,j,1)) - call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_u(I), eos_S_to_ppt*S_u(I), & - eos_RL2_T2_to_Pa*pres_u(I), drho_dT_u(I), drho_dS_u(I)) - drho_dT_u(I) = (eos_kg_m3_to_R*eos_C_to_degC) * drho_dT_u(I) - drho_dS_u(I) = (eos_kg_m3_to_R*eos_S_to_ppt) * drho_dS_u(I) - enddo + ! In this branch use_EOS is necessarily true. The density derivatives use the device-callable + ! elementwise dispatcher (v-points too), fused into the work loop and scalarized. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(pres_us,T_us,S_us,dT_us,dS_us,drdiB,G_scale) +#else + !$OMP parallel do default(shared) private(pres_us,T_us,S_us,dT_us,dS_us,drdiB,G_scale) +#endif + do j=js,je ; do I=is-1,ie + pres_us = 0.5*(pres(i,j,1) + pres(i+1,j,1)) + T_us = 0.5*(T(i,j,1) + T(i+1,j,1)) + S_us = 0.5*(S(i,j,1) + S(i+1,j,1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_us, eos_S_to_ppt*S_us, & + eos_RL2_T2_to_Pa*pres_us, dT_us, dS_us) + dT_us = (eos_kg_m3_to_R*eos_C_to_degC) * dT_us + dS_us = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_us + + uhD(I,j,1) = -uhtot(I,j) + + G_scale = GV%g_Earth * GV%H_to_Z + drdiB = dT_us * (T(i+1,j,1)-T(i,j,1)) + dS_us * (S(i+1,j,1)-S(i,j,1)) +#ifndef __NVCOMPILER_OPENMP_GPU + if (allocated(tv%SpV_avg)) then + G_scale = GV%H_to_RZ * GV%g_Earth * & + ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i+1,j,1)+hn_2) * tv%SpV_avg(i+1,j,1)) ) / & + ( (h(i,j,1) + h(i+1,j,1)) + 2.0*hn_2 ) ) endif - do I=is-1,ie - uhD(I,j,1) = -uhtot(I,j) +#endif + if (CS%use_GM_work_bug) then + Work_u(I,j) = Work_u(I,j) + G_scale * & + ( (uhD(I,j,1) * drdiB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) + else + Work_u(I,j) = Work_u(I,j) - G_scale * & + ( (uhD(I,j,1) * drdiB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) + endif + enddo ; enddo - G_scale = GV%g_Earth * GV%H_to_Z - if (use_EOS) then - drdiB = drho_dT_u(I) * (T(i+1,j,1)-T(i,j,1)) + & - drho_dS_u(I) * (S(i+1,j,1)-S(i,j,1)) - if (allocated(tv%SpV_avg)) then - G_scale = GV%H_to_RZ * GV%g_Earth * & - ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i+1,j,1)+hn_2) * tv%SpV_avg(i+1,j,1)) ) / & - ( (h(i,j,1) + h(i+1,j,1)) + 2.0*hn_2 ) ) - endif - endif - if (CS%use_GM_work_bug) then - Work_u(I,j) = Work_u(I,j) + G_scale * & - ( (uhD(I,j,1) * drdiB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) - else - Work_u(I,j) = Work_u(I,j) - G_scale * & - ( (uhD(I,j,1) * drdiB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) - endif - enddo - enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(pres_vs,T_vs,S_vs,dT_vs,dS_vs,drdjB,G_scale) +#else + !$OMP parallel do default(shared) private(pres_vs,T_vs,S_vs,dT_vs,dS_vs,drdjB,G_scale) +#endif + do J=js-1,je ; do i=is,ie + pres_vs = 0.5*(pres(i,j,1) + pres(i,j+1,1)) + T_vs = 0.5*(T(i,j,1) + T(i,j+1,1)) + S_vs = 0.5*(S(i,j,1) + S(i,j+1,1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_vs, eos_S_to_ppt*S_vs, & + eos_RL2_T2_to_Pa*pres_vs, dT_vs, dS_vs) + dT_vs = (eos_kg_m3_to_R*eos_C_to_degC) * dT_vs + dS_vs = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_vs + + vhD(i,J,1) = -vhtot(i,J) + + G_scale = GV%g_Earth * GV%H_to_Z + drdjB = dT_vs * (T(i,j+1,1)-T(i,j,1)) + dS_vs * (S(i,j+1,1)-S(i,j,1)) +#ifndef __NVCOMPILER_OPENMP_GPU + if (allocated(tv%SpV_avg)) then + G_scale = GV%H_to_RZ * GV%g_Earth * & + ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i,j+1,1)+hn_2) * tv%SpV_avg(i,j+1,1)) ) / & + ( (h(i,j,1) + h(i,j+1,1)) + 2.0*hn_2 ) ) + endif +#endif + Work_v(i,J) = Work_v(i,J) - G_scale * & + ( (vhD(i,J,1) * drdjB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i,j+1,1) + e(i,j+1,2))) ) + enddo ; enddo + endif - EOSdom_v(:) = EOS_domain(G%HI) - !$OMP parallel do default(shared) private(pres_v,T_v,S_v,drho_dT_v,drho_dS_v,drdjB,G_scale) - do J=js-1,je - if (use_EOS) then - do i=is,ie - pres_v(i) = 0.5*(pres(i,j,1) + pres(i,j+1,1)) - T_v(i) = 0.5*(T(i,j,1) + T(i,j+1,1)) - S_v(i) = 0.5*(S(i,j,1) + S(i,j+1,1)) + ! Assemble the GM work / energy source (writing the hoisted local arrays, not the pointer/ + ! allocatable components). + if (find_work) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(Work_h) +#endif + do j=js,je ; do i=is,ie + ! Note that the units of Work_v and Work_u are [R Z L4 T-3 ~> W], while Work_h is in [R Z L2 T-3 ~> W m-2]. + Work_h = 0.5 * G%IareaT(i,j) * & + ((Work_u(I-1,j) + Work_u(I,j)) + (Work_v(i,J-1) + Work_v(i,J))) + if (have_GMwork) GMwork_loc(i,j) = Work_h + if (.not. CS%GM_src_alt) then ; if (have_GM_src) then + GM_src_loc(i,j) = GM_src_loc(i,j) + Work_h + endif ; endif +#ifndef __NVCOMPILER_OPENMP_GPU + if (skeb_use_gm) then + skeb_gm_work(i,j) = STOCH%skeb_gm_coef * Work_h + skeb_ebt_norm2(i,j) = 0.0 + do k=1,nz + skeb_ebt_norm2(i,j) = skeb_ebt_norm2(i,j) + h(i,j,k) * VarMix%ebt_struct(i,j,k)**2 enddo - call calculate_density_derivs(T_v, S_v, pres_v, drho_dT_v, drho_dS_v, & - tv%eqn_of_state, EOSdom_v) + skeb_ebt_norm2(i,j) = GV%H_to_RZ * (skeb_ebt_norm2(i,j) + h_neglect) endif - do i=is,ie - vhD(i,J,1) = -vhtot(i,J) - - G_scale = GV%g_Earth * GV%H_to_Z - if (use_EOS) then - drdjB = drho_dT_v(i) * (T(i,j+1,1)-T(i,j,1)) + & - drho_dS_v(i) * (S(i,j+1,1)-S(i,j,1)) - if (allocated(tv%SpV_avg)) then - G_scale = GV%H_to_RZ * GV%g_Earth * & - ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i,j+1,1)+hn_2) * tv%SpV_avg(i,j+1,1)) ) / & - ( (h(i,j,1) + h(i,j+1,1)) + 2.0*hn_2 ) ) - endif - endif - Work_v(i,J) = Work_v(i,J) - G_scale * & - ( (vhD(i,J,1) * drdjB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i,j+1,1) + e(i,j+1,2))) ) - enddo - enddo +#endif + enddo ; enddo endif - if (find_work) then ; do j=js,je ; do i=is,ie - ! Note that the units of Work_v and Work_u are [R Z L4 T-3 ~> W], while Work_h is in [R Z L2 T-3 ~> W m-2]. - Work_h = 0.5 * G%IareaT(i,j) * & - ((Work_u(I-1,j) + Work_u(I,j)) + (Work_v(i,J-1) + Work_v(i,J))) - if (allocated(CS%GMwork)) CS%GMwork(i,j) = Work_h - if (.not. CS%GM_src_alt) then ; if (allocated(MEKE%GM_src)) then - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + Work_h - endif ; endif - if (skeb_use_gm) then - skeb_gm_work(i,j) = STOCH%skeb_gm_coef * Work_h - skeb_ebt_norm2(i,j) = 0.0 - do k=1,nz - skeb_ebt_norm2(i,j) = skeb_ebt_norm2(i,j) + h(i,j,k) * VarMix%ebt_struct(i,j,k)**2 - enddo - skeb_ebt_norm2(i,j) = GV%H_to_RZ * (skeb_ebt_norm2(i,j) + h_neglect) - endif - enddo ; enddo ; endif - +#ifndef __NVCOMPILER_OPENMP_GPU if (skeb_use_gm) then ! This block spreads the GM work down through the column using the ebt vertical structure, squared. ! Note the sign convention. @@ -1659,38 +1855,71 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV VarMix%ebt_struct(i,j,k)**2 / skeb_ebt_norm2(i,j) enddo ; enddo ; enddo endif +#endif - if (find_work .and. CS%GM_src_alt) then ; if (allocated(MEKE%GM_src)) then + if (find_work .and. CS%GM_src_alt) then ; if (have_GM_src) then + ! Accumulate the S^2 N^2 kappa energy release down each column (a serial-K sum into GM_src_loc). if (CS%MEKE_src_answer_date >= 20240601) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(PE_release_h) +#endif do j=js,je ; do i=is,ie ; do k=nz,1,-1 PE_release_h = -0.25 * GV%H_to_RZ * & ( ((KH_u(I,j,k)*(Slope_x_PE(I,j,k)**2) * hN2_x_PE(I,j,k)) + & (Kh_u(I-1,j,k)*(Slope_x_PE(I-1,j,k)**2) * hN2_x_PE(I-1,j,k))) + & ((Kh_v(i,J,k)*(Slope_y_PE(i,J,k)**2) * hN2_y_PE(i,J,k)) + & (Kh_v(i,J-1,k)*(Slope_y_PE(i,J-1,k)**2) * hN2_y_PE(i,J-1,k))) ) - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + PE_release_h + GM_src_loc(i,j) = GM_src_loc(i,j) + PE_release_h enddo ; enddo ; enddo else +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(PE_release_h) +#endif do j=js,je ; do i=is,ie ; do k=nz,1,-1 PE_release_h = -0.25 * GV%H_to_RZ * & ((KH_u(I,j,k)*(Slope_x_PE(I,j,k)**2) * hN2_x_PE(I,j,k)) + & (Kh_u(I-1,j,k)*(Slope_x_PE(I-1,j,k)**2) * hN2_x_PE(I-1,j,k)) + & (Kh_v(i,J,k)*(Slope_y_PE(i,J,k)**2) * hN2_y_PE(i,J,k)) + & (Kh_v(i,J-1,k)*(Slope_y_PE(i,J-1,k)**2) * hN2_y_PE(i,J-1,k))) - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + PE_release_h + GM_src_loc(i,j) = GM_src_loc(i,j) + PE_release_h enddo ; enddo ; enddo endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%debug) then - call hchksum(MEKE%GM_src, 'MEKE%GM_src', G%HI, unscale=US%RZ3_T3_to_W_m2*US%L_to_Z**2) + call hchksum(GM_src_loc, 'MEKE%GM_src', G%HI, unscale=US%RZ3_T3_to_W_m2*US%L_to_Z**2) call uvchksum("KH_[uv]", Kh_u, Kh_v, G%HI, unscale=US%L_to_m**2*US%s_to_T, & scalar_pair=.true.) call uvchksum("Slope_[xy]_PE", Slope_x_PE, Slope_y_PE, G%HI, unscale=US%Z_to_L) call uvchksum("hN2_[xy]_PE", hN2_x_PE, hN2_y_PE, G%HI, unscale=GV%H_to_mks*US%L_to_Z**2*US%s_to_T**2, & scalar_pair=.true.) endif +#endif endif ; endif + ! ------------------------------------------------------------------------------------------ + ! Close the device data region: copy the outputs back and release every mapping (release, not + ! delete — a per-call delete forces the refcount to zero and would destroy any outer/persistent + ! mapping of the same storage; KNOWLEDGE §8), then copy the hoisted GM work back to the host. + ! ------------------------------------------------------------------------------------------ +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target exit data map(from: uhD, vhD, GM_src_loc, GMwork_loc) & + !$omp map(release: e, dz, T, S, Kh_u, Kh_v, int_slope_u, int_slope_v, CS, p_surf_loc, & + !$omp Slope_x_PE, Slope_y_PE, hN2_x_PE, hN2_y_PE, pres, h_avail, h_frac, & + !$omp h_avail_rsum, uhtot, vhtot, Work_u, Work_v, drdi_u, drdj_v, & + !$omp drdkDe_u, drdkDe_v, dzN2_u, dzN2_v, Sfn_unlim_u, Sfn_unlim_v, & + !$omp slope2_Ratio_u, slope2_Ratio_v) & + !$omp map(release: h) + if (present_slope_x) then + !$omp target exit data map(release: slope_x) + endif + if (present_slope_y) then + !$omp target exit data map(release: slope_y) + endif +#endif + if (have_GM_src) then ; do j=js,je ; do i=is,ie ; MEKE%GM_src(i,j) = GM_src_loc(i,j) ; enddo ; enddo ; endif + if (have_GMwork) then ; do j=js,je ; do i=is,ie ; CS%GMwork(i,j) = GMwork_loc(i,j) ; enddo ; enddo ; endif + if (CS%id_slope_x > 0) call post_data(CS%id_slope_x, CS%diagSlopeX, CS%diag) if (CS%id_slope_y > 0) call post_data(CS%id_slope_y, CS%diagSlopeY, CS%diag) if (CS%id_sfn_x > 0) call post_data(CS%id_sfn_x, diag_sfn_x, CS%diag) From 726842da6ee0ec89a70728ac1f6ea565c33a2e4b Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 19:44:56 +1000 Subject: [PATCH 39/45] ePBL: document the device-vs-host transcendental bitwise-repro sites (TODO for Marshall) Comment-only. ePBL is the one offloaded routine on this branch that is not yet GPU/CPU bit-for-bit: the device run differs from host at ~1e-13 in the energy diagnostic (T/S/mass bit-identical) solely from device-vs-host libm/libdevice divergence in transcendental ops in its on-device kernels, which have no deterministic drop-in (unlike the x**(1/3) cube roots, already handled via cuberoot() on the answer_date>=20240101 path). Add a consolidated, greppable `TODO(gpu-bitwise-repro)` audit at the module level enumerating the live (benchmark_ALE) transcendental sites for the real matched-transcendental fix (being done by Marshall -- not duplicated here): find_PE_chg exp() (core, likely dominant), exp_decay_TKE_adjust exp(), find_mstar mstar_N log() (OM4 Ekman), mstar_Langmuir **(-1.33) (arbitrary power), and the MixLen pow(x,1.0) (trivial identity stopgap noted but not applied). MKE_src exp() is inert (MKE_TO_TKE_EFFIC=0). No numerical change, so the ePBL-device baseline and the no-new-diff offload gate are unaffected; an nvfortran-CPU reference build is staged at ocean_only/cpu_build for the GPU==CPU one-source verification once Marshall's fix lands. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_energetic_PBL.F90 | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index a319c196ae..7ae7e0aa9b 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -333,6 +333,27 @@ module MOM_energetic_PBL !$omp declare target(find_mstar, mstar_Langmuir) !$omp declare target(ePBL_column) +! TODO(gpu-bitwise-repro): ePBL is the ONE routine on feat/port-ePBL that is not yet GPU/CPU +! bit-for-bit. The device run differs from host at ~1e-13 in the energy diagnostic (T/S/mass are +! bit-identical) SOLELY because of device-vs-host libm/libdevice divergence in transcendental ops +! in this routine's on-device kernels -- there is no deterministic drop-in for these (unlike the +! x**(1/3) cube roots, which the answer_date>=20240101 path already handles via cuberoot()). The +! real fix is a matched-transcendental / bitwise-repro implementation being done by Marshall; do NOT +! duplicate it here. The live (benchmark_ALE) device-vs-host transcendental sites, for Marshall's +! pass / for the eventual GPU==CPU verify against ocean_only/cpu_build: +! - find_PE_chg / find_PE_chg_orig : exp() in the PE-change integrals (core, per interface x +! iteration -- the likely dominant contributor). +! - exp_decay_TKE_adjust (~:1533) : exp(-h*Idecay_len_TKE) TKE decay. +! - find_mstar mstar_N (~:3753) : log(...) Ekman-limit term (EPBL_MSTAR_SCHEME=OM4). +! - mstar_Langmuir (~:3875,:3878) : Convect_Langmuir_Number**LT_enhance_exp (LT_ENHANCE_EXP=-1.33, +! an ARBITRARY real power -- no cuberoot/nth_root form; = exp(y*log x) under the hood). +! - MixLen_shape (~:1492) : (...)**MixLenExponent with MixLenExponent=1.0 == pow(x,1.0); +! a trivial local stopgap would be an `==1.0` identity fast-path next to the existing `==2.0` +! one, but it alone does NOT restore bitwise (the exp/log above dominate), so left for Marshall. +! (MKE_src exp() sites ~:1755+ are multiplied by MKE_TO_TKE_EFFIC=0.0 in benchmark_ALE -> harmless.) +! See the [[gpu-transcendental-bitwise-plan]] auto-memory. Until this lands, ePBL offload increments +! are gated on "no NEW diff vs the ePBL-device baseline", not bit-for-bit vs the golden. + !> GPU port: fixes the per-column private scratch in ePBL_column to a compile-time-constant size on !! GPU builds, so each device thread gets stack (local memory) arrays rather than runtime-sized !! device-heap automatics (NVFORTRAN-W-0155 / KNOWLEDGE row 21). Checked against GV%ke in From 2e9062dd80767ef24ae31b2d2ae917667226821f Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 19:46:34 +1000 Subject: [PATCH 40/45] wave_interface: flag the LF17 Langmuir on-device transcendental sites (TODO for Marshall) Comment-only. The LF17 Langmuir chain (get_StokesSL_LiFoxKemper / one_minus_exp_x / ust_2_u10_coare3p5) runs on-device when called from ePBL_column and contributes device-vs-host transcendental divergence to the ~1e-13 ePBL energy floor via LA -> mstar_LT. Add a greppable TODO(gpu-bitwise-repro) enumerating the live sites (exp in one_minus_exp_x; erfc in get_StokesSL r5 large-root branch; atan(1.0) for PI, usually folded; sqrt is exact) for Marshall's matched-transcendental pass, cross-referencing the ePBL audit. No numerical change. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/user/MOM_wave_interface.F90 | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index fa377a338d..664851d422 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -303,6 +303,14 @@ module MOM_wave_interface !! (USE_LA_LI2016) statistical-wave path is compiled (the wave-model branches are host-only). !$omp declare target(get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) +! TODO(gpu-bitwise-repro): this LF17 Langmuir chain runs on-device (called from ePBL_column) and +! adds device-vs-host libm/libdevice transcendental divergence that feeds the ~1e-13 ePBL energy +! floor (LA -> mstar_LT). Live sites (benchmark_ALE, WAVE_INTERFACE_ANSWER_DATE=99991231): exp() in +! one_minus_exp_x; erfc() in get_StokesSL_LiFoxKemper r5 (the large root_2kz branch); PI=4*atan(1.0) +! (a compile-time constant, usually folded). sqrt() is IEEE-exact (fine). No deterministic drop-in; +! part of Marshall's matched-transcendental / bitwise-repro pass -- see the audit in +! MOM_energetic_PBL.F90 and the [[gpu-transcendental-bitwise-plan]] auto-memory. Do NOT duplicate. + contains !> Initializes parameters related to MOM_wave_interface From 3c1ed41e45fb54cbf0fb8343fbf5a219c34f5303 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 20:11:56 +1000 Subject: [PATCH 41/45] MOM_intrinsic_functions: add exp_reprod, a bit-reproducible device-callable exp() The intrinsic exp() lowers to host libm on the CPU and CUDA libdevice on the GPU, whose last-bit rounding differs, so any do concurrent / omp target kernel that calls exp() is not bit-for-bit CPU==GPU -- the dominant source of the ~1e-13 ePBL energy-budget divergence on feat/port-ePBL (see the TODO(gpu-bitwise-repro) audit in MOM_energetic_PBL.F90). exp_reprod avoids the library call: Cody-Waite range reduction (x = k*ln2 + r with a two-part ln2 so k*ln2_hi is near-exact), a degree-12 Horner polynomial for exp(r) whose reciprocal-factorial coefficients are compile-time constant-folded (identical host/device), and scale(.,k) for 2**k. Every op is + - * / (plus nint/scale, exact), all bit-identical host vs device under -Mnofma -- so the result is reproducible by construction, the same way cuberoot/nth_root replace x**(1/n). declare target; compiles to device code (symbol confirmed); scale/nint verified device-callable. Accuracy: max relative error 3.1e-16 (~1.4 ULP) vs the intrinsic over x in [-70, 20]; unit test Test_exp_reprod added (tol 1e-14). Like cuberoot it is NOT bit-identical to intrinsic exp, so adopting it in a kernel changes answers and needs a golden regeneration. Nothing calls it yet (bitwise-neutral to commit); offered as the exp building block for the transcendental-repro pass. Log/erfc/arbitrary-pow still need the same treatment. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/framework/MOM_intrinsic_functions.F90 | 72 ++++++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) diff --git a/src/framework/MOM_intrinsic_functions.F90 b/src/framework/MOM_intrinsic_functions.F90 index de668f56bd..dc0e68da45 100644 --- a/src/framework/MOM_intrinsic_functions.F90 +++ b/src/framework/MOM_intrinsic_functions.F90 @@ -11,7 +11,7 @@ module MOM_intrinsic_functions implicit none ; private -public :: invcosh, cuberoot, nth_root +public :: invcosh, cuberoot, nth_root, exp_reprod public :: intrinsic_functions_unit_tests ! Floating point model, if bit layout from high to low is (sign, exp, frac) @@ -117,6 +117,47 @@ elemental function cuberoot(x) result(root) end function cuberoot +!> Bit-reproducible exponential, exp(x), suitable for evaluation inside `!$omp target` / +!! `do concurrent` offloaded regions. +!! +!! The intrinsic `exp()` lowers to host libm on the CPU and CUDA libdevice on the GPU, whose +!! last-bit rounding differs -- so a `do concurrent`/`omp target` kernel that calls `exp()` is not +!! bit-for-bit CPU==GPU (the class of ~1e-13 divergence seen in the ePBL energy budget). This routine +!! avoids the library call entirely: Cody-Waite range reduction x = k*ln2 + r (|r| <= ln2/2) using a +!! two-part ln2 so that k*ln2_hi is (near-)exact, a degree-12 Taylor/Horner polynomial for exp(r) whose +!! reciprocal-factorial coefficients are compile-time constant-folded (hence identical on host and +!! device), and `scale(.,k)` for the 2**k factor. Every operation is +,-,*,/ (plus `nint`/`scale`, +!! which are exact), all of which are bit-identical host vs device under `-Mnofma`, so the result is +!! reproducible by construction. Accuracy is ~1.4 ULP (max relative error 3.1e-16 vs the intrinsic +!! over x in [-70, 20]); it is NOT bit-identical to the intrinsic exp (like cuberoot vs x**(1/3), it +!! changes answers and requires a reference/golden regeneration when adopted). +elemental function exp_reprod(x) result(ex) + !$omp declare target + real, intent(in) :: x !< The argument of the exponential [nondim or arbitrary] + real :: ex !< The reproducible exponential of x [same as exp(x)] + + ! Cody-Waite split of ln(2): ln2 = ln2_hi + ln2_lo, with ln2_hi chosen so k*ln2_hi is ~exact. + real, parameter :: invln2 = 1.44269504088896338700 ! 1/ln(2) [nondim] + real, parameter :: ln2_hi = 0.693147180369123816490 ! High part of ln(2) [nondim] + real, parameter :: ln2_lo = 1.90821492927058770002e-10 ! Low part of ln(2) [nondim] + ! Reciprocal factorials 1/2! .. 1/12! (compile-time constant-folded -> identical host/device). + real, parameter :: c2 = 1.0/2.0, c3 = 1.0/6.0, c4 = 1.0/24.0 + real, parameter :: c5 = 1.0/120.0, c6 = 1.0/720.0, c7 = 1.0/5040.0 + real, parameter :: c8 = 1.0/40320.0, c9 = 1.0/362880.0, c10 = 1.0/3628800.0 + real, parameter :: c11 = 1.0/39916800.0, c12 = 1.0/479001600.0 + real :: r ! The reduced argument, x - k*ln2, in [-ln2/2, ln2/2] [nondim] + real :: p ! The polynomial estimate of exp(r) [nondim] + integer :: k ! The integer number of factors of 2 in exp(x) [nondim] + + k = nint(x*invln2) + r = (x - real(k)*ln2_hi) - real(k)*ln2_lo + ! Horner form of 1 + r + r^2/2! + ... + r^12/12! + p = 1.0 + r*(1.0 + r*(c2 + r*(c3 + r*(c4 + r*(c5 + r*(c6 + r*(c7 + & + r*(c8 + r*(c9 + r*(c10 + r*(c11 + r*c12))))))))))) + ex = scale(p, k) +end function exp_reprod + + !> Bit-stable n-th root of x for x in (0, +inf) and integer n >= 1, suitable !! for evaluation inside `!$omp target` / `do concurrent` offloaded regions. !! @@ -298,6 +339,16 @@ function intrinsic_functions_unit_tests(verbose) result(fail) fail = fail .or. Test_cuberoot(v, testval) testval = (-2.908 * (1.414213562373 + 1.2345678901234e-5*n)) * testval enddo + + v = verbose + fail = fail .or. Test_exp_reprod(v, 0.0) + fail = fail .or. Test_exp_reprod(v, 1.0) + fail = fail .or. Test_exp_reprod(v, -3.7) + fail = fail .or. Test_exp_reprod(v, 20.0) + v = .false. + do n=-700,200 + fail = fail .or. Test_exp_reprod(v, 0.1*n) + enddo end function intrinsic_functions_unit_tests !> True if the cube of cuberoot(val) does not closely match val. False otherwise. @@ -318,4 +369,23 @@ logical function Test_cuberoot(verbose, val) endif end function Test_cuberoot +!> True if exp_reprod(val) does not closely match the intrinsic exp(val). False otherwise. +logical function Test_exp_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test [nondim] + ! Local variables + real :: e_ref ! The intrinsic exponential of val [nondim] + real :: relerr ! The relative difference between exp_reprod(val) and exp(val) [nondim] + + e_ref = exp(val) + relerr = abs(exp_reprod(val) - e_ref) / e_ref + Test_exp_reprod = (relerr > 1.0e-14) + + if (Test_exp_reprod) then + write(stdout, '("For val = ",ES22.15,", exp_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", exp_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_exp_reprod + end module MOM_intrinsic_functions From edd2468aaf925d34f7de8570001d49294d6db81c Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 20:26:10 +1000 Subject: [PATCH 42/45] MOM_intrinsic_functions: add log_reprod, a bit-reproducible device-callable log() Companion to exp_reprod: a device-callable natural log built from only +,-,*,/ (plus the exact exponent/fraction bit intrinsics), so it is bit-identical host vs device under -Mnofma. x = m*2**k via exponent/fraction; m reduced to [sqrt(1/2), sqrt(2)); log(m) = 2*(s + s^3/3 + ... ), s=(m-1)/(m+1), |s|<=0.172, as a degree-21 Horner poly in s^2 with compile-folded reciprocal-odd coeffs; then + k*ln2. Max relative error 3.7e-16 (~1.7 ULP) vs intrinsic log over x in [1e-30, 1e30]; unit test Test_log_reprod added (tol 1e-13). declare target; device code + symbol confirmed. Together exp_reprod + log_reprod make arbitrary real powers reproducible via x**y = exp_reprod(y*log_reprod(x)) -- covering the ePBL Langmuir Convect_Langmuir_Number**(-1.33) and the MixLen **1.0. Bitwise-neutral commit (nothing calls it yet). Building block for the transcendental-repro pass; erfc and site wiring + golden regen still to come. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/framework/MOM_intrinsic_functions.F90 | 68 ++++++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/src/framework/MOM_intrinsic_functions.F90 b/src/framework/MOM_intrinsic_functions.F90 index dc0e68da45..70fec8103b 100644 --- a/src/framework/MOM_intrinsic_functions.F90 +++ b/src/framework/MOM_intrinsic_functions.F90 @@ -11,7 +11,7 @@ module MOM_intrinsic_functions implicit none ; private -public :: invcosh, cuberoot, nth_root, exp_reprod +public :: invcosh, cuberoot, nth_root, exp_reprod, log_reprod public :: intrinsic_functions_unit_tests ! Floating point model, if bit layout from high to low is (sign, exp, frac) @@ -158,6 +158,43 @@ elemental function exp_reprod(x) result(ex) end function exp_reprod +!> Bit-reproducible natural logarithm, log(x) for x > 0, suitable for evaluation inside +!! `!$omp target` / `do concurrent` offloaded regions (companion to exp_reprod; together they make +!! arbitrary real powers reproducible via x**y = exp_reprod(y*log_reprod(x))). +!! +!! Same rationale as exp_reprod: the intrinsic log() differs host libm vs CUDA libdevice in the last +!! bit. This routine uses the exponent/fraction split x = m * 2**k (exact bit intrinsics), reduces the +!! mantissa to m in [sqrt(1/2), sqrt(2)), and evaluates log(m) = 2*(s + s^3/3 + s^5/5 + ...) with +!! s = (m-1)/(m+1) (|s| <= 0.172) as a Horner polynomial in s^2 to degree 21, then adds k*ln2. Every +!! operation is +,-,*,/ (plus exponent/fraction, exact), bit-identical host vs device under -Mnofma. +!! Accuracy ~1.7 ULP (max relative error 3.7e-16 vs the intrinsic over x in [1e-30, 1e30]). Requires +!! x > 0. Like exp_reprod it is not bit-identical to intrinsic log (changes answers on adoption). +elemental function log_reprod(x) result(lx) + !$omp declare target + real, intent(in) :: x !< The argument of the logarithm, x > 0 [nondim or arbitrary] + real :: lx !< The reproducible natural logarithm of x [nondim] + + real, parameter :: ln2 = 0.69314718055994530942 ! ln(2) [nondim] + real, parameter :: sqrt2_2 = 0.70710678118654752440 ! sqrt(1/2), the mantissa reduction threshold [nondim] + ! Reciprocal odd integers 1/3 .. 1/21 for the atanh series (compile-folded -> identical host/device). + real, parameter :: a3=1.0/3.0, a5=1.0/5.0, a7=1.0/7.0, a9=1.0/9.0 + real, parameter :: a11=1.0/11.0, a13=1.0/13.0, a15=1.0/15.0, a17=1.0/17.0 + real, parameter :: a19=1.0/19.0, a21=1.0/21.0 + real :: m ! The mantissa of x, reduced to [sqrt(1/2), sqrt(2)) [nondim] + real :: s ! (m-1)/(m+1), the atanh-series argument, |s| <= 0.172 [nondim] + real :: s2 ! s*s [nondim] + real :: poly ! The polynomial estimate of log(m) [nondim] + integer :: k ! The binary exponent of x [nondim] + + k = exponent(x) ; m = fraction(x) ! x = m * 2**k, m in [0.5, 1) + if (m < sqrt2_2) then ; m = m + m ; k = k - 1 ; endif ! recenter m to [sqrt(1/2), sqrt(2)) + s = (m - 1.0) / (m + 1.0) ; s2 = s*s + poly = 2.0*s*(1.0 + s2*(a3 + s2*(a5 + s2*(a7 + s2*(a9 + s2*(a11 + s2*(a13 + & + s2*(a15 + s2*(a17 + s2*(a19 + s2*a21)))))))))) + lx = poly + real(k)*ln2 +end function log_reprod + + !> Bit-stable n-th root of x for x in (0, +inf) and integer n >= 1, suitable !! for evaluation inside `!$omp target` / `do concurrent` offloaded regions. !! @@ -349,6 +386,15 @@ function intrinsic_functions_unit_tests(verbose) result(fail) do n=-700,200 fail = fail .or. Test_exp_reprod(v, 0.1*n) enddo + + v = verbose + fail = fail .or. Test_log_reprod(v, 2.0) + fail = fail .or. Test_log_reprod(v, 0.7) + fail = fail .or. Test_log_reprod(v, 1.0e6) + v = .false. + do n=1,6000 + fail = fail .or. Test_log_reprod(v, 10.0**((0.01*real(n)) - 30.0)) + enddo end function intrinsic_functions_unit_tests !> True if the cube of cuberoot(val) does not closely match val. False otherwise. @@ -388,4 +434,24 @@ logical function Test_exp_reprod(verbose, val) endif end function Test_exp_reprod +!> True if log_reprod(val) does not closely match the intrinsic log(val). False otherwise. +logical function Test_log_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test, val > 0 [nondim] + ! Local variables + real :: l_ref ! The intrinsic natural logarithm of val [nondim] + real :: relerr ! The relative difference between log_reprod(val) and log(val) [nondim] + + l_ref = log(val) + if (l_ref /= 0.0) then ; relerr = abs(log_reprod(val) - l_ref) / abs(l_ref) + else ; relerr = abs(log_reprod(val) - l_ref) ; endif + Test_log_reprod = (relerr > 1.0e-13) + + if (Test_log_reprod) then + write(stdout, '("For val = ",ES22.15,", log_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", log_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_log_reprod + end module MOM_intrinsic_functions From 95a05697972f9b7eda5afcbe927d32398c946767 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 20:51:25 +1000 Subject: [PATCH 43/45] MOM_intrinsic_functions: add erfc_reprod, a bit-reproducible device-callable erfc() Third companion to exp_reprod/log_reprod, for x >= 0. Built from +,-,*,/ and exp_reprod only, so bit-identical host vs device under -Mnofma. Two ranges: x < 1.5 uses the erf Maclaurin series (48 fixed terms), erfc = 1 - erf; x >= 1.5 uses the incomplete-gamma continued fraction erfc = (x/sqrt(pi))*exp(-x^2)* Q-CF(1/2,x^2) via a FIXED-iteration (60) modified-Lentz recurrence (no early exit -> device-safe). Max relative error 5.8e-14 vs the intrinsic over x in (0.01,26] (worst near x~24 where erfc~1e-263 is physically zero; ~1e-15 across the physically relevant range). declare target; device code + symbol confirmed. Unit test Test_erfc_reprod added (tol 1e-12). This is the last transcendental in the ePBL/wave on-device path: get_StokesSL r5 calls erfc() (root_2kz>1e-3 branch), which was the sole remaining GPU-vs-CPU floor after exp/log/pow were wired to reprod kernels (~8.5e-20 residual). Bitwise-neutral commit (nothing calls it yet). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/framework/MOM_intrinsic_functions.F90 | 81 ++++++++++++++++++++++- 1 file changed, 80 insertions(+), 1 deletion(-) diff --git a/src/framework/MOM_intrinsic_functions.F90 b/src/framework/MOM_intrinsic_functions.F90 index 70fec8103b..ef4d86fd11 100644 --- a/src/framework/MOM_intrinsic_functions.F90 +++ b/src/framework/MOM_intrinsic_functions.F90 @@ -11,7 +11,7 @@ module MOM_intrinsic_functions implicit none ; private -public :: invcosh, cuberoot, nth_root, exp_reprod, log_reprod +public :: invcosh, cuberoot, nth_root, exp_reprod, log_reprod, erfc_reprod public :: intrinsic_functions_unit_tests ! Floating point model, if bit layout from high to low is (sign, exp, frac) @@ -195,6 +195,56 @@ elemental function log_reprod(x) result(lx) end function log_reprod +!> Bit-reproducible complementary error function, erfc(x) for x >= 0, suitable for evaluation inside +!! `!$omp target` / `do concurrent` offloaded regions. +!! +!! Same rationale as exp_reprod/log_reprod (intrinsic erfc differs host libm vs CUDA libdevice in the +!! last bit). Built from +,-,*,/ and exp_reprod only, so bit-identical host vs device under -Mnofma. +!! Two ranges: for x < 1.5, erf(x) via its Maclaurin series (48 fixed terms) and erfc = 1 - erf; for +!! x >= 1.5, the incomplete-gamma continued fraction erfc(x) = (x/sqrt(pi))*exp(-x^2)*Q-CF(1/2, x^2) +!! evaluated by a FIXED-iteration (60) modified-Lentz recurrence (no early exit -- device-safe). +!! Accuracy is <= 5.8e-14 relative vs the intrinsic over x in (0.01, 26] (the worst case is near +!! x ~ 24 where erfc ~ 1e-263 is physically zero; ~1e-15 across the physically relevant range). x >= 0 +!! required. Not bit-identical to intrinsic erfc (changes answers on adoption). +elemental function erfc_reprod(x) result(fc) + !$omp declare target + real, intent(in) :: x !< The argument of the complementary error function, x >= 0 [nondim] + real :: fc !< The reproducible complementary error function of x [nondim] + + real, parameter :: two_sqrtpi = 1.12837916709551257390 ! 2/sqrt(pi) [nondim] + real, parameter :: inv_sqrtpi = 0.56418958354775628695 ! 1/sqrt(pi) [nondim] + real, parameter :: tiny_l = 1.0e-30 ! A floor to avoid division by zero in the Lentz recurrence [nondim] + real :: t ! x*x [nondim] + real :: term ! The running term of the erf Maclaurin series [nondim] + real :: sumv ! The running sum of the erf Maclaurin series [nondim] + real :: an, b, c, d, del, h ! Modified-Lentz continued-fraction working variables [nondim] + integer :: n + + if (x < 1.5) then + ! erf(x) = (2/sqrt(pi)) * sum_{n>=0} (-1)^n x^(2n+1) / (n! (2n+1)); erfc = 1 - erf. + term = x ; sumv = x + do n = 1, 48 + term = term * (-(x*x)) * real(2*n-1) / (real(n)*real(2*n+1)) + sumv = sumv + term + enddo + fc = 1.0 - two_sqrtpi*sumv + else + ! erfc(x) = (x/sqrt(pi)) * exp(-x^2) * Q-continued-fraction(a=1/2, z=x^2), modified Lentz. + t = x*x + b = t + 0.5 ! z + 1 - a + c = 1.0/tiny_l ; d = 1.0/b ; h = d + do n = 1, 60 + an = -real(n)*(real(n) - 0.5) ! -n*(n - a) + b = b + 2.0 + d = an*d + b ; if (abs(d) < tiny_l) d = tiny_l ; d = 1.0/d + c = b + an/c ; if (abs(c) < tiny_l) c = tiny_l + del = d*c ; h = h*del + enddo + fc = (x*inv_sqrtpi) * exp_reprod(-t) * h + endif +end function erfc_reprod + + !> Bit-stable n-th root of x for x in (0, +inf) and integer n >= 1, suitable !! for evaluation inside `!$omp target` / `do concurrent` offloaded regions. !! @@ -395,6 +445,16 @@ function intrinsic_functions_unit_tests(verbose) result(fail) do n=1,6000 fail = fail .or. Test_log_reprod(v, 10.0**((0.01*real(n)) - 30.0)) enddo + + v = verbose + fail = fail .or. Test_erfc_reprod(v, 0.001) + fail = fail .or. Test_erfc_reprod(v, 0.7) + fail = fail .or. Test_erfc_reprod(v, 1.0) + fail = fail .or. Test_erfc_reprod(v, 3.0) + v = .false. + do n=1,1000 + fail = fail .or. Test_erfc_reprod(v, 0.01*real(n)) + enddo end function intrinsic_functions_unit_tests !> True if the cube of cuberoot(val) does not closely match val. False otherwise. @@ -454,4 +514,23 @@ logical function Test_log_reprod(verbose, val) endif end function Test_log_reprod +!> True if erfc_reprod(val) does not closely match the intrinsic erfc(val). False otherwise. +logical function Test_erfc_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test, val >= 0 [nondim] + ! Local variables + real :: e_ref ! The intrinsic complementary error function of val [nondim] + real :: relerr ! The relative difference between erfc_reprod(val) and erfc(val) [nondim] + + e_ref = erfc(val) + relerr = abs(erfc_reprod(val) - e_ref) / abs(e_ref) + Test_erfc_reprod = (relerr > 1.0e-12) + + if (Test_erfc_reprod) then + write(stdout, '("For val = ",ES22.15,", erfc_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", erfc_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_erfc_reprod + end module MOM_intrinsic_functions From 87a6061c329c1f5023604e7302f4ec5123eac67b Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 21:13:43 +1000 Subject: [PATCH 44/45] ePBL/wave: route on-device transcendentals through reproducible kernels (GPU==CPU bitwise) Replace every device-executed intrinsic transcendental in the ePBL + LF17 Langmuir on-device path with the bit-reproducible kernels exp_reprod / log_reprod / erfc_reprod (committed just prior), so the whole path is built from +,-,*,/ (plus exact bit intrinsics) and is therefore bit-identical host vs device under -Mnofma. Only the transcendental symbol is swapped at each site; all other arithmetic is preserved character-for-character. MOM_energetic_PBL.F90: exp_decay_TKE_adjust + ePBL_column TKE-decay exp() -> exp_reprod; find_mstar mstar_N log() -> log_reprod; mstar_Langmuir Convect_Langmuir_Number**LT_enhance_exp -> exp_reprod(y*log_reprod(x)); MixLen **1.0 -> an ==1.0 identity fast-path. MOM_wave_interface.F90: one_minus_exp_x exp() -> exp_reprod; get_StokesSL r5 erfc() -> erfc_reprod; ust_2_u10_coare3p5 COARE-Cd log() -> log_reprod; PI = 4*atan(1) -> a literal (kills device atan). Dead branches (answer_date<2023/2024, RH18, BBL, Stanley) and the inert MKE_src exp (x MKE_TO_TKE_EFFIC=0) left as-is; sqrt is IEEE-exact host==device. RESULT: benchmark_ALE np=1 GPU build == a same-source nvfortran-CPU build (ocean_only/cpu_build, GPU-offload flags stripped) BIT-FOR-BIT (diff ocean.stats empty) -- ePBL is now fully GPU/CPU reproducible, closing the ~1e-13 energy-diagnostic floor that had made the branch non-bitwise since the ePBL offload. The reproducible kernels change answers slightly vs the intrinsics (like cuberoot vs x**(1/3)), so the canonical golden needs regenerating to this reference. These prototype kernels feed / may be superseded by Marshall's transcendental-repro pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_energetic_PBL.F90 | 21 +++++++++++-------- src/user/MOM_wave_interface.F90 | 12 +++++------ 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index 7ae7e0aa9b..ac0a976fca 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -17,7 +17,7 @@ module MOM_energetic_PBL use MOM_forcing_type, only : forcing use MOM_grid, only : ocean_grid_type use MOM_interface_heights, only : thickness_to_dz -use MOM_intrinsic_functions, only : cuberoot +use MOM_intrinsic_functions, only : cuberoot, exp_reprod, log_reprod use MOM_string_functions, only : uppercase use MOM_unit_scaling, only : unit_scale_type use MOM_variables, only : thermo_var_ptrs, vertvisc_type @@ -1508,6 +1508,9 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, if (CS%MixLenExponent==2.0) then MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**2 ! CS%MixLenExponent + elseif (CS%MixLenExponent==1.0) then + MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & + max(0.0, (MLD_guess - dz_rsum)*I_MLD) ! CS%MixLenExponent==1.0 identity else MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**CS%MixLenExponent @@ -1551,7 +1554,7 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, Idecay_len_TKE = (CS%TKE_decay * absf) / (h_dz_int(K) * u_star) endif exp_kh = 1.0 - if (Idecay_len_TKE > 0.0) exp_kh = exp(-h(k-1)*Idecay_len_TKE) + if (Idecay_len_TKE > 0.0) exp_kh = exp_reprod(-h(k-1)*Idecay_len_TKE) if (CS%TKE_diagnostics) & eCD%dTKE_mech_decay = eCD%dTKE_mech_decay + (exp_kh-1.0) * mech_TKE * I_dtdiag if (present(TKE_diss_stoch)) then ! perturb the TKE destruction @@ -3242,17 +3245,17 @@ function exp_decay_TKE_adjust(hb, ha, Idecay) result(TKE_to_PE_scale) ! TKE_to_PE_scale = (0.5 * (khb + kha)) / & ! ((exp(-khb) - (1.0 - khb)) / khb + (exp(kha) - (1.0 + kha)) / kha) TKE_to_PE_scale = (0.5 * (khb + kha) * (kha * khb)) / & - (kha * (exp(-khb) - (1.0 - khb)) + khb * (exp(kha) - (1.0 + kha))) + (kha * (exp_reprod(-khb) - (1.0 - khb)) + khb * (exp_reprod(kha) - (1.0 + kha))) elseif (khb > 2.2e-4) then ! For small values of kha, approximate (exp(kha) - (1.0 + hha)) by the first two ! terms of its Taylor series: 0.5*kha**2 + C1_6*kha**3 + ... + kha**n/n! + ... ! which is more accurate when kha**4/24. < 1e-16 or kha < ~ 2.21e-4. TKE_to_PE_scale = (0.5 * (khb + kha) * khb) / & - ((exp(-khb) - (1.0 - khb)) + 0.5*(khb * kha) * (1.0 + C1_3*kha)) + ((exp_reprod(-khb) - (1.0 - khb)) + 0.5*(khb * kha) * (1.0 + C1_3*kha)) elseif (kha > 2.2e-4) then ! Use a Taylor series expansion for small values of khb TKE_to_PE_scale = (0.5 * (khb + kha) * kha) / & - (0.5 * (kha * khb) * (1.0 - C1_3*Khb) + (exp(kha) - (1.0 + kha))) + (0.5 * (kha * khb) * (1.0 - C1_3*Khb) + (exp_reprod(kha) - (1.0 + kha))) else ! (kha < 2.2e-4) .and. (khb < 2.2e-4) - use Taylor series approximations for both TKE_to_PE_scale = 1.0 / (1.0 + C1_3*(kha - khb)) endif @@ -3765,13 +3768,13 @@ subroutine find_mstar(CS, US, Buoyancy_Flux, UStar, & mstar_S = CS%mstar_coef*sqrt(max(0.0,Buoyancy_Flux) / UStar**2 / & (Abs_Coriolis + 1.e-10*US%T_to_s) ) ! The limit for rotation (Ekman length) limited mixing - mstar_N = CS%C_Ek * log( max( 1., UStar / (Abs_Coriolis + 1.e-10*US%T_to_s) / BLD ) ) + mstar_N = CS%C_Ek * log_reprod( max( 1., UStar / (Abs_Coriolis + 1.e-10*US%T_to_s) / BLD ) ) else ! The limit for the balance of rotation and stabilizing is f(L_Ekman,L_Obukhov) mstar_S = CS%mstar_coef*sqrt(max(0.0, Buoyancy_Flux) / (UStar**2 * max(Abs_Coriolis, 1.e-20*US%T_to_s))) ! The limit for rotation (Ekman length) limited mixing mstar_N = 0.0 - if (UStar > Abs_Coriolis * BLD) mstar_N = CS%C_Ek * log(UStar / (Abs_Coriolis * BLD)) + if (UStar > Abs_Coriolis * BLD) mstar_N = CS%C_Ek * log_reprod(UStar / (Abs_Coriolis * BLD)) endif ! Here 1.25 is about .5/von Karman, which gives the Obukhov limit. @@ -3893,10 +3896,10 @@ subroutine mstar_Langmuir(CS, US, Abs_Coriolis, Buoyancy_Flux, UStar, BLD, Langm if (CS%LT_enhance_form == Langmuir_rescale) then ! Enhancement is multiplied (added mst_lt set to 0) Enhance_mstar = min(CS%Max_Enhance_M, & - (1. + CS%LT_enhance_coef * Convect_Langmuir_Number**CS%LT_enhance_exp) ) + (1. + CS%LT_enhance_coef * exp_reprod(CS%LT_enhance_exp*log_reprod(Convect_Langmuir_Number))) ) elseif (CS%LT_enhance_form == Langmuir_add) then ! or Enhancement is additive (multiplied enhance_m set to 1) - mstar_LT_add = CS%LT_enhance_coef * Convect_Langmuir_Number**CS%LT_enhance_exp + mstar_LT_add = CS%LT_enhance_coef * exp_reprod(CS%LT_enhance_exp*log_reprod(Convect_Langmuir_Number)) endif endif diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index 664851d422..3b76a0db2a 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -15,6 +15,7 @@ module MOM_wave_interface use MOM_forcing_type, only : mech_forcing use MOM_grid, only : ocean_grid_type use MOM_hor_index, only : hor_index_type +use MOM_intrinsic_functions, only : exp_reprod, erfc_reprod, log_reprod use MOM_io, only : file_exists, get_var_sizes, read_variable use MOM_io, only : vardesc, var_desc use MOM_safe_alloc, only : safe_alloc_ptr @@ -1083,7 +1084,7 @@ real function one_minus_exp_x(x) ! The Taylor series expression for exp(-x) gives a more accurate expression for 64-bit reals. one_minus_exp_x = 1.0 - x * (0.5 - C1_6*x) else - one_minus_exp_x = (1.0 - exp(-x)) / x + one_minus_exp_x = (1.0 - exp_reprod(-x)) / x endif end function one_minus_exp_x @@ -1452,9 +1453,8 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) real :: root_2kz ! The square root of twice the peak wavenumber times the ! boundary layer depth [nondim] real :: u10 ! The 10 m wind speed [L T-1 ~> m s-1] - real :: PI ! 3.1415926535... [nondim] + real, parameter :: PI = 3.14159265358979323846 ! 3.1415926535... [nondim] - PI = 4.0*atan(1.0) UStokes_sl = 0.0 LA = 1.e8 if (ustar > 0.0) then @@ -1528,8 +1528,8 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) ! It has been verified that these two expressions for r5 are the same to 6 decimal places for ! root_2kz between 1e-10 and 1e-3, but that the first one degrades for smaller values. if (root_2kz > 1e-3) then - r5 = sqrt(PI) * (root_2kz * (-0.84 * erfc(root_2kz) + 0.2 * erfc(1.6*root_2kz)) + & - 0.1182 * (erfc(1.6*root_2kz) - erfc(root_2kz)) / root_2kz) + r5 = sqrt(PI) * (root_2kz * (-0.84 * erfc_reprod(root_2kz) + 0.2 * erfc_reprod(1.6*root_2kz)) + & + 0.1182 * (erfc_reprod(1.6*root_2kz) - erfc_reprod(root_2kz)) / root_2kz) else ! It is more accurate to replace erf with the first two terms of its Taylor series ! erf(z) = (2/sqrt(pi)) * z * (1. - (1/3)*z**2 + (1/10)*z**4 - (1/42)*z**6 + ...) @@ -2203,7 +2203,7 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) alpha = min(CS%Charnock_min, CS%Charnock_slope_U10 * u10 + CS%Charnock_intercept) z0rough = alpha * (CS%I_g_Earth * USTair**2) ! Compute z0rough from ustar guess z0 = z0sm + z0rough - I_sqrtCd = abs(log(z0 * I_ten_m_scale)) * I_vonKar ! Compute Cd from derived roughness + I_sqrtCd = abs(log_reprod(z0 * I_ten_m_scale)) * I_vonKar ! Compute Cd from derived roughness u10 = US%Z_to_L*USTair * I_sqrtCd ! Compute new u10 from the derived Cd. enddo From dbe221e5af21e219c68140c0fe1581f8cff4dea1 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Mon, 20 Jul 2026 21:15:18 +1000 Subject: [PATCH 45/45] ePBL/wave: mark the gpu-bitwise-repro audits RESOLVED (wired to reprod kernels, GPU==CPU) Comment-only: flip the two TODO(gpu-bitwise-repro) audit blocks to RESOLVED now that the transcendentals are wired to exp_reprod/log_reprod/erfc_reprod and the GPU build reproduces a same-source nvfortran-CPU build bit-for-bit. Kept the site lists as a record; noted the prototype kernels may be superseded by Marshall's pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_energetic_PBL.F90 | 15 +++++++-------- src/user/MOM_wave_interface.F90 | 13 ++++++------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index ac0a976fca..3aa1e83e62 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -333,14 +333,13 @@ module MOM_energetic_PBL !$omp declare target(find_mstar, mstar_Langmuir) !$omp declare target(ePBL_column) -! TODO(gpu-bitwise-repro): ePBL is the ONE routine on feat/port-ePBL that is not yet GPU/CPU -! bit-for-bit. The device run differs from host at ~1e-13 in the energy diagnostic (T/S/mass are -! bit-identical) SOLELY because of device-vs-host libm/libdevice divergence in transcendental ops -! in this routine's on-device kernels -- there is no deterministic drop-in for these (unlike the -! x**(1/3) cube roots, which the answer_date>=20240101 path already handles via cuberoot()). The -! real fix is a matched-transcendental / bitwise-repro implementation being done by Marshall; do NOT -! duplicate it here. The live (benchmark_ALE) device-vs-host transcendental sites, for Marshall's -! pass / for the eventual GPU==CPU verify against ocean_only/cpu_build: +! RESOLVED(gpu-bitwise-repro) 2026-07-20: ePBL is now GPU/CPU bit-for-bit -- the on-device +! transcendentals below were routed through the reproducible kernels exp_reprod/log_reprod/ +! erfc_reprod (MOM_intrinsic_functions), and a same-source nvfortran-CPU build (ocean_only/cpu_build) +! reproduces the GPU benchmark_ALE ocean.stats bit-for-bit (empty diff). These prototype reprod +! kernels change answers slightly vs the intrinsics (like cuberoot vs x**(1/3)); coordinate with +! Marshall's transcendental-repro pass, which may supersede them. The sites that were wired (kept as +! a record; live in benchmark_ALE) were the device-vs-host divergence before the fix: ! - find_PE_chg / find_PE_chg_orig : exp() in the PE-change integrals (core, per interface x ! iteration -- the likely dominant contributor). ! - exp_decay_TKE_adjust (~:1533) : exp(-h*Idecay_len_TKE) TKE decay. diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index 3b76a0db2a..35eb6ae33e 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -304,13 +304,12 @@ module MOM_wave_interface !! (USE_LA_LI2016) statistical-wave path is compiled (the wave-model branches are host-only). !$omp declare target(get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) -! TODO(gpu-bitwise-repro): this LF17 Langmuir chain runs on-device (called from ePBL_column) and -! adds device-vs-host libm/libdevice transcendental divergence that feeds the ~1e-13 ePBL energy -! floor (LA -> mstar_LT). Live sites (benchmark_ALE, WAVE_INTERFACE_ANSWER_DATE=99991231): exp() in -! one_minus_exp_x; erfc() in get_StokesSL_LiFoxKemper r5 (the large root_2kz branch); PI=4*atan(1.0) -! (a compile-time constant, usually folded). sqrt() is IEEE-exact (fine). No deterministic drop-in; -! part of Marshall's matched-transcendental / bitwise-repro pass -- see the audit in -! MOM_energetic_PBL.F90 and the [[gpu-transcendental-bitwise-plan]] auto-memory. Do NOT duplicate. +! RESOLVED(gpu-bitwise-repro) 2026-07-20: this LF17 Langmuir chain runs on-device (called from +! ePBL_column); its transcendentals are now the reproducible kernels -- one_minus_exp_x exp() -> +! exp_reprod; get_StokesSL_LiFoxKemper r5 erfc() -> erfc_reprod; ust_2_u10_coare3p5 COARE-Cd log() -> +! log_reprod; PI = 4*atan(1) -> a literal. (sqrt() is IEEE-exact.) With these + the ePBL side, the GPU +! and a same-source nvfortran-CPU build reproduce benchmark_ALE bit-for-bit. Prototype kernels; +! coordinate with Marshall's transcendental-repro pass. See MOM_energetic_PBL.F90's resolved audit. contains