diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 3f3c79db5a..cc45403599 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -30,6 +30,15 @@ module MOM_mixed_layer_restrat implicit none ; private #include +#include "do_concurrent_compat.h" + +#ifdef __NVCOMPILER_OPENMP_GPU +integer, parameter :: default_nkblock = 0 !< Default k block size for the mixed layer density integral [nondim] +integer, parameter :: default_njblock = 0 !< Default j block size for the mixed layer density integral [nondim] +#else +integer, parameter :: default_nkblock = 1 !< Default k block size for the mixed layer density integral [nondim] +integer, parameter :: default_njblock = 1 !< Default j block size for the mixed layer density integral [nondim] +#endif public mixedlayer_restrat public mixedlayer_restrat_init @@ -111,6 +120,12 @@ module MOM_mixed_layer_restrat !! units [H2 T-1 ~> m2 s-1 or kg2 m-4 s-1] logical :: MLD_grid !< If true, read a spacially varying field for MLD_decaying_Tfilt logical :: Cr_grid !< If true, read a spacially varying field for Cr + integer :: nkblock = default_nkblock !< The number of layers whose density is evaluated in one call to the + !! equation of state in the Bodner mixed-layer buoyancy integral, or 0 + !! to do the whole column in a single block [nondim]. + integer :: njblock = default_njblock !< The number of rows whose density is evaluated in one call to the + !! equation of state in the Bodner mixed-layer buoyancy integral, or 0 + !! to do the whole j compute domain in a single block [nondim]. real, dimension(:,:), allocatable :: & MLD_filtered, & !< Time-filtered MLD [H ~> m or kg m-2] @@ -716,7 +731,7 @@ subroutine mixedlayer_restrat_OM4(h, uhtr, vhtr, tv, forces, dt, h_MLD, VarMix, end subroutine mixedlayer_restrat_OM4 !> Stream function shape as a function of non-dimensional position within mixed-layer [nondim] -real function mu(sigma, dh) +pure real function mu(sigma, dh) real, intent(in) :: sigma !< Fractional position within mixed layer [nondim] !! z=0 is surface, z=-1 is the bottom of the mixed layer real, intent(in) :: dh !< Non-dimensional distance over which to extend stream @@ -796,11 +811,19 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! mode [Z T-1 ~> m s-1] real :: covTS(SZI_(G)) ! SGS TS covariance in Stanley param; currently 0 [C S ~> degC ppt] real :: varS(SZI_(G)) ! SGS S variance in Stanley param; currently 0 [S2 ~> ppt2] - real :: dmu(SZK_(GV)) ! Change in mu(z) across layer k [nondim] real :: Rml_int(SZI_(G)) ! Potential density integrated through the mixed layer [R H ~> kg m-2 or kg2 m-5] real :: SpV_ml(SZI_(G)) ! Specific volume evaluated at the surface pressure [R-1 ~> m3 kg-1] real :: SpV_int(SZI_(G)) ! Specific volume integrated through the mixed layer [H R-1 ~> m4 kg-1 or m] real :: rho_ml(SZI_(G)) ! Potential density relative to the surface [R ~> kg m-3] + real :: rho_blk(SZI_(G), merge(G%jed-G%jsd+1, CS%njblock, CS%njblock==0), & + merge(GV%ke, CS%nkblock, CS%nkblock==0)) + ! Potential density relative to the surface for a tile [R ~> kg m-3] + real :: p_blk(SZI_(G), merge(G%jed-G%jsd+1, CS%njblock, CS%njblock==0), & + merge(GV%ke, CS%nkblock, CS%nkblock==0)) + ! A pressure of 0 for a tile [R L2 T-2 ~> Pa] + real :: Rml_blk(SZI_(G), merge(G%jed-G%jsd+1, CS%njblock, CS%njblock==0)) + ! Density integrated through the mixed layer, for the rows of a tile, + ! carried across the k blocks of a column [R H ~> kg m-2 or kg2 m-5] real :: p0(SZI_(G)) ! A pressure of 0 [R L2 T-2 ~> Pa] real :: g_Rho0 ! G_Earth/Rho0 times a thickness conversion factor ! [L2 H-1 T-2 R-1 ~> m4 s-2 kg-1 or m7 s-2 kg-2] @@ -830,8 +853,13 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d real :: m2_s2_to_Z2_T2 ! Conversion factors to restore scaling after a term is raised to a ! fractional power [Z2 s2 T-2 m-2 ~> 1] real, parameter :: two_thirds = 2./3. ! [nondim] + real :: dmu ! Change in mu(z) across a layer [nondim] logical :: line_is_empty, keep_going integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state + integer :: EOSdom3(3,2) ! The (i,j,k) computational domain for the blocked equation of state calls + integer :: nkblock, njblock ! The number of layers and rows in each tile of the density integral [nondim] + integer :: kstart, kend ! The first and last layer of the tile being worked on [nondim] + integer :: jstart, jend ! The first and last row of the tile being worked on [nondim] integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke @@ -879,37 +907,48 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif + !$omp target enter data map(to: U_star_2d, h_MLD) + !$omp target enter data map(alloc: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) + !$omp target enter data map(alloc: vol_dt_avail, uhml, vhml) + ! Apply time filter to h_MLD (to remove diurnal cycle) to obtain "little h". ! "little h" is representative of the active mixing layer depth, used in B22 formula (eq 27). - do j=js-1,je+1 ; do i=is-1,ie+1 + do concurrent (j=js-1:je+1, i=is-1:ie+1) little_h(i,j) = rmean2ts(h_MLD(i,j), CS%MLD_filtered(i,j), & CS%BLD_growing_Tfilt, CS%BLD_decaying_Tfilt, dt) CS%MLD_filtered(i,j) = little_h(i,j) - enddo ; enddo + enddo ! Calculate "big H", representative of the mixed layer depth, used in B22 formula (eq 27). + ! not offloaded, send back to the host if (CS%MLD_grid) then + !$omp target update from(little_h, CS%MLD_filtered_slow) do j=js-1,je+1 ; do i=is-1,ie+1 big_H(i,j) = rmean2ts(little_h(i,j), CS%MLD_filtered_slow(i,j), & CS%MLD_growing_Tfilt, CS%MLD_Tfilt_space(i,j), dt) enddo ; enddo + !$omp target update to(big_H) elseif (CS%Bodner_detect_MLD) then + !$omp target update from(CS%MLD_filtered_slow) call detect_mld(h, tv, MLD, G, GV, CS) do j=js-1,je+1 ; do i=is-1,ie+1 big_H(i,j) = rmean2ts(MLD(i,j), CS%MLD_filtered_slow(i,j), & CS%MLD_growing_Tfilt, CS%MLD_decaying_Tfilt, dt) enddo ; enddo + !$omp target update to(big_H) else - do j=js-1,je+1 ; do i=is-1,ie+1 + do concurrent (j=js-1:je+1, i=is-1:ie+1) big_H(i,j) = rmean2ts(little_h(i,j), CS%MLD_filtered_slow(i,j), & CS%MLD_growing_Tfilt, CS%MLD_decaying_Tfilt, dt) - enddo ; enddo + enddo endif - do j=js-1,je+1 ; do i=is-1,ie+1 + do concurrent (j=js-1:je+1, i=is-1:ie+1) CS%MLD_filtered_slow(i,j) = big_H(i,j) - enddo ; enddo + enddo ! Estimate w'u' at h-points, with a floor to avoid division by zero later. + ! needed here, otehrwe death + !$omp target enter data map(to: BLD, bflux) if (allocated(tv%SpV_avg) .and. .not.(GV%Boussinesq .or. GV%semi_Boussinesq)) then do j=js-1,je+1 ; do i=is-1,ie+1 ! This expression differs by a factor of 1. / (Rho_0 * SpV_avg) compared with the other @@ -926,6 +965,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! bflux (such as SpV_avg**-2/3 combining with other terms in bflux to give the thermal ! expansion coefficient) and because the specific volume does vary within the mixed layer. enddo ; enddo + !$omp target update to(wpup) elseif (CS%answer_date < 20240201) then Z3_T3_to_m3_s3 = (US%Z_to_m * US%s_to_T)**3 m2_s2_to_Z2_T2 = (US%m_to_Z * US%T_to_s)**2 @@ -935,22 +975,27 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d wpup(i,j) = max(m2_s2_to_Z2_T2 * (Z3_T3_to_m3_s3 * ( CS%mstar * u_star3 + CS%nstar * w_star3 ) )**two_thirds, & CS%min_wstar2) * US%Z_to_L * GV%Z_to_H ! In [L H T-2 ~> m2 s-2 or kg m-1 s-2] enddo ; enddo + !$omp target update to(wpup) else - do j=js-1,je+1 ; do i=is-1,ie+1 - w_star3 = max(0., -bflux(i,j)) * BLD(i,j) ! In [Z3 T-3 ~> m3 s-3] - wpup(i,j) = max( (cuberoot(CS%mstar * U_star_2d(i,j)**3 + CS%nstar * w_star3))**2, CS%min_wstar2 ) & + do concurrent (j=js-1:je+1, i=is-1:ie+1) + wpup(i,j) = max( (cuberoot(CS%mstar * U_star_2d(i,j)**3 & + + CS%nstar * (max(0., -bflux(i,j)) * BLD(i,j))))**2, CS%min_wstar2 ) & * US%Z_to_L * GV%Z_to_H ! In [L H T-2 ~> m2 s-2 or kg m-1 s-2] - enddo ; enddo + enddo endif + !$omp target exit data map(release: BLD, bflux) + ! We filter w'u' with the same time scales used for "little h" - do j=js-1,je+1 ; do i=is-1,ie+1 + do concurrent (j=js-1:je+1, i=is-1:ie+1) wpup(i,j) = rmean2ts(wpup(i,j), CS%wpup_filtered(i,j), & CS%BLD_growing_Tfilt, CS%BLD_decaying_Tfilt, dt) CS%wpup_filtered(i,j) = wpup(i,j) - enddo ; enddo + enddo if (CS%id_lfbod > 0) then + ! This diagnostic is evaluated on the host, so it needs the device's "little h". + !$omp target update from(little_h) do j=js-1,je+1 ; do i=is-1,ie+1 ! Calculate front length used in B22 formula (eq 24). w_star3 = max(0., -bflux(i,j)) * BLD(i,j) @@ -980,6 +1025,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d endif if (CS%debug) then + !$omp target update from(little_h, big_H, wpup, CS%MLD_filtered, CS%MLD_filtered_slow) call hchksum(little_h,'mle_Bodner: little_h', G%HI, haloshift=1, unscale=GV%H_to_mks) call hchksum(big_H,'mle_Bodner: big_H', G%HI, haloshift=1, unscale=GV%H_to_mks) call hchksum(CS%MLD_filtered,'mle_Bodner: MLD_filtered 2', & @@ -994,62 +1040,128 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! in-situ density would contain the MLD gradient (through the pressure dependence). p0(:) = 0.0 EOSdom(:) = EOS_domain(G%HI, halo=1) - !$OMP parallel & - !$OMP default(shared) & - !$OMP private(i, j, k, keep_going, line_is_empty, dh, & - !$OMP grid_dsd, absf, h_sml, h_big, grd_b, r_wpup, psi_mag, IhTot, & - !$OMP sigint, muzb, muza, hAtVel, Rml_int, SpV_int) - !$OMP do - do j=js-1,je+1 - rho_ml(:) = 0.0 ; SpV_ml(:) = 0.0 - do i=is-1,ie+1 - htot(i,j) = 0.0 ; Rml_int(i) = 0.0 ; SpV_int(i) = 0.0 + do concurrent (k=1:nz, j=js-1:je+1, i=is-1:ie+1) + vol_dt_avail(i,j,k) = max(I4dt*G%areaT(i,j)*(h(i,j,k)-GV%Angstrom_H), 0.0) + enddo + + if ((GV%Boussinesq .or. GV%semi_Boussinesq) .and. .not.CS%use_Stanley_ML) then + ! Walk the mixed layer in tiles of njblock rows by nkblock layers: evaluate the density for a + ! tile, then integrate it. The columns are independent of one another, so the integral runs as + ! one concurrent loop over (i,j) with the sequential k-dependence kept inside each column, and + ! the running totals for the tile's rows carry across its k blocks. + nkblock = merge(nz, CS%nkblock, CS%nkblock==0) + njblock = merge((je+1) - (js-1) + 1, CS%njblock, CS%njblock==0) + EOSdom3(1,:) = EOS_domain(G%HI, halo=1) + + ! The tile scratch never leaves the device, and no other branch uses it. + !$omp target enter data map(alloc: rho_blk, p_blk, Rml_blk) + + if (nkblock < nz) then + ! The early-exit test below compares htot against big_H on the host, and big_H has been + ! computed on the device, so fetch it once. With a single k block that test never runs. + !$omp target update from(big_H) + endif + + do concurrent (k=1:nkblock, j=1:njblock, i=is-1:ie+1) + p_blk(i,j,k) = 0.0 enddo - keep_going = .true. - do k=1,nz + + do jstart=js-1,je+1,njblock + jend = min(jstart+njblock-1, je+1) + EOSdom3(2,:) = [1, jend-jstart+1] + + do concurrent (j=jstart:jend, i=is-1:ie+1) + htot(i,j) = 0.0 ; Rml_blk(i,j-jstart+1) = 0.0 + enddo + + keep_going = .true. + do kstart=1,nz,nkblock ; if (keep_going) then + kend = min(kstart+nkblock-1, nz) + EOSdom3(3,:) = [1, kend-kstart+1] + call calculate_density(tv%T(:,jstart:jend,kstart:kend), tv%S(:,jstart:jend,kstart:kend), & + p_blk, rho_blk, tv%eqn_of_state, EOSdom3) + + do concurrent (j=jstart:jend, i=is-1:ie+1) DO_LOCALITY(local(k, dh)) + do k=kstart,kend + if (htot(i,j) < big_H(i,j)) then + dh = min( h(i,j,k), big_H(i,j) - htot(i,j) ) + ! Rml_blk is in [R H ~> kg m-2] + Rml_blk(i,j-jstart+1) = Rml_blk(i,j-jstart+1) + dh*rho_blk(i,j-jstart+1,k-kstart+1) + htot(i,j) = htot(i,j) + dh + endif + enddo + enddo + + if (nkblock < nz) then + ! Stop calling the equation of state once every column in the tile has been filled to + ! "big H". With a single k block there is nothing left to skip, and this test would drag + ! htot back to the host for nothing. + !$omp target update from(htot) + keep_going = .false. + do j=jstart,jend ; do i=is-1,ie+1 + if (htot(i,j) < big_H(i,j)) keep_going = .true. + enddo ; enddo + endif + endif ; enddo + + do concurrent (j=jstart:jend, i=is-1:ie+1) + ! Buoy_av has units (L2 H-1 T-2 R-1) * (R H) * H-1 = [L2 H-1 T-2 ~> m s-2 or m4 kg-1 s-2] + buoy_av(i,j) = -( g_Rho0 * Rml_blk(i,j-jstart+1) ) / (htot(i,j) + h_neglect) + enddo + enddo + + !$omp target exit data map(release: rho_blk, p_blk, Rml_blk) + else + ! The Stanley and non-Boussinesq variants of the equation of state have no array-of-columns + ! interface to block over, so they retain the original row-at-a-time form on the host. + !$OMP parallel do default(shared) & + !$OMP private(i, k, keep_going, line_is_empty, dh, rho_ml, SpV_ml, Rml_int, SpV_int) + do j=js-1,je+1 + rho_ml(:) = 0.0 ; SpV_ml(:) = 0.0 do i=is-1,ie+1 - vol_dt_avail(i,j,k) = max(I4dt*G%areaT(i,j)*(h(i,j,k)-GV%Angstrom_H),0.0) + htot(i,j) = 0.0 ; Rml_int(i) = 0.0 ; SpV_int(i) = 0.0 enddo - if (keep_going) then - if (GV%Boussinesq .or. GV%semi_Boussinesq) then - if (CS%use_Stanley_ML) then + keep_going = .true. + do k=1,nz + if (keep_going) then + if (GV%Boussinesq .or. GV%semi_Boussinesq) then call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p0, tv%varT(:,j,k), covTS, varS, & rho_ml, tv%eqn_of_state, EOSdom) else - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p0, rho_ml, tv%eqn_of_state, EOSdom) + call calculate_spec_vol(tv%T(:,j,k), tv%S(:,j,k), p0, SpV_ml, tv%eqn_of_state, EOSdom) endif - else - call calculate_spec_vol(tv%T(:,j,k), tv%S(:,j,k), p0, SpV_ml, tv%eqn_of_state, EOSdom) + line_is_empty = .true. + do i=is-1,ie+1 + if (htot(i,j) < big_H(i,j)) then + dh = min( h(i,j,k), big_H(i,j) - htot(i,j) ) + Rml_int(i) = Rml_int(i) + dh*rho_ml(i) ! Rml_int has units of [R H ~> kg m-2] + SpV_int(i) = SpV_int(i) + dh*SpV_ml(i) ! SpV_int has units of [H R-1 ~> m4 kg-1 or m] + htot(i,j) = htot(i,j) + dh + line_is_empty = .false. + endif + enddo + if (line_is_empty) keep_going=.false. endif - line_is_empty = .true. + enddo + + if (GV%Boussinesq .or. GV%semi_Boussinesq) then do i=is-1,ie+1 - if (htot(i,j) < big_H(i,j)) then - dh = min( h(i,j,k), big_H(i,j) - htot(i,j) ) - Rml_int(i) = Rml_int(i) + dh*rho_ml(i) ! Rml_int has units of [R H ~> kg m-2] - SpV_int(i) = SpV_int(i) + dh*SpV_ml(i) ! SpV_int has units of [H R-1 ~> m4 kg-1 or m] - htot(i,j) = htot(i,j) + dh - line_is_empty = .false. - endif + ! Buoy_av has units (L2 H-1 T-2 R-1) * (R H) * H-1 = [L2 H-1 T-2 ~> m s-2 or m4 kg-1 s-2] + buoy_av(i,j) = -( g_Rho0 * Rml_int(i) ) / (htot(i,j) + h_neglect) + enddo + else + do i=is-1,ie+1 + ! Buoy_av has units (R L2 H-1 T-2) * (R-1 H) * H-1 = [L2 H-1 T-2 ~> m s-2 or m4 kg-1 s-2] + buoy_av(i,j) = (GV%H_to_RZ*GV%g_Earth * SpV_int(i)) / (htot(i,j) + h_neglect) enddo - if (line_is_empty) keep_going=.false. endif enddo - - if (GV%Boussinesq .or. GV%semi_Boussinesq) then - do i=is-1,ie+1 - ! Buoy_av has units (L2 H-1 T-2 R-1) * (R H) * H-1 = [L2 H-1 T-2 ~> m s-2 or m4 kg-1 s-2] - buoy_av(i,j) = -( g_Rho0 * Rml_int(i) ) / (htot(i,j) + h_neglect) - enddo - else - do i=is-1,ie+1 - ! Buoy_av has units (R L2 H-1 T-2) * (R-1 H) * H-1 = [L2 H-1 T-2 ~> m s-2 or m4 kg-1 s-2] - buoy_av(i,j) = (GV%H_to_RZ*GV%g_Earth * SpV_int(i)) / (htot(i,j) + h_neglect) - enddo - endif - enddo + !$omp target update to(htot, buoy_av) + endif if (CS%debug) then + !$omp target update from(htot, vol_dt_avail, buoy_av) call hchksum(htot,'mle_Bodner: htot', G%HI, haloshift=1, unscale=GV%H_to_mks) call hchksum(vol_dt_avail,'mle_Bodner: vol_dt_avail', G%HI, haloshift=1, & unscale=US%L_to_m**2*GV%H_to_mks*US%s_to_T) @@ -1057,8 +1169,8 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d endif ! U - Component - !$OMP do - do j=js,je ; do I=is-1,ie + do concurrent (j=js:je, I=is-1:ie) & + DO_LOCALITY(local(k,dmu,grid_dsd,absf,h_sml,h_big,grd_b,r_wpup,psi_mag,IhTot,sigint,muzb,muza,hAtVel)) if (G%OBCmaskCu(I,j) > 0.) then grid_dsd = sqrt(0.5*( G%dxCu(I,j)**2 + G%dyCu(I,j)**2 )) * G%dyCu(I,j) ! [L2 ~> m2] absf = 0.5*(abs(G%CoriolisBu(I,J-1)) + abs(G%CoriolisBu(I,J))) ! [T-1 ~> s-1] @@ -1079,27 +1191,29 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d muza = muzb ! mu(z/MLD) for upper interface [nondim] hAtVel = 0.5*(h(i,j,k) + h(i+1,j,k)) ! Thickness at velocity point [H ~> m or kg m-2] sigint = sigint - (hAtVel * IhTot) ! z/H for lower interface [nondim] - muzb = mu(sigint, CS%MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] - dmu(k) = muza - muzb ! Change in mu(z) across layer [nondim] - ! dmu(k)*psi_mag is the transport in this layer [L2 H T-1 ~> m3 s-1] + muzb = mu(sigint, CS%MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] + dmu = muza - muzb ! Change in mu(z) across layer [nondim] + uhml(I,j,k) = dmu ! Stash dmu in uhml: the columns run concurrently, so there is nowhere + ! else to keep a per-column profile. It is scaled by psi_mag below. + ! dmu*psi_mag is the transport in this layer [L2 H T-1 ~> m3 s-1] ! Limit magnitude (psi_mag) if it would violate CFL - if (dmu(k)*psi_mag > 0.0) then - if (dmu(k)*psi_mag > vol_dt_avail(i,j,k)) psi_mag = vol_dt_avail(i,j,k) / dmu(k) - elseif (dmu(k)*psi_mag < 0.0) then - if (-dmu(k)*psi_mag > vol_dt_avail(i+1,j,k)) psi_mag = -vol_dt_avail(i+1,j,k) / dmu(k) + if (dmu*psi_mag > 0.0) then + if (dmu*psi_mag > vol_dt_avail(i,j,k)) psi_mag = vol_dt_avail(i,j,k) / dmu + elseif (dmu*psi_mag < 0.0) then + if (-dmu*psi_mag > vol_dt_avail(i+1,j,k)) psi_mag = -vol_dt_avail(i+1,j,k) / dmu endif enddo ! These loops cannot be fused because psi_mag applies to the whole column do k=1,nz - uhml(I,j,k) = dmu(k) * psi_mag ! [L2 H T-1 ~> m3 s-1 or kg s-1] + uhml(I,j,k) = uhml(I,j,k) * psi_mag ! [L2 H T-1 ~> m3 s-1 or kg s-1] uhtr(I,j,k) = uhtr(I,j,k) + uhml(I,j,k) * dt ! [L2 H ~> m3 or kg] enddo uDml_diag(I,j) = psi_mag - enddo ; enddo + enddo ! V- component - !$OMP do - do J=js-1,je ; do i=is,ie + do concurrent (J=js-1:je, i=is:ie) & + DO_LOCALITY(local(k,dmu,grid_dsd,absf,h_sml,h_big,grd_b,r_wpup,psi_mag,IhTot,sigint,muzb,muza,hAtVel)) if (G%OBCmaskCv(i,J) > 0.) then grid_dsd = sqrt(0.5*( G%dxCv(i,J)**2 + G%dyCv(i,J)**2 )) * G%dxCv(i,J) ! [L2 ~> m2] absf = 0.5*(abs(G%CoriolisBu(I-1,J)) + abs(G%CoriolisBu(I,J))) ! [T-1 ~> s-1] @@ -1120,30 +1234,44 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d muza = muzb ! mu(z/MLD) for upper interface [nondim] hAtVel = 0.5*(h(i,j,k) + h(i,j+1,k)) ! Thickness at velocity point [H ~> m or kg m-2] sigint = sigint - (hAtVel * IhTot) ! z/H for lower interface [nondim] - muzb = mu(sigint, CS%MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] - dmu(k) = muza - muzb ! Change in mu(z) across layer [nondim] - ! dmu(k)*psi_mag is the transport in this layer [L2 H T-1 ~> m3 s-1 or kg s-1] + muzb = mu(sigint, CS%MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] + dmu = muza - muzb ! Change in mu(z) across layer [nondim] + vhml(i,J,k) = dmu ! Stash dmu in vhml, as for uhml above; scaled by psi_mag below. + ! dmu*psi_mag is the transport in this layer [L2 H T-1 ~> m3 s-1 or kg s-1] ! Limit magnitude (psi_mag) if it would violate CFL - if (dmu(k)*psi_mag > 0.0) then - if (dmu(k)*psi_mag > vol_dt_avail(i,j,k)) psi_mag = vol_dt_avail(i,j,k) / dmu(k) - elseif (dmu(k)*psi_mag < 0.0) then - if (-dmu(k)*psi_mag > vol_dt_avail(i,j+1,k)) psi_mag = -vol_dt_avail(i,j+1,k) / dmu(k) + if (dmu*psi_mag > 0.0) then + if (dmu*psi_mag > vol_dt_avail(i,j,k)) psi_mag = vol_dt_avail(i,j,k) / dmu + elseif (dmu*psi_mag < 0.0) then + if (-dmu*psi_mag > vol_dt_avail(i,j+1,k)) psi_mag = -vol_dt_avail(i,j+1,k) / dmu endif enddo ! These loops cannot be fused because psi_mag applies to the whole column do k=1,nz - vhml(i,J,k) = dmu(k) * psi_mag ! [L2 H T-1 ~> m3 s-1 or kg s-1] + vhml(i,J,k) = vhml(i,J,k) * psi_mag ! [L2 H T-1 ~> m3 s-1 or kg s-1] vhtr(i,J,k) = vhtr(i,J,k) + vhml(i,J,k) * dt ! [L2 H ~> m3 or kg] enddo vDml_diag(i,J) = psi_mag - enddo ; enddo + enddo - !$OMP do - do j=js,je ; do k=1,nz ; do i=is,ie + do concurrent (j=js:je, k=1:nz, i=is:ie) h(i,j,k) = h(i,j,k) - dt*G%IareaT(i,j) * & ((uhml(I,j,k) - uhml(I-1,j,k)) + (vhml(i,J,k) - vhml(i,J-1,k))) - enddo ; enddo ; enddo - !$OMP end parallel + enddo + + ! h, uhtr and vhtr stay mapped for the whole run, but step_MOM_dynamics pushes the host copies + ! back to the device after this call, so the values just computed here have to reach the host. + !$omp target update from(h, uhtr, vhtr) + + ! uhml and vhml are otherwise device-only; retrieve them only when they are being diagnosed. + if (CS%id_uhml > 0) then + !$omp target update from(uhml) + endif + if (CS%id_vhml > 0) then + !$omp target update from(vhml) + endif + + !$omp target exit data map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) + !$omp target exit data map(release: vol_dt_avail, uhml, vhml, U_star_2d, h_MLD) if (CS%id_uhml > 0 .or. CS%id_vhml > 0) & ! Remapped uhml and vhml require east/north halo updates of h @@ -1735,6 +1863,14 @@ logical function mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, CS, call get_param(param_file, mdl, "USE_STANLEY_TVAR", CS%use_Stanley_ML, & "If true, turn on Stanley SGS T variance parameterization "// & "in ML restrat code.", default=.false.) + call get_param(param_file, mdl, "MLE_NKBLOCK", CS%nkblock, & + "Vertical block size for the mixed layer density integral. "//& + "0 processes all layers in a single block.", & + default=default_nkblock) + call get_param(param_file, mdl, "MLE_NJBLOCK", CS%njblock, & + "j-tile size for the mixed layer density integral. "//& + "0 uses the full j compute domain.", & + default=default_njblock) call get_param(param_file, mdl, "USE_CR_GRID", CS%Cr_grid, & "If true, read in a spatially varying Cr field.", default=.false.) call get_param(param_file, mdl, "USE_MLD_GRID", CS%MLD_grid, & @@ -1955,6 +2091,13 @@ logical function mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, CS, if (allocated(CS%MLD_filtered_slow)) call pass_var(CS%MLD_filtered_slow, G%domain) if (allocated(CS%wpup_filtered)) call pass_var(CS%wpup_filtered, G%domain) + if (CS%use_Bodner) then + ! very important! + !$omp target update to(CS) + !$omp target enter data map(to: CS%Cr_space) + !$omp target enter data map(to: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) + endif + end function mixedlayer_restrat_init !> Allocate and register fields in the mixed layer restratification structure for restarts