From a039d7f2b804d41ea8e4ab87d83f10a40b97579c Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Wed, 15 Jul 2026 09:51:29 +1000 Subject: [PATCH 01/15] MOM_cpu_clock_infra: emit NVTX ranges for MOM6 cpu clocks Record each clock's name as cpu_clock_id() registers it, then open an NVTX range in cpu_clock_begin() and close it in cpu_clock_end(). Every existing cpu_clock_id() name becomes a named range in an nsys timeline with no call-site changes. The range is opened before, and closed after, the mpp clock so that it encloses it. Begin and end test the same condition, so starts and ends stay balanced for handles that were never named or that fall outside the table. The code is guarded behind -DMOM_USE_NVTX, which is undefined by default: without it the preprocessed source is byte-identical to before and no extra library is linked. Profiling builds need -DMOM_USE_NVTX and -cudalib=nvtx. Applied to both the FMS1 and FMS2 infra shims, which are byte-identical to each other. This also covers the halo and message-passing clocks, which MOM_domain_infra begins and ends through this module directly rather than through the MOM_cpu_clock framework wrapper. Based on Edward Yang's ae67665d3 ("add nvtx markers to clocks", branch benchmark_ALE_nvtx_clocks), which introduced the name-table approach and the range nesting; this adds the build guard so non-NVHPC builds are unaffected. Co-authored-by: Edward Yang Co-Authored-By: Claude Opus 4.8 (1M context) --- config_src/infra/FMS1/MOM_cpu_clock_infra.F90 | 34 +++++++++++++++++++ config_src/infra/FMS2/MOM_cpu_clock_infra.F90 | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 b/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 index aeca65b863..f76c38fe10 100644 --- a/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 +++ b/config_src/infra/FMS1/MOM_cpu_clock_infra.F90 @@ -5,6 +5,12 @@ !> Wraps the MPP cpu clock functions !! !! The functions and constants should be accessed via mom_cpu_clock +!! +!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu +!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline +!! with no call-site changes. It requires nvfortran and the NVTX library +!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles +!! exactly as before and links no extra library. module MOM_cpu_clock_infra ! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA +#ifdef MOM_USE_NVTX +use nvtx, only : nvtxStartRange, nvtxEndRange +#endif implicit none ; private +#ifdef MOM_USE_NVTX +!> The largest clock handle for which an NVTX range name is retained. +integer, parameter :: MAX_NVTX_CLOCKS = 4096 +!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry +!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test +!! the same condition, so starts and ends stay balanced for handles that were never named +!! or that fall outside the table. +character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = "" +#endif + ! Public entities public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra subroutine cpu_clock_begin(id) integer, intent(in) :: id !< Handle for clock +#ifdef MOM_USE_NVTX + ! Opened before, and closed after, the mpp clock so the NVTX range encloses it. + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id))) + endif +#endif call mpp_clock_begin(id) end subroutine cpu_clock_begin @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id) integer, intent(in) :: id !< Handle for clock call mpp_clock_end(id) +#ifdef MOM_USE_NVTX + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange + endif +#endif end subroutine cpu_clock_end @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain) endif cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain) +#ifdef MOM_USE_NVTX + if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) & + nvtx_clock_names(cpu_clock_id) = name +#endif end function cpu_clock_id end module MOM_cpu_clock_infra diff --git a/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 b/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 index aeca65b863..f76c38fe10 100644 --- a/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 +++ b/config_src/infra/FMS2/MOM_cpu_clock_infra.F90 @@ -5,6 +5,12 @@ !> Wraps the MPP cpu clock functions !! !! The functions and constants should be accessed via mom_cpu_clock +!! +!! Compiling with -DMOM_USE_NVTX additionally emits an NVTX range around every MOM6 cpu +!! clock, so each existing cpu_clock_id() name becomes a named range in an nsys timeline +!! with no call-site changes. It requires nvfortran and the NVTX library +!! (-DMOM_USE_NVTX ... -cudalib=nvtx). Undefined by default: a normal build compiles +!! exactly as before and links no extra library. module MOM_cpu_clock_infra ! These interfaces and constants from MPP/FMS will not be directly exposed outside of this module @@ -18,9 +24,22 @@ module MOM_cpu_clock_infra use mpp_mod, only : MPP_CLOCK_ROUTINE => CLOCK_ROUTINE use mpp_mod, only : MPP_CLOCK_LOOP => CLOCK_LOOP use mpp_mod, only : MPP_CLOCK_INFRA => CLOCK_INFRA +#ifdef MOM_USE_NVTX +use nvtx, only : nvtxStartRange, nvtxEndRange +#endif implicit none ; private +#ifdef MOM_USE_NVTX +!> The largest clock handle for which an NVTX range name is retained. +integer, parameter :: MAX_NVTX_CLOCKS = 4096 +!> The NVTX range name for each clock handle, recorded by cpu_clock_id(). An empty entry +!! means no range is emitted for that handle. cpu_clock_begin() and cpu_clock_end() test +!! the same condition, so starts and ends stay balanced for handles that were never named +!! or that fall outside the table. +character(len=64), dimension(MAX_NVTX_CLOCKS) :: nvtx_clock_names = "" +#endif + ! Public entities public :: cpu_clock_id, cpu_clock_begin, cpu_clock_end public :: CLOCK_COMPONENT, CLOCK_SUBCOMPONENT, CLOCK_MODULE_DRIVER, CLOCK_MODULE @@ -60,6 +79,12 @@ module MOM_cpu_clock_infra subroutine cpu_clock_begin(id) integer, intent(in) :: id !< Handle for clock +#ifdef MOM_USE_NVTX + ! Opened before, and closed after, the mpp clock so the NVTX range encloses it. + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxStartRange(trim(nvtx_clock_names(id))) + endif +#endif call mpp_clock_begin(id) end subroutine cpu_clock_begin @@ -69,6 +94,11 @@ subroutine cpu_clock_end(id) integer, intent(in) :: id !< Handle for clock call mpp_clock_end(id) +#ifdef MOM_USE_NVTX + if (id > 0 .and. id <= MAX_NVTX_CLOCKS) then + if (len_trim(nvtx_clock_names(id)) > 0) call nvtxEndRange + endif +#endif end subroutine cpu_clock_end @@ -96,6 +126,10 @@ integer function cpu_clock_id(name, sync, grain) endif cpu_clock_id = mpp_clock_id(name, flags=clock_flags, grain=grain) +#ifdef MOM_USE_NVTX + if (cpu_clock_id > 0 .and. cpu_clock_id <= MAX_NVTX_CLOCKS) & + nvtx_clock_names(cpu_clock_id) = name +#endif end function cpu_clock_id end module MOM_cpu_clock_infra From c145f07d4bb829f3b185f174f05448de435fcf7a Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 12:56:49 +1000 Subject: [PATCH 02/15] kappa_shear: offload h_at_u/h_at_v vertex interpolation to GPU First increment of the MOM_kappa_shear GPU port. Converts the two thickness-to-vertex interpolation loops in Calc_kappa_shear_vertex (both branches of vertex_shear_OBC_bug) from !$OMP parallel do to do concurrent, offloading them to the device. Residency: - h is host-authoritative in the (host-only) diabatic stack and is mutated on host during diabatic (applyBoundaryFluxesInOut), so the globally-resident device copy (CS%h) is refreshed with update to(h) before the kernels; a plain map(to:) does not refresh a present object. - G%mask2dCu/Cv/T are already device-resident (initialize_MOM). - h_at_u/h_at_v are device workspace (map alloc), copied back with update from for the still-host per-column solver below. enter/exit data are mirrored and refcount-neutral (release, not delete). The per-column solver (kappa_shear_column/find_kappa_tke) stays on host for now; later increments move it to the device, at which point the update from(h_at_u,h_at_v) round-trip is removed. Verification: - benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). - nsys --stats=true confirms both do concurrent loops run on the GPU (calc_kappa_shear_vertex_{550,553}_gpu, 12 launches each). - Reviewed; no code defect found. The >=2-rank bitwise run remains the mandatory pre-merge gate and is still pending (single-GPU node). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 1cff115fe5..3f4a301800 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -538,33 +538,37 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Convert layer thicknesses into geometric thickness in height units. call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=1) + ! --- GPU port increment 1: h_at_u/h_at_v interpolation offloaded to device. + ! h is host-authoritative in the (host-only) diabatic stack; refresh the device copy. + ! G%mask2dCu/Cv/T are already device-resident (mapped in initialize_MOM). h_at_u/h_at_v + ! are device workspace, copied back for the (still host) per-column loop below. + !$omp target enter data map(to: h) + !$omp target update to(h) + !$omp target enter data map(alloc: h_at_u, h_at_v) + if (CS%vertex_shear_OBC_bug) then - !$OMP parallel do default(shared) - do k=1,nz - do j=JsB,JeB+1 ; do I=IsB,IeB - h_at_u(I,j,k) = G%mask2dCu(I,j) * (h(i,j,k) + h(i+1,j,k)) * 0.5 - enddo ; enddo - do J=JsB,JeB ; do i=IsB,IeB+1 - h_at_v(i,J,k) = G%mask2dCv(i,J) * (h(i,j,k) + h(i,j+1,k)) * 0.5 - enddo ; enddo + do concurrent (k=1:nz, j=JsB:JeB+1, I=IsB:IeB) + h_at_u(I,j,k) = G%mask2dCu(I,j) * (h(i,j,k) + h(i+1,j,k)) * 0.5 + enddo + do concurrent (k=1:nz, J=JsB:JeB, i=IsB:IeB+1) + h_at_v(i,J,k) = G%mask2dCv(i,J) * (h(i,j,k) + h(i,j+1,k)) * 0.5 enddo else ! Because G%mask2dCu(I,j) is zero if either G%mask2dT(i,j) or G%mask2dT(i+1,j) except at OBC ! faces, the following form give equivalent answers to those above unless OBCs are in use, ! although the former is clearly less complicated and costly. - !$OMP parallel do default(shared) - do k=1,nz - do j=JsB,JeB+1 ; do I=IsB,IeB - h_at_u(I,j,k) = G%mask2dCu(I,j) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j) * h(i+1,j,k)) / & - (G%mask2dT(i,j) + G%mask2dT(i+1,j) + 1.0e-36) - enddo ; enddo - do J=JsB,JeB ; do i=IsB,IeB+1 - h_at_v(i,J,k) = G%mask2dCv(i,J) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) / & - (G%mask2dT(i,j) + G%mask2dT(i,j+1) + 1.0e-36) - enddo ; enddo + do concurrent (k=1:nz, j=JsB:JeB+1, I=IsB:IeB) + h_at_u(I,j,k) = G%mask2dCu(I,j) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j) * h(i+1,j,k)) / & + (G%mask2dT(i,j) + G%mask2dT(i+1,j) + 1.0e-36) + enddo + do concurrent (k=1:nz, J=JsB:JeB, i=IsB:IeB+1) + h_at_v(i,J,k) = G%mask2dCv(i,J) * (G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) / & + (G%mask2dT(i,j) + G%mask2dT(i,j+1) + 1.0e-36) enddo endif + !$omp target update from(h_at_u, h_at_v) + !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & @@ -860,6 +864,10 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) + ! --- GPU port increment 1: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: h_at_u, h_at_v) + !$omp target exit data map(release: h) + end subroutine Calc_kappa_shear_vertex From fbcfd11a5b8dc643308404922a9e885d1444a5b3 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 13:30:54 +1000 Subject: [PATCH 03/15] MOM_EOS: add device-callable density-derivs dispatcher (declare target) Groundwork for whole-column GPU kernels (kappa_shear next) that must compute density derivatives from INSIDE a device region, where the polymorphic EOS%type%... dispatch is not usable. - Expose the existing this-free _loc elementals calculate_density_derivs_elem_{Roquet_rho,buggy_Wright}_loc as public and !$omp declare target. - Add calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_dT, drho_dS): a declare-target subroutine that select-case's on the integer form and calls the matching _loc kernel, returning mks derivatives. Unit rescaling and any `scale` factor remain the caller's responsibility (as in calculate_density_derivs_1d). Forms without a device _loc kernel hit the default branch; a device-using module must FATAL at init on a GPU build before reaching it. This is the reusable, form-general enabler chosen over a Roquet-scoped inline; it is aligned with doc 06 6.4's select-case direction but adds only the device-callable point path (not the full polymorphism removal). No functional change: nothing calls the dispatcher yet. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). Dispatcher is exercised (and Fable-reviewed) in the follow-up that wires kappa_shear to it. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/equation_of_state/MOM_EOS.F90 | 30 ++++++++++++++++++++ src/equation_of_state/MOM_EOS_Roquet_rho.F90 | 4 +++ src/equation_of_state/MOM_EOS_Wright.F90 | 4 +++ 3 files changed, 38 insertions(+) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 724a4ca28b..75d1cee684 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -9,6 +9,7 @@ module MOM_EOS use MOM_EOS_linear, only : linear_EOS, avg_spec_vol_linear use MOM_EOS_linear, only : int_density_dz_linear, int_spec_vol_dp_linear use MOM_EOS_Wright, only : buggy_Wright_EOS, avg_spec_vol_buggy_Wright +use MOM_EOS_Wright, only : calculate_density_derivs_elem_buggy_Wright_loc use MOM_EOS_Wright, only : int_density_dz_wright, int_spec_vol_dp_wright use MOM_EOS_Wright_full, only : Wright_full_EOS, avg_spec_vol_Wright_full use MOM_EOS_Wright_full, only : int_density_dz_wright_full, int_spec_vol_dp_wright_full @@ -17,6 +18,7 @@ module MOM_EOS use MOM_EOS_Jackett06, only : Jackett06_EOS use MOM_EOS_UNESCO, only : UNESCO_EOS use MOM_EOS_Roquet_rho, only : Roquet_rho_EOS +use MOM_EOS_Roquet_rho, only : calculate_density_derivs_elem_Roquet_rho_loc use MOM_EOS_Roquet_SpV, only : Roquet_SpV_EOS use MOM_EOS_TEOS10, only : TEOS10_EOS use MOM_EOS_TEOS10, only : gsw_sp_from_sr, gsw_pt_from_ct, gsw_sr_from_sp, gsw_ct_from_pt @@ -46,6 +48,7 @@ module MOM_EOS public calculate_density_elem public calculate_density public calculate_density_derivs +public calculate_density_derivs_elem_loc public calculate_density_second_derivs public calculate_spec_vol public calculate_specific_vol_derivs @@ -992,6 +995,33 @@ subroutine calculate_density_derivs_1d(T, S, pressure, drho_dT, drho_dS, EOS, do end subroutine calculate_density_derivs_1d +!> Device-callable dispatcher for density derivatives at a single point, in mks units, +!! selecting the equation-of-state form at runtime by integer id (no polymorphic dispatch) +!! so it can be called from inside a do concurrent / target region by whole-column GPU +!! kernels. Unit rescaling (EOS%*_to_* factors) and any `scale` factor are the caller's +!! responsibility, exactly as in calculate_density_derivs_1d. Forms without a device-callable +!! _loc kernel are not handled here; a device-using module must FATAL at init on a GPU build +!! before reaching this with an unsupported form. +subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_dT, drho_dS) + integer, intent(in) :: form_of_EOS !< The equation of state form (EOS_ROQUET_RHO, EOS_WRIGHT, ...) + real, intent(in) :: T !< Temperature in the EOS kernel's mks units [degC] + real, intent(in) :: S !< Salinity in the EOS kernel's mks units [ppt or g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + real, intent(out) :: drho_dT !< Partial derivative of density wrt temperature [kg m-3 degC-1] + real, intent(out) :: drho_dS !< Partial derivative of density wrt salinity [kg m-3 ppt-1] + !$omp declare target + + select case (form_of_EOS) + case (EOS_ROQUET_RHO) + call calculate_density_derivs_elem_Roquet_rho_loc(T, S, pressure, drho_dT, drho_dS) + case (EOS_WRIGHT) + call calculate_density_derivs_elem_buggy_Wright_loc(T, S, pressure, drho_dT, drho_dS) + case default + drho_dT = 0.0 ; drho_dS = 0.0 + end select + +end subroutine calculate_density_derivs_elem_loc + !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs. subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom) diff --git a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 index 73ea02a400..98ddb0723d 100644 --- a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 +++ b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 @@ -10,6 +10,10 @@ module MOM_EOS_Roquet_rho implicit none ; private public Roquet_rho_EOS +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels +! can compute density derivatives in-region without polymorphic dispatch. +public calculate_density_derivs_elem_Roquet_rho_loc +!$omp declare target(calculate_density_derivs_elem_Roquet_rho_loc) real, parameter :: Pa2kb = 1.e-8 !< Conversion factor between Pa and kbar [kbar Pa-1] !>@{ Parameters in the Roquet_rho (Roquet density) equation of state diff --git a/src/equation_of_state/MOM_EOS_Wright.F90 b/src/equation_of_state/MOM_EOS_Wright.F90 index 748ca12e69..22632ebdbc 100644 --- a/src/equation_of_state/MOM_EOS_Wright.F90 +++ b/src/equation_of_state/MOM_EOS_Wright.F90 @@ -12,6 +12,10 @@ module MOM_EOS_Wright implicit none ; private public buggy_Wright_EOS +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels +! can compute density derivatives in-region without polymorphic dispatch. +public calculate_density_derivs_elem_buggy_Wright_loc +!$omp declare target(calculate_density_derivs_elem_buggy_Wright_loc) public int_density_dz_wright, int_spec_vol_dp_wright public avg_spec_vol_buggy_Wright public set_params_buggy_Wright From e15c35bf60416ea6e36b8b6308eb3c49cdc8ac06 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:06:04 +1000 Subject: [PATCH 04/15] kappa_shear: compute density derivs via device-callable EOS dispatcher on GPU builds Wires kappa_shear_column's Boussinesq density-derivative computation to the device-callable path added in 47ac931cb, so the column solver can later run inside a target region (the polymorphic calculate_density_derivs interface is not device-callable). - MOM_EOS: add get_EOS_form_and_scaling(EOS, form, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) so callers can reproduce, host-side, the unit rescaling that calculate_density_derivs_1d applies (EOS_type components are private). - kappa_shear_column: under #ifdef __NVCOMPILER_OPENMP_GPU, replace the calculate_density_derivs array call with an inline K=2..nzc loop calling calculate_density_derivs_elem_loc + the exact scale math of calculate_density_derivs_1d(dom=(/2,nzc/), scale=-g_R0). Non-GPU builds keep the original polymorphic call verbatim for all EOS forms. On GPU builds only ROQUET_RHO and WRIGHT have device _loc kernels; others FATAL. Verification: benchmark_ALE (ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1) with the GPU path active on host. Reviewed; bitwise equivalence confirmed operation-by-operation, verdict ship. Follow-ups required by the offload step (increment 3), per review: hoist get_EOS_form_and_scaling + the form FATAL out of kappa_shear_column to the driver (both are host-only and cannot run in a device region), passing the form/scaling in; and FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). Untested-by-benchmark but code-identical: the non-unity unit-rescale branch and the WRIGHT dispatcher case. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/equation_of_state/MOM_EOS.F90 | 21 ++++++++++ .../vertical/MOM_kappa_shear.F90 | 40 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 75d1cee684..427e81592c 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -49,6 +49,7 @@ module MOM_EOS public calculate_density public calculate_density_derivs public calculate_density_derivs_elem_loc +public get_EOS_form_and_scaling public calculate_density_second_derivs public calculate_spec_vol public calculate_specific_vol_derivs @@ -1022,6 +1023,26 @@ subroutine calculate_density_derivs_elem_loc(form_of_EOS, T, S, pressure, drho_d end subroutine calculate_density_derivs_elem_loc +!> Return the equation-of-state form id and the unit-rescaling factors held in an EOS_type. +!! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion +!! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without +!! needing access to the private components of EOS_type. +subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) + type(EOS_type), intent(in) :: EOS !< Equation of state structure + integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) + real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal density unit R [R m3 kg-1 ~> 1] + real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] + real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] + real, intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1] + + form_of_EOS = EOS%form_of_EOS + kg_m3_to_R = EOS%kg_m3_to_R + C_to_degC = EOS%C_to_degC + S_to_ppt = EOS%S_to_ppt + RL2_T2_to_Pa = EOS%RL2_T2_to_Pa + +end subroutine get_EOS_form_and_scaling + !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs. subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 3f4a301800..9c0e20fbda 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -19,6 +19,8 @@ module MOM_kappa_shear use MOM_verticalGrid, only : verticalGrid_type use MOM_EOS, only : calculate_density_derivs use MOM_EOS, only : calculate_density, calculate_specific_vol_derivs +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT implicit none ; private @@ -978,6 +980,12 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real :: gR0 ! A conversion factor from H to pressure, Rho_0 times g in Boussinesq ! mode, or just g when non-Boussinesq [R L2 T-2 H-1 ~> kg m-2 s-2 or m s-2]. real :: g_R0 ! g_R0 is a rescaled version of g/Rho [Z R-1 T-2 ~> m4 kg-1 s-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + ! Locals for the device-callable EOS derivs path (reproduces calculate_density_derivs_1d). + integer :: eos_form ! The equation-of-state form id. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. + real :: eos_rho_scale, dRdT_scale, dRdS_scale ! Output rescaling factors [various]. +#endif real :: Norm ! A factor that normalizes two weights to 1 [H-2 ~> m-2 or m4 kg-2]. real :: tol_dksrc ! Tolerance for the change in the kappa source within an iteration ! relative to the local source [nondim]. This must be greater than 1. @@ -1141,8 +1149,40 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la Sal_int(K) = 0.5*(Sal(k-1) + Sal(k)) enddo if (GV%Boussinesq .or. GV%semi_Boussinesq) then +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable EOS path: dispatch density derivatives by form id (host-resolved + ! scaling), reproducing calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) + ! bit-for-bit. Needed because the polymorphic calculate_density_derivs interface is + ! not callable from inside the (soon-to-be-offloaded) column device region. + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if ((eos_RL2_T2_to_Pa == 1.0) .and. (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0)) then + do K=2,nzc + call calculate_density_derivs_elem_loc(eos_form, T_int(K), Sal_int(K), pressure(K), & + dbuoy_dT(K), dbuoy_dS(K)) + enddo + else + do K=2,nzc + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_int(K), & + eos_S_to_ppt*Sal_int(K), eos_RL2_T2_to_Pa*pressure(K), dbuoy_dT(K), dbuoy_dS(K)) + enddo + endif + eos_rho_scale = eos_kg_m3_to_R * (-g_R0) + dRdT_scale = eos_rho_scale * eos_C_to_degC + dRdS_scale = eos_rho_scale * eos_S_to_ppt + if ((dRdT_scale /= 1.0) .or. (dRdS_scale /= 1.0)) then + do K=2,nzc + dbuoy_dT(K) = dRdT_scale * dbuoy_dT(K) + dbuoy_dS(K) = dRdS_scale * dbuoy_dS(K) + enddo + endif +#else call calculate_density_derivs(T_int, Sal_int, pressure, dbuoy_dT, dbuoy_dS, & tv%eqn_of_state, (/2,nzc/), scale=-g_R0 ) +#endif else ! These should perhaps be combined into a single call to calculate the thermal expansion ! and haline contraction coefficients? From af578cbe97e9d8337214c045fdf8f6c9b283eed2 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:34:12 +1000 Subject: [PATCH 05/15] kappa_shear: hoist EOS form/scaling resolution to the drivers (offload prep) Prerequisite for offloading the column solver (increment 3): get_EOS_form_and_scaling and MOM_error are host-only and cannot run inside a device region, so move them out of kappa_shear_column up into both drivers (Calculate_kappa_shear and Calc_kappa_shear_vertex). - Resolve eos_form + the four unit-scaling factors once per driver call (guarded by use_temperature), and thread them into kappa_shear_column as new intent(in) args. - Move the "unsupported EOS form on GPU" FATAL to the drivers, and add a FATAL on non-Boussinesq on GPU builds (that branch still uses polymorphic EOS calls). - Add the five hoisted scalars to the shared() clause of both drivers' !$OMP parallel do (they are loop-invariant and read-only in the region; default(private) would leave them uninitialized inside the loop -- caught by the bitwise check). No functional change: same values, computed once per call instead of once per column. Verification: benchmark_ALE (vertex driver, ROQUET_RHO, Boussinesq) ocean.stats bit-identical to the dev/gpu reference (np=1). Reviewed; verdict ship, bitwise safety confirmed by inspection for both drivers. Not runtime-exercised by the benchmark (inspection-only): the h-point driver, the non-unity unit-scaling branch. Notes for the offload step: the driver FATAL guards execution not compilation, so declare target on kappa_shear_column will require #ifdef-excluding the non-Boussinesq branch; CS (pointer dummy) and tv (pointer-laden) must be demoted/pre-extracted before declare target. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 70 +++++++++++++++---- 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 9c0e20fbda..914a21750e 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -214,11 +214,29 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for ! interpolating back to the original index space [nondim]. integer :: is, ie, js, je, i, j, k, nz, nzc + integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke use_temperature = associated(tv%T) + ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor + ! and MOM_error are not device-callable) so they can be passed into the column solver. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_temperature) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if (.not. (GV%Boussinesq .or. GV%semi_Boussinesq)) call MOM_error(FATAL, & + "kappa_shear GPU build: the non-Boussinesq density-derivs path is not device-callable.") +#endif + endif + k0dt = dt*CS%kappa_0 dz_massless = 0.1*sqrt((US%Z_to_m*GV%m_to_H)*k0dt) @@ -229,6 +247,7 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,tv,G,GV,US, & !$OMP CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do j=js,je @@ -329,7 +348,8 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, & - tv, CS, GV, US) + tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. @@ -520,6 +540,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ real :: H_tiny ! A sub-roundoff thickness to use in the denominator when calculating ! thickness-weighted averages [H ~> m or kg m-2] integer :: IsB, IeB, JsB, JeB, i, j, k, nz, nzc + integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. ! Diagnostics that should be deleted? isB = G%isc-1 ; ieB = G%iecB ; jsB = G%jsc-1 ; jeB = G%jecB ; nz = GV%ke @@ -532,6 +554,23 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ use_temperature = associated(tv%T) + ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor + ! and MOM_error are not device-callable) so they can be passed into the column solver, which + ! will run inside a device region. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_temperature) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & + eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + if (.not. (GV%Boussinesq .or. GV%semi_Boussinesq)) call MOM_error(FATAL, & + "kappa_shear GPU build: the non-Boussinesq density-derivs path is not device-callable.") +#endif + endif + k0dt = dt*CS%kappa_0 dz_massless = 0.1*sqrt((US%Z_to_m*GV%m_to_H)*k0dt) I_Prandtl = 0.0 ; if (CS%Prandtl_turb > 0.0) I_Prandtl = 1.0 / CS%Prandtl_turb @@ -575,6 +614,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do J=JsB,JeB @@ -708,7 +748,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & - tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US) + tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) ! call cpu_clock_begin(Id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then @@ -876,7 +917,8 @@ end subroutine Calc_kappa_shear_vertex !> This subroutine calculates shear-driven diffusivity and TKE in a single column subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_lay, & u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, tke_avg, N2_init, S2_init, & - N2_mean, S2_mean, tv, CS, GV, US ) + N2_mean, S2_mean, tv, CS, GV, US, & + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ) type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. real, dimension(SZK_(GV)+1), & intent(inout) :: kappa !< The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -918,6 +960,12 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la type(Kappa_shear_CS), pointer :: CS !< The control structure returned by a previous !! call to kappa_shear_init. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + integer, intent(in) :: eos_form !< The equation-of-state form id, resolved + !! host-side for the device-callable EOS derivs path. + real, intent(in) :: eos_kg_m3_to_R !< EOS factor converting kg m-3 to R [R m3 kg-1 ~> 1] + real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] + real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] + real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure to Pa [Pa T2 R-1 L-2 ~> 1] ! Local variables real, dimension(nzc) :: & @@ -981,9 +1029,7 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la ! mode, or just g when non-Boussinesq [R L2 T-2 H-1 ~> kg m-2 s-2 or m s-2]. real :: g_R0 ! g_R0 is a rescaled version of g/Rho [Z R-1 T-2 ~> m4 kg-1 s-2]. #ifdef __NVCOMPILER_OPENMP_GPU - ! Locals for the device-callable EOS derivs path (reproduces calculate_density_derivs_1d). - integer :: eos_form ! The equation-of-state form id. - real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. + ! Locals for the device-callable EOS derivs path (form + scaling are passed in as arguments). real :: eos_rho_scale, dRdT_scale, dRdS_scale ! Output rescaling factors [various]. #endif real :: Norm ! A factor that normalizes two weights to 1 [H-2 ~> m-2 or m4 kg-2]. @@ -1150,15 +1196,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la enddo if (GV%Boussinesq .or. GV%semi_Boussinesq) then #ifdef __NVCOMPILER_OPENMP_GPU - ! Device-callable EOS path: dispatch density derivatives by form id (host-resolved - ! scaling), reproducing calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) - ! bit-for-bit. Needed because the polymorphic calculate_density_derivs interface is - ! not callable from inside the (soon-to-be-offloaded) column device region. - call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, & - eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) - if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & - "kappa_shear GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & - "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") + ! Device-callable EOS path: dispatch density derivatives by form id (form + unit scaling + ! resolved host-side in the driver and passed in), reproducing + ! calculate_density_derivs_1d(..., dom=(/2,nzc/), scale=-g_R0) bit-for-bit. if ((eos_RL2_T2_to_Pa == 1.0) .and. (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0)) then do K=2,nzc call calculate_density_derivs_elem_loc(eos_form, T_int(K), Sal_int(K), pressure(K), & From 5660f98076e75762deb2e2f3e7d89a1e7255e7d1 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 14:55:47 +1000 Subject: [PATCH 06/15] kappa_shear: make the column solver device-callable (declare target) Prepares the per-column solver for offload: the three workers (kappa_shear_column, find_kappa_tke, calculate_projected_state) are now !$omp declare target, so the driver column loop can run inside a target region in the next step. Device code is generated for all three (confirmed by -Minfo). Execution is still on the host (the column loop is still !$OMP parallel do), so this is bitwise-neutral. To make them device-compilable: - Demote CS from `pointer` to `type(Kappa_shear_CS), intent(in)` in kappa_shear_column and find_kappa_tke (the type is all scalars bar a diag pointer that the column math never touches). - Hoist use_temperature to a kappa_shear_column argument (resolved in the drivers) so the device path no longer reads the tv%T pointer via associated(); tv is then untouched on GPU builds. - #ifndef __NVCOMPILER_OPENMP_GPU around the non-Boussinesq branch's polymorphic calculate_specific_vol_derivs/calculate_density calls so they are excluded from device compilation (unreachable on GPU anyway -- the driver FATALs on non-Boussinesq). Notably the 26 early exit/return sites in the two solvers compile for the device without error on NVHPC 26.5 (the 25.11 exit-under-parallel miscompile appears fixed); their runtime correctness on device is gated by the bitwise check in the offload step, not pre-emptive rewrites. Verification: benchmark_ALE ocean.stats bit-identical to the dev/gpu reference (np=1). Next step flips the driver !$OMP parallel do to a target teams loop and adds residency; the solver then runs on the GPU. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../vertical/MOM_kappa_shear.F90 | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 914a21750e..a5f7d916c3 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -131,6 +131,10 @@ module MOM_kappa_shear ! integer :: id_clock_project, id_clock_KQ, id_clock_avg, id_clock_setup +! The per-column solver and its helpers are device-callable so the driver column loop +! can run inside a target region (GPU port increment 3). +!$omp declare target(kappa_shear_column, find_kappa_tke, calculate_projected_state) + contains !> Subroutine for calculating shear-driven diffusivity and TKE in tracer columns @@ -349,7 +353,8 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, & tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature) ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. @@ -749,7 +754,8 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, & tke_avg, N2_init, S2_init, N2_mean, S2_mean, tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature) ! call cpu_clock_begin(Id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then @@ -918,7 +924,8 @@ end subroutine Calc_kappa_shear_vertex subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_lay, & u0xdz, v0xdz, T0xdz, S0xdz, kappa_avg, tke_avg, N2_init, S2_init, & N2_mean, S2_mean, tv, CS, GV, US, & - eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ) + eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, & + use_temperature ) type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. real, dimension(SZK_(GV)+1), & intent(inout) :: kappa !< The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -957,8 +964,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la type(thermo_var_ptrs), intent(in) :: tv !< A structure containing pointers to any !! available thermodynamic fields. Absent fields !! have NULL ptrs. - type(Kappa_shear_CS), pointer :: CS !< The control structure returned by a previous - !! call to kappa_shear_init. + type(Kappa_shear_CS), intent(in) :: CS !< The control structure returned by a previous + !! call to kappa_shear_init. Plain (not pointer) so + !! the routine is device-callable. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type integer, intent(in) :: eos_form !< The equation-of-state form id, resolved !! host-side for the device-callable EOS derivs path. @@ -966,6 +974,9 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure to Pa [Pa T2 R-1 L-2 ~> 1] + logical, intent(in) :: use_temperature !< If true, temperature and salinity are + !! state variables (resolved host-side to avoid a device + !! read of the tv%T pointer). ! Local variables real, dimension(nzc) :: & @@ -1055,8 +1066,6 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real :: k0dt ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1]. real :: I_lz_rescale_sqr ! The inverse of a rescaling factor for L2_bdry (Lz) squared [nondim]. logical :: valid_dt ! If true, all levels so far exhibit acceptably small changes in k_src. - logical :: use_temperature ! If true, temperature and salinity have been - ! allocated and are being used as state variables. integer :: ks_kappa, ke_kappa ! The k-range with nonzero kappas. integer :: dt_refinements ! The number of 2-fold refinements that will be used ! to estimate the maximum permitted time step. I.e., @@ -1083,7 +1092,6 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la endif tol2 = 2.0*CS%kappa_tol_err dt_refinements = 5 ! Selected so that 1/2^dt_refinements < 1-tol_dksrc_low - use_temperature = .false. ; if (associated(tv%T)) use_temperature = .true. ! Set up Idz as the inverse of layer thicknesses. @@ -1224,6 +1232,7 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la tv%eqn_of_state, (/2,nzc/), scale=-g_R0 ) #endif else +#ifndef __NVCOMPILER_OPENMP_GPU ! These should perhaps be combined into a single call to calculate the thermal expansion ! and haline contraction coefficients? call calculate_specific_vol_derivs(T_int, Sal_int, pressure, dSpV_dT, dSpV_dS, & @@ -1233,6 +1242,11 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la dbuoy_dT(K) = GV%g_Earth_Z_T2 * (rho_int(K) * dSpV_dT(K)) dbuoy_dS(K) = GV%g_Earth_Z_T2 * (rho_int(K) * dSpV_dS(K)) enddo +#else + ! The non-Boussinesq density-derivs path uses the polymorphic EOS interface, which is not + ! device-callable, and is excluded from device compilation. It is unreachable on GPU builds: + ! the driver FATALs on non-Boussinesq before the column solver runs. +#endif endif elseif (GV%Boussinesq .or. GV%semi_Boussinesq) then do K=1,nzc+1 ; dbuoy_dT(K) = -g_R0 ; dbuoy_dS(K) = 0.0 ; enddo @@ -1612,7 +1626,8 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b !! boundaries [H-1 Z-1 ~> m-2 or m kg-1]. real, dimension(nz), intent(in) :: Idz !< The inverse grid spacing of layers [Z-1 ~> m-1]. real, intent(in) :: f2 !< The squared Coriolis parameter [T-2 ~> s-2]. - type(Kappa_shear_CS), pointer :: CS !< A pointer to this module's control structure. + type(Kappa_shear_CS), intent(in) :: CS !< This module's control structure (plain, not pointer, + !! so the routine is device-callable). type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type real, dimension(nz+1), intent(inout) :: K_Q !< The shear-driven diapycnal diffusivity divided by From 9bf8a8940919e0200574f3862c9770a3af91f861 Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:07:26 +1000 Subject: [PATCH 07/15] kappa_shear: offload the vertex slab interpolation to GPU (3-D slabs) Promote the per-J 2-D vertex slabs (u_2d..rho_2d) in Calc_kappa_shear_vertex to 3-D arrays (u_slab..rho_slab) computed in device do-concurrent passes ahead of the column loop. The loop bodies are verbatim except that I_hwt is inlined as a reciprocal multiply and the use_temperature branch is hoisted so each pass is purely elementwise; answers are bitwise-identical (verified np=1 against repro_ocean.stats on benchmark_ALE, V100/NVHPC 26.5). The key residency fix relative to the earlier reverted attempt: u_in/v_in and T_in/S_in (tv%T, tv%S) are already device-resident (mapped in MOM.F90), so a map(to:) on them does not copy - and the diabatic stack modifies T/S on the host before set_diffusivity, leaving the device copies stale. They now get an explicit "target update to" refresh, which resolves the NaN-diffusivity crash the first attempt hit. dz_3d is a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are now consumed on the device by the slab passes, so increment 1's copy-back is dropped; the slabs are copied back for the (still host) per-column solver loop, which reads them as X_slab(I,J,k). Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 150 ++++++++++-------- 1 file changed, 86 insertions(+), 64 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index a5f7d916c3..8855164d4c 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -494,11 +494,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ h_at_u ! A mask-weighted thickness interpolated to u-points [H ~> m or kg m-2] real, dimension(SZI_(G),SZJB_(G),SZK_(GV)) :: & h_at_v ! A mask-weighted thickness interpolated to v-points [H ~> m or kg m-2] - real, dimension(SZIB_(G),SZK_(GV)) :: & - h_2d, & ! A 2-D version of h interpolated to vertices [H ~> m or kg m-2]. - dz_2d, & ! Vertical distance between interface heights [Z ~> m]. - u_2d, v_2d, & ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1]. - T_2d, S_2d, rho_2d ! 2-D versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. + real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)) :: & + h_slab, & ! A version of h interpolated to vertices [H ~> m or kg m-2]. + dz_slab, & ! Vertical distance between interface heights at vertices [Z ~> m]. + u_slab, v_slab, & ! Versions of u_in and v_in interpolated to vertices [L T-1 ~> m s-1]. + T_slab, S_slab, rho_slab ! Vertex versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. real, dimension(SZIB_(G),SZK_(GV)+1) :: & kappa_2d ! 2-D slice of kappa_vert [H Z T-1 ~> m2 s-1 or Pa s] real, dimension(SZIB_(G),SZK_(GV)+1) :: & @@ -587,7 +587,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! --- GPU port increment 1: h_at_u/h_at_v interpolation offloaded to device. ! h is host-authoritative in the (host-only) diabatic stack; refresh the device copy. ! G%mask2dCu/Cv/T are already device-resident (mapped in initialize_MOM). h_at_u/h_at_v - ! are device workspace, copied back for the (still host) per-column loop below. + ! are device workspace, consumed on the device by the slab interpolation below. !$omp target enter data map(to: h) !$omp target update to(h) !$omp target enter data map(alloc: h_at_u, h_at_v) @@ -613,52 +613,72 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - !$omp target update from(h_at_u, h_at_v) + ! --- GPU port increment 2: the per-J 2-D vertex slabs are promoted to 3-D arrays computed on + ! the device in one pass over J before the column loop. The loop bodies are verbatim from the + ! former per-J loops, except that I_hwt is inlined as a reciprocal multiply (bitwise-identical) + ! and the temperature/salinity branch is hoisted out of the loop so each body is purely + ! elementwise. u_in/v_in (the dycore u,v) and T_in/S_in (tv%T, tv%S) are already device- + ! resident but host-authoritative at this point in the (host-only) diabatic stack, so they + ! need an explicit refresh — a map(to:) on an already-present object does NOT copy. dz_3d is + ! a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are consumed on the + ! device here, so increment 1's copy-back is no longer needed. + !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) + !$omp target update to(u_in, v_in, T_in, S_in) + !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) + + ! Interpolate the various quantities to the corners, using masks. + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + u_slab(I,J,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & + ( (h_at_u(I,j,k) + h_at_u(I,j+1,k)) + H_tiny ) + v_slab(I,J,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & + ( (h_at_v(i,J,k) + h_at_v(i+1,J,k)) + H_tiny ) + + h_slab(I,J,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) + dz_slab(I,J,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * dz_3d(i+1,j,k) + G%mask2dT(i,j+1) * dz_3d(i,j+1,k)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) +! h_slab(I,J,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) +! h_slab(I,J,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & +! ((h(i+1,j,k)**2) + (h(i,j+1,k)**2))) * I_hwt + enddo + if (use_temperature) then + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + T_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & + G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * T_in(i+1,j+1,k))) + & + (G%mask2dT(i+1,j) * (h(i+1,j,k) * T_in(i+1,j,k)) + & + G%mask2dT(i,j+1) * (h(i,j+1,k) * T_in(i,j+1,k))) ) * & + (1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & + GV%H_subroundoff)) + S_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & + G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * S_in(i+1,j+1,k))) + & + (G%mask2dT(i+1,j) * (h(i+1,j,k) * S_in(i+1,j,k)) + & + G%mask2dT(i,j+1) * (h(i,j+1,k) * S_in(i,j+1,k))) ) * & + (1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & + GV%H_subroundoff)) + enddo + else + do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) + rho_slab(I,J,k) = GV%Rlay(k) + enddo + endif + ! The slabs feed the (still host) per-column loop below; copy them back. + !$omp target update from(u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & + !$OMP u_slab,v_slab,T_slab,S_slab,h_slab,dz_slab,rho_slab, & !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) do J=JsB,JeB - ! Interpolate the various quantities to the corners, using masks. - do k=1,nz ; do I=IsB,IeB - u_2d(I,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & - ( (h_at_u(I,j,k) + h_at_u(I,j+1,k)) + H_tiny ) - v_2d(I,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & - ( (h_at_v(i,J,k) + h_at_v(i+1,J,k)) + H_tiny ) - - I_hwt = 1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & - GV%H_subroundoff) - if (use_temperature) then - T_2d(I,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & - G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * T_in(i+1,j+1,k))) + & - (G%mask2dT(i+1,j) * (h(i+1,j,k) * T_in(i+1,j,k)) + & - G%mask2dT(i,j+1) * (h(i,j+1,k) * T_in(i,j+1,k))) ) * I_hwt - S_2d(I,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & - G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * S_in(i+1,j+1,k))) + & - (G%mask2dT(i+1,j) * (h(i+1,j,k) * S_in(i+1,j,k)) + & - G%mask2dT(i,j+1) * (h(i,j+1,k) * S_in(i,j+1,k))) ) * I_hwt - endif - h_2d(I,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) - dz_2d(I,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & - (G%mask2dT(i+1,j) * dz_3d(i+1,j,k) + G%mask2dT(i,j+1) * dz_3d(i,j+1,k)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) -! h_2d(I,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) -! h_2d(I,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & -! ((h(i+1,j,k)**2) + (h(i,j+1,k)**2))) * I_hwt - enddo ; enddo - if (.not.use_temperature) then ; do k=1,nz ; do I=IsB,IeB - rho_2d(I,k) = GV%Rlay(k) - enddo ; enddo ; endif - !--------------------------------------- ! Work on each column. !--------------------------------------- @@ -675,25 +695,25 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ T0xdz(k) = 0.0 ; S0xdz(k) = 0.0 ! Add a new layer if this one has mass. -! if ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless)) nzc = nzc+1 +! if ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 if ((k>CS%nkml) .and. (h_lay(nzc) > 0.0) .and. & - (h_2d(I,k) > dz_massless)) nzc = nzc+1 + (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 ! Only merge clusters of massless layers. ! if ((h_lay(nzc) > dz_massless) .or. & -! ((h_lay(nzc) > 0.0) .and. (h_2d(I,k) > dz_massless))) nzc = nzc+1 +! ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless))) nzc = nzc+1 kc(k) = nzc - h_lay(nzc) = h_lay(nzc) + h_2d(I,k) - dz_lay(nzc) = dz_lay(nzc) + dz_2d(I,k) - u0xdz(nzc) = u0xdz(nzc) + u_2d(I,k)*h_2d(I,k) - v0xdz(nzc) = v0xdz(nzc) + v_2d(I,k)*h_2d(I,k) + h_lay(nzc) = h_lay(nzc) + h_slab(I,J,k) + dz_lay(nzc) = dz_lay(nzc) + dz_slab(I,J,k) + u0xdz(nzc) = u0xdz(nzc) + u_slab(I,J,k)*h_slab(I,J,k) + v0xdz(nzc) = v0xdz(nzc) + v_slab(I,J,k)*h_slab(I,J,k) if (use_temperature) then - T0xdz(nzc) = T0xdz(nzc) + T_2d(I,k)*h_2d(I,k) - S0xdz(nzc) = S0xdz(nzc) + S_2d(I,k)*h_2d(I,k) + T0xdz(nzc) = T0xdz(nzc) + T_slab(I,J,k)*h_slab(I,J,k) + S0xdz(nzc) = S0xdz(nzc) + S_slab(I,J,k)*h_slab(I,J,k) else - T0xdz(nzc) = T0xdz(nzc) + rho_2d(I,k)*h_2d(I,k) - S0xdz(nzc) = S0xdz(nzc) + rho_2d(I,k)*h_2d(I,k) + T0xdz(nzc) = T0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) + S0xdz(nzc) = S0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) endif enddo kc(nz+1) = nzc+1 @@ -703,28 +723,28 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Now determine kf, the fractional weight of interface kc when ! interpolating between interfaces kc and kc+1. - kf(1) = 0.0 ; dz_in_lay = h_2d(I,1) + kf(1) = 0.0 ; dz_in_lay = h_slab(I,J,1) do k=2,nz if (kc(k) > kc(k-1)) then - kf(k) = 0.0 ; dz_in_lay = h_2d(I,k) + kf(k) = 0.0 ; dz_in_lay = h_slab(I,J,k) else - kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(I,k) + kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_slab(I,J,k) endif enddo kf(nz+1) = 0.0 else do k=1,nz - h_lay(k) = h_2d(I,k) - dz_lay(k) = dz_2d(I,k) - u0xdz(k) = u_2d(I,k)*h_lay(k) ; v0xdz(k) = v_2d(I,k)*h_lay(k) + h_lay(k) = h_slab(I,J,k) + dz_lay(k) = dz_slab(I,J,k) + u0xdz(k) = u_slab(I,J,k)*h_lay(k) ; v0xdz(k) = v_slab(I,J,k)*h_lay(k) enddo if (use_temperature) then do k=1,nz - T0xdz(k) = T_2d(I,k)*h_lay(k) ; S0xdz(k) = S_2d(I,k)*h_lay(k) + T0xdz(k) = T_slab(I,J,k)*h_lay(k) ; S0xdz(k) = S_slab(I,J,k)*h_lay(k) enddo else do k=1,nz - T0xdz(k) = rho_2d(I,k)*h_lay(k) ; S0xdz(k) = rho_2d(I,k)*h_lay(k) + T0xdz(k) = rho_slab(I,J,k)*h_lay(k) ; S0xdz(k) = rho_slab(I,J,k)*h_lay(k) enddo endif nzc = nz @@ -817,7 +837,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Store the 2-d slices back in the 3-d arrays for restarts or interpolation back to tracer points. if (CS%VS_ThicknessMean) then do K=1,nz+1 ; do I=IsB,IeB - h_vert(I,J,k) = h_2d(I,k) + h_vert(I,J,k) = h_slab(I,J,k) enddo ; enddo endif if (CS%VS_viscosity_bug) then @@ -913,7 +933,9 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) - ! --- GPU port increment 1: mirror the enter-data above (balance discipline). + ! --- GPU port increments 1+2: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) + !$omp target exit data map(release: u_in, v_in, T_in, S_in, dz_3d) !$omp target exit data map(release: h_at_u, h_at_v) !$omp target exit data map(release: h) From 685037c9b7fb709b98d0269b7af0f6a6b94c83ca Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:19:38 +1000 Subject: [PATCH 08/15] kappa_shear: run the vertex column solver on the GPU Convert Calc_kappa_shear_vertex's per-column loop from a host OMP parallel do over J into an "omp target teams loop collapse(2)" over (J,I), calling the declare-target kappa_shear_column on the device. Bitwise-identical (np=1, benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5), which also demonstrates that the solver's 26 early exit/return sites are runtime-correct on 26.5. Restructuring, all bitwise-neutral: - The per-J staging slices kappa_2d/tke_2d become device-resident 3-D staging arrays kappa_3d/tke_3d (map(alloc:)), written per column and consumed by new elementwise do-concurrent write-back passes for kappa_vertex/tke_io/kv_io. - The surface pressure is interpolated to vertices on the host ahead of the region (surface_pres_2d, verbatim expressions), keeping the possibly-null p_surf pointer and the CS%psurf_bug branch out of device code. - The h_vert fill stays on the host (it is only consumed by the host tracer-point averaging when VS_ThicknessMean=true), fed by a guarded update from(h_slab). - All per-column scratch is private at the loop; loop-invariant scalars are firstprivate; CS gets a per-call shallow map(to:) in the driver (all scalars plus a diag pointer that device code never dereferences). kappa_vertex, tke_io and kv_io are mapped to: (not alloc:) so host-set halo/initial values survive the full-array update from. diag_N2/S2_init/mean are mapped after their conditional host zeroing and copied back under one guard matching their consumers (post_data/Bchksum). Known limitation, to be fixed next: the solver's automatic arrays (sized nzc or SZK_(GV)) are device-heap allocated, so runs currently need NVCOMPILER_ACC_CUDA_HEAPSIZE=1073741824 (the default heap is too small: "DEVICE FORTRAN AUTO ALLOCATION FAILED"). Fixed-size locals will remove this. Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 148 +++++++++++------- 1 file changed, 94 insertions(+), 54 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 8855164d4c..6a7738939a 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -499,10 +499,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ dz_slab, & ! Vertical distance between interface heights at vertices [Z ~> m]. u_slab, v_slab, & ! Versions of u_in and v_in interpolated to vertices [L T-1 ~> m s-1]. T_slab, S_slab, rho_slab ! Vertex versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. - real, dimension(SZIB_(G),SZK_(GV)+1) :: & - kappa_2d ! 2-D slice of kappa_vert [H Z T-1 ~> m2 s-1 or Pa s] - real, dimension(SZIB_(G),SZK_(GV)+1) :: & - tke_2d ! 2-D version tke_io [Z2 T-2 ~> m2 s-2]. + real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1) :: & + kappa_3d, & ! Device staging array for the columns' averaged kappa [H Z T-1 ~> m2 s-1 or Pa s] + tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. + real, dimension(SZIB_(G),SZJB_(G)) :: & + surface_pres_2d ! The surface pressure interpolated to vertices [R L2 T-2 ~> Pa]. real, dimension(SZK_(GV)) :: & Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] @@ -527,7 +528,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ real :: dz_in_lay ! The running sum of the thickness in a layer [H ~> m or kg m-2] real :: k0dt ! The background diffusivity times the timestep [H Z ~> m2 or kg m-1] real :: dz_massless ! A layer thickness that is considered massless [H ~> m or kg m-2] - real :: I_hwt ! The inverse of the sum of the adjacent masked thickness weights [H-1 ~> m-1 or m2 kg-1] real :: I_htot ! The inverse of the sum of the thicknesses at adjacent vertices [H-1 ~> m-1 or m2 kg-1] real :: I_Prandtl ! The inverse of the turbulent Prandtl number [nondim]. logical :: use_temperature ! If true, temperature and salinity have been @@ -581,6 +581,28 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ I_Prandtl = 0.0 ; if (CS%Prandtl_turb > 0.0) I_Prandtl = 1.0 / CS%Prandtl_turb H_tiny = 0.5 * GV%H_subroundoff + ! GPU port increment 3: interpolate the surface pressure to the vertices on the host, ahead + ! of the device column region (p_surf is a possibly-unassociated pointer, which is simpler to + ! keep off the device). The expressions are verbatim from the former in-column code, but are + ! evaluated for every vertex instead of only ocean vertices; the extra values are never read. + surface_pres_2d(:,:) = 0.0 + if (associated(p_surf)) then + if (CS%psurf_bug) then + ! This is wrong because it is averaging values from land in some places. + do J=JsB,JeB ; do I=IsB,IeB + surface_pres_2d(I,J) = 0.25 * ((p_surf(i,j) + p_surf(i+1,j+1)) + & + (p_surf(i+1,j) + p_surf(i,j+1))) + enddo ; enddo + else + do J=JsB,JeB ; do I=IsB,IeB + surface_pres_2d(I,J) = ((G%mask2dT(i,j) * p_surf(i,j) + G%mask2dT(i+1,j+1) * p_surf(i+1,j+1)) + & + (G%mask2dT(i+1,j) * p_surf(i+1,j) + G%mask2dT(i,j+1) * p_surf(i,j+1)) ) / & + ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) + enddo ; enddo + endif + endif + ! Convert layer thicknesses into geometric thickness in height units. call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=1) @@ -668,20 +690,31 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - ! The slabs feed the (still host) per-column loop below; copy them back. - !$omp target update from(u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) - - !$OMP parallel do default(private) shared(jsB,jeB,isB,ieB,nz,h,u_in,v_in,T_in,S_in,h_at_u,h_at_v,dz_3d,H_tiny, & - !$OMP use_temperature,tv,G,GV,US,CS,kappa_io, & - !$OMP u_slab,v_slab,T_slab,S_slab,h_slab,dz_slab,rho_slab, & - !$OMP dz_massless,k0dt,p_surf,dt,tke_io,kv_io,kappa_vertex,h_vert,I_Prandtl, & - !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa, & - !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) - do J=JsB,JeB + ! --- GPU port increment 3: run the per-column solver on the device. The columns iterate as + ! a target teams loop collapsed over (J,I); every piece of per-column scratch is private (the + ! declare-target solver's own locals are automatically private per device thread). CS is all + ! scalars plus a diag pointer that is never dereferenced in the device code, so a per-call + ! shallow map(to:) suffices. kappa_vertex/tke_io/kv_io are mapped to: (not alloc) so that the + ! host-set values - kv_io is intent(inout), kappa_vertex is zeroed on the host - survive the + ! full-array update from below. The diag_* arrays are mapped after their (conditional) host + ! zeroing for the same reason. kappa_3d/tke_3d are device-only staging for what was the per-J + ! kappa_2d/tke_2d, written per column and consumed by the write-back passes below. + !$omp target enter data map(to: CS) + !$omp target enter data map(to: surface_pres_2d) + !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) + !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target enter data map(alloc: kappa_3d, tke_3d) !--------------------------------------- ! Work on each column. !--------------------------------------- + !$omp target teams loop collapse(2) & + !$omp private(nzc, kc, kf, Idz, h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, dz_in_lay, & + !$omp f2, surface_pres, kappa, tke, kappa_avg, tke_avg, N2_init, S2_init, & + !$omp N2_mean, S2_mean, k) & + !$omp firstprivate(nz, dt, k0dt, dz_massless, use_temperature, eos_form, & + !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + do J=JsB,JeB do I=IsB,IeB ; if ((G%mask2dCu(I,j) + G%mask2dCu(I,j+1)) + & (G%mask2dCv(i,J) + G%mask2dCv(i+1,J)) > 0.0) then ! call cpu_clock_begin(Id_clock_setup) @@ -752,19 +785,7 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ endif f2 = G%Coriolis2Bu(I,J) - surface_pres = 0.0 - if (associated(p_surf)) then - if (CS%psurf_bug) then - ! This is wrong because it is averaging values from land in some places. - surface_pres = 0.25 * ((p_surf(i,j) + p_surf(i+1,j+1)) + & - (p_surf(i+1,j) + p_surf(i,j+1))) - else - surface_pres = ((G%mask2dT(i,j) * p_surf(i,j) + G%mask2dT(i+1,j+1) * p_surf(i+1,j+1)) + & - (G%mask2dT(i+1,j) * p_surf(i+1,j) + G%mask2dT(i,j+1) * p_surf(i,j+1)) ) / & - ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & - (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) - endif - endif + surface_pres = surface_pres_2d(I,J) ! ---------------------------------------------------- ! Set the initial guess for kappa, here defined at interfaces. @@ -780,11 +801,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_2d(I,K) = kappa_avg(K) + kappa_3d(I,J,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_2d(I,K) = tke(K) + tke_3d(I,J,K) = tke(K) else - tke_2d(I,K) = tke_avg(K) + tke_3d(I,J,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -802,11 +823,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ else do K=1,nz+1 if (kf(K) == 0.0) then - kappa_2d(I,K) = kappa_avg(kc(K)) - tke_2d(I,K) = tke_avg(kc(K)) + kappa_3d(I,J,K) = kappa_avg(kc(K)) + tke_3d(I,J,K) = tke_avg(kc(K)) else - kappa_2d(I,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) - tke_2d(I,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) + kappa_3d(I,J,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) + tke_3d(I,J,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) endif enddo do K=1,nz+1 @@ -830,30 +851,44 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! call cpu_clock_end(Id_clock_setup) else ! Land points, still inside the i-loop. do K=1,nz+1 - kappa_2d(I,K) = 0.0 ; tke_2d(I,K) = 0.0 + kappa_3d(I,J,K) = 0.0 ; tke_3d(I,J,K) = 0.0 enddo endif ; enddo ! i-loop + enddo ! end of J-loop - ! Store the 2-d slices back in the 3-d arrays for restarts or interpolation back to tracer points. - if (CS%VS_ThicknessMean) then + ! Store the columns' results back in the 3-d arrays for restarts or interpolation back to + ! tracer points. h_vert is only used when VS_ThicknessMean is true, and only by the (host) + ! tracer-point averaging below, so its fill stays a host loop fed by a guarded copy-back of + ! the device-computed h_slab. + if (CS%VS_ThicknessMean) then + !$omp target update from(h_slab) + do J=JsB,JeB do K=1,nz+1 ; do I=IsB,IeB h_vert(I,J,k) = h_slab(I,J,k) enddo ; enddo - endif - if (CS%VS_viscosity_bug) then - do K=1,nz+1 ; do I=IsB,IeB - kappa_vertex(I,J,K) = kappa_2d(I,K) - tke_io(I,J,K) = G%mask2dBu(I,J) * tke_2d(I,K) - kv_io(I,J,K) = ( G%mask2dBu(I,J) * kappa_vertex(I,J,K) ) * CS%Prandtl_turb - enddo ; enddo - else - do K=1,nz+1 ; do I=IsB,IeB - kappa_vertex(I,J,K) = kappa_2d(I,K) - tke_io(I,J,K) = tke_2d(I,K) - kv_io(I,J,K) = kappa_vertex(I,J,K) * CS%Prandtl_turb - enddo ; enddo - endif - enddo ! end of J-loop + enddo + endif + if (CS%VS_viscosity_bug) then + do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) + kappa_vertex(I,J,K) = kappa_3d(I,J,K) + tke_io(I,J,K) = G%mask2dBu(I,J) * tke_3d(I,J,K) + kv_io(I,J,K) = ( G%mask2dBu(I,J) * kappa_vertex(I,J,K) ) * CS%Prandtl_turb + enddo + else + do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) + kappa_vertex(I,J,K) = kappa_3d(I,J,K) + tke_io(I,J,K) = tke_3d(I,J,K) + kv_io(I,J,K) = kappa_vertex(I,J,K) * CS%Prandtl_turb + enddo + endif + + ! The vertex-to-tracer-point averaging below, the checksums and post_data are all still on + ! the host. The diag_* transfer is guarded by the same conditions as their consumers. + !$omp target update from(kappa_vertex, tke_io, kv_io) + if ((CS%id_N2_init>0) .or. (CS%id_S2_init>0) .or. (CS%id_N2_mean>0) .or. (CS%id_S2_mean>0) & + .or. CS%debug) then + !$omp target update from(diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + endif ! Set the diffusivities in tracer columns from the values at vertices. @@ -933,7 +968,12 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) - ! --- GPU port increments 1+2: mirror the enter-data above (balance discipline). + ! --- GPU port increments 1-3: mirror the enter-data above (balance discipline). + !$omp target exit data map(release: kappa_3d, tke_3d) + !$omp target exit data map(release: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target exit data map(release: kappa_vertex, tke_io, kv_io) + !$omp target exit data map(release: surface_pres_2d) + !$omp target exit data map(release: CS) !$omp target exit data map(release: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) !$omp target exit data map(release: u_in, v_in, T_in, S_in, dz_3d) !$omp target exit data map(release: h_at_u, h_at_v) From 2ba3995e4cfc594ef913178bee61fcd9a3efc25d Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:31:50 +1000 Subject: [PATCH 09/15] kappa_shear: fixed-size column arrays in GPU builds (drop device-heap allocs) The vertex driver's per-column private scratch and the column routines' automatic locals are sized by GV%ke/nzc, which is not a compile-time constant, so nvfortran allocated them from the device heap on every column call (NVFORTRAN-W-0155). That exhausted the default device heap ("DEVICE FORTRAN AUTO ALLOCATION FAILED" unless NVCOMPILER_ACC_CUDA_HEAPSIZE was raised) and serialized columns on the device allocator (~2x total model slowdown). In GPU builds (#ifdef __NVCOMPILER_OPENMP_GPU) these arrays now use a compile-time ceiling, the new module parameter GPU_nk_max=128, so each device thread gets stack ("local memory") arrays; kappa_shear_init fatals if GV%ke > GPU_nk_max when the scheme is active. CPU builds keep the exact-size declarations. No code queries these arrays' sizes and all accesses are indexed to nzc/nz, so the enlargement is bitwise-neutral: verified np=1 against repro_ocean.stats at the default device heap, wall time back to the pre-offload baseline (98s -> 45s), and the W-0155 warnings are gone. Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 6a7738939a..c891e41903 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -135,6 +135,13 @@ module MOM_kappa_shear ! can run inside a target region (GPU port increment 3). !$omp declare target(kappa_shear_column, find_kappa_tke, calculate_projected_state) +!> A compile-time ceiling on the number of layers in GPU builds, used to give the +!! device-executed column routines and the driver's per-column private scratch fixed-size +!! (stack) arrays instead of per-call device-heap automatic allocations, which exhaust the +!! default device heap and serialize on the device allocator. Checked against GV%ke in +!! kappa_shear_init. Unused in CPU builds, where the declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 + contains !> Subroutine for calculating shear-driven diffusivity and TKE in tracer columns @@ -504,7 +511,14 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. real, dimension(SZIB_(G),SZJB_(G)) :: & surface_pres_2d ! The surface pressure interpolated to vertices [R L2 T-2 ~> Pa]. + ! In GPU builds the per-column private scratch has compile-time-constant sizes so each + ! device thread gets stack ("local memory") arrays; runtime-sized privates are device-heap + ! allocated per column, which exhausts the default heap and serializes on the allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] dz_lay, & ! The geometric layer thickness in height units [Z ~> m] @@ -512,7 +526,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ v0xdz, & ! The initial meridional velocity times dz [H L T-1 ~> m2 s-1 or kg m-1 s-1] T0xdz, & ! The initial temperature times dz [C H ~> degC m or degC kg m-2] S0xdz ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif kappa, & ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s] tke, & ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2]. kappa_avg, & ! The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -533,10 +551,18 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ logical :: use_temperature ! If true, temperature and salinity have been ! allocated and are being used as state variables. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max+1) :: kc ! The index map between the original +#else integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original +#endif ! interfaces and the interfaces with massless layers ! merged into nearby massive layers. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: kf ! The fractional weight of interface kc+1 for +#else real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for +#endif ! interpolating back to the original index space [nondim]. real :: h_SW, h_SE, h_NW, h_NE ! Thicknesses at adjacent vertices [H ~> m or kg m-2] real :: mks_to_HZ_T ! A factor used to restore dimensional scaling after the geometric mean @@ -1041,7 +1067,14 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la !! read of the tv%T pointer). ! Local variables + ! In GPU builds these locals have compile-time-constant sizes so each device call uses + ! stack ("local memory") arrays; runtime-sized automatics are device-heap allocated per + ! call, which exhausts the default heap and serializes on the device allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(nzc) :: & +#endif u, & ! The zonal velocity after a timestep of mixing [L T-1 ~> m s-1]. v, & ! The meridional velocity after a timestep of mixing [L T-1 ~> m s-1]. Idz, & ! The inverse of the distance between TKE points [Z-1 ~> m-1]. @@ -1050,7 +1083,11 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la u_test, v_test, & ! Temporary velocities [L T-1 ~> m s-1]. T_test, S_test ! Temporary temperatures [C ~> degC] and salinities [S ~> ppt]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nzc+1) :: & +#endif N2, & ! The squared buoyancy frequency at an interface [T-2 ~> s-2]. h_Int, & ! The extent of a finite-volume space surrounding an interface, ! as used in calculating kappa and TKE [H ~> m or kg m-2] @@ -1574,7 +1611,11 @@ subroutine calculate_projected_state(kappa, u0, v0, T0, S0, dt, nz, dz, I_dz_int !! diffusivity. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: c1 ! A tridiagonal variable [nondim] +#else real, dimension(nz+1) :: c1 ! A tridiagonal variable [nondim] +#endif real :: a_a, a_b ! Tridiagonal coupling coefficients [H ~> m or kg m-2] real :: b1, b1nz_0 ! Tridiagonal variables [H-1 ~> m-1 or m2 kg-1] real :: bd1 ! A term in the denominator of b1 [H ~> m or kg m-2] @@ -1707,10 +1748,18 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b ! This subroutine calculates new, consistent estimates of TKE and kappa. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(nz) :: & +#endif aQ, & ! aQ is the coupling between adjacent interfaces in the TKE equations [H T-1 ~> m s-1 or kg m-2 s-1] dQdz ! Half the partial derivative of TKE with depth [Z T-2 ~> m s-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nz+1) :: & +#endif dK, & ! The change in kappa [H Z T-1 ~> m2 s-1 or Pa s]. dQ, & ! The change in TKE [Z2 T-2 ~> m2 s-2]. cQ, cK, & ! cQ and cK are the upward influences in the tridiagonal and @@ -1788,7 +1837,11 @@ subroutine find_kappa_tke(N2, S2, kappa_in, Idz, h_Int, dz_Int, dz_h_Int, I_L2_b logical, parameter :: debug_soln = .false. real :: K_err_lin ! The imbalance in the K equation [H T-1 ~> m s-1 or kg m-2 s-1] real :: Q_err_lin ! The imbalance in the Q equation [H Z T-3 ~> m2 s-3 or kg m-1 s-3] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(nz+1) :: & +#endif I_Ld2_debug, & ! A separate version of I_Ld2 for debugging [H-1 Z-1 ~> m-2 or m kg-1]. kappa_prev, & ! The value of kappa at the start of the current iteration [H Z T-1 ~> m2 s-1 or Pa s] TKE_prev ! The value of TKE at the start of the current iteration [Z2 T-2 ~> m2 s-2]. @@ -2495,6 +2548,13 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) 's-2', conversion=US%s_to_T**2) endif +#ifdef __NVCOMPILER_OPENMP_GPU + ! The device-executed column routines use fixed-size local arrays in GPU builds. + if (kappa_shear_init .and. (GV%ke > GPU_nk_max)) call MOM_error(FATAL, & + "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max; "//& + "increase GPU_nk_max in MOM_kappa_shear.F90.") +#endif + end function kappa_shear_init !> This function indicates to other modules whether the Jackson et al shear mixing From 6530045cb088d0475e05bffdabc5a11b0f5179ba Mon Sep 17 00:00:00 2001 From: Jorge Galvez Vallejo Date: Fri, 17 Jul 2026 16:53:14 +1000 Subject: [PATCH 10/15] kappa_shear: review fixes - refresh device tke_io/kv_io, correct T/S residency Review of the column-offload series found one latent halo bug: kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident, mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update, so this routine's map(to: kv_io) copies nothing and the full-array "update from" wrote stale pre-pass_var device halos over the corrected host halos - silently resurrecting the restart non-reproducibility in non-symmetric mode that the pass_var was added to fix. Invisible to a cold-start np=1 run. Fixed with an explicit "update to(tke_io, kv_io)" after the maps (tke_io defensively, so the routine does not depend on it staying unmapped elsewhere). Also corrects the increment-2 residency comment and drops the redundant "update to(T_in, S_in)": the actuals are the caller's convection-filtered T_f/S_f, fresh host locals each call, so their map(to:) already copies (they are not tv%T/tv%S as the comment claimed). This also establishes that the earlier reverted slab attempt's NaN was not T/S staleness; the operative fix was restructuring the interpolation (hoisting the use_temperature branch out of the do concurrent). And the GPU_nk_max FATAL message now notes it covers the shared tracer-point column routines too. Bit-identical at np=1 (benchmark_ALE vs repro_ocean.stats, V100/NVHPC 26.5). Co-Authored-By: Claude Fable 5 --- .../vertical/MOM_kappa_shear.F90 | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index c891e41903..7ca8143fba 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -665,13 +665,14 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! the device in one pass over J before the column loop. The loop bodies are verbatim from the ! former per-J loops, except that I_hwt is inlined as a reciprocal multiply (bitwise-identical) ! and the temperature/salinity branch is hoisted out of the loop so each body is purely - ! elementwise. u_in/v_in (the dycore u,v) and T_in/S_in (tv%T, tv%S) are already device- - ! resident but host-authoritative at this point in the (host-only) diabatic stack, so they - ! need an explicit refresh — a map(to:) on an already-present object does NOT copy. dz_3d is - ! a fresh local each call, so its map(to:) does copy. h_at_u/h_at_v are consumed on the - ! device here, so increment 1's copy-back is no longer needed. + ! elementwise. u_in/v_in (the dycore u,v) are already device-resident but host- + ! authoritative at this point in the (host-only) diabatic stack, so they need an explicit + ! refresh — a map(to:) on an already-present object does NOT copy. T_in/S_in (the caller's + ! convection-filtered T_f/S_f) and dz_3d are fresh host locals each call, so their map(to:) + ! does copy. h_at_u/h_at_v are consumed on the device here, so increment 1's copy-back is + ! no longer needed. !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) - !$omp target update to(u_in, v_in, T_in, S_in) + !$omp target update to(u_in, v_in) !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) ! Interpolate the various quantities to the corners, using masks. @@ -728,6 +729,13 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ !$omp target enter data map(to: CS) !$omp target enter data map(to: surface_pres_2d) !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) + ! kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident + ! (mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update), so the + ! map(to:) above does not copy it; without this refresh, the full-array update from below + ! would overwrite the corrected host halos with stale device values after a restart in + ! non-symmetric mode. tke_io is not currently mapped elsewhere, but is refreshed too so + ! this routine does not silently depend on that staying true. + !$omp target update to(tke_io, kv_io) !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) !$omp target enter data map(alloc: kappa_3d, tke_3d) @@ -2551,8 +2559,10 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) #ifdef __NVCOMPILER_OPENMP_GPU ! The device-executed column routines use fixed-size local arrays in GPU builds. if (kappa_shear_init .and. (GV%ke > GPU_nk_max)) call MOM_error(FATAL, & - "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max; "//& - "increase GPU_nk_max in MOM_kappa_shear.F90.") + "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max because the "//& + "column routines use fixed-size local arrays on the device (this applies to the "//& + "tracer-point scheme too, which shares those routines); increase GPU_nk_max in "//& + "MOM_kappa_shear.F90 or use a CPU build.") #endif end function kappa_shear_init From eaf17628330ee36b5ba7258c7d61e5cd9a5a4633 Mon Sep 17 00:00:00 2001 From: Jorge Date: Sun, 26 Jul 2026 21:21:56 -0500 Subject: [PATCH 11/15] remove port notes --- .../vertical/MOM_kappa_shear.F90 | 42 +------------------ 1 file changed, 1 insertion(+), 41 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 7ca8143fba..5dea402789 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -132,7 +132,7 @@ module MOM_kappa_shear ! integer :: id_clock_project, id_clock_KQ, id_clock_avg, id_clock_setup ! The per-column solver and its helpers are device-callable so the driver column loop -! can run inside a target region (GPU port increment 3). +! can run inside a target region . !$omp declare target(kappa_shear_column, find_kappa_tke, calculate_projected_state) !> A compile-time ceiling on the number of layers in GPU builds, used to give the @@ -232,8 +232,6 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & use_temperature = associated(tv%T) - ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor - ! and MOM_error are not device-callable) so they can be passed into the column solver. eos_form = -1 eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 if (use_temperature) then @@ -585,9 +583,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ use_temperature = associated(tv%T) - ! GPU port increment 3a: resolve the EOS form + unit scaling once on the host (the accessor - ! and MOM_error are not device-callable) so they can be passed into the column solver, which - ! will run inside a device region. eos_form = -1 eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 if (use_temperature) then @@ -607,10 +602,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ I_Prandtl = 0.0 ; if (CS%Prandtl_turb > 0.0) I_Prandtl = 1.0 / CS%Prandtl_turb H_tiny = 0.5 * GV%H_subroundoff - ! GPU port increment 3: interpolate the surface pressure to the vertices on the host, ahead - ! of the device column region (p_surf is a possibly-unassociated pointer, which is simpler to - ! keep off the device). The expressions are verbatim from the former in-column code, but are - ! evaluated for every vertex instead of only ocean vertices; the extra values are never read. surface_pres_2d(:,:) = 0.0 if (associated(p_surf)) then if (CS%psurf_bug) then @@ -632,10 +623,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Convert layer thicknesses into geometric thickness in height units. call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=1) - ! --- GPU port increment 1: h_at_u/h_at_v interpolation offloaded to device. - ! h is host-authoritative in the (host-only) diabatic stack; refresh the device copy. - ! G%mask2dCu/Cv/T are already device-resident (mapped in initialize_MOM). h_at_u/h_at_v - ! are device workspace, consumed on the device by the slab interpolation below. !$omp target enter data map(to: h) !$omp target update to(h) !$omp target enter data map(alloc: h_at_u, h_at_v) @@ -661,16 +648,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - ! --- GPU port increment 2: the per-J 2-D vertex slabs are promoted to 3-D arrays computed on - ! the device in one pass over J before the column loop. The loop bodies are verbatim from the - ! former per-J loops, except that I_hwt is inlined as a reciprocal multiply (bitwise-identical) - ! and the temperature/salinity branch is hoisted out of the loop so each body is purely - ! elementwise. u_in/v_in (the dycore u,v) are already device-resident but host- - ! authoritative at this point in the (host-only) diabatic stack, so they need an explicit - ! refresh — a map(to:) on an already-present object does NOT copy. T_in/S_in (the caller's - ! convection-filtered T_f/S_f) and dz_3d are fresh host locals each call, so their map(to:) - ! does copy. h_at_u/h_at_v are consumed on the device here, so increment 1's copy-back is - ! no longer needed. !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) !$omp target update to(u_in, v_in) !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) @@ -717,24 +694,9 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ enddo endif - ! --- GPU port increment 3: run the per-column solver on the device. The columns iterate as - ! a target teams loop collapsed over (J,I); every piece of per-column scratch is private (the - ! declare-target solver's own locals are automatically private per device thread). CS is all - ! scalars plus a diag pointer that is never dereferenced in the device code, so a per-call - ! shallow map(to:) suffices. kappa_vertex/tke_io/kv_io are mapped to: (not alloc) so that the - ! host-set values - kv_io is intent(inout), kappa_vertex is zeroed on the host - survive the - ! full-array update from below. The diag_* arrays are mapped after their (conditional) host - ! zeroing for the same reason. kappa_3d/tke_3d are device-only staging for what was the per-J - ! kappa_2d/tke_2d, written per column and consumed by the write-back passes below. !$omp target enter data map(to: CS) !$omp target enter data map(to: surface_pres_2d) !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) - ! kv_io's actual argument (visc%Kv_shear_Bu) is already persistently device-resident - ! (mapped in set_visc_init BEFORE the restart-reproducibility pass_var halo update), so the - ! map(to:) above does not copy it; without this refresh, the full-array update from below - ! would overwrite the corrected host halos with stale device values after a restart in - ! non-symmetric mode. tke_io is not currently mapped elsewhere, but is refreshed too so - ! this routine does not silently depend on that staying true. !$omp target update to(tke_io, kv_io) !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) !$omp target enter data map(alloc: kappa_3d, tke_3d) @@ -1002,7 +964,6 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ if (CS%id_N2_mean > 0) call post_data(CS%id_N2_mean, diag_N2_mean, CS%diag) if (CS%id_S2_mean > 0) call post_data(CS%id_S2_mean, diag_S2_mean, CS%diag) - ! --- GPU port increments 1-3: mirror the enter-data above (balance discipline). !$omp target exit data map(release: kappa_3d, tke_3d) !$omp target exit data map(release: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) !$omp target exit data map(release: kappa_vertex, tke_io, kv_io) @@ -2557,7 +2518,6 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) endif #ifdef __NVCOMPILER_OPENMP_GPU - ! The device-executed column routines use fixed-size local arrays in GPU builds. if (kappa_shear_init .and. (GV%ke > GPU_nk_max)) call MOM_error(FATAL, & "kappa_shear_init: GPU builds of kappa_shear require GV%ke <= GPU_nk_max because the "//& "column routines use fixed-size local arrays on the device (this applies to the "//& From 3797a4dc24a930c9ebe047f7665077e33f18a652 Mon Sep 17 00:00:00 2001 From: Jorge Date: Sun, 26 Jul 2026 22:23:27 -0500 Subject: [PATCH 12/15] fix doxygfen --- src/parameterizations/vertical/MOM_kappa_shear.F90 | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 5dea402789..b20d4d1e38 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -1030,7 +1030,8 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real, intent(in) :: eos_kg_m3_to_R !< EOS factor converting kg m-3 to R [R m3 kg-1 ~> 1] real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] - real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure to Pa [Pa T2 R-1 L-2 ~> 1] + real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure + !! to Pa [Pa T2 R-1 L-2 ~> 1] logical, intent(in) :: use_temperature !< If true, temperature and salinity are !! state variables (resolved host-side to avoid a device !! read of the tv%T pointer). From 659644eeb25723a2a9e847ac6b80dfd45badd25a Mon Sep 17 00:00:00 2001 From: Jorge Date: Sun, 26 Jul 2026 22:47:22 -0500 Subject: [PATCH 13/15] doxygen fix --- src/equation_of_state/MOM_EOS.F90 | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 427e81592c..502d0640db 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -1030,7 +1030,8 @@ end subroutine calculate_density_derivs_elem_loc subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) type(EOS_type), intent(in) :: EOS !< Equation of state structure integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) - real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal density unit R [R m3 kg-1 ~> 1] + real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal + !! density unit R [R m3 kg-1 ~> 1] real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] real, intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1] From dd8e3e31e0e1347b0378cfeacec4bb5d0380e86a Mon Sep 17 00:00:00 2001 From: Jorge Date: Sun, 26 Jul 2026 22:53:25 -0500 Subject: [PATCH 14/15] trailing white space --- src/equation_of_state/MOM_EOS.F90 | 2 +- src/parameterizations/vertical/MOM_kappa_shear.F90 | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 502d0640db..370c64860d 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -1030,7 +1030,7 @@ end subroutine calculate_density_derivs_elem_loc subroutine get_EOS_form_and_scaling(EOS, form_of_EOS, kg_m3_to_R, C_to_degC, S_to_ppt, RL2_T2_to_Pa) type(EOS_type), intent(in) :: EOS !< Equation of state structure integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) - real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal + real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal !! density unit R [R m3 kg-1 ~> 1] real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index b20d4d1e38..e648a8be91 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -1030,7 +1030,7 @@ subroutine kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, hlay, dz_la real, intent(in) :: eos_kg_m3_to_R !< EOS factor converting kg m-3 to R [R m3 kg-1 ~> 1] real, intent(in) :: eos_C_to_degC !< EOS factor converting temperature to degC [degC C-1 ~> 1] real, intent(in) :: eos_S_to_ppt !< EOS factor converting salinity to ppt [ppt S-1 ~> 1] - real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure + real, intent(in) :: eos_RL2_T2_to_Pa !< EOS factor converting pressure !! to Pa [Pa T2 R-1 L-2 ~> 1] logical, intent(in) :: use_temperature !< If true, temperature and salinity are !! state variables (resolved host-side to avoid a device From e9cc067fbbeb6aca142b9c01490659642e413a07 Mon Sep 17 00:00:00 2001 From: Jorge Date: Mon, 27 Jul 2026 05:38:21 -0500 Subject: [PATCH 15/15] block over j --- .../vertical/MOM_kappa_shear.F90 | 167 +++++++++++------- 1 file changed, 103 insertions(+), 64 deletions(-) diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index e648a8be91..0ca70e9fe4 100644 --- a/src/parameterizations/vertical/MOM_kappa_shear.F90 +++ b/src/parameterizations/vertical/MOM_kappa_shear.F90 @@ -25,6 +25,7 @@ module MOM_kappa_shear implicit none ; private #include +#include "do_concurrent_compat.h" public Calculate_kappa_shear, Calc_kappa_shear_vertex, kappa_shear_init public kappa_shear_is_used, kappa_shear_at_vertex @@ -84,6 +85,8 @@ module MOM_kappa_shear !! iteration. The bug causes under-corrections when dz > 1m. logical :: KS_at_vertex !< If true, do the calculations of the shear-driven mixing !! at the cell vertices (i.e., the vorticity points). + integer :: njblock !< The J-direction block size used for the vertex scratch slabs + !! in Calc_kappa_shear_vertex [nondim]. logical :: eliminate_massless !< If true, massless layers are merged with neighboring !! massive layers in this calculation. ! I can think of no good reason why this should be false. - RWH @@ -499,12 +502,18 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ h_at_u ! A mask-weighted thickness interpolated to u-points [H ~> m or kg m-2] real, dimension(SZI_(G),SZJB_(G),SZK_(GV)) :: & h_at_v ! A mask-weighted thickness interpolated to v-points [H ~> m or kg m-2] - real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)) :: & + ! The vertex scratch slabs below are blocked in J: their second extent is the block size, not + ! the full J range, and they are indexed by the block-local jj = J - Jstart + 1. The full J + ! extent at vertices is JeB-JsB+1 = G%JecB-(G%jsc-1)+1; CS%njblock==0 selects it, which collapses + ! the block loop to a single pass over the whole domain (the GPU default). A block size of 1 + ! (the CPU default) reduces these to one vertex row at a time, so the scratch stays cache + ! resident as it did before these were promoted to full 3-D arrays for offload. + real, dimension(SZIB_(G),merge(G%JecB-G%jsc+2,CS%njblock,CS%njblock==0),SZK_(GV)) :: & h_slab, & ! A version of h interpolated to vertices [H ~> m or kg m-2]. dz_slab, & ! Vertical distance between interface heights at vertices [Z ~> m]. u_slab, v_slab, & ! Versions of u_in and v_in interpolated to vertices [L T-1 ~> m s-1]. T_slab, S_slab, rho_slab ! Vertex versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. - real, dimension(SZIB_(G),SZJB_(G),SZK_(GV)+1) :: & + real, dimension(SZIB_(G),merge(G%JecB-G%jsc+2,CS%njblock,CS%njblock==0),SZK_(GV)+1) :: & kappa_3d, & ! Device staging array for the columns' averaged kappa [H Z T-1 ~> m2 s-1 or Pa s] tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. real, dimension(SZIB_(G),SZJB_(G)) :: & @@ -569,11 +578,13 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ real :: H_tiny ! A sub-roundoff thickness to use in the denominator when calculating ! thickness-weighted averages [H ~> m or kg m-2] integer :: IsB, IeB, JsB, JeB, i, j, k, nz, nzc + integer :: Jstart, Jend, njblock, jj ! J-blocking bounds, block size and block-local J index. integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa ! EOS unit-rescaling factors. ! Diagnostics that should be deleted? isB = G%isc-1 ; ieB = G%iecB ; jsB = G%jsc-1 ; jeB = G%jecB ; nz = GV%ke + njblock = merge(JeB-JsB+1, CS%njblock, CS%njblock==0) if ((CS%id_N2_init>0) .or. CS%debug) diag_N2_init(:,:,:) = 0.0 if ((CS%id_S2_init>0) .or. CS%debug) diag_S2_init(:,:,:) = 0.0 @@ -651,36 +662,49 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ !$omp target enter data map(to: u_in, v_in, T_in, S_in, dz_3d) !$omp target update to(u_in, v_in) !$omp target enter data map(alloc: u_slab, v_slab, T_slab, S_slab, h_slab, dz_slab, rho_slab) + !$omp target enter data map(to: CS) + !$omp target enter data map(to: surface_pres_2d) + !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) + !$omp target update to(tke_io, kv_io) + !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target enter data map(alloc: kappa_3d, tke_3d) + + ! The scratch slabs are allocated on the device once, at the block size, and reused by every + ! block; only the loop bounds below change from one block to the next. + do Jstart=JsB,JeB,njblock + Jend = min(Jstart+njblock-1, JeB) ! Interpolate the various quantities to the corners, using masks. - do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) - u_slab(I,J,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & + do concurrent (k=1:nz, J=Jstart:Jend, I=IsB:IeB) DO_LOCALITY(local(jj)) + jj = J - Jstart + 1 + u_slab(I,jj,k) = ( (u_in(I,j,k) * h_at_u(I,j,k)) + (u_in(I,j+1,k) * h_at_u(I,j+1,k)) ) / & ( (h_at_u(I,j,k) + h_at_u(I,j+1,k)) + H_tiny ) - v_slab(I,J,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & + v_slab(I,jj,k) = ( (v_in(i,J,k) * h_at_v(i,J,k)) + (v_in(i+1,J,k) * h_at_v(i+1,J,k)) ) / & ( (h_at_v(i,J,k) + h_at_v(i+1,J,k)) + H_tiny ) - h_slab(I,J,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & + h_slab(I,jj,k) = ((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k)) ) / & ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) - dz_slab(I,J,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & + dz_slab(I,jj,k) = ((G%mask2dT(i,j) * dz_3d(i,j,k) + G%mask2dT(i+1,j+1) * dz_3d(i+1,j+1,k)) + & (G%mask2dT(i+1,j) * dz_3d(i+1,j,k) + G%mask2dT(i,j+1) * dz_3d(i,j+1,k)) ) / & ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + & (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) + 1.0e-36 ) -! h_slab(I,J,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) -! h_slab(I,J,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & +! h_slab(I,jj,k) = 0.25*((h(i,j,k) + h(i+1,j+1,k)) + (h(i+1,j,k) + h(i,j+1,k))) +! h_slab(I,jj,k) = (((h(i,j,k)**2) + (h(i+1,j+1,k)**2)) + & ! ((h(i+1,j,k)**2) + (h(i,j+1,k)**2))) * I_hwt enddo if (use_temperature) then - do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) - T_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & + do concurrent (k=1:nz, J=Jstart:Jend, I=IsB:IeB) DO_LOCALITY(local(jj)) + jj = J - Jstart + 1 + T_slab(I,jj,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * T_in(i,j,k)) + & G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * T_in(i+1,j+1,k))) + & (G%mask2dT(i+1,j) * (h(i+1,j,k) * T_in(i+1,j,k)) + & G%mask2dT(i,j+1) * (h(i,j+1,k) * T_in(i,j+1,k))) ) * & (1.0 / (((G%mask2dT(i,j) * h(i,j,k) + G%mask2dT(i+1,j+1) * h(i+1,j+1,k)) + & (G%mask2dT(i+1,j) * h(i+1,j,k) + G%mask2dT(i,j+1) * h(i,j+1,k))) + & GV%H_subroundoff)) - S_slab(I,J,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & + S_slab(I,jj,k) = ( (G%mask2dT(i,j) * (h(i,j,k) * S_in(i,j,k)) + & G%mask2dT(i+1,j+1) * (h(i+1,j+1,k) * S_in(i+1,j+1,k))) + & (G%mask2dT(i+1,j) * (h(i+1,j,k) * S_in(i+1,j,k)) + & G%mask2dT(i,j+1) * (h(i,j+1,k) * S_in(i,j+1,k))) ) * & @@ -689,29 +713,23 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ GV%H_subroundoff)) enddo else - do concurrent (k=1:nz, J=JsB:JeB, I=IsB:IeB) - rho_slab(I,J,k) = GV%Rlay(k) + do concurrent (k=1:nz, J=Jstart:Jend, I=IsB:IeB) DO_LOCALITY(local(jj)) + jj = J - Jstart + 1 + rho_slab(I,jj,k) = GV%Rlay(k) enddo endif - !$omp target enter data map(to: CS) - !$omp target enter data map(to: surface_pres_2d) - !$omp target enter data map(to: kappa_vertex, tke_io, kv_io) - !$omp target update to(tke_io, kv_io) - !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) - !$omp target enter data map(alloc: kappa_3d, tke_3d) - !--------------------------------------- ! Work on each column. !--------------------------------------- !$omp target teams loop collapse(2) & !$omp private(nzc, kc, kf, Idz, h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, dz_in_lay, & !$omp f2, surface_pres, kappa, tke, kappa_avg, tke_avg, N2_init, S2_init, & - !$omp N2_mean, S2_mean, k) & + !$omp N2_mean, S2_mean, k, jj) & !$omp firstprivate(nz, dt, k0dt, dz_massless, use_temperature, eos_form, & - !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) - do J=JsB,JeB - do I=IsB,IeB ; if ((G%mask2dCu(I,j) + G%mask2dCu(I,j+1)) + & + !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, Jstart) + do J=Jstart,Jend + do I=IsB,IeB ; jj = J - Jstart + 1 ; if ((G%mask2dCu(I,j) + G%mask2dCu(I,j+1)) + & (G%mask2dCv(i,J) + G%mask2dCv(i+1,J)) > 0.0) then ! call cpu_clock_begin(Id_clock_setup) ! Store a transposed version of the initial arrays. @@ -724,25 +742,25 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ T0xdz(k) = 0.0 ; S0xdz(k) = 0.0 ! Add a new layer if this one has mass. -! if ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 +! if ((h_lay(nzc) > 0.0) .and. (h_slab(I,jj,k) > dz_massless)) nzc = nzc+1 if ((k>CS%nkml) .and. (h_lay(nzc) > 0.0) .and. & - (h_slab(I,J,k) > dz_massless)) nzc = nzc+1 + (h_slab(I,jj,k) > dz_massless)) nzc = nzc+1 ! Only merge clusters of massless layers. ! if ((h_lay(nzc) > dz_massless) .or. & -! ((h_lay(nzc) > 0.0) .and. (h_slab(I,J,k) > dz_massless))) nzc = nzc+1 +! ((h_lay(nzc) > 0.0) .and. (h_slab(I,jj,k) > dz_massless))) nzc = nzc+1 kc(k) = nzc - h_lay(nzc) = h_lay(nzc) + h_slab(I,J,k) - dz_lay(nzc) = dz_lay(nzc) + dz_slab(I,J,k) - u0xdz(nzc) = u0xdz(nzc) + u_slab(I,J,k)*h_slab(I,J,k) - v0xdz(nzc) = v0xdz(nzc) + v_slab(I,J,k)*h_slab(I,J,k) + h_lay(nzc) = h_lay(nzc) + h_slab(I,jj,k) + dz_lay(nzc) = dz_lay(nzc) + dz_slab(I,jj,k) + u0xdz(nzc) = u0xdz(nzc) + u_slab(I,jj,k)*h_slab(I,jj,k) + v0xdz(nzc) = v0xdz(nzc) + v_slab(I,jj,k)*h_slab(I,jj,k) if (use_temperature) then - T0xdz(nzc) = T0xdz(nzc) + T_slab(I,J,k)*h_slab(I,J,k) - S0xdz(nzc) = S0xdz(nzc) + S_slab(I,J,k)*h_slab(I,J,k) + T0xdz(nzc) = T0xdz(nzc) + T_slab(I,jj,k)*h_slab(I,jj,k) + S0xdz(nzc) = S0xdz(nzc) + S_slab(I,jj,k)*h_slab(I,jj,k) else - T0xdz(nzc) = T0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) - S0xdz(nzc) = S0xdz(nzc) + rho_slab(I,J,k)*h_slab(I,J,k) + T0xdz(nzc) = T0xdz(nzc) + rho_slab(I,jj,k)*h_slab(I,jj,k) + S0xdz(nzc) = S0xdz(nzc) + rho_slab(I,jj,k)*h_slab(I,jj,k) endif enddo kc(nz+1) = nzc+1 @@ -752,28 +770,28 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Now determine kf, the fractional weight of interface kc when ! interpolating between interfaces kc and kc+1. - kf(1) = 0.0 ; dz_in_lay = h_slab(I,J,1) + kf(1) = 0.0 ; dz_in_lay = h_slab(I,jj,1) do k=2,nz if (kc(k) > kc(k-1)) then - kf(k) = 0.0 ; dz_in_lay = h_slab(I,J,k) + kf(k) = 0.0 ; dz_in_lay = h_slab(I,jj,k) else - kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_slab(I,J,k) + kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_slab(I,jj,k) endif enddo kf(nz+1) = 0.0 else do k=1,nz - h_lay(k) = h_slab(I,J,k) - dz_lay(k) = dz_slab(I,J,k) - u0xdz(k) = u_slab(I,J,k)*h_lay(k) ; v0xdz(k) = v_slab(I,J,k)*h_lay(k) + h_lay(k) = h_slab(I,jj,k) + dz_lay(k) = dz_slab(I,jj,k) + u0xdz(k) = u_slab(I,jj,k)*h_lay(k) ; v0xdz(k) = v_slab(I,jj,k)*h_lay(k) enddo if (use_temperature) then do k=1,nz - T0xdz(k) = T_slab(I,J,k)*h_lay(k) ; S0xdz(k) = S_slab(I,J,k)*h_lay(k) + T0xdz(k) = T_slab(I,jj,k)*h_lay(k) ; S0xdz(k) = S_slab(I,jj,k)*h_lay(k) enddo else do k=1,nz - T0xdz(k) = rho_slab(I,J,k)*h_lay(k) ; S0xdz(k) = rho_slab(I,J,k)*h_lay(k) + T0xdz(k) = rho_slab(I,jj,k)*h_lay(k) ; S0xdz(k) = rho_slab(I,jj,k)*h_lay(k) enddo endif nzc = nz @@ -797,11 +815,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_3d(I,J,K) = kappa_avg(K) + kappa_3d(I,jj,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_3d(I,J,K) = tke(K) + tke_3d(I,jj,K) = tke(K) else - tke_3d(I,J,K) = tke_avg(K) + tke_3d(I,jj,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -819,11 +837,11 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ else do K=1,nz+1 if (kf(K) == 0.0) then - kappa_3d(I,J,K) = kappa_avg(kc(K)) - tke_3d(I,J,K) = tke_avg(kc(K)) + kappa_3d(I,jj,K) = kappa_avg(kc(K)) + tke_3d(I,jj,K) = tke_avg(kc(K)) else - kappa_3d(I,J,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) - tke_3d(I,J,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) + kappa_3d(I,jj,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) + tke_3d(I,jj,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) endif enddo do K=1,nz+1 @@ -847,37 +865,43 @@ subroutine Calc_kappa_shear_vertex(u_in, v_in, h, T_in, S_in, tv, p_surf, kappa_ ! call cpu_clock_end(Id_clock_setup) else ! Land points, still inside the i-loop. do K=1,nz+1 - kappa_3d(I,J,K) = 0.0 ; tke_3d(I,J,K) = 0.0 + kappa_3d(I,jj,K) = 0.0 ; tke_3d(I,jj,K) = 0.0 enddo endif ; enddo ! i-loop enddo ! end of J-loop - ! Store the columns' results back in the 3-d arrays for restarts or interpolation back to - ! tracer points. h_vert is only used when VS_ThicknessMean is true, and only by the (host) - ! tracer-point averaging below, so its fill stays a host loop fed by a guarded copy-back of - ! the device-computed h_slab. + ! Store this block's columns back into the full-J 3-d arrays, for restarts or interpolation + ! back to tracer points. These have to happen inside the block loop because the slabs they + ! read are only valid for the block that just ran. h_vert is only used when VS_ThicknessMean + ! is true, and only by the (host) tracer-point averaging below, so its fill stays a host loop + ! fed by a guarded copy-back of the device-computed h_slab. if (CS%VS_ThicknessMean) then !$omp target update from(h_slab) - do J=JsB,JeB + do J=Jstart,Jend + jj = J - Jstart + 1 do K=1,nz+1 ; do I=IsB,IeB - h_vert(I,J,k) = h_slab(I,J,k) + h_vert(I,J,k) = h_slab(I,jj,k) enddo ; enddo enddo endif if (CS%VS_viscosity_bug) then - do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) - kappa_vertex(I,J,K) = kappa_3d(I,J,K) - tke_io(I,J,K) = G%mask2dBu(I,J) * tke_3d(I,J,K) + do concurrent (K=1:nz+1, J=Jstart:Jend, I=IsB:IeB) DO_LOCALITY(local(jj)) + jj = J - Jstart + 1 + kappa_vertex(I,J,K) = kappa_3d(I,jj,K) + tke_io(I,J,K) = G%mask2dBu(I,J) * tke_3d(I,jj,K) kv_io(I,J,K) = ( G%mask2dBu(I,J) * kappa_vertex(I,J,K) ) * CS%Prandtl_turb enddo else - do concurrent (K=1:nz+1, J=JsB:JeB, I=IsB:IeB) - kappa_vertex(I,J,K) = kappa_3d(I,J,K) - tke_io(I,J,K) = tke_3d(I,J,K) + do concurrent (K=1:nz+1, J=Jstart:Jend, I=IsB:IeB) DO_LOCALITY(local(jj)) + jj = J - Jstart + 1 + kappa_vertex(I,J,K) = kappa_3d(I,jj,K) + tke_io(I,J,K) = tke_3d(I,jj,K) kv_io(I,J,K) = kappa_vertex(I,J,K) * CS%Prandtl_turb enddo endif + enddo ! end of Jstart block loop. + ! The vertex-to-tracer-point averaging below, the checksums and post_data are all still on ! the host. The diag_* transfer is guarded by the same conditions as their consumers. !$omp target update from(kappa_vertex, tke_io, kv_io) @@ -2284,6 +2308,11 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) ! This include declares and sets the variable "version". # include "version_variable.h" character(len=40) :: mdl = "MOM_kappa_shear" ! This module's name. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, parameter :: default_njblock = 0 +#else + integer, parameter :: default_njblock = 1 +#endif if (associated(CS)) then call MOM_error(WARNING, "kappa_shear_init called with an associated "// & @@ -2453,6 +2482,16 @@ function kappa_shear_init(Time, G, GV, US, param_file, diag, CS) "at the corner if VERTEX_SHEAR=True. Otherwise mask out any land points in "//& "the average.", default=.false., do_not_log=(just_read .or. (.not.CS%KS_at_vertex))) + call get_param(param_file, mdl, "KAPPA_SHEAR_NJBLOCK", CS%njblock, & + "The J-direction block size used for the vertex scratch slabs in the "//& + "shear-driven mixing calculations when VERTEX_SHEAR=True. The default 0 "//& + "setting dynamically uses the full J extent, which maximizes the parallelism "//& + "available to a GPU but makes the scratch working set too large for a CPU "//& + "cache; the default 1 setting on CPU builds restores row-at-a-time reuse.", & + default=default_njblock, layoutParam=.true., & + do_not_log=(just_read .or. (.not.CS%KS_at_vertex))) + if (CS%njblock < 0) call MOM_error(FATAL, "KAPPA_SHEAR_NJBLOCK must be >= 0.") + call get_param(param_file, mdl, "KAPPA_SHEAR_ITER_BUG", CS%dKdQ_iteration_bug, & "If true, use an older, dimensionally inconsistent estimate of the "//& "derivative of diffusivity with energy in the Newton's method iteration. "//&