From 774039e6f1187d713b3ce48c197618827660a144 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 19:21:24 -0500 Subject: [PATCH 01/11] Offload the Bodner MLE restratification to the GPU Rewrite of mixedlayer_restrat_Bodner as do concurrent loops inside a single persistent device data region, replacing the earlier naive port. * The mixed-layer density integral evaluates the equation of state a block of MLE_NKBLOCK layers at a time rather than hoisting it to a whole-column 3D call. The block scratch is (SZI,SZJ,nkblock), so the CPU default of 1 keeps it to a single slab; the GPU default of 0 puts the column in one block. The row-at-a-time early exit becomes a slab-wide one, which only changes work that the htot < big_H guard already discarded. * No shadow copies of CS/GV/US scalars or of the input fields. The kernels read CS% and tv% directly, with those parents mapped alongside their components as nvfortran requires once any component is mapped. * One target data region covers the whole calculation, so nothing round-trips between phases. Arrays that only one branch or diagnostic touches (the block scratch, uhml/vhml, the host-side big_H and wpup variants, the chksums) are mapped or synced inside that branch. * dmu is stashed in uhml/vhml instead of a per-column array, which do concurrent cannot give each iteration without a per-thread allocation. mu is now pure so it can be called from a concurrent loop. Co-Authored-By: Claude Opus 5 (1M context) --- .../lateral/MOM_mixed_layer_restrat.F90 | 277 +++++++++++++----- 1 file changed, 198 insertions(+), 79 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 3f3c79db5a..cfaa8b798b 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -30,6 +30,13 @@ 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] +#else +integer, parameter :: default_nkblock = 1 !< Default k block size for the mixed layer density integral [nondim] +#endif public mixedlayer_restrat public mixedlayer_restrat_init @@ -111,6 +118,9 @@ 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]. real, dimension(:,:), allocatable :: & MLD_filtered, & !< Time-filtered MLD [H ~> m or kg m-2] @@ -716,7 +726,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 +806,18 @@ 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 :: Rml_int(SZI_(G),SZJ_(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] + ! The density integral works through the column nkblock layers at a time, so that the equation + ! of state is handed enough points to be worth offloading while these scratch arrays stay small. + ! MLE_NKBLOCK=1 (the CPU default) makes them a single 2D slab; 0 puts the whole column in one block. + real :: rho_blk(SZI_(G),SZJ_(G),merge(GV%ke, CS%nkblock, CS%nkblock==0)) + ! Potential density relative to the surface for a block of layers [R ~> kg m-3] + real :: p_blk(SZI_(G),SZJ_(G),merge(GV%ke, CS%nkblock, CS%nkblock==0)) + ! A pressure of 0 for a block of layers [R L2 T-2 ~> Pa] 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 +847,12 @@ 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 ! The number of layers in each block of the density integral [nondim] + integer :: kstart, kend ! The first and last layer of the block 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,35 +900,56 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif + ! Everything from here to the diagnostics is evaluated on the device inside a single data region, + ! so no field crosses the bus between phases. Two rules shape the map clauses: + ! - A derived type is mapped alongside its components, parent first. nvfortran turns off implicit + ! derived-type copyin for any region that maps a component, and initialize_MOM maps the G% + ! metrics, so every kernel in MOM6 is already in that regime. G, GV and US are mapped there and + ! stay present, which is why CS and tv are the only parents that need mapping here. + ! - h, uhtr and vhtr are mapped for the whole run by initialize_MOM and are used present. + ! Fields that only one branch or one diagnostic touches are mapped within that branch instead. + !$omp target data & + !$omp map(to: CS, tv) & + !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d) & + !$omp map(tofrom: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) & + !$omp map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) & + !$omp map(alloc: Rml_int, 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). + ! The two spatially-varying variants below are not offloaded, so they push their result to the + ! device for the kernels that follow. 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. if (allocated(tv%SpV_avg) .and. .not.(GV%Boussinesq .or. GV%semi_Boussinesq)) then @@ -926,6 +968,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 +978,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 ) & + ! w_star3 = max(0., -bflux(i,j)) * BLD(i,j), in [Z3 T-3 ~> m3 s-3], is written out in place + ! here rather than held in a scalar that would have to be made local to the concurrent loop. + 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 ! 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 +1028,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 +1043,113 @@ 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 + ! Evaluate the density for a block of nkblock layers at a time and then integrate that block + ! through the mixed layer. 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. + nkblock = merge(nz, CS%nkblock, CS%nkblock==0) + EOSdom3(1,:) = EOS_domain(G%HI, halo=1) + EOSdom3(2,:) = [(js-1) - (G%jsd-1), (je+1) - (G%jsd-1)] + + ! The block scratch never leaves the device, and no other branch uses it. + !$omp target data map(alloc: rho_blk, p_blk) + + do concurrent (k=1:nkblock, j=js-1:je+1, i=is-1:ie+1) + p_blk(i,j,k) = 0.0 enddo + do concurrent (j=js-1:je+1, i=is-1:ie+1) + htot(i,j) = 0.0 ; Rml_int(i,j) = 0.0 + enddo + keep_going = .true. - do k=1,nz - 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) - enddo - if (keep_going) then - if (GV%Boussinesq .or. GV%semi_Boussinesq) then - if (CS%use_Stanley_ML) 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) - endif - else - call calculate_spec_vol(tv%T(:,j,k), tv%S(:,j,k), p0, SpV_ml, tv%eqn_of_state, EOSdom) - endif - line_is_empty = .true. - do i=is-1,ie+1 + 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(:,:,kstart:kend), tv%S(:,:,kstart:kend), p_blk, rho_blk, & + tv%eqn_of_state, EOSdom3) + + do concurrent (j=js-1:je+1, 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_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] + Rml_int(i,j) = Rml_int(i,j) + dh*rho_blk(i,j,k-kstart+1) ! Rml_int is in [R H ~> kg m-2] htot(i,j) = htot(i,j) + dh - line_is_empty = .false. endif enddo - if (line_is_empty) keep_going=.false. + enddo + + if (nkblock < nz) then + ! Stop calling the equation of state once every column has been filled to "big H". With a + ! single block there is nothing left to skip, and this test would drag htot back to the host. + !$omp target update from(htot) + keep_going = .false. + do j=js-1,je+1 ; 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=js-1:je+1, 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,j) ) / (htot(i,j) + h_neglect) enddo - if (GV%Boussinesq .or. GV%semi_Boussinesq) then + !$omp end target data + 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, SpV_int) + do j=js-1,je+1 + rho_ml(:) = 0.0 ; SpV_ml(:) = 0.0 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) + htot(i,j) = 0.0 ; Rml_int(i,j) = 0.0 ; SpV_int(i) = 0.0 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) + 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_spec_vol(tv%T(:,j,k), tv%S(:,j,k), p0, SpV_ml, tv%eqn_of_state, EOSdom) + endif + 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,j) = Rml_int(i,j) + 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 enddo - 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,j) ) / (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 +1157,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] @@ -1080,26 +1180,28 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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] + 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] @@ -1121,29 +1223,42 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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] + 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 end target data if (CS%id_uhml > 0 .or. CS%id_vhml > 0) & ! Remapped uhml and vhml require east/north halo updates of h @@ -1735,6 +1850,10 @@ 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, "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, & From 1ddc9ff44af40cebb17dcb545ae395d25034a4f4 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 19:50:05 -0500 Subject: [PATCH 02/11] Map the PBL pointer dummies explicitly h_MLD, BLD and bflux were left to implicit copyin, which the mapped-component switch disables along with derived-type copyin, so the kernels read uninitialized device memory and little_h/wpup came back NaN. BLD and bflux are mapped around the w'u' branches that dereference them; h_MLD joins the outer region. Co-Authored-By: Claude Opus 5 (1M context) --- src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index cfaa8b798b..ea8100e0f0 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -907,10 +907,13 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! metrics, so every kernel in MOM6 is already in that regime. G, GV and US are mapped there and ! stay present, which is why CS and tv are the only parents that need mapping here. ! - h, uhtr and vhtr are mapped for the whole run by initialize_MOM and are used present. + ! - Nothing is left to implicit copyin, including the pointer dummies from the PBL scheme: the + ! same switch that disables derived-type copyin leaves them reading uninitialized device + ! memory, which shows up as NaN rather than a fault. ! Fields that only one branch or one diagnostic touches are mapped within that branch instead. !$omp target data & !$omp map(to: CS, tv) & - !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d) & + !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & !$omp map(tofrom: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) & !$omp map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) & !$omp map(alloc: Rml_int, vol_dt_avail, uhml, vhml) @@ -952,6 +955,8 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo ! Estimate w'u' at h-points, with a floor to avoid division by zero later. + ! BLD and bflux are dereferenced only by these branches, so they are mapped only around them. + !$omp target 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 @@ -988,6 +993,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d * US%Z_to_L * GV%Z_to_H ! In [L H T-2 ~> m2 s-2 or kg m-1 s-2] enddo endif + !$omp end target data ! We filter w'u' with the same time scales used for "little h" do concurrent (j=js-1:je+1, i=is-1:ie+1) From ff83fa5773be58fa85a16b81aeec726083ebba48 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 20:50:21 -0500 Subject: [PATCH 03/11] Map CS and tv in an enclosing region so their scalars reach the device Mapping a derived-type component makes nvfortran create a parent entry that supersedes an explicit map(to: CS) in the same directive. The component descriptors still attach -- CS%Cr_space read correctly -- but the struct body is never copied, so CS%mstar, CS%nstar and CS%min_wstar2 silently read as zero. wpup then collapsed to exactly zero, r_wpup went infinite, and psi_mag came out 0*Inf = NaN. Mapping the parents in their own enclosing region copies the struct body. benchmark_ALE now reproduces repro_ocean.stats bitwise. Co-Authored-By: Claude Opus 5 (1M context) --- .../lateral/MOM_mixed_layer_restrat.F90 | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index ea8100e0f0..d3e23b9579 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -901,18 +901,21 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d endif ! Everything from here to the diagnostics is evaluated on the device inside a single data region, - ! so no field crosses the bus between phases. Two rules shape the map clauses: - ! - A derived type is mapped alongside its components, parent first. nvfortran turns off implicit - ! derived-type copyin for any region that maps a component, and initialize_MOM maps the G% - ! metrics, so every kernel in MOM6 is already in that regime. G, GV and US are mapped there and - ! stay present, which is why CS and tv are the only parents that need mapping here. + ! so no field crosses the bus between phases. Three rules shape the map clauses: + ! - A derived type must be mapped for its components to be readable at all: nvfortran turns off + ! implicit derived-type copyin for any region that maps a component, and initialize_MOM maps the + ! G% metrics, so every kernel in MOM6 is already in that regime. G, GV and US are mapped there + ! and stay present, which is why CS and tv are the only parents to map here. + ! - The parent must be mapped in an OUTER region, not in the same directive as its components. + ! Mapping a component implicitly creates a parent entry that supersedes an explicit map(to: CS) + ! in the same directive: the component descriptors attach, but the struct body is never copied, + ! so CS%mstar and friends silently read as zero while CS%Cr_space reads correctly. ! - h, uhtr and vhtr are mapped for the whole run by initialize_MOM and are used present. - ! - Nothing is left to implicit copyin, including the pointer dummies from the PBL scheme: the - ! same switch that disables derived-type copyin leaves them reading uninitialized device - ! memory, which shows up as NaN rather than a fault. + ! Nothing is left to implicit copyin, including the pointer dummies from the PBL scheme, which + ! would otherwise read uninitialized device memory and show up as NaN rather than as a fault. ! Fields that only one branch or one diagnostic touches are mapped within that branch instead. + !$omp target data map(always, to: CS, tv) !$omp target data & - !$omp map(to: CS, tv) & !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & !$omp map(tofrom: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) & !$omp map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) & @@ -1264,6 +1267,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d !$omp target update from(vhml) endif + !$omp end target data !$omp end target data if (CS%id_uhml > 0 .or. CS%id_vhml > 0) & From f0e8c87a3fdd1d92801a19f5e5c5cf23de6f2f78 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 21:58:16 -0500 Subject: [PATCH 04/11] Tile the density integral in j as well as k With MLE_NKBLOCK=1 the equation-of-state scratch was still a full (SZI,SZJ) slab, where upstream used a single row rho_ml(SZI). That loses the locality the original had, where the EOS fills a row and the integral consumes it straight back out of cache. Add MLE_NJBLOCK alongside MLE_NKBLOCK, following PGF_PLM_NJBLOCK: 1 and 1 on the CPU reduce the scratch to the one row upstream used, 0 and 0 on the GPU make it the whole halo-1 j domain by the whole column, which is a single kernel. Rml_int goes back to being the 1D host array upstream declares, with the tiled path carrying its running totals in Rml_blk. The early-exit test compares htot against big_H on the host, so big_H is fetched once when it can actually run. Without that it read a stale big_H, concluded every column was full and stopped calling the equation of state early, which truncated the integral -- benchmark_ALE energy moved in the 7th digit at NJBLOCK=NKBLOCK=1 while the untiled default stayed bitwise. benchmark_ALE reproduces repro_ocean.stats bitwise at both 0/0 and 1/1. Co-Authored-By: Claude Opus 5 (1M context) --- .../lateral/MOM_mixed_layer_restrat.F90 | 139 +++++++++++------- 1 file changed, 85 insertions(+), 54 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index d3e23b9579..00c1629033 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -34,8 +34,10 @@ module MOM_mixed_layer_restrat #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 @@ -121,6 +123,9 @@ module MOM_mixed_layer_restrat 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] @@ -806,18 +811,24 @@ 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 :: Rml_int(SZI_(G),SZJ_(G)) ! Potential density integrated through the mixed layer - ! [R H ~> kg m-2 or kg2 m-5] + 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] - ! The density integral works through the column nkblock layers at a time, so that the equation - ! of state is handed enough points to be worth offloading while these scratch arrays stay small. - ! MLE_NKBLOCK=1 (the CPU default) makes them a single 2D slab; 0 puts the whole column in one block. - real :: rho_blk(SZI_(G),SZJ_(G),merge(GV%ke, CS%nkblock, CS%nkblock==0)) - ! Potential density relative to the surface for a block of layers [R ~> kg m-3] - real :: p_blk(SZI_(G),SZJ_(G),merge(GV%ke, CS%nkblock, CS%nkblock==0)) - ! A pressure of 0 for a block of layers [R L2 T-2 ~> Pa] + ! The density integral walks the mixed layer in tiles of njblock rows by nkblock layers, so that + ! the equation of state is handed enough points to be worth offloading without the scratch growing + ! past what a cache can hold. MLE_NJBLOCK=MLE_NKBLOCK=1 (the CPU defaults) reduce these to the one + ! row upstream uses, produced and consumed in the same tile; 0 and 0 (the GPU defaults) make them + ! the whole halo-1 j domain by the whole column, which is one kernel launch. + 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] @@ -851,8 +862,9 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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 ! The number of layers in each block of the density integral [nondim] - integer :: kstart, kend ! The first and last layer of the block being worked on [nondim] + 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 @@ -919,7 +931,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & !$omp map(tofrom: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) & !$omp map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) & - !$omp map(alloc: Rml_int, vol_dt_avail, uhml, vhml) + !$omp 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). @@ -1058,54 +1070,69 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo if ((GV%Boussinesq .or. GV%semi_Boussinesq) .and. .not.CS%use_Stanley_ML) then - ! Evaluate the density for a block of nkblock layers at a time and then integrate that block - ! through the mixed layer. 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. + ! 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) - EOSdom3(2,:) = [(js-1) - (G%jsd-1), (je+1) - (G%jsd-1)] - ! The block scratch never leaves the device, and no other branch uses it. - !$omp target data map(alloc: rho_blk, p_blk) + ! The tile scratch never leaves the device, and no other branch uses it. + !$omp target 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=js-1:je+1, i=is-1:ie+1) + do concurrent (k=1:nkblock, j=1:njblock, i=is-1:ie+1) p_blk(i,j,k) = 0.0 enddo - do concurrent (j=js-1:je+1, i=is-1:ie+1) - htot(i,j) = 0.0 ; Rml_int(i,j) = 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(:,:,kstart:kend), tv%S(:,:,kstart:kend), p_blk, rho_blk, & - tv%eqn_of_state, EOSdom3) - - do concurrent (j=js-1:je+1, 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_int(i,j) = Rml_int(i,j) + dh*rho_blk(i,j,k-kstart+1) ! Rml_int is in [R H ~> kg m-2] - htot(i,j) = htot(i,j) + dh - endif - enddo + 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 - if (nkblock < nz) then - ! Stop calling the equation of state once every column has been filled to "big H". With a - ! single block there is nothing left to skip, and this test would drag htot back to the host. - !$omp target update from(htot) - keep_going = .false. - do j=js-1,je+1 ; do i=is-1,ie+1 - if (htot(i,j) < big_H(i,j)) keep_going = .true. - enddo ; enddo - endif - endif ; 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 - do concurrent (j=js-1:je+1, 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,j) ) / (htot(i,j) + h_neglect) + 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 end target data @@ -1113,11 +1140,11 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! 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, SpV_int) + !$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 - htot(i,j) = 0.0 ; Rml_int(i,j) = 0.0 ; SpV_int(i) = 0.0 + htot(i,j) = 0.0 ; Rml_int(i) = 0.0 ; SpV_int(i) = 0.0 enddo keep_going = .true. do k=1,nz @@ -1132,7 +1159,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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,j) = Rml_int(i,j) + dh*rho_ml(i) ! Rml_int has units of [R H ~> kg m-2] + 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. @@ -1145,7 +1172,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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,j) ) / (htot(i,j) + h_neglect) + buoy_av(i,j) = -( g_Rho0 * Rml_int(i) ) / (htot(i,j) + h_neglect) enddo else do i=is-1,ie+1 @@ -1864,6 +1891,10 @@ logical function mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, CS, "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, & From 11e34d0b07adc66318216b275b7a0c5fac26bac4 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 22:20:05 -0500 Subject: [PATCH 05/11] Trim the map-clause commentary to the load-bearing line Co-Authored-By: Claude Opus 5 (1M context) --- .../lateral/MOM_mixed_layer_restrat.F90 | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 00c1629033..147ec94383 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -912,20 +912,8 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif - ! Everything from here to the diagnostics is evaluated on the device inside a single data region, - ! so no field crosses the bus between phases. Three rules shape the map clauses: - ! - A derived type must be mapped for its components to be readable at all: nvfortran turns off - ! implicit derived-type copyin for any region that maps a component, and initialize_MOM maps the - ! G% metrics, so every kernel in MOM6 is already in that regime. G, GV and US are mapped there - ! and stay present, which is why CS and tv are the only parents to map here. - ! - The parent must be mapped in an OUTER region, not in the same directive as its components. - ! Mapping a component implicitly creates a parent entry that supersedes an explicit map(to: CS) - ! in the same directive: the component descriptors attach, but the struct body is never copied, - ! so CS%mstar and friends silently read as zero while CS%Cr_space reads correctly. - ! - h, uhtr and vhtr are mapped for the whole run by initialize_MOM and are used present. - ! Nothing is left to implicit copyin, including the pointer dummies from the PBL scheme, which - ! would otherwise read uninitialized device memory and show up as NaN rather than as a fault. - ! Fields that only one branch or one diagnostic touches are mapped within that branch instead. + ! CS and tv are mapped in their own region: a parent mapped in the same directive as its + ! components is superseded by them, and its scalars then read as zero on the device. !$omp target data map(always, to: CS, tv) !$omp target data & !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & From 7aef4a7faebda5f025766f0a4ab68313d48e9b3e Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 00:04:33 -0500 Subject: [PATCH 06/11] Record that always is load-bearing on the CS/tv parent map Dropping it reproduces the zeroed-scalar NaN even with the parent in its own enclosing region. Co-Authored-By: Claude Opus 5 (1M context) --- src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 147ec94383..e86c6e3f29 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -912,8 +912,9 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif - ! CS and tv are mapped in their own region: a parent mapped in the same directive as its - ! components is superseded by them, and its scalars then read as zero on the device. + ! CS and tv are mapped in their own region, and "always" is load-bearing: a parent mapped in the + ! same directive as its components, or mapped without "always", is superseded by them and its + ! scalars then read as zero on the device. Dropping either gives NaN, not a fault. !$omp target data map(always, to: CS, tv) !$omp target data & !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & From 1487e7e15c32c11d2f22fdf5e6715ce39e33a48d Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 01:44:13 -0500 Subject: [PATCH 07/11] Map the CS arrays at init and hoist its scalars, dropping map(always) The previous form leaned on `map(always, to: CS, tv)` in an enclosing structured region, a construct used nowhere else in MOM6. Replace it with what the rest of the GPU port does: unstructured `target enter data`, the control structure's array components mapped once in the init routine as hor_visc_init and friends do, and no structured `target data` regions anywhere. CS's scalars cannot be mapped alongside its arrays. Mapping any component of a derived type costs that type its implicit copyin, after which its scalars read as zero on the device while its array components still read correctly -- silent, and it surfaces only as a NaN much further downstream. So the eight scalars the kernels need are hoisted to plain locals, which is the practice the rest of the port already follows. Leaving the arrays unmapped instead is not an option: that faults with an illegal address on the first CS%MLD_filtered_slow read. benchmark_ALE reproduces repro_ocean.stats bitwise at MLE_NJBLOCK/NKBLOCK of both 0/0 and 1/1. Co-Authored-By: Claude Opus 5 (1M context) --- src/core/MOM.F90 | 1 + .../lateral/MOM_mixed_layer_restrat.F90 | 63 ++++++++++++------- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/src/core/MOM.F90 b/src/core/MOM.F90 index e15ac57acf..6dee72f714 100644 --- a/src/core/MOM.F90 +++ b/src/core/MOM.F90 @@ -3778,6 +3778,7 @@ subroutine initialize_MOM(Time, Time_init, param_file, dirs, CS, & CS%mixedlayer_restrat = mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, & CS%mixedlayer_restrat_CSp, restart_CSp) + !$omp target enter data map(alloc: CS%mixedlayer_restrat_CSp) if (GV%Boussinesq .and. associated(CS%visc%h_ML)) then ! This is here to allow for a transition of restart files between model versions. diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index e86c6e3f29..6d19947e55 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -859,6 +859,13 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! 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] + ! Scalars of CS read by the kernels below. They are copied to plain locals because the kernels + ! also use CS's array components, and mapping those costs CS its implicit copyin on the device. + real :: tau_bgrow, tau_bdecay ! Filter timescales for the boundary layer depth [T ~> s] + real :: tau_mgrow, tau_mdecay ! Filter timescales for the mixed layer depth [T ~> s] + real :: l_mstar, l_nstar ! Coefficients of u*^3 and w*^3 in the momentum flux [nondim] + real :: l_min_wstar2 ! Floor on the vertical momentum flux [Z2 T-2 ~> m2 s-2] + real :: l_MLE_tail_dh ! Fractional depth by which the stream function is extended [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 @@ -912,21 +919,24 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif - ! CS and tv are mapped in their own region, and "always" is load-bearing: a parent mapped in the - ! same directive as its components, or mapped without "always", is superseded by them and its - ! scalars then read as zero on the device. Dropping either gives NaN, not a fault. - !$omp target data map(always, to: CS, tv) - !$omp target data & - !$omp map(to: CS%Cr_space, tv%T, tv%S, U_star_2d, h_MLD) & - !$omp map(tofrom: CS%MLD_filtered, CS%MLD_filtered_slow, CS%wpup_filtered) & - !$omp map(from: little_h, big_H, wpup, htot, buoy_av, uDml_diag, vDml_diag) & - !$omp map(alloc: vol_dt_avail, uhml, vhml) + ! The whole calculation below runs on the device. Only plain local arrays and the PBL pointer + ! dummies are mapped: nothing of CS or tv is, because mapping a derived-type component costs the + ! kernels implicit copyin of that whole type and its scalars then read as zero. Left alone, CS% + ! and tv% are copied in, and CS's filtered fields copied back, by the loops that use them. + !$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) + + tau_bgrow = CS%BLD_growing_Tfilt ; tau_bdecay = CS%BLD_decaying_Tfilt + tau_mgrow = CS%MLD_growing_Tfilt ; tau_mdecay = CS%MLD_decaying_Tfilt + l_mstar = CS%mstar ; l_nstar = CS%nstar ; l_min_wstar2 = CS%min_wstar2 + l_MLE_tail_dh = CS%MLE_tail_dh ! 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 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) + tau_bgrow, tau_bdecay, dt) CS%MLD_filtered(i,j) = little_h(i,j) enddo @@ -951,7 +961,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d else 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) + tau_mgrow, tau_mdecay, dt) enddo endif do concurrent (j=js-1:je+1, i=is-1:ie+1) @@ -960,7 +970,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! Estimate w'u' at h-points, with a floor to avoid division by zero later. ! BLD and bflux are dereferenced only by these branches, so they are mapped only around them. - !$omp target data map(to: BLD, bflux) + !$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 @@ -992,17 +1002,18 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! w_star3 = max(0., -bflux(i,j)) * BLD(i,j), in [Z3 T-3 ~> m3 s-3], is written out in place ! here rather than held in a scalar that would have to be made local to the concurrent loop. 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 ) & + wpup(i,j) = max( (cuberoot(l_mstar * U_star_2d(i,j)**3 & + + l_nstar * (max(0., -bflux(i,j)) * BLD(i,j))))**2, l_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 endif - !$omp end target data + !$omp target exit data map(release: BLD, bflux) + ! We filter w'u' with the same time scales used for "little h" 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) + tau_bgrow, tau_bdecay, dt) CS%wpup_filtered(i,j) = wpup(i,j) enddo @@ -1068,7 +1079,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d EOSdom3(1,:) = EOS_domain(G%HI, halo=1) ! The tile scratch never leaves the device, and no other branch uses it. - !$omp target data map(alloc: rho_blk, p_blk, Rml_blk) + !$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 @@ -1124,7 +1135,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo enddo - !$omp end target data + !$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. @@ -1204,7 +1215,7 @@ 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] + muzb = mu(sigint, l_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. @@ -1247,7 +1258,7 @@ 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] + muzb = mu(sigint, l_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] @@ -1283,8 +1294,8 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d !$omp target update from(vhml) endif - !$omp end target data - !$omp end target data + !$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 @@ -2104,6 +2115,14 @@ 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) + ! Map the array components the Bodner kernels use, once, as hor_visc_init and friends do. Their + ! scalars cannot come along: mapping any component of a derived type costs that type its implicit + ! copyin, and its scalars then read as zero on the device, so those are hoisted in the routine. + if (CS%use_Bodner) then + !$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 From 05fbd465b72acbe1a5e2cfff2a4a15b9a1c70abd Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 01:46:32 -0500 Subject: [PATCH 08/11] Drop leftover scaffolding from MOM.F90 Co-Authored-By: Claude Opus 5 (1M context) --- src/core/MOM.F90 | 1 - 1 file changed, 1 deletion(-) diff --git a/src/core/MOM.F90 b/src/core/MOM.F90 index 6dee72f714..e15ac57acf 100644 --- a/src/core/MOM.F90 +++ b/src/core/MOM.F90 @@ -3778,7 +3778,6 @@ subroutine initialize_MOM(Time, Time_init, param_file, dirs, CS, & CS%mixedlayer_restrat = mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, & CS%mixedlayer_restrat_CSp, restart_CSp) - !$omp target enter data map(alloc: CS%mixedlayer_restrat_CSp) if (GV%Boussinesq .and. associated(CS%visc%h_ML)) then ! This is here to allow for a transition of restart files between model versions. From ff5c22a5bb6785c2b72e2a6fe4b47cc8ff8e4a3e Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 05:02:51 -0500 Subject: [PATCH 09/11] Read CS scalars directly in the Bodner kernels via an init-time update The solo driver puts the whole MOM_control_struct on the device with `enter data map(alloc: MOM_CSp)`: the address range is claimed but its contents are never copied. Any later map(to:) of anything inside that range -- including the implicit copyin every kernel gets -- is then a present-table no-op, so kernels read zeros for the scalars of any control structure held as an inline member. Mapped array components were unaffected because their maps copy the heap data and attach the descriptor into the parent's device block; scalars have no attach equivalent. This was the mechanism behind the l_*/tau_* hoists. `target update` always copies regardless of presence, so issue `target update to(CS)` in mixedlayer_restrat_init once the parameters are set, BEFORE the component maps whose attach then fixes this struct's array descriptors in the device copy. With that, delete the eight hoisted locals and read CS% directly in the kernels, like every other module in the tree. Verified bitwise against repro_ocean.stats at both MLE_NJBLOCK/NKBLOCK 0/0 and 1/1 (overrides confirmed applied in MOM_parameter_doc.all). Co-Authored-By: Claude Opus 5 (1M context) --- .../lateral/MOM_mixed_layer_restrat.F90 | 36 ++++++++----------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 6d19947e55..94462b41e7 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -859,13 +859,6 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! 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] - ! Scalars of CS read by the kernels below. They are copied to plain locals because the kernels - ! also use CS's array components, and mapping those costs CS its implicit copyin on the device. - real :: tau_bgrow, tau_bdecay ! Filter timescales for the boundary layer depth [T ~> s] - real :: tau_mgrow, tau_mdecay ! Filter timescales for the mixed layer depth [T ~> s] - real :: l_mstar, l_nstar ! Coefficients of u*^3 and w*^3 in the momentum flux [nondim] - real :: l_min_wstar2 ! Floor on the vertical momentum flux [Z2 T-2 ~> m2 s-2] - real :: l_MLE_tail_dh ! Fractional depth by which the stream function is extended [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 @@ -927,16 +920,11 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d !$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) - tau_bgrow = CS%BLD_growing_Tfilt ; tau_bdecay = CS%BLD_decaying_Tfilt - tau_mgrow = CS%MLD_growing_Tfilt ; tau_mdecay = CS%MLD_decaying_Tfilt - l_mstar = CS%mstar ; l_nstar = CS%nstar ; l_min_wstar2 = CS%min_wstar2 - l_MLE_tail_dh = CS%MLE_tail_dh - ! 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 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), & - tau_bgrow, tau_bdecay, dt) + CS%BLD_growing_Tfilt, CS%BLD_decaying_Tfilt, dt) CS%MLD_filtered(i,j) = little_h(i,j) enddo @@ -961,7 +949,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d else 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), & - tau_mgrow, tau_mdecay, dt) + CS%MLD_growing_Tfilt, CS%MLD_decaying_Tfilt, dt) enddo endif do concurrent (j=js-1:je+1, i=is-1:ie+1) @@ -1002,8 +990,8 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! w_star3 = max(0., -bflux(i,j)) * BLD(i,j), in [Z3 T-3 ~> m3 s-3], is written out in place ! here rather than held in a scalar that would have to be made local to the concurrent loop. do concurrent (j=js-1:je+1, i=is-1:ie+1) - wpup(i,j) = max( (cuberoot(l_mstar * U_star_2d(i,j)**3 & - + l_nstar * (max(0., -bflux(i,j)) * BLD(i,j))))**2, l_min_wstar2 ) & + 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 endif @@ -1013,7 +1001,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! We filter w'u' with the same time scales used for "little h" do concurrent (j=js-1:je+1, i=is-1:ie+1) wpup(i,j) = rmean2ts(wpup(i,j), CS%wpup_filtered(i,j), & - tau_bgrow, tau_bdecay, dt) + CS%BLD_growing_Tfilt, CS%BLD_decaying_Tfilt, dt) CS%wpup_filtered(i,j) = wpup(i,j) enddo @@ -1215,7 +1203,7 @@ 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, l_MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] + 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. @@ -1258,7 +1246,7 @@ 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, l_MLE_tail_dh) ! mu(z/MLD) for lower interface [nondim] + 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] @@ -2115,10 +2103,14 @@ 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) - ! Map the array components the Bodner kernels use, once, as hor_visc_init and friends do. Their - ! scalars cannot come along: mapping any component of a derived type costs that type its implicit - ! copyin, and its scalars then read as zero on the device, so those are hoisted in the routine. + ! This CS is an inline member of MOM_control_struct, whose whole address range the solo + ! driver puts on the device with `enter data map(alloc: MOM_CSp)` -- allocated, never + ! copied. Any later map(to:) of this CS (implicit or explicit) is therefore a present- + ! table no-op and kernels read zeros for its scalars. `target update` always copies, so + ! issue it here, once the parameters above are set and BEFORE the component maps below, + ! whose attach then fixes this struct's array descriptors in the device copy. if (CS%use_Bodner) then + !$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 From 7a2a4849b1775d3e3193ae8340370eb15b58e1b4 Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 05:16:38 -0500 Subject: [PATCH 10/11] fix fix --- .../lateral/MOM_mixed_layer_restrat.F90 | 27 ++++--------------- 1 file changed, 5 insertions(+), 22 deletions(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 94462b41e7..32969c9272 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -815,11 +815,6 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d 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] - ! The density integral walks the mixed layer in tiles of njblock rows by nkblock layers, so that - ! the equation of state is handed enough points to be worth offloading without the scratch growing - ! past what a cache can hold. MLE_NJBLOCK=MLE_NKBLOCK=1 (the CPU defaults) reduce these to the one - ! row upstream uses, produced and consumed in the same tile; 0 and 0 (the GPU defaults) make them - ! the whole halo-1 j domain by the whole column, which is one kernel launch. 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] @@ -912,10 +907,6 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d G%HI, haloshift=1, unscale=GV%H_to_mks) endif - ! The whole calculation below runs on the device. Only plain local arrays and the PBL pointer - ! dummies are mapped: nothing of CS or tv is, because mapping a derived-type component costs the - ! kernels implicit copyin of that whole type and its scalars then read as zero. Left alone, CS% - ! and tv% are copied in, and CS's filtered fields copied back, by the loops that use them. !$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) @@ -929,8 +920,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo ! Calculate "big H", representative of the mixed layer depth, used in B22 formula (eq 27). - ! The two spatially-varying variants below are not offloaded, so they push their result to the - ! device for the kernels that follow. + ! 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 @@ -957,7 +947,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo ! Estimate w'u' at h-points, with a floor to avoid division by zero later. - ! BLD and bflux are dereferenced only by these branches, so they are mapped only around them. + ! 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 @@ -987,8 +977,6 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d enddo ; enddo !$omp target update to(wpup) else - ! w_star3 = max(0., -bflux(i,j)) * BLD(i,j), in [Z3 T-3 ~> m3 s-3], is written out in place - ! here rather than held in a scalar that would have to be made local to the concurrent loop. 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 ) & @@ -1182,7 +1170,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! U - Component 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)) + 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] @@ -1225,7 +1213,7 @@ subroutine mixedlayer_restrat_Bodner(CS, G, GV, US, h, uhtr, vhtr, tv, forces, d ! V- component 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)) + 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] @@ -2103,13 +2091,8 @@ 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) - ! This CS is an inline member of MOM_control_struct, whose whole address range the solo - ! driver puts on the device with `enter data map(alloc: MOM_CSp)` -- allocated, never - ! copied. Any later map(to:) of this CS (implicit or explicit) is therefore a present- - ! table no-op and kernels read zeros for its scalars. `target update` always copies, so - ! issue it here, once the parameters above are set and BEFORE the component maps below, - ! whose attach then fixes this struct's array descriptors in the device copy. 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) From c1f26dd11a15504ae33faa91c36d81c36992f2d2 Mon Sep 17 00:00:00 2001 From: Jorge Date: Tue, 28 Jul 2026 05:48:04 -0500 Subject: [PATCH 11/11] fix doxygne --- src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 index 32969c9272..cc45403599 100644 --- a/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 +++ b/src/parameterizations/lateral/MOM_mixed_layer_restrat.F90 @@ -2092,7 +2092,7 @@ logical function mixedlayer_restrat_init(Time, G, GV, US, param_file, diag, CS, if (allocated(CS%wpup_filtered)) call pass_var(CS%wpup_filtered, G%domain) if (CS%use_Bodner) then - ! very important! + ! 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)