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 diff --git a/src/ALE/MOM_regridding.F90 b/src/ALE/MOM_regridding.F90 index 59ec8d4d7d..e489bc69c0 100644 --- a/src/ALE/MOM_regridding.F90 +++ b/src/ALE/MOM_regridding.F90 @@ -32,6 +32,7 @@ module MOM_regridding use coord_zlike, only : zlike_CS use coord_zlike, only : init_coord_zlike, set_zlike_params, build_zstar_column, end_coord_zlike +use coord_zlike, only : build_zstar_column_loc, get_zlike_coord_res use coord_sigma, only : sigma_CS use coord_sigma, only : init_coord_sigma, set_sigma_params, build_sigma_column, end_coord_sigma use coord_rho, only : init_coord_rho, rho_CS, set_rho_params, build_rho_column, end_coord_rho @@ -48,6 +49,13 @@ module MOM_regridding #include +!> On GPU builds this fixes the per-column private scratch in the regrid column loops +!! (build_zstar_grid's zNew/zOld) to a compile-time-constant size, so each device thread gets +!! stack (local memory) arrays instead of runtime-sized device-heap allocations, which nvfortran +!! cannot place in local memory (NVFORTRAN-W-0155). Checked against CS%nk / GV%ke in +!! build_zstar_grid. Unused in CPU builds, where the declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 + ! A note on unit descriptions in comments: MOM6 uses units that can be rescaled for dimensional ! consistency testing. These are noted in comments with units like Z, H, L, and T, along with ! their mks counterparts with notation like "a velocity [Z T-1 ~> m s-1]". If the units @@ -1465,6 +1473,28 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) real, dimension(CS%nk+1), intent(in) :: z_new !< New grid position before filtering [H ~> m or kg m-2] real, dimension(CS%nk+1), intent(inout) :: dz_g !< Change in interface positions including !! the effects of filtering [H ~> m or kg m-2] + + call filtered_grid_motion_loc( CS%nk, CS%depth_of_time_filter_shallow, CS%depth_of_time_filter_deep, & + CS%old_grid_weight, nk, z_old, z_new, dz_g ) + +end subroutine filtered_grid_motion + +!> Element-wise, control-structure-free form of filtered_grid_motion: it takes the (host-resolved) +!! target level count and filter parameters explicitly so it can be made device-callable, avoiding a +!! dereference of regridding_CS inside a device region. The arithmetic is verbatim. +subroutine filtered_grid_motion_loc( nk_tgt, filt_shallow, filt_deep, old_grid_wt, nk, z_old, z_new, dz_g ) + integer, intent(in) :: nk_tgt !< Number of cells in the target grid (was CS%nk) + real, intent(in) :: filt_shallow !< Depth at which the shallow filter + !! timescale applies [H ~> m or kg m-2] (was CS%depth_of_time_filter_shallow) + real, intent(in) :: filt_deep !< Depth at which the deep filter timescale + !! applies [H ~> m or kg m-2] (was CS%depth_of_time_filter_deep) + real, intent(in) :: old_grid_wt !< Weight given to the old grid when + !! filtering [nondim] (was CS%old_grid_weight) + integer, intent(in) :: nk !< Number of cells in source grid + real, dimension(nk+1), intent(in) :: z_old !< Old grid position [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(in) :: z_new !< New grid position before filtering [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(inout) :: dz_g !< Change in interface positions including + !! the effects of filtering [H ~> m or kg m-2] ! Local variables real :: sgn ! The sign convention for downward [nondim]. real :: dz_tgt ! The target grid movement of the unfiltered grid [H ~> m or kg m-2] @@ -1486,24 +1516,37 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) real :: Int_zd ! A depth integral of the weights in [H ~> m or kg m-2] real :: dInt_zs_zd ! The depth integral of the weights between the deep and shallow depths in [H ~> m or kg m-2] ! For debugging: +#ifndef __NVCOMPILER_OPENMP_GPU real, dimension(nk+1) :: z_act ! The final grid positions after the filtered movement [H ~> m or kg m-2] ! real, dimension(nk+1) :: ddz_g_s, ddz_g_d logical :: debug = .false. +#endif + logical :: massless ! True for a massless column, in which case no grid motion is applied. integer :: k - - if ((z_old(nk+1) - z_old(1)) * (z_new(CS%nk+1) - z_new(1)) < 0.0) then + !$omp declare target + + ! GPU port: the massless-column early return is expressed as an if-guard (device loops cannot + ! return early), and the "should never happen" sign/tangling MOM_error asserts are compiled out + ! of device builds (kept on the host). + massless = .false. ; sgn = 1.0 + if ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) < 0.0) then +#ifndef __NVCOMPILER_OPENMP_GPU call MOM_error(FATAL, "filtered_grid_motion: z_old and z_new use different sign conventions.") - elseif ((z_old(nk+1) - z_old(1)) * (z_new(CS%nk+1) - z_new(1)) == 0.0) then - ! This is a massless column, so do nothing and return. - do k=1,CS%nk+1 ; dz_g(k) = 0.0 ; enddo ; return - elseif ((z_old(nk+1) - z_old(1)) + (z_new(CS%nk+1) - z_new(1)) > 0.0) then +#endif + elseif ((z_old(nk+1) - z_old(1)) * (z_new(nk_tgt+1) - z_new(1)) == 0.0) then + ! This is a massless column, so do nothing (the filtering block below is skipped). + massless = .true. ; do k=1,nk_tgt+1 ; dz_g(k) = 0.0 ; enddo + elseif ((z_old(nk+1) - z_old(1)) + (z_new(nk_tgt+1) - z_new(1)) > 0.0) then sgn = 1.0 else sgn = -1.0 endif + if (.not. massless) then + +#ifndef __NVCOMPILER_OPENMP_GPU if (debug) then - do k=2,CS%nk+1 + do k=2,nk_tgt+1 if (sgn*(z_new(k)-z_new(k-1)) < -5e-16*(abs(z_new(k))+abs(z_new(k-1))) ) & call MOM_error(FATAL, "filtered_grid_motion: z_new is tangled.") enddo @@ -1513,10 +1556,11 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) enddo ! ddz_g_s(:) = 0.0 ; ddz_g_d(:) = 0.0 endif +#endif - zs = CS%depth_of_time_filter_shallow - zd = CS%depth_of_time_filter_deep - wtd = 1.0 - CS%old_grid_weight + zs = filt_shallow + zd = filt_deep + wtd = 1.0 - old_grid_wt Iwtd = 1.0 / wtd dzwt = (zd - zs) @@ -1526,7 +1570,7 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) dz_g(1) = 0.0 z_old_k = z_old(1) - do k = 2,CS%nk+1 + do k = 2,nk_tgt+1 if (k<=nk+1) z_old_k = z_old(k) ! This allows for virtual z_old interface at bottom of the model ! zr1 is positive and increases with depth, and dz_tgt is positive downward. dz_tgt = sgn*(z_new(k) - z_old_k) @@ -1593,19 +1637,23 @@ subroutine filtered_grid_motion( CS, nk, z_old, z_new, dz_g ) enddo !dz_g(CS%nk+1) = 0.0 +#ifndef __NVCOMPILER_OPENMP_GPU if (debug) then z_old_k = z_old(1) - do k=1,CS%nk+1 + do k=1,nk_tgt+1 if (k<=nk+1) z_old_k = z_old(k) ! This allows for virtual z_old interface at bottom of the model z_act(k) = z_old_k + dz_g(k) enddo - do k=2,CS%nk+1 + do k=2,nk_tgt+1 if (sgn*((z_act(k))-z_act(k-1)) < -1e-15*(abs(z_act(k))+abs(z_act(k-1))) ) & call MOM_error(FATAL, "filtered_grid_motion: z_output is tangled.") enddo endif +#endif -end subroutine filtered_grid_motion + endif ! .not. massless + +end subroutine filtered_grid_motion_loc !> Builds a z*-coordinate grid with partial steps (Adcroft and Campin, 2004). !! z* is defined as @@ -1634,8 +1682,21 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ #ifdef __DO_SAFETY_CHECKS__ real :: dh ! The larger of the total column thickness or bathymetric depth [H ~> m or kg m-2] #endif + ! GPU port: per-column private scratch gets compile-time-constant sizes on GPU builds so each + ! device thread uses stack (local memory) arrays rather than runtime-sized device-heap allocations. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: zOld ! Previous coordinate interface heights [H ~> m or kg m-2] + real, dimension(GPU_nk_max+1) :: zNew ! New coordinate interface heights [H ~> m or kg m-2] +#else real, dimension(SZK_(GV)+1) :: zOld ! Previous coordinate interface heights [H ~> m or kg m-2] real, dimension(CS%nk+1) :: zNew ! New coordinate interface heights [H ~> m or kg m-2] +#endif + ! GPU port: the regridding_CS scalars + coordinateResolution are hoisted into plain locals so the + ! device column loop never dereferences the nested (pointer + allocatable) control structure. + real, dimension(CS%nk) :: coordRes ! A local copy of CS%zlike_CS%coordinateResolution [Z ~> m] + real :: min_thick, filt_shallow, filt_deep, old_grid_wt, z_scale_l ! Hoisted CS scalars + integer :: nk_tgt ! A local copy of CS%nk (target level count) + logical :: do_filter, use_adjust ! Hoisted CS control flags integer :: i, j, k, nz logical :: ice_shelf @@ -1643,20 +1704,35 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ minThickness = CS%min_thickness ice_shelf = present(frac_shelf_h) - !$OMP parallel do default(none) shared(G,GV,dzInterface,CS,nz,h,frac_shelf_h, & - !$OMP ice_shelf,minThickness,zScale,nom_depth_H) & - !$OMP private(nominalDepth,totalThickness, & -#ifdef __DO_SAFETY_CHECKS__ - !$OMP dh, & + ! Hoist the control-structure reads onto the host; the device column loop reads only plain locals. + nk_tgt = CS%nk + min_thick = CS%min_thickness + filt_shallow = CS%depth_of_time_filter_shallow + filt_deep = CS%depth_of_time_filter_deep + old_grid_wt = CS%old_grid_weight + do_filter = CS%use_depth_based_time_filter .or. CS%old_grid_weight > 0. + use_adjust = CS%use_adjust_interface_motion + z_scale_l = 1.0 ; if (present(zScale)) z_scale_l = zScale + call get_zlike_coord_res(CS%zlike_CS, coordRes) +#ifdef __NVCOMPILER_OPENMP_GPU + if (nk_tgt > GPU_nk_max .or. nz > GPU_nk_max) call MOM_error(FATAL, & + "build_zstar_grid: GPU builds require CS%nk and GV%ke <= GPU_nk_max; increase GPU_nk_max "//& + "in MOM_regridding.F90.") #endif - !$OMP zNew,zOld) + + ! GPU port: run the per-column regrid on the device. h is host-authoritative in the (host) ALE + ! stack -> refresh; nom_depth_H/coordRes are fresh host locals; dzInterface is device-written and + ! copied back for the (host) remap that follows. + !$omp target update to(h) + !$omp target enter data map(to: nom_depth_H, coordRes) map(alloc: dzInterface) + !$omp target teams loop collapse(2) & + !$omp private(nominalDepth, totalThickness, zNew, zOld, k) & + !$omp firstprivate(nz, nk_tgt, min_thick, filt_shallow, filt_deep, old_grid_wt, & + !$omp do_filter, use_adjust, ice_shelf, z_scale_l) do j = G%jsc-1,G%jec+1 do i = G%isc-1,G%iec+1 - if (G%mask2dT(i,j)==0.) then - dzInterface(i,j,:) = 0. - cycle - endif + if (G%mask2dT(i,j)/=0.) then ! Local depth (positive downward) nominalDepth = nom_depth_H(i,j) @@ -1677,21 +1753,21 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ if (ice_shelf) then if (frac_shelf_h(i,j) > 0.) then ! under ice shelf - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, zNew, & - z_rigid_top=totalThickness-nominalDepth, & - eta_orig=zOld(1), zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, zNew, & + z_scale_l, totalThickness-nominalDepth, zOld(1)) else - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, & - zNew, zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, & + zNew, z_scale_l) endif else - call build_zstar_column(CS%zlike_CS, nominalDepth, totalThickness, & - zNew, zScale=zScale) + call build_zstar_column_loc(nk_tgt, min_thick, coordRes, nominalDepth, totalThickness, & + zNew, z_scale_l) endif ! Calculate the final change in grid position after blending new and old grids - if (CS%use_depth_based_time_filter .or. CS%old_grid_weight>0.) & - call filtered_grid_motion(CS, nz, zOld, zNew, dzInterface(i,j,:)) + if (do_filter) & + call filtered_grid_motion_loc(nk_tgt, filt_shallow, filt_deep, old_grid_wt, nz, zOld, zNew, & + dzInterface(i,j,:)) #ifdef __DO_SAFETY_CHECKS__ dh = max(nominalDepth,totalThickness) @@ -1716,11 +1792,18 @@ subroutine build_zstar_grid( CS, G, GV, h, nom_depth_H, dzInterface, frac_shelf_ endif #endif - if (CS%use_adjust_interface_motion) call adjust_interface_motion( CS, nz, h(i,j,:), dzInterface(i,j,:) ) + if (use_adjust) call adjust_interface_motion_loc( nk_tgt, min_thick, nz, h(i,j,:), dzInterface(i,j,:) ) + + else + dzInterface(i,j,:) = 0. + endif ! G%mask2dT(i,j) /= 0. enddo enddo + !$omp target update from(dzInterface) + !$omp target exit data map(release: nom_depth_H, coordRes, dzInterface) + end subroutine build_zstar_grid !------------------------------------------------------------------------------ @@ -2164,6 +2247,21 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) real, dimension(nk), intent(in) :: h_old !< Layer thicknesses on the old grid [H ~> m or kg m-2] real, dimension(CS%nk+1), intent(inout) :: dz_int !< Interface movements, adjusted to keep the thicknesses !! thicker than their minimum value [H ~> m or kg m-2] + + call adjust_interface_motion_loc( CS%nk, CS%min_thickness, nk, h_old, dz_int ) + +end subroutine adjust_interface_motion + +!> Element-wise, control-structure-free form of adjust_interface_motion: it takes the (host-resolved) +!! target level count and minimum thickness explicitly so it can be made device-callable. Verbatim. +subroutine adjust_interface_motion_loc( nk_tgt, min_thick, nk, h_old, dz_int ) + integer, intent(in) :: nk_tgt !< Number of cells in the target grid (was CS%nk) + real, intent(in) :: min_thick !< Minimum allowed thickness [H ~> m or kg m-2] + !! (was CS%min_thickness) + integer, intent(in) :: nk !< Number of layers in h_old + real, dimension(nk), intent(in) :: h_old !< Layer thicknesses on the old grid [H ~> m or kg m-2] + real, dimension(nk_tgt+1), intent(inout) :: dz_int !< Interface movements, adjusted to keep the thicknesses + !! thicker than their minimum value [H ~> m or kg m-2] ! Local variables real :: h_new ! A layer thickness on the new grid [H ~> m or kg m-2] real :: eps ! A tiny relative thickness [nondim] @@ -2171,14 +2269,16 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) real :: h_err ! An error tolerance that use used to flag unacceptably large negative layer thicknesses ! that can not be explained by roundoff errors [H ~> m or kg m-2] integer :: k + !$omp declare target eps = 1. ; eps = epsilon(eps) h_total = 0. ; h_err = 0. - do k = 1, min(CS%nk,nk) + do k = 1, min(nk_tgt,nk) h_total = h_total + h_old(k) h_err = h_err + max( h_old(k), abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) +#ifndef __NVCOMPILER_OPENMP_GPU if (h_new < -3.0*h_err) then write(0,*) 'h<0 at k=',k,'h_old=',h_old(k), & 'wup=',dz_int(k),'wdn=',dz_int(k+1),'dw_dz=',dz_int(k) - dz_int(k+1), & @@ -2186,11 +2286,13 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) call MOM_error( FATAL, 'MOM_regridding: adjust_interface_motion() - '//& 'implied h<0 is larger than roundoff!') endif +#endif enddo - if (CS%nk>nk) then - do k = nk+1, CS%nk + if (nk_tgt>nk) then + do k = nk+1, nk_tgt h_err = h_err + max( abs(dz_int(k)), abs(dz_int(k+1)) )*eps h_new = ( dz_int(k) - dz_int(k+1) ) +#ifndef __NVCOMPILER_OPENMP_GPU if (h_new < -3.0*h_err) then write(0,*) 'h<0 at k=',k,'h_old was empty',& 'wup=',dz_int(k),'wdn=',dz_int(k+1),'dw_dz=',dz_int(k) - dz_int(k+1), & @@ -2198,16 +2300,18 @@ subroutine adjust_interface_motion( CS, nk, h_old, dz_int ) call MOM_error( FATAL, 'MOM_regridding: adjust_interface_motion() - '//& 'implied h<0 is larger than roundoff!') endif +#endif enddo endif - do k = min(CS%nk,nk),2,-1 + do k = min(nk_tgt,nk),2,-1 h_new = h_old(k) + ( dz_int(k) - dz_int(k+1) ) - if (h_new Builds a z* coordinate with a minimum thickness +!> Builds a z* coordinate with a minimum thickness (thin wrapper preserving the CS-taking API; +!! the body lives in the device-callable build_zstar_column_loc, GPU port increment 1). subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & z_rigid_top, eta_orig, zScale) type(zlike_CS), intent(in) :: CS !< Coordinate control structure @@ -78,17 +80,56 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & real, optional, intent(in) :: zScale !< Scaling factor from the target coordinate resolution !! in Z to desired units for zInterface, perhaps Z_to_H, !! often [nondim] or [H Z-1 ~> 1 or kg m-3] + real :: z_scale_w ! zScale with its default resolved on the host + + z_scale_w = 1.0 ; if (present(zScale)) z_scale_w = zScale + call build_zstar_column_loc(CS%nk, CS%min_thickness, CS%coordinateResolution, depth, & + total_thickness, zInterface, z_scale_w, z_rigid_top, eta_orig) + +end subroutine build_zstar_column + +!> Copy the target coordinate resolution out of a zlike_CS (a private component), for callers that +!! need it as a plain array -- e.g. to hoist it onto a device region without dereferencing the CS. +subroutine get_zlike_coord_res(CS, coordinateResolution) + type(zlike_CS), intent(in) :: CS !< Coordinate control structure + real, dimension(:), intent(out) :: coordinateResolution !< Target coordinate resolution [Z ~> m] + integer :: k + do k=1,CS%nk ; coordinateResolution(k) = CS%coordinateResolution(k) ; enddo +end subroutine get_zlike_coord_res + +!> Builds a z* coordinate with a minimum thickness. This is the element-wise, control-structure-free +!! form of build_zstar_column: it takes the (host-resolved) scalar and array parameters explicitly so +!! it can be made device-callable, avoiding a dereference of the nested (pointer + allocatable) zlike_CS +!! inside a device region. The arithmetic is verbatim from the former build_zstar_column body. +subroutine build_zstar_column_loc(nk, min_thick, coordinateResolution, depth, total_thickness, & + zInterface, z_scale, z_rigid_top, eta_orig) + integer, intent(in) :: nk !< Number of levels to be generated + real, intent(in) :: min_thick !< Minimum thickness allowed for layers, in the + !! same units as depth [Z ~> m] or [H ~> m or kg m-2] + real, dimension(nk), intent(in) :: coordinateResolution !< Target coordinate resolution [Z ~> m] + real, intent(in) :: depth !< Depth of ocean bottom (positive downward in the + !! output units), units may be [Z ~> m] or [H ~> m or kg m-2] + real, intent(in) :: total_thickness !< Column thickness (positive definite in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, dimension(nk+1), intent(inout) :: zInterface !< Absolute positions of interfaces (in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, intent(in) :: z_scale !< Scaling factor from the target coordinate resolution + !! in Z to desired units for zInterface (was optional zScale; + !! required here so device call sites need no keyword args) + real, optional, intent(in) :: z_rigid_top !< The height of a rigid top (positive upward in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] + real, optional, intent(in) :: eta_orig !< The actual original height of the top (in the same + !! units as depth) [Z ~> m] or [H ~> m or kg m-2] ! Local variables real :: eta ! Free surface height [Z ~> m] or [H ~> m or kg m-2] real :: stretching ! A stretching factor for the coordinate [nondim] - real :: dh, min_thickness, z0_top, z_star, z_scale ! Thicknesses or heights [Z ~> m] or [H ~> m or kg m-2] + real :: dh, min_thickness, z0_top, z_star ! Thicknesses or heights [Z ~> m] or [H ~> m or kg m-2] integer :: k logical :: new_zstar_def - - z_scale = 1.0 ; if (present(zScale)) z_scale = zScale + !$omp declare target new_zstar_def = .false. - min_thickness = min( CS%min_thickness, total_thickness/real(CS%nk) ) + min_thickness = min( min_thick, total_thickness/real(nk) ) z0_top = 0. if (present(z_rigid_top)) then z0_top = z_rigid_top @@ -112,37 +153,37 @@ subroutine build_zstar_column(CS, depth, total_thickness, zInterface, & ! z_star is the notional z* coordinate in absence of upper/lower topography z_star = 0. ! z*=0 at the free-surface zInterface(1) = eta ! The actual position of the top of the column - do k = 2,CS%nk - z_star = z_star - CS%coordinateResolution(k-1)*z_scale + do k = 2,nk + z_star = z_star - coordinateResolution(k-1)*z_scale ! This ensures that z is below a rigid upper surface (ice shelf bottom) zInterface(k) = min( eta + stretching * ( z_star - z0_top ), z0_top ) ! This ensures that the layer in inflated zInterface(k) = min( zInterface(k), zInterface(k-1) - min_thickness ) ! This ensures that z is above or at the topography - zInterface(k) = max( zInterface(k), -depth + real(CS%nk+1-k) * min_thickness ) + zInterface(k) = max( zInterface(k), -depth + real(nk+1-k) * min_thickness ) enddo - zInterface(CS%nk+1) = -depth + zInterface(nk+1) = -depth else ! Integrate down from the top for a notional new grid, ignoring topography ! The starting position is offset by z0_top which, if z0_top<0, will place ! interfaces above the rigid boundary. zInterface(1) = eta - do k = 1,CS%nk - dh = stretching * CS%coordinateResolution(k)*z_scale ! Notional grid spacing + do k = 1,nk + dh = stretching * coordinateResolution(k)*z_scale ! Notional grid spacing zInterface(k+1) = zInterface(k) - dh enddo ! Integrating up from the bottom adjusting interface position to accommodate ! inflating layers without disturbing the interface above - zInterface(CS%nk+1) = -depth - do k = CS%nk,1,-1 + zInterface(nk+1) = -depth + do k = nk,1,-1 if ( zInterface(k) < (zInterface(k+1) + min_thickness) ) then zInterface(k) = zInterface(k+1) + min_thickness endif enddo endif -end subroutine build_zstar_column +end subroutine build_zstar_column_loc end module coord_zlike diff --git a/src/core/MOM_PressureForce_FV.F90 b/src/core/MOM_PressureForce_FV.F90 index 1dcb872597..ed136a802c 100644 --- a/src/core/MOM_PressureForce_FV.F90 +++ b/src/core/MOM_PressureForce_FV.F90 @@ -19,6 +19,7 @@ module MOM_PressureForce_FV use MOM_variables, only : thermo_var_ptrs, accel_diag_ptrs use MOM_verticalGrid, only : verticalGrid_type use MOM_EOS, only : calculate_density, calculate_spec_vol, EOS_domain +use MOM_EOS, only : get_EOS_form_and_scaling use MOM_density_integrals, only : int_density_dz, int_specific_vol_dp use MOM_density_integrals, only : int_density_dz_generic_plm, int_density_dz_generic_ppm use MOM_density_integrals, only : int_spec_vol_dp_generic_plm @@ -1101,6 +1102,12 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, integer, dimension(2) :: EOSdom_u ! The i-computational domain for the equation of state at u-velocity points integer, dimension(2) :: EOSdom_v ! The i-computational domain for the equation of state at v-velocity points integer :: EOSdom2d(2,2) ! The 2D compute domain for the equation of state + logical :: idz_on_device ! If true, int_density_dz_generic_plm leaves its outputs (dpa, intz_dpa, + ! intx_dpa, inty_dpa) device-resident, so the caller neither pushes e to the + ! host for it nor pulls the integrals back (GPU builds, supported config). + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, eos_R_to_kg_m3 + ! EOS unit-rescaling factors, used only to test eos_unity for idz_on_device. + integer :: eos_form_pf ! EOS form id (from get_EOS_form_and_scaling; not otherwise used here). integer :: is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, nkmb integer :: i, j, k, m @@ -1295,15 +1302,43 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, ! Calculate 4 integrals through the layer that are required in the ! subsequent calculation. + + ! Determine whether int_density_dz_generic_plm will leave its outputs (dpa, intz_dpa, intx_dpa, + ! inty_dpa) on the device: GPU build, ALE PLM reconstruction (scheme 1 or 3), non-Stanley, + ! accurate anomaly form, and unscaled EOS units. This mirrors that routine's internal offload + ! gate so the caller can drop the host round-trip for exactly that case. It is k-independent. + idz_on_device = .false. +#ifdef __NVCOMPILER_OPENMP_GPU + if (use_EOS .and. use_ALE .and. (CS%Recon_Scheme == 1 .or. CS%Recon_Scheme == 3) .and. & + (.not. CS%use_stanley_pgf) .and. (.not. CS%use_inaccurate_pgf_rho_anom)) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form_pf, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa, R_to_kg_m3=eos_R_to_kg_m3) + idz_on_device = (eos_RL2_T2_to_Pa == 1.0) .and. (eos_R_to_kg_m3 == 1.0) .and. & + (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0) + endif +#endif + !$omp target enter data map(alloc: dpa, intx_dpa, inty_dpa, intz_dpa) if (use_EOS) then - !$omp target update from(e) if( (use_ALE .and. CS%Recon_Scheme > 0) .or. & + ! e is pulled to the host only for host-side consumers: a host-run int_density (when not + ! idz_on_device -- e.g. Recon scheme 2, Stanley, non-unity units), the MassWt diagnostics, the + ! correction/reset-intxpa reference-interface searches, or the SAL/tides "new answers" blocks + ! (which read e(:,:,1) and read-modify e(:,:,K) on the host). When idz_on_device would otherwise + ! suppress the pull, those SAL/tides host reads would see stale e, so keep the pull for them. + !$omp target update from(e) if( (use_ALE .and. CS%Recon_Scheme > 0 .and. (.not. idz_on_device)) .or. & + !$omp CS%correction_intxpa .or. CS%reset_intxpa_integral .or. & + !$omp (CS%calculate_SAL .and. CS%tides_answer_date>20250131) .or. & + !$omp (CS%tides .and. CS%tides_answer_date>20250131) .or. & !$omp (CS%id_MassWt_u > 0) .or. (CS%id_MassWt_v > 0)) ! transfer tv_tmp%* only if int_density_dz is called !$omp target enter data map(to: tv_tmp, tv_tmp%T, tv_tmp%S) & !$omp if(.not.(use_ALE .and. CS%Recon_Scheme > 0)) + ! The layer edge T/S reconstructions are read on the device by the (offloaded) generic_plm + ! density integrals; map them once here (they are the same for every layer k). + !$omp target enter data map(to: T_t, T_b, S_t, S_b) if(use_ALE .and. CS%Recon_Scheme > 0) + ! The following routine computes the integrals that are needed to ! calculate the pressure gradient force. Linear profiles for T and S are ! assumed when regridding is activated. Otherwise, the previous version @@ -1327,8 +1362,13 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, MassWghtInterp=CS%MassWghtInterp, Z_0p=Z_0p, & MassWghtInterpVanOnly=CS%MassWghtInterpVanOnly, h_nv=dz_nonvanished) endif - ! defensive update - not sure if it works - !$omp target update to(dpa, intx_dpa, inty_dpa, intz_dpa) + ! Push the host-computed integrals to the device only when int_density ran on the host + ! (Recon scheme 2 / Stanley / non-unity units). When idz_on_device the generic_plm kernels + ! already wrote them on the device, so this copy is both unnecessary and would clobber them + ! with a stale host copy. + if (.not. idz_on_device) then + !$omp target update to(dpa, intx_dpa, inty_dpa, intz_dpa) + endif else call int_density_dz(tv_tmp%T(:,:,k), tv_tmp%S(:,:,k), e(:,:,K), e(:,:,K+1), & rho_ref, rho0_int_density, GV%g_Earth, G%HI, tv%eqn_of_state, US, dpa(:,:,k), & @@ -1346,6 +1386,7 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, G%HI, MassWt_u(:,:,k), MassWt_v(:,:,k), & MassWghtInterpVanOnly=CS%MassWghtInterpVanOnly, h_nv=CS%h_nonvanished) enddo + !$omp target exit data map(release: T_t, T_b, S_t, S_b) if(use_ALE .and. CS%Recon_Scheme > 0) !$omp target exit data map(release: tv_tmp, tv_tmp%T, tv_tmp%S) & !$omp if(.not.(use_ALE .and. CS%Recon_Scheme > 0)) else @@ -1949,6 +1990,12 @@ subroutine PressureForce_FV_Bouss(h, tv, PFu, PFv, G, GV, US, CS, ALE_CSp, ADp, !$omp target exit data map(delete: Z_0p) if (use_EOS) + ! When idz_on_device, intx_dpa/inty_dpa were written on the device and are not otherwise copied + ! back. The SAL/tides gradient diagnostics below read them on the host, so pull them back before + ! the device mapping is deleted -- guarded to exactly that host-consuming configuration. + !$omp target update from(intx_dpa, inty_dpa) & + !$omp if(idz_on_device .and. (.not. CS%bq_sal_tides) .and. (CS%calculate_SAL .or. CS%tides)) + !$omp target exit data & !$omp map(delete: pa, dpa) & !$omp map(delete: intx_pa, inty_pa, intx_dpa, inty_dpa, intz_dpa) diff --git a/src/core/MOM_density_integrals.F90 b/src/core/MOM_density_integrals.F90 index bb398d9f00..c039d44cfb 100644 --- a/src/core/MOM_density_integrals.F90 +++ b/src/core/MOM_density_integrals.F90 @@ -10,6 +10,8 @@ module MOM_density_integrals use MOM_EOS, only : analytic_int_density_dz use MOM_EOS, only : analytic_int_specific_vol_dp use MOM_EOS, only : calculate_density +use MOM_EOS, only : calculate_density_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_EOS, only : calculate_spec_vol use MOM_EOS, only : calculate_specific_vol_derivs use MOM_EOS, only : average_specific_vol @@ -532,6 +534,21 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & integer, dimension(2) :: EOSdom_h5 ! The 5-point h-point i-computational domain for the equation of state integer, dimension(2) :: EOSdom_q15 ! The 3x5-point q-point i-computational domain for the equation of state integer, dimension(2) :: EOSdom_h15 ! The 3x5-point h-point i-computational domain for the equation of state + integer :: eos_form ! The equation-of-state form id, resolved host-side for the device-callable density path. + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa, eos_R_to_kg_m3 + ! EOS unit-rescaling factors, resolved host-side. + logical :: eos_unity ! True when all EOS unit-rescaling factors are 1, matching the no-rescale + ! fast path of calculate_density_1d. + logical :: do_intz ! present(intz_dpa), hoisted host-side for use inside the device region. + logical :: offload_phase1 ! True on GPU builds when phase 1 (vertical integrals) runs on device. + logical :: offload_phase2, offload_phase3 ! True on GPU builds when phases 2/3 (x/y integrals) run on device. + integer :: n15 ! A subgrid-point loop index over the 15 q-point slots in a column. + real :: T5l(5), S5l(5), p5l(5), r5l(5) ! Per-(i,j) size-5 quadrature locals for the device phase-1 + ! kernel (compile-time sized to avoid device auto-allocation). + real :: dz_l ! A per-(i,j) layer thickness for the device phase-1 kernel [Z ~> m]. + real :: T15l(15), S15l(15), p15l(15), r15l(15) ! Per-(I,j)/(J,i) size-15 q-point locals for the + ! device phase-2/3 kernels (compile-time sized). + real :: dz_x_l(5), dz_y_l(5) ! Per-cell subgrid layer thicknesses for the device x/y kernels [Z ~> m]. integer :: Isq, Ieq, Jsq, Jeq, i, j, m, n, pos Isq = HI%IscB ; Ieq = HI%IecB ; Jsq = HI%JscB ; Jeq = HI%JecB @@ -569,6 +586,42 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & use_varS = associated(tv%varS) endif + ! GPU port increment B: resolve the EOS form + unit scaling once on the host (the accessor is + ! not device-callable) so the non-Stanley density integrals can be evaluated point-by-point via + ! the device-callable dispatcher calculate_density_elem_loc, instead of the polymorphic + ! calculate_density interface, when this is (later) offloaded. eos_unity mirrors + ! calculate_density_1d's no-rescale fast path. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 + eos_RL2_T2_to_Pa = 1.0 ; eos_R_to_kg_m3 = 1.0 + call get_EOS_form_and_scaling(EOS, eos_form, eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, & + eos_RL2_T2_to_Pa, R_to_kg_m3=eos_R_to_kg_m3) + eos_unity = (eos_RL2_T2_to_Pa == 1.0) .and. (eos_R_to_kg_m3 == 1.0) .and. & + (eos_C_to_degC == 1.0) .and. (eos_S_to_ppt == 1.0) + do_intz = present(intz_dpa) ! Hoisted host-side so the device phase-1 region needs no present(). + ! Phase 1 runs on device only on GPU builds and only for the device-supported configuration + ! (non-Stanley, unscaled EOS units, accurate anomaly form); everything else keeps the host path. + offload_phase1 = .false. ; offload_phase2 = .false. ; offload_phase3 = .false. +#ifdef __NVCOMPILER_OPENMP_GPU + offload_phase1 = (.not. use_stanley_eos) .and. eos_unity .and. use_rho_ref + ! These must track offload_phase1: the caller (MOM_PressureForce_FV) mirrors this gate to decide + ! idz_on_device and no longer copies the outputs back to the host, so decoupling a phase would + ! leave its host-computed integral (intx/inty_dpa) stranded on the host while the caller reads a + ! stale device copy. Do not set them independently now that the device->host round-trip is gone. + offload_phase2 = offload_phase1 + offload_phase3 = offload_phase1 + ! On GPU builds the non-Stanley density integrals below are evaluated with the device-callable + ! dispatcher, which only covers ROQUET_RHO (in-situ density and anomaly) and buggy_Wright + ! (in-situ density only). FATAL host-side, before any device region, on an unsupported config. + if (.not. use_stanley_eos) then + if (.not. ((eos_form == EOS_ROQUET_RHO) .or. & + ((eos_form == EOS_WRIGHT) .and. (.not. use_rho_ref)))) call MOM_error(FATAL, & + "int_density_dz_generic_plm GPU build: no device-callable density kernel for this "// & + "EQN_OF_STATE / anomaly mode (only ROQUET_RHO, and buggy_Wright without rho_ref, are "// & + "supported); use a CPU build or add a _loc kernel.") + endif +#endif + T25(:) = 0. TS5(:) = 0. S25(:) = 0. @@ -586,7 +639,39 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & EOSdom_q15(1) = 1 ; EOSdom_q15(2) = 15*(Ieq-Isq+1) EOSdom_h15(1) = 1 ; EOSdom_h15(2) = 15*(HI%iec-HI%isc+1) + ! z0pres is read by all offloaded phases (1-3); map it once here and release it after phase 3. + if (offload_phase1) then + !$omp target enter data map(to: z0pres) + endif + ! 1. Compute vertical integrals + if (offload_phase1) then + ! Offload phase 1 (vertical integrals -> dpa, intz_dpa). Per-(i,j) size-5 quadrature locals + ! reproduce the host loop below point-for-point; dpa/intz_dpa are independent across (i,j). + ! Residency: e, dpa, intz_dpa are already device-resident (mapped by the caller); S_t/S_b/T_t/T_b + ! are mapped by the caller before the k-loop; z0pres is mapped just above (bracketing phases 1-3). + !$omp target teams loop collapse(2) & + !$omp private(T5l, S5l, p5l, r5l, rho_anom, dz_l, n) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, do_intz, k, eos_form, wt_t, wt_b) + do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1 + dz_l = e(i,j,K) - e(i,j,K+1) + do n=1,5 + p5l(n) = -GxRho*((e(i,j,K) - z0pres(i,j)) - 0.25*real(n-1)*dz_l) + S5l(n) = wt_t(n) * S_t(i,j,k) + wt_b(n) * S_b(i,j,k) + T5l(n) = wt_t(n) * T_t(i,j,k) + wt_b(n) * T_b(i,j,k) + enddo + do n=1,5 + r5l(n) = calculate_density_elem_loc(eos_form, T5l(n), S5l(n), p5l(n), use_rho_ref, rho_ref) + enddo + ! Boole's rule for the pressure anomaly change (use_rho_ref accurate form). + rho_anom = C1_90*(7.0*(r5l(1)+r5l(5)) + 32.0*(r5l(2)+r5l(4)) + 12.0*r5l(3)) + dpa(i,j) = G_e*dz_l*rho_anom + if (do_intz) intz_dpa(i,j) = 0.5*G_e*dz_l**2 * & + (rho_anom - C1_90*(16.0*(r5l(4)-r5l(2)) + 7.0*(r5l(5)-r5l(1))) ) + enddo ; enddo + ! dpa/intz_dpa stay on the device; phases 2 & 3 (also offloaded) read dpa there, and a single + ! consolidated copy-back to the host happens after phase 3. + else do j=Jsq,Jeq+1 do i = Isq,Ieq+1 dz(i) = e(i,j,K) - e(i,j,K+1) @@ -603,12 +688,32 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & if (use_Stanley_eos) then call calculate_density(T5, S5, p5, T25, TS5, S25, r5, EOS, EOSdom_h5, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d(T5,S5,p5,r5,EOS,EOSdom_h5 + ! [,rho_ref]) bit-for-bit: same _loc kernel per point, same unit rescaling. The eos_unity + ! branch is hoisted outside the point loop so the loop body is branch-free (offload-ready). + ! Iterating (i,n) covers exactly the EOSdom_h5 index range that the array call works on. + if (eos_unity) then + do i=Isq,Ieq+1 ; do n=1,5 + r5(i*5+n) = calculate_density_elem_loc(eos_form, T5(i*5+n), S5(i*5+n), p5(i*5+n), & + use_rho_ref, rho_ref) + enddo ; enddo + else + do i=Isq,Ieq+1 ; do n=1,5 + r5(i*5+n) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T5(i*5+n), eos_S_to_ppt*S5(i*5+n), eos_RL2_T2_to_Pa*p5(i*5+n), & + use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif + if (.not. use_rho_ref) u5(:) = r5(:) - rho_ref +#else if (use_rho_ref) then call calculate_density(T5, S5, p5, r5, EOS, EOSdom_h5, rho_ref=rho_ref) else call calculate_density(T5, S5, p5, r5, EOS, EOSdom_h5) u5(:) = r5(:) - rho_ref endif +#endif endif if (use_rho_ref) then @@ -638,9 +743,76 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & enddo endif enddo ! end loops on j + endif ! offload_phase1 ! 2. Compute horizontal integrals in the x direction - if (present(intx_dpa)) then ; do j=HI%jsc,HI%jec + if (present(intx_dpa)) then + if (offload_phase2) then + ! Offload phase 2 (x-integrals -> intx_dpa). Per-(I,j) size-15 q-point locals reproduce the host + ! loop below point-for-point; intx_dpa is independent across (I,j). dpa was written on the device + ! by phase 1 and is read here without a host round-trip. Both hWght branches are carried verbatim. + !$omp target teams loop collapse(2) & + !$omp private(T15l, S15l, p15l, r15l, dz_x_l, intz, hWght, hWghtTop, hL, hR, iDenom, & + !$omp Ttl, Ttr, Tbl, Tbr, Stl, Str, Sbl, Sbr, w_left, w_right, m, n, n15, pos) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, k, eos_form, massWeightToggle, & + !$omp TopWeightToggle, massWeightNVonlyToggle, h_nonvanished, dz_subroundoff, wt_t, wt_b) + do j=HI%jsc,HI%jec ; do I=Isq,Ieq + hWght = massWeightToggle * & + max(0., -bathyT(i,j)-e(i+1,j,K), -bathyT(i+1,j)-e(i,j,K)) + hWghtTop = TopWeightToggle * & + max(0., e(i+1,j,K+1)-e(i,j,1), e(i,j,K+1)-e(i+1,j,1)) + hWght = max(hWght, hWghtTop) + if (((e(i,j,K) - e(i,j,K+1)) > h_nonvanished) .and. ((e(i+1,j,K) - e(i+1,j,K+1)) > h_nonvanished)) then + hWght = massWeightNVonlyToggle * hWght + endif + if (hWght > 0.) then + hL = (e(i,j,K) - e(i,j,K+1)) + dz_subroundoff + hR = (e(i+1,j,K) - e(i+1,j,K+1)) + dz_subroundoff + hWght = hWght * ( (hL-hR)/(hL+hR) )**2 + iDenom = 1./( hWght*(hR + hL) + hL*hR ) + Ttl = ( (hWght*hR)*T_t(i+1,j,k) + (hWght*hL + hR*hL)*T_t(i,j,k) ) * iDenom + Ttr = ( (hWght*hL)*T_t(i,j,k) + (hWght*hR + hR*hL)*T_t(i+1,j,k) ) * iDenom + Tbl = ( (hWght*hR)*T_b(i+1,j,k) + (hWght*hL + hR*hL)*T_b(i,j,k) ) * iDenom + Tbr = ( (hWght*hL)*T_b(i,j,k) + (hWght*hR + hR*hL)*T_b(i+1,j,k) ) * iDenom + Stl = ( (hWght*hR)*S_t(i+1,j,k) + (hWght*hL + hR*hL)*S_t(i,j,k) ) * iDenom + Str = ( (hWght*hL)*S_t(i,j,k) + (hWght*hR + hR*hL)*S_t(i+1,j,k) ) * iDenom + Sbl = ( (hWght*hR)*S_b(i+1,j,k) + (hWght*hL + hR*hL)*S_b(i,j,k) ) * iDenom + Sbr = ( (hWght*hL)*S_b(i,j,k) + (hWght*hR + hR*hL)*S_b(i+1,j,k) ) * iDenom + else + Ttl = T_t(i,j,k) ; Tbl = T_b(i,j,k) ; Ttr = T_t(i+1,j,k) ; Tbr = T_b(i+1,j,k) + Stl = S_t(i,j,k) ; Sbl = S_b(i,j,k) ; Str = S_t(i+1,j,k) ; Sbr = S_b(i+1,j,k) + endif + do m=2,4 + w_left = wt_t(m) ; w_right = wt_b(m) + dz_x_l(m) = (w_left*(e(i,j,K) - e(i,j,K+1))) + (w_right*(e(i+1,j,K) - e(i+1,j,K+1))) + pos = (m-2)*5 + T15l(pos+1) = (w_left*Ttl) + (w_right*Ttr) + T15l(pos+5) = (w_left*Tbl) + (w_right*Tbr) + S15l(pos+1) = (w_left*Stl) + (w_right*Str) + S15l(pos+5) = (w_left*Sbl) + (w_right*Sbr) + p15l(pos+1) = -GxRho * ((w_left*(e(i,j,K)-z0pres(i,j))) + (w_right*(e(i+1,j,K)-z0pres(i+1,j)))) + do n=2,5 + p15l(pos+n) = p15l(pos+n-1) + GxRho*0.25*dz_x_l(m) + enddo + do n=2,4 + S15l(pos+n) = wt_t(n) * S15l(pos+1) + wt_b(n) * S15l(pos+5) + T15l(pos+n) = wt_t(n) * T15l(pos+1) + wt_b(n) * T15l(pos+5) + enddo + enddo + do n15=1,15 + r15l(n15) = calculate_density_elem_loc(eos_form, T15l(n15), S15l(n15), p15l(n15), use_rho_ref, rho_ref) + enddo + intz(1) = dpa(i,j) ; intz(5) = dpa(i+1,j) + do m=2,4 + pos = (m-2)*5 + intz(m) = (G_e*dz_x_l(m)*( C1_90*(7.0*(r15l(pos+1)+r15l(pos+5)) + 32.0*(r15l(pos+2)+r15l(pos+4)) + & + 12.0*r15l(pos+3)) )) + enddo + intx_dpa(I,j) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & + 12.0*intz(3)) + enddo ; enddo + else + do j=HI%jsc,HI%jec do I=Isq,Ieq ! Corner values of T and S ! hWght is the distance measure by which the cell is violation of @@ -720,11 +892,28 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & if (use_stanley_eos) then call calculate_density(T15, S15, p15, T215, TS15, S215, r15, EOS, EOSdom_q15, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d(T15,S15,p15,r15,EOS,EOSdom_q15 + ! [,rho_ref]) bit-for-bit. Iterating (I,n15) covers exactly the EOSdom_q15 index range. + if (eos_unity) then + do I=Isq,Ieq ; do n15=1,15 + r15(I*15+n15) = calculate_density_elem_loc(eos_form, T15(I*15+n15), S15(I*15+n15), & + p15(I*15+n15), use_rho_ref, rho_ref) + enddo ; enddo + else + do I=Isq,Ieq ; do n15=1,15 + r15(I*15+n15) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T15(I*15+n15), eos_S_to_ppt*S15(I*15+n15), & + eos_RL2_T2_to_Pa*p15(I*15+n15), use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif +#else if (use_rho_ref) then call calculate_density(T15, S15, p15, r15, EOS, EOSdom_q15, rho_ref=rho_ref) else call calculate_density(T15, S15, p15, r15, EOS, EOSdom_q15) endif +#endif endif do I=Isq,Ieq @@ -748,10 +937,79 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & intx_dpa(I,j) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & 12.0*intz(3)) enddo - enddo ; endif + enddo + endif ! offload_phase2 + endif ! present(intx_dpa) ! 3. Compute horizontal integrals in the y direction - if (present(inty_dpa)) then ; do J=Jsq,Jeq + if (present(inty_dpa)) then + if (offload_phase3) then + ! Offload phase 3 (y-integrals -> inty_dpa). Per-(J,i) size-15 q-point locals reproduce the host + ! loop below point-for-point; inty_dpa is independent across (J,i). dpa (device, from phase 1) is + ! read here without a host round-trip. Both hWght branches are carried verbatim. + !$omp target teams loop collapse(2) & + !$omp private(T15l, S15l, p15l, r15l, dz_y_l, intz, hWght, hWghtTop, hL, hR, iDenom, & + !$omp Ttl, Ttr, Tbl, Tbr, Stl, Str, Sbl, Sbr, w_left, w_right, m, n, n15, pos) & + !$omp firstprivate(GxRho, G_e, rho_ref, use_rho_ref, k, eos_form, massWeightToggle, & + !$omp TopWeightToggle, massWeightNVonlyToggle, h_nonvanished, dz_subroundoff, wt_t, wt_b) + do J=Jsq,Jeq ; do i=HI%isc,HI%iec + hWght = massWeightToggle * & + max(0., -bathyT(i,j)-e(i,j+1,K), -bathyT(i,j+1)-e(i,j,K)) + hWghtTop = TopWeightToggle * & + max(0., e(i,j+1,K+1)-e(i,j,1), e(i,j,K+1)-e(i,j+1,1)) + hWght = max(hWght, hWghtTop) + if (((e(i,j,K) - e(i,j,K+1)) > h_nonvanished) .and. ((e(i,j+1,K) - e(i,j+1,K+1)) > h_nonvanished)) then + hWght = massWeightNVonlyToggle * hWght + endif + if (hWght > 0.) then + hL = (e(i,j,K) - e(i,j,K+1)) + dz_subroundoff + hR = (e(i,j+1,K) - e(i,j+1,K+1)) + dz_subroundoff + hWght = hWght * ( (hL-hR)/(hL+hR) )**2 + iDenom = 1./( hWght*(hR + hL) + hL*hR ) + Ttl = ( (hWght*hR)*T_t(i,j+1,k) + (hWght*hL + hR*hL)*T_t(i,j,k) ) * iDenom + Ttr = ( (hWght*hL)*T_t(i,j,k) + (hWght*hR + hR*hL)*T_t(i,j+1,k) ) * iDenom + Tbl = ( (hWght*hR)*T_b(i,j+1,k) + (hWght*hL + hR*hL)*T_b(i,j,k) ) * iDenom + Tbr = ( (hWght*hL)*T_b(i,j,k) + (hWght*hR + hR*hL)*T_b(i,j+1,k) ) * iDenom + Stl = ( (hWght*hR)*S_t(i,j+1,k) + (hWght*hL + hR*hL)*S_t(i,j,k) ) * iDenom + Str = ( (hWght*hL)*S_t(i,j,k) + (hWght*hR + hR*hL)*S_t(i,j+1,k) ) * iDenom + Sbl = ( (hWght*hR)*S_b(i,j+1,k) + (hWght*hL + hR*hL)*S_b(i,j,k) ) * iDenom + Sbr = ( (hWght*hL)*S_b(i,j,k) + (hWght*hR + hR*hL)*S_b(i,j+1,k) ) * iDenom + else + Ttl = T_t(i,j,k) ; Tbl = T_b(i,j,k) ; Ttr = T_t(i,j+1,k) ; Tbr = T_b(i,j+1,k) + Stl = S_t(i,j,k) ; Sbl = S_b(i,j,k) ; Str = S_t(i,j+1,k) ; Sbr = S_b(i,j+1,k) + endif + do m=2,4 + w_left = wt_t(m) ; w_right = wt_b(m) + dz_y_l(m) = (w_left*(e(i,j,K) - e(i,j,K+1))) + (w_right*(e(i,j+1,K) - e(i,j+1,K+1))) + pos = (m-2)*5 + T15l(pos+1) = (w_left*Ttl) + (w_right*Ttr) + T15l(pos+5) = (w_left*Tbl) + (w_right*Tbr) + S15l(pos+1) = (w_left*Stl) + (w_right*Str) + S15l(pos+5) = (w_left*Sbl) + (w_right*Sbr) + p15l(pos+1) = -GxRho * ((w_left*(e(i,j,K)-z0pres(i,j))) + (w_right*(e(i,j+1,K)-z0pres(i,j+1)))) + do n=2,5 + p15l(pos+n) = p15l(pos+n-1) + GxRho*0.25*dz_y_l(m) + enddo + do n=2,4 + S15l(pos+n) = wt_t(n) * S15l(pos+1) + wt_b(n) * S15l(pos+5) + T15l(pos+n) = wt_t(n) * T15l(pos+1) + wt_b(n) * T15l(pos+5) + enddo + enddo + do n15=1,15 + r15l(n15) = calculate_density_elem_loc(eos_form, T15l(n15), S15l(n15), p15l(n15), use_rho_ref, rho_ref) + enddo + intz(1) = dpa(i,j) ; intz(5) = dpa(i,j+1) + do m=2,4 + pos = (m-2)*5 + intz(m) = (G_e*dz_y_l(m)*( C1_90*(7.0*(r15l(pos+1)+r15l(pos+5)) + & + 32.0*(r15l(pos+2)+r15l(pos+4)) + & + 12.0*r15l(pos+3)) )) + enddo + inty_dpa(i,J) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & + 12.0*intz(3)) + enddo ; enddo + else + do J=Jsq,Jeq do i=HI%isc,HI%iec ! Corner values of T and S ! hWght is the distance measure by which the cell is violation of @@ -834,6 +1092,22 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & T215(15*HI%isc+1:), TS15(15*HI%isc+1:), S215(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15, rho_ref=rho_ref) else +#ifdef __NVCOMPILER_OPENMP_GPU + ! Device-callable element path reproducing calculate_density_1d over the h15 domain + ! bit-for-bit. Iterating (i,n15) covers exactly the r15(15*HI%isc+1:)/EOSdom_h15 range. + if (eos_unity) then + do i=HI%isc,HI%iec ; do n15=1,15 + r15(i*15+n15) = calculate_density_elem_loc(eos_form, T15(i*15+n15), S15(i*15+n15), & + p15(i*15+n15), use_rho_ref, rho_ref) + enddo ; enddo + else + do i=HI%isc,HI%iec ; do n15=1,15 + r15(i*15+n15) = eos_kg_m3_to_R * calculate_density_elem_loc(eos_form, & + eos_C_to_degC*T15(i*15+n15), eos_S_to_ppt*S15(i*15+n15), & + eos_RL2_T2_to_Pa*p15(i*15+n15), use_rho_ref, eos_R_to_kg_m3*rho_ref) + enddo ; enddo + endif +#else if (use_rho_ref) then call calculate_density(T15(15*HI%isc+1:), S15(15*HI%isc+1:), p15(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15, rho_ref=rho_ref) @@ -841,6 +1115,7 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & call calculate_density(T15(15*HI%isc+1:), S15(15*HI%isc+1:), p15(15*HI%isc+1:), & r15(15*HI%isc+1:), EOS, EOSdom_h15) endif +#endif endif do i=HI%isc,HI%iec @@ -866,7 +1141,16 @@ subroutine int_density_dz_generic_plm(k, tv, T_t, T_b, S_t, S_b, e, rho_ref, & inty_dpa(i,J) = C1_90*(7.0*(intz(1)+intz(5)) + 32.0*(intz(2)+intz(4)) + & 12.0*intz(3)) enddo - enddo ; endif + enddo + endif ! offload_phase3 + endif ! present(inty_dpa) + + ! The outputs dpa/intz_dpa/intx_dpa/inty_dpa are left device-resident: the caller consumes them + ! entirely in its own device regions (interface-pressure accumulation and the PFu/PFv loops), so + ! there is no device->host copy-back here. Only the per-call z0pres scratch mapping is released. + if (offload_phase1) then + !$omp target exit data map(release: z0pres) + endif end subroutine int_density_dz_generic_plm diff --git a/src/core/MOM_interface_heights.F90 b/src/core/MOM_interface_heights.F90 index c70c5e8f44..861b73a4c8 100644 --- a/src/core/MOM_interface_heights.F90 +++ b/src/core/MOM_interface_heights.F90 @@ -41,6 +41,11 @@ module MOM_interface_heights module procedure thickness_to_dz_3d, thickness_to_dz_jslice end interface thickness_to_dz +!> Computes the density of the near-bottom water in the ocean. +interface find_rho_bottom + module procedure find_rho_bottom_1d, find_rho_bottom_2d +end interface find_rho_bottom + contains !> Calculates the change in height across layers, using the appropriate form for @@ -486,7 +491,7 @@ end subroutine find_col_mass !> Determine the in situ density averaged over a specified distance from the bottom, !! calculating it as the inverse of the mass-weighted average specific volume. -subroutine find_rho_bottom(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h_bot, k_bot) +subroutine find_rho_bottom_1d(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h_bot, k_bot) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -663,8 +668,205 @@ subroutine find_rho_bottom(G, GV, US, tv, h, dz, pres_int, dz_avg, j, Rho_bot, h enddo endif -end subroutine find_rho_bottom +end subroutine find_rho_bottom_1d + +!> Computes the density of the near-bottom water for a j-block of rows, using pre-computed +!! blocked arrays with an explicit j-block index dimension. +subroutine find_rho_bottom_2d(G, GV, US, tv, h, dz, pres_int, dz_avg, jstart, jend, nj, Rho_bot, h_bot, k_bot) + type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure + type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + type(thermo_var_ptrs), intent(in) :: tv !< Structure containing pointers to any available + !! thermodynamic fields. + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), & + intent(in) :: h !< Layer thicknesses [H ~> m or kg m-2] + integer, intent(in) :: nj !< Number of j-rows in the block + real, dimension(SZI_(G),SZK_(GV),nj), & + intent(in) :: dz !< Height change across layers [Z ~> m] + real, dimension(SZI_(G),SZK_(GV)+1,nj), & + intent(in) :: pres_int !< Pressure at each interface [R L2 T-2 ~> Pa] + real, dimension(SZI_(G),nj), intent(in) :: dz_avg !< The vertical distance over which to average [Z ~> m] + integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block + real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. + real, dimension(SZI_(G),nj), intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2] + integer, dimension(SZI_(G),nj), intent(out) :: k_bot !< Bottom boundary layer top layer index + + ! Local variables + real :: hb(SZI_(G),nj) ! Running sum of the thickness in the bottom boundary layer [H ~> m or kg m-2] + real :: SpV_h_bot(SZI_(G),nj) ! Running sum of the specific volume times thickness in the bottom + ! boundary layer [H R-1 ~> m4 kg-1 or m] + real :: dz_bbl_rem(SZI_(G),nj) ! Vertical extent of the boundary layer that has yet to be accounted + ! for [Z ~> m] + real :: h_bbl_frac(SZI_(G),nj) ! Thickness of the fractional layer that makes up the top of the + ! boundary layer [H ~> m or kg m-2] + real :: T_bbl(SZI_(G),nj) ! Temperature of the fractional layer that makes up the top of the + ! boundary layer [C ~> degC] + real :: S_bbl(SZI_(G),nj) ! Salinity of the fractional layer that makes up the top of the + ! boundary layer [S ~> ppt] + real :: P_bbl(SZI_(G),nj) ! Pressure the top of the boundary layer [R L2 T-2 ~> Pa] + real :: dp(SZI_(G),nj) ! Pressure change across the fractional layer that makes up the top + ! of the boundary layer [R L2 T-2 ~> Pa] + real :: SpV_bbl(SZI_(G),nj) ! In situ specific volume of the fractional layer that makes up the + ! top of the boundary layer [R-1 ~> m3 kg-1] + real :: frac_in ! The fraction of a layer that is within the bottom boundary layer [nondim] + logical :: do_i(SZI_(G),nj), do_any + logical :: use_EOS + integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state + integer :: i, j, k, is, ie, nz, jj + + is = G%isc ; ie = G%iec ; nz = GV%ke + + use_EOS = associated(tv%T) .and. associated(tv%S) .and. associated(tv%eqn_of_state) + + if (GV%Boussinesq .or. GV%semi_Boussinesq .or. .not.allocated(tv%SpV_avg)) then + ! Obtain bottom boundary layer thickness and index of top layer + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + rho_bot(i,jj) = GV%Rho0 + hb(i,jj) = 0.0 ; h_bot(i,jj) = 0.0 ; k_bot(i,jj) = nz + dz_bbl_rem(i,jj) = G%mask2dT(i,j) * max(0.0, dz_avg(i,jj)) + do_i(i,jj) = .true. + if (G%mask2dT(i,j) <= 0.0) then + h_bbl_frac(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + enddo + enddo + + do j=jstart,jend ; jj = j - jstart + 1 ; do k=nz,1,-1 + do_any = .false. + + do i=is,ie ; if (do_i(i,jj)) then + if (dz(i,k,jj) < dz_bbl_rem(i,jj)) then + ! This layer is fully within the averaging depth. + dz_bbl_rem(i,jj) = dz_bbl_rem(i,jj) - dz(i,k,jj) + hb(i,jj) = hb(i,jj) + h(i,j,k) + k_bot(i,jj) = k + do_any = .true. + else + if (dz(i,k,jj) > 0.0) then + frac_in = dz_bbl_rem(i,jj) / dz(i,k,jj) + if (frac_in >= 0.5) k_bot(i,jj) = k ! update bbl top index if >= 50% of layer + else + frac_in = 0.0 + endif + h_bbl_frac(i,jj) = frac_in * h(i,j,k) + dz_bbl_rem(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + endif ; enddo + enddo + if (.not.do_any) exit + enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + ! The nominal bottom boundary layer is thicker than the water column, but layer 1 is + ! already included in the averages. These values are set so that the call to find + ! the layer-average specific volume will behave sensibly. + h_bbl_frac(i,jj) = 0.0 + endif ; enddo + enddo + + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (hb(i,jj) + h_bbl_frac(i,jj) < GV%H_subroundoff) h_bbl_frac(i,jj) = GV%H_subroundoff + h_bot(i,jj) = hb(i,jj) + h_bbl_frac(i,jj) + enddo + enddo + + else + ! Check that SpV_avg has been set. + if (tv%valid_SpV_halo < 0) call MOM_error(FATAL, & + "find_rho_bottom called in fully non-Boussinesq mode with invalid values of SpV_avg.") + + ! Set the bottom density to the inverse of the in situ specific volume averaged over the + ! specified distance, with care taken to avoid having compressibility lead to an imprint + ! of the layer thicknesses on this density. + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + hb(i,jj) = 0.0 ; SpV_h_bot(i,jj) = 0.0 ; h_bot(i,jj) = 0.0 ; k_bot(i,jj) = nz + dz_bbl_rem(i,jj) = G%mask2dT(i,j) * max(0.0, dz_avg(i,jj)) + do_i(i,jj) = .true. + if (G%mask2dT(i,j) <= 0.0) then + ! Set acceptable values for calling the equation of state over land. + T_bbl(i,jj) = 0.0 ; S_bbl(i,jj) = 0.0 ; dp(i,jj) = 0.0 ; P_bbl(i,jj) = 0.0 + SpV_bbl(i,jj) = 1.0 ! This value is arbitrary, provided it is non-zero. + h_bbl_frac(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + enddo + enddo + + do k=nz,1,-1 + do_any = .false. + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + if (dz(i,k,jj) < dz_bbl_rem(i,jj)) then + ! This layer is fully within the averaging depth. + SpV_h_bot(i,jj) = SpV_h_bot(i,jj) + h(i,j,k) * tv%SpV_avg(i,j,k) + dz_bbl_rem(i,jj) = dz_bbl_rem(i,jj) - dz(i,k,jj) + hb(i,jj) = hb(i,jj) + h(i,j,k) + k_bot(i,jj) = k + do_any = .true. + else + if (dz(i,k,jj) > 0.0) then + frac_in = dz_bbl_rem(i,jj) / dz(i,k,jj) + if (frac_in >= 0.5) k_bot(i,jj) = k ! update bbl top index if >= 50% of layer + else + frac_in = 0.0 + endif + if (use_EOS) then + ! Store the properties of this layer to determine the average + ! specific volume of the portion that is within the BBL. + T_bbl(i,jj) = tv%T(i,j,k) ; S_bbl(i,jj) = tv%S(i,j,k) + dp(i,jj) = frac_in * (GV%g_Earth*GV%H_to_RZ * h(i,j,k)) + P_bbl(i,jj) = pres_int(i,K,jj) + (1.0-frac_in) * (GV%g_Earth*GV%H_to_RZ * h(i,j,k)) + else + SpV_bbl(i,jj) = tv%SpV_avg(i,j,k) + endif + h_bbl_frac(i,jj) = frac_in * h(i,j,k) + dz_bbl_rem(i,jj) = 0.0 + do_i(i,jj) = .false. + endif + endif ; enddo + enddo + if (.not.do_any) exit + enddo + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; if (do_i(i,jj)) then + ! The nominal bottom boundary layer is thicker than the water column, but layer 1 is + ! already included in the averages. These values are set so that the call to find + ! the layer-average specific volume will behave sensibly. + if (use_EOS) then + T_bbl(i,jj) = tv%T(i,j,1) ; S_bbl(i,jj) = tv%S(i,j,1) + dp(i,jj) = 0.0 + P_bbl(i,jj) = pres_int(i,1,jj) + else + SpV_bbl(i,jj) = tv%SpV_avg(i,j,1) + endif + h_bbl_frac(i,jj) = 0.0 + endif ; enddo + enddo + + if (use_EOS) then + ! Find the average specific volume of the fractional layer atop the BBL. + EOSdom(:) = EOS_domain(G%HI) + do j=jstart,jend ; jj = j - jstart + 1 + call average_specific_vol(T_bbl(:,jj), S_bbl(:,jj), P_bbl(:,jj), dp(:,jj), SpV_bbl(:,jj), tv%eqn_of_state, EOSdom) + enddo + endif + + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (hb(i,jj) + h_bbl_frac(i,jj) < GV%H_subroundoff) h_bbl_frac(i,jj) = GV%H_subroundoff + rho_bot(i,jj) = G%mask2dT(i,j) * (hb(i,jj) + h_bbl_frac(i,jj)) / (SpV_h_bot(i,jj) + h_bbl_frac(i,jj)*SpV_bbl(i,jj)) + h_bot(i,jj) = hb(i,jj) + h_bbl_frac(i,jj) + enddo + enddo + endif +end subroutine find_rho_bottom_2d !> Converts thickness from geometric height units to thickness units, perhaps via an !! inversion of the integral of the density in pressure using variables stored in diff --git a/src/equation_of_state/MOM_EOS.F90 b/src/equation_of_state/MOM_EOS.F90 index 754f293b69..39fa0a86bb 100644 --- a/src/equation_of_state/MOM_EOS.F90 +++ b/src/equation_of_state/MOM_EOS.F90 @@ -9,6 +9,8 @@ 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 : density_elem_buggy_Wright_loc use MOM_EOS_Wright, only : int_density_dz_wright, int_spec_vol_dp_wright use MOM_EOS_Wright_full, only : Wright_full_EOS, avg_spec_vol_Wright_full use MOM_EOS_Wright_full, only : int_density_dz_wright_full, int_spec_vol_dp_wright_full @@ -17,6 +19,8 @@ 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_rho, only : density_elem_Roquet_rho_loc, density_anomaly_elem_Roquet_rho_loc use MOM_EOS_Roquet_SpV, only : Roquet_SpV_EOS use MOM_EOS_TEOS10, only : TEOS10_EOS use MOM_EOS_TEOS10, only : gsw_sp_from_sr, gsw_pt_from_ct, gsw_sr_from_sp, gsw_ct_from_pt @@ -46,6 +50,9 @@ module MOM_EOS public calculate_density_elem public calculate_density public calculate_density_derivs +public calculate_density_derivs_elem_loc +public calculate_density_elem_loc +public get_EOS_form_and_scaling public calculate_density_second_derivs public calculate_spec_vol public calculate_specific_vol_derivs @@ -893,6 +900,94 @@ 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 + +!> Device-callable dispatcher for in-situ density (or its anomaly relative to rho_ref) at a single +!! point, in mks units, selecting the equation-of-state form at runtime by integer id (no +!! polymorphic dispatch) so it can be called from inside a do concurrent / target region by GPU +!! kernels (e.g. the finite-volume pressure-gradient density integrals). Unit rescaling +!! (EOS%*_to_* factors) and any `scale` factor are the caller's responsibility, exactly as in +!! calculate_density_elem. If use_rho_ref is true, the density anomaly relative to rho_ref (all in +!! mks units) is returned via the form's device-callable anomaly kernel. Combinations without a +!! device-callable _loc kernel -- an unsupported form, or the anomaly branch of a form that has no +!! anomaly kernel (e.g. buggy_Wright) -- are not handled here and return 0; a device-using module +!! must FATAL at init on a GPU build before reaching this with such a combination. +real function calculate_density_elem_loc(form_of_EOS, T, S, pressure, use_rho_ref, rho_ref) + integer, intent(in) :: form_of_EOS !< The equation of state form (EOS_ROQUET_RHO, EOS_WRIGHT, ...) + real, intent(in) :: T !< Temperature in the EOS kernel's mks units [degC] + real, intent(in) :: S !< Salinity in the EOS kernel's mks units [ppt or g kg-1] + real, intent(in) :: pressure !< Pressure [Pa] + logical, intent(in) :: use_rho_ref !< If true, return the density anomaly relative to rho_ref + real, intent(in) :: rho_ref !< A reference density [kg m-3], subtracted when use_rho_ref is true + !$omp declare target + + calculate_density_elem_loc = 0.0 + select case (form_of_EOS) + case (EOS_ROQUET_RHO) + if (use_rho_ref) then + calculate_density_elem_loc = density_anomaly_elem_Roquet_rho_loc(T, S, pressure, rho_ref) + else + calculate_density_elem_loc = density_elem_Roquet_rho_loc(T, S, pressure) + endif + case (EOS_WRIGHT) + ! buggy_Wright has a device-callable in-situ density kernel but no anomaly kernel, so the + ! anomaly branch is unsupported and returns 0 (host FATALs at init on GPU builds). + if (.not. use_rho_ref) & + calculate_density_elem_loc = density_elem_buggy_Wright_loc(T, S, pressure) + case default + calculate_density_elem_loc = 0.0 + end select + +end function calculate_density_elem_loc + +!> Return the equation-of-state form id and the unit-rescaling factors held in an EOS_type. +!! Lets a caller (e.g. a whole-column GPU kernel) reproduce, host-side, the unit conversion +!! and rescaling that calculate_density_derivs_1d applies around the mks _loc kernels, without +!! 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, & + R_to_kg_m3) + type(EOS_type), intent(in) :: EOS !< Equation of state structure + integer, intent(out) :: form_of_EOS !< The equation of state form id (EOS_ROQUET_RHO, ...) + real, intent(out) :: kg_m3_to_R !< Factor converting kg m-3 to the internal density unit R [R m3 kg-1 ~> 1] + real, intent(out) :: C_to_degC !< Factor converting the temperature unit to degC [degC C-1 ~> 1] + real, intent(out) :: S_to_ppt !< Factor converting the salinity unit to ppt [ppt S-1 ~> 1] + real, intent(out) :: RL2_T2_to_Pa !< Factor converting the pressure unit to Pa [Pa T2 R-1 L-2 ~> 1] + real, optional, intent(out) :: R_to_kg_m3 !< Factor converting the internal density unit R to kg m-3 [kg R-1 m-3 ~> 1] + + form_of_EOS = EOS%form_of_EOS + kg_m3_to_R = EOS%kg_m3_to_R + C_to_degC = EOS%C_to_degC + S_to_ppt = EOS%S_to_ppt + RL2_T2_to_Pa = EOS%RL2_T2_to_Pa + if (present(R_to_kg_m3)) R_to_kg_m3 = EOS%R_to_kg_m3 + +end subroutine get_EOS_form_and_scaling + !> Calls the appropriate subroutine to calculate density derivatives for 1-D array inputs. subroutine calculate_density_derivs_2d(T, S, pressure, drho_dT, drho_dS, EOS, dom) diff --git a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 index 86110d6aa1..6a4f3cc7a2 100644 --- a/src/equation_of_state/MOM_EOS_Roquet_rho.F90 +++ b/src/equation_of_state/MOM_EOS_Roquet_rho.F90 @@ -10,6 +10,15 @@ 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) +! Exposed as device-callable (declare target) elementals so whole-column GPU kernels can compute +! in-situ density and its anomaly in-region without polymorphic dispatch. +public density_elem_Roquet_rho_loc, density_anomaly_elem_Roquet_rho_loc +!$omp declare target(density_elem_Roquet_rho_loc) +!$omp declare target(density_anomaly_elem_Roquet_rho_loc) real, parameter :: Pa2kb = 1.e-8 !< Conversion factor between Pa and kbar [kbar Pa-1] !>@{ Parameters in the Roquet_rho (Roquet density) equation of state diff --git a/src/equation_of_state/MOM_EOS_Wright.F90 b/src/equation_of_state/MOM_EOS_Wright.F90 index c2861c451d..08c85858d6 100644 --- a/src/equation_of_state/MOM_EOS_Wright.F90 +++ b/src/equation_of_state/MOM_EOS_Wright.F90 @@ -12,6 +12,14 @@ 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) +! Exposed as a device-callable (declare target) elemental so whole-column GPU kernels can compute +! in-situ density in-region without polymorphic dispatch. (buggy_Wright has no anomaly _loc kernel.) +public density_elem_buggy_Wright_loc +!$omp declare target(density_elem_buggy_Wright_loc) public int_density_dz_wright, int_spec_vol_dp_wright public avg_spec_vol_buggy_Wright public set_params_buggy_Wright diff --git a/src/framework/MOM_intrinsic_functions.F90 b/src/framework/MOM_intrinsic_functions.F90 index de668f56bd..ef4d86fd11 100644 --- a/src/framework/MOM_intrinsic_functions.F90 +++ b/src/framework/MOM_intrinsic_functions.F90 @@ -11,7 +11,7 @@ module MOM_intrinsic_functions implicit none ; private -public :: invcosh, cuberoot, nth_root +public :: invcosh, cuberoot, nth_root, exp_reprod, log_reprod, erfc_reprod public :: intrinsic_functions_unit_tests ! Floating point model, if bit layout from high to low is (sign, exp, frac) @@ -117,6 +117,134 @@ elemental function cuberoot(x) result(root) end function cuberoot +!> Bit-reproducible exponential, exp(x), suitable for evaluation inside `!$omp target` / +!! `do concurrent` offloaded regions. +!! +!! The intrinsic `exp()` lowers to host libm on the CPU and CUDA libdevice on the GPU, whose +!! last-bit rounding differs -- so a `do concurrent`/`omp target` kernel that calls `exp()` is not +!! bit-for-bit CPU==GPU (the class of ~1e-13 divergence seen in the ePBL energy budget). This routine +!! avoids the library call entirely: Cody-Waite range reduction x = k*ln2 + r (|r| <= ln2/2) using a +!! two-part ln2 so that k*ln2_hi is (near-)exact, a degree-12 Taylor/Horner polynomial for exp(r) whose +!! reciprocal-factorial coefficients are compile-time constant-folded (hence identical on host and +!! device), and `scale(.,k)` for the 2**k factor. Every operation is +,-,*,/ (plus `nint`/`scale`, +!! which are exact), all of which are bit-identical host vs device under `-Mnofma`, so the result is +!! reproducible by construction. Accuracy is ~1.4 ULP (max relative error 3.1e-16 vs the intrinsic +!! over x in [-70, 20]); it is NOT bit-identical to the intrinsic exp (like cuberoot vs x**(1/3), it +!! changes answers and requires a reference/golden regeneration when adopted). +elemental function exp_reprod(x) result(ex) + !$omp declare target + real, intent(in) :: x !< The argument of the exponential [nondim or arbitrary] + real :: ex !< The reproducible exponential of x [same as exp(x)] + + ! Cody-Waite split of ln(2): ln2 = ln2_hi + ln2_lo, with ln2_hi chosen so k*ln2_hi is ~exact. + real, parameter :: invln2 = 1.44269504088896338700 ! 1/ln(2) [nondim] + real, parameter :: ln2_hi = 0.693147180369123816490 ! High part of ln(2) [nondim] + real, parameter :: ln2_lo = 1.90821492927058770002e-10 ! Low part of ln(2) [nondim] + ! Reciprocal factorials 1/2! .. 1/12! (compile-time constant-folded -> identical host/device). + real, parameter :: c2 = 1.0/2.0, c3 = 1.0/6.0, c4 = 1.0/24.0 + real, parameter :: c5 = 1.0/120.0, c6 = 1.0/720.0, c7 = 1.0/5040.0 + real, parameter :: c8 = 1.0/40320.0, c9 = 1.0/362880.0, c10 = 1.0/3628800.0 + real, parameter :: c11 = 1.0/39916800.0, c12 = 1.0/479001600.0 + real :: r ! The reduced argument, x - k*ln2, in [-ln2/2, ln2/2] [nondim] + real :: p ! The polynomial estimate of exp(r) [nondim] + integer :: k ! The integer number of factors of 2 in exp(x) [nondim] + + k = nint(x*invln2) + r = (x - real(k)*ln2_hi) - real(k)*ln2_lo + ! Horner form of 1 + r + r^2/2! + ... + r^12/12! + p = 1.0 + r*(1.0 + r*(c2 + r*(c3 + r*(c4 + r*(c5 + r*(c6 + r*(c7 + & + r*(c8 + r*(c9 + r*(c10 + r*(c11 + r*c12))))))))))) + ex = scale(p, k) +end function exp_reprod + + +!> Bit-reproducible natural logarithm, log(x) for x > 0, suitable for evaluation inside +!! `!$omp target` / `do concurrent` offloaded regions (companion to exp_reprod; together they make +!! arbitrary real powers reproducible via x**y = exp_reprod(y*log_reprod(x))). +!! +!! Same rationale as exp_reprod: the intrinsic log() differs host libm vs CUDA libdevice in the last +!! bit. This routine uses the exponent/fraction split x = m * 2**k (exact bit intrinsics), reduces the +!! mantissa to m in [sqrt(1/2), sqrt(2)), and evaluates log(m) = 2*(s + s^3/3 + s^5/5 + ...) with +!! s = (m-1)/(m+1) (|s| <= 0.172) as a Horner polynomial in s^2 to degree 21, then adds k*ln2. Every +!! operation is +,-,*,/ (plus exponent/fraction, exact), bit-identical host vs device under -Mnofma. +!! Accuracy ~1.7 ULP (max relative error 3.7e-16 vs the intrinsic over x in [1e-30, 1e30]). Requires +!! x > 0. Like exp_reprod it is not bit-identical to intrinsic log (changes answers on adoption). +elemental function log_reprod(x) result(lx) + !$omp declare target + real, intent(in) :: x !< The argument of the logarithm, x > 0 [nondim or arbitrary] + real :: lx !< The reproducible natural logarithm of x [nondim] + + real, parameter :: ln2 = 0.69314718055994530942 ! ln(2) [nondim] + real, parameter :: sqrt2_2 = 0.70710678118654752440 ! sqrt(1/2), the mantissa reduction threshold [nondim] + ! Reciprocal odd integers 1/3 .. 1/21 for the atanh series (compile-folded -> identical host/device). + real, parameter :: a3=1.0/3.0, a5=1.0/5.0, a7=1.0/7.0, a9=1.0/9.0 + real, parameter :: a11=1.0/11.0, a13=1.0/13.0, a15=1.0/15.0, a17=1.0/17.0 + real, parameter :: a19=1.0/19.0, a21=1.0/21.0 + real :: m ! The mantissa of x, reduced to [sqrt(1/2), sqrt(2)) [nondim] + real :: s ! (m-1)/(m+1), the atanh-series argument, |s| <= 0.172 [nondim] + real :: s2 ! s*s [nondim] + real :: poly ! The polynomial estimate of log(m) [nondim] + integer :: k ! The binary exponent of x [nondim] + + k = exponent(x) ; m = fraction(x) ! x = m * 2**k, m in [0.5, 1) + if (m < sqrt2_2) then ; m = m + m ; k = k - 1 ; endif ! recenter m to [sqrt(1/2), sqrt(2)) + s = (m - 1.0) / (m + 1.0) ; s2 = s*s + poly = 2.0*s*(1.0 + s2*(a3 + s2*(a5 + s2*(a7 + s2*(a9 + s2*(a11 + s2*(a13 + & + s2*(a15 + s2*(a17 + s2*(a19 + s2*a21)))))))))) + lx = poly + real(k)*ln2 +end function log_reprod + + +!> Bit-reproducible complementary error function, erfc(x) for x >= 0, suitable for evaluation inside +!! `!$omp target` / `do concurrent` offloaded regions. +!! +!! Same rationale as exp_reprod/log_reprod (intrinsic erfc differs host libm vs CUDA libdevice in the +!! last bit). Built from +,-,*,/ and exp_reprod only, so bit-identical host vs device under -Mnofma. +!! Two ranges: for x < 1.5, erf(x) via its Maclaurin series (48 fixed terms) and erfc = 1 - erf; for +!! x >= 1.5, the incomplete-gamma continued fraction erfc(x) = (x/sqrt(pi))*exp(-x^2)*Q-CF(1/2, x^2) +!! evaluated by a FIXED-iteration (60) modified-Lentz recurrence (no early exit -- device-safe). +!! Accuracy is <= 5.8e-14 relative vs the intrinsic over x in (0.01, 26] (the worst case is near +!! x ~ 24 where erfc ~ 1e-263 is physically zero; ~1e-15 across the physically relevant range). x >= 0 +!! required. Not bit-identical to intrinsic erfc (changes answers on adoption). +elemental function erfc_reprod(x) result(fc) + !$omp declare target + real, intent(in) :: x !< The argument of the complementary error function, x >= 0 [nondim] + real :: fc !< The reproducible complementary error function of x [nondim] + + real, parameter :: two_sqrtpi = 1.12837916709551257390 ! 2/sqrt(pi) [nondim] + real, parameter :: inv_sqrtpi = 0.56418958354775628695 ! 1/sqrt(pi) [nondim] + real, parameter :: tiny_l = 1.0e-30 ! A floor to avoid division by zero in the Lentz recurrence [nondim] + real :: t ! x*x [nondim] + real :: term ! The running term of the erf Maclaurin series [nondim] + real :: sumv ! The running sum of the erf Maclaurin series [nondim] + real :: an, b, c, d, del, h ! Modified-Lentz continued-fraction working variables [nondim] + integer :: n + + if (x < 1.5) then + ! erf(x) = (2/sqrt(pi)) * sum_{n>=0} (-1)^n x^(2n+1) / (n! (2n+1)); erfc = 1 - erf. + term = x ; sumv = x + do n = 1, 48 + term = term * (-(x*x)) * real(2*n-1) / (real(n)*real(2*n+1)) + sumv = sumv + term + enddo + fc = 1.0 - two_sqrtpi*sumv + else + ! erfc(x) = (x/sqrt(pi)) * exp(-x^2) * Q-continued-fraction(a=1/2, z=x^2), modified Lentz. + t = x*x + b = t + 0.5 ! z + 1 - a + c = 1.0/tiny_l ; d = 1.0/b ; h = d + do n = 1, 60 + an = -real(n)*(real(n) - 0.5) ! -n*(n - a) + b = b + 2.0 + d = an*d + b ; if (abs(d) < tiny_l) d = tiny_l ; d = 1.0/d + c = b + an/c ; if (abs(c) < tiny_l) c = tiny_l + del = d*c ; h = h*del + enddo + fc = (x*inv_sqrtpi) * exp_reprod(-t) * h + endif +end function erfc_reprod + + !> Bit-stable n-th root of x for x in (0, +inf) and integer n >= 1, suitable !! for evaluation inside `!$omp target` / `do concurrent` offloaded regions. !! @@ -298,6 +426,35 @@ function intrinsic_functions_unit_tests(verbose) result(fail) fail = fail .or. Test_cuberoot(v, testval) testval = (-2.908 * (1.414213562373 + 1.2345678901234e-5*n)) * testval enddo + + v = verbose + fail = fail .or. Test_exp_reprod(v, 0.0) + fail = fail .or. Test_exp_reprod(v, 1.0) + fail = fail .or. Test_exp_reprod(v, -3.7) + fail = fail .or. Test_exp_reprod(v, 20.0) + v = .false. + do n=-700,200 + fail = fail .or. Test_exp_reprod(v, 0.1*n) + enddo + + v = verbose + fail = fail .or. Test_log_reprod(v, 2.0) + fail = fail .or. Test_log_reprod(v, 0.7) + fail = fail .or. Test_log_reprod(v, 1.0e6) + v = .false. + do n=1,6000 + fail = fail .or. Test_log_reprod(v, 10.0**((0.01*real(n)) - 30.0)) + enddo + + v = verbose + fail = fail .or. Test_erfc_reprod(v, 0.001) + fail = fail .or. Test_erfc_reprod(v, 0.7) + fail = fail .or. Test_erfc_reprod(v, 1.0) + fail = fail .or. Test_erfc_reprod(v, 3.0) + v = .false. + do n=1,1000 + fail = fail .or. Test_erfc_reprod(v, 0.01*real(n)) + enddo end function intrinsic_functions_unit_tests !> True if the cube of cuberoot(val) does not closely match val. False otherwise. @@ -318,4 +475,62 @@ logical function Test_cuberoot(verbose, val) endif end function Test_cuberoot +!> True if exp_reprod(val) does not closely match the intrinsic exp(val). False otherwise. +logical function Test_exp_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test [nondim] + ! Local variables + real :: e_ref ! The intrinsic exponential of val [nondim] + real :: relerr ! The relative difference between exp_reprod(val) and exp(val) [nondim] + + e_ref = exp(val) + relerr = abs(exp_reprod(val) - e_ref) / e_ref + Test_exp_reprod = (relerr > 1.0e-14) + + if (Test_exp_reprod) then + write(stdout, '("For val = ",ES22.15,", exp_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", exp_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_exp_reprod + +!> True if log_reprod(val) does not closely match the intrinsic log(val). False otherwise. +logical function Test_log_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test, val > 0 [nondim] + ! Local variables + real :: l_ref ! The intrinsic natural logarithm of val [nondim] + real :: relerr ! The relative difference between log_reprod(val) and log(val) [nondim] + + l_ref = log(val) + if (l_ref /= 0.0) then ; relerr = abs(log_reprod(val) - l_ref) / abs(l_ref) + else ; relerr = abs(log_reprod(val) - l_ref) ; endif + Test_log_reprod = (relerr > 1.0e-13) + + if (Test_log_reprod) then + write(stdout, '("For val = ",ES22.15,", log_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", log_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_log_reprod + +!> True if erfc_reprod(val) does not closely match the intrinsic erfc(val). False otherwise. +logical function Test_erfc_reprod(verbose, val) + logical, intent(in) :: verbose !< If true, write results to stdout + real, intent(in) :: val !< The real value to test, val >= 0 [nondim] + ! Local variables + real :: e_ref ! The intrinsic complementary error function of val [nondim] + real :: relerr ! The relative difference between erfc_reprod(val) and erfc(val) [nondim] + + e_ref = erfc(val) + relerr = abs(erfc_reprod(val) - e_ref) / abs(e_ref) + Test_erfc_reprod = (relerr > 1.0e-12) + + if (Test_erfc_reprod) then + write(stdout, '("For val = ",ES22.15,", erfc_reprod relative error = ",ES9.2," <-- FAIL")') val, relerr + elseif (verbose) then + write(stdout, '("For val = ",ES22.15,", erfc_reprod relative error = ",ES9.2)') val, relerr + endif +end function Test_erfc_reprod + end module MOM_intrinsic_functions diff --git a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 index 15cb9bc9ea..0cc0382581 100644 --- a/src/parameterizations/lateral/MOM_thickness_diffuse.F90 +++ b/src/parameterizations/lateral/MOM_thickness_diffuse.F90 @@ -13,6 +13,8 @@ module MOM_thickness_diffuse use MOM_error_handler, only : MOM_error, FATAL, WARNING, is_root_pe use MOM_EOS, only : calculate_density, calculate_density_derivs, EOS_domain use MOM_EOS, only : calculate_density_second_derivs +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_file_parser, only : get_param, log_version, param_file_type use MOM_grid, only : ocean_grid_type use MOM_io, only : MOM_read_data, slasher @@ -706,6 +708,12 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drho_dT_dT_hr ! The second derivative of density with temperature at h (+1) points [R C-2 ~> kg m-3 degC-2] real :: uhtot(SZIB_(G),SZJ_(G)) ! The vertical sum of uhD [H L2 T-1 ~> m3 s-1 or kg s-1]. real :: vhtot(SZI_(G),SZJB_(G)) ! The vertical sum of vhD [H L2 T-1 ~> m3 s-1 or kg s-1]. + real :: pres_us, T_us, S_us ! Scalar (per-iteration) pressure, temperature and salinity on the + ! interface at a u-point, for the fused device-callable EOS path. + real :: dT_us, dS_us ! Scalar density derivatives with T and S at a u-point. + real :: pres_vs, T_vs, S_vs ! Scalar (per-iteration) pressure, temperature and salinity on the + ! interface at a v-point, for the fused device-callable EOS path. + real :: dT_vs, dS_vs ! Scalar density derivatives with T and S at a v-point. real, dimension(SZIB_(G)) :: & T_u, & ! Temperature on the interface at the u-point [C ~> degC]. S_u, & ! Salinity on the interface at the u-point [S ~> ppt]. @@ -733,11 +741,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV real :: drdjA, drdjB ! Along layer meridional potential density gradients in the layers above (A) ! and below (B) the interface times the grid spacing [R ~> kg m-3]. real :: drdkL, drdkR ! Vertical density differences across an interface [R ~> kg m-3]. - real :: drdi_u(SZIB_(G),SZK_(GV)) ! Copy of drdi at u-points [R ~> kg m-3]. - real :: drdj_v(SZI_(G),SZK_(GV)) ! Copy of drdj at v-points [R ~> kg m-3]. - real :: drdkDe_u(SZIB_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at u-points + real :: drdi_u(SZIB_(G),SZJ_(G),SZK_(GV)) ! Copy of drdi at u-points [R ~> kg m-3]. + real :: drdj_v(SZI_(G),SZJB_(G),SZK_(GV)) ! Copy of drdj at v-points [R ~> kg m-3]. + real :: drdkDe_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at u-points ! [Z R ~> kg m-2]. - real :: drdkDe_v(SZI_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at v-points + real :: drdkDe_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Lateral difference of product of drdk and e at v-points ! [Z R ~> kg m-2]. real :: hg2A, hg2B, hg2L, hg2R ! Squares of geometric mean thicknesses [H2 ~> m2 or kg2 m-4]. real :: haA, haB, haL, haR ! Arithmetic mean thicknesses [H ~> m or kg m-2]. @@ -751,16 +759,16 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV real :: dz_harm ! Harmonic mean layer vertical extent [Z ~> m]. real :: c2_dz_u(SZIB_(G),SZK_(GV)+1) ! Wave speed squared divided by dz at u-points [L2 Z-1 T-2 ~> m s-2] real :: c2_dz_v(SZI_(G),SZK_(GV)+1) ! Wave speed squared divided by dz at v-points [L2 Z-1 T-2 ~> m s-2] - real :: dzN2_u(SZIB_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above u-points times + real :: dzN2_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above u-points times ! rescaling factors from vertical to horizontal distances [L2 Z-1 T-2 ~> m s-2] - real :: dzN2_v(SZI_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above v-points times + real :: dzN2_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Vertical extent times N2 at interfaces above v-points times ! rescaling factors from vertical to horizontal distances [L2 Z-1 T-2 ~> m s-2] real :: Sfn_est ! A preliminary estimate (before limiting) of the overturning ! streamfunction [H L2 T-1 ~> m3 s-1 or kg s-1]. - real :: Sfn_unlim_u(SZIB_(G),SZK_(GV)+1) ! Volume streamfunction for u-points [Z L2 T-1 ~> m3 s-1] - real :: Sfn_unlim_v(SZI_(G),SZK_(GV)+1) ! Volume streamfunction for v-points [Z L2 T-1 ~> m3 s-1] - real :: slope2_Ratio_u(SZIB_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] - real :: slope2_Ratio_v(SZI_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] + real :: Sfn_unlim_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! Volume streamfunction for u-points [Z L2 T-1 ~> m3 s-1] + real :: Sfn_unlim_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! Volume streamfunction for v-points [Z L2 T-1 ~> m3 s-1] + real :: slope2_Ratio_u(SZIB_(G),SZJ_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] + real :: slope2_Ratio_v(SZI_(G),SZJB_(G),SZK_(GV)+1) ! The ratio of the slope squared to slope_max squared [nondim] real :: Sfn_in_h ! The overturning streamfunction [H L2 T-1 ~> m3 s-1 or kg s-1] (note that ! the units are different from other Sfn vars). real :: Sfn_safe ! The streamfunction that goes linearly back to 0 at the surface @@ -784,6 +792,11 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV integer :: nk_linear ! The number of layers over which the streamfunction goes to 0. real :: G_rho0 ! g/Rho0 [L2 R-1 Z-1 T-2 ~> m4 kg-1 s-2]. real :: Rho_avg ! The in situ density averaged to an interface [R ~> kg m-3] + integer :: eos_form ! The equation-of-state form code, for the device-callable EOS dispatcher. + real :: eos_kg_m3_to_R ! Factor converting the EOS kernel's density to model units [R m3 kg-1 ~> 1] + real :: eos_C_to_degC ! Factor converting model temperature to the EOS kernel's degC [degC C-1 ~> 1] + real :: eos_S_to_ppt ! Factor converting model salinity to the EOS kernel's ppt [ppt S-1 ~> 1] + real :: eos_RL2_T2_to_Pa ! Factor converting model pressure to Pa [Pa T2 R-1 L-2 ~> 1] real :: N2_floor ! A floor for N2 to avoid degeneracy in the elliptic solver ! times unit conversion factors [L2 Z-2 T-2 ~> s-2] real :: N2_unlim ! An unlimited estimate of the buoyancy frequency @@ -801,6 +814,17 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! applying limiters [H L2 T-1 ~> m3 s-1 or kg s-1] real, allocatable :: skeb_gm_work(:,:) ! Temp array to hold GM work for SKEB real, allocatable :: skeb_ebt_norm2(:,:) ! Used to normalize EBT for SKEB + real :: GM_src_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of MEKE%GM_src [R Z L2 T-3 ~> W m-2]. + ! MEKE%GM_src is an allocatable component of a pointer/derived + ! type, so it is copied to/from this plain array around the + ! device region rather than deep-mapped (the pointer trap). + real :: GMwork_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of CS%GMwork [R Z L2 T-3 ~> W m-2]. + real :: p_surf_loc(SZI_(G),SZJ_(G)) ! Hoisted plain-array copy of the surface pressure tv%p_surf + ! [R L2 T-2 ~> Pa] (0 where tv%p_surf is not associated), so the + ! device reads a plain array rather than a pointer component of tv. + logical :: have_GM_src ! True if MEKE%GM_src is allocated (evaluated on the host). + logical :: have_GMwork ! True if CS%GMwork is allocated (evaluated on the host). + logical :: have_p_surf ! True if tv%p_surf is associated (evaluated on the host). logical :: present_slope_x, present_slope_y, calc_derivatives integer, dimension(2) :: EOSdom_u ! The shifted I-computational domain to use for equation of @@ -857,37 +881,98 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (CS%use_FGNV_streamfn .and. .not. associated(cg1)) call MOM_error(FATAL, & "cg1 must be associated when using FGNV streamfunction.") - !$OMP parallel default(shared) + ! --------------------------------------------------------------------------------------------- + ! GPU offload of thickness_diffuse_full. One target data region spans the whole compute below. + ! Notes on residency: tv%T/tv%S are NOT mapped — the density inputs on device are the host-filled + ! local arrays T and S (from vert_fill_TS), copied in fresh with map(to:). MEKE%GM_src and + ! CS%GMwork are allocatable components of a pointer/derived type, so they are hoisted to the plain + ! local arrays GM_src_loc/GMwork_loc (mapped, written on device, copied back) rather than + ! deep-mapped (the pointer-in-derived-type trap). The persistent, diabatic-mutated h is refreshed + ! with update to(h) (map-to-present is a no-copy). + ! --------------------------------------------------------------------------------------------- +#ifdef __NVCOMPILER_OPENMP_GPU + if (use_stanley) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: USE_STANLEY_GM is not supported on device.") + if (CS%use_FGNV_streamfn) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the FGNV streamfunction is not supported on device.") + if (skeb_use_gm) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: SKEB GM work is not supported on device.") + if (allocated(tv%SpV_avg)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the non-Boussinesq SpV_avg path is not supported on device.") + if ((CS%id_slope_x>0) .or. (CS%id_slope_y>0) .or. (CS%id_sfn_x>0) .or. (CS%id_sfn_y>0) .or. & + (CS%id_sfn_unlim_x>0) .or. (CS%id_sfn_unlim_y>0)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: the streamfunction/slope diagnostics are not supported on device.") +#endif + + have_p_surf = associated(tv%p_surf) + have_GM_src = allocated(MEKE%GM_src) + have_GMwork = allocated(CS%GMwork) + + ! Hoist MEKE%GM_src / CS%GMwork to plain local arrays (copied back after the region). + GM_src_loc(:,:) = 0.0 ; GMwork_loc(:,:) = 0.0 + if (have_GM_src) then ; do j=js,je ; do i=is,ie ; GM_src_loc(i,j) = MEKE%GM_src(i,j) ; enddo ; enddo ; endif + if (have_GMwork) then ; do j=js,je ; do i=is,ie ; GMwork_loc(i,j) = CS%GMwork(i,j) ; enddo ; enddo ; endif + + ! Hoist the surface pressure onto a plain array so the device never dereferences the tv%p_surf pointer. + p_surf_loc(:,:) = 0.0 + if (have_p_surf) then + do j=js-1,je+1 ; do i=is-1,ie+1 ; p_surf_loc(i,j) = tv%p_surf(i,j) ; enddo ; enddo + endif + +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target enter data map(to: h, e, dz, T, S, Kh_u, Kh_v, int_slope_u, int_slope_v, CS, & + !$omp Slope_x_PE, Slope_y_PE, hN2_x_PE, hN2_y_PE, GM_src_loc, GMwork_loc, p_surf_loc) & + !$omp map(alloc: pres, h_avail, h_frac, h_avail_rsum, uhtot, vhtot, Work_u, Work_v, & + !$omp uhD, vhD, drdi_u, drdj_v, drdkDe_u, drdkDe_v, dzN2_u, dzN2_v, & + !$omp Sfn_unlim_u, Sfn_unlim_v, slope2_Ratio_u, slope2_Ratio_v) + if (present_slope_x) then + !$omp target enter data map(to: slope_x) + endif + if (present_slope_y) then + !$omp target enter data map(to: slope_y) + endif + ! h is persistently mapped and mutated on the host by diabatic, so refresh it (map-to-present + ! does not copy). + !$omp target update to(h) +#endif + ! Find the maximum and minimum permitted streamfunction. - !$OMP do +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js-1,je+1 ; do i=is-1,ie+1 h_avail_rsum(i,j,1) = 0.0 - pres(i,j,1) = 0.0 - if (associated(tv%p_surf)) then ; pres(i,j,1) = tv%p_surf(i,j) ; endif - + pres(i,j,1) = p_surf_loc(i,j) h_avail(i,j,1) = max(I4dt*G%areaT(i,j)*(h(i,j,1)-GV%Angstrom_H),0.0) h_avail_rsum(i,j,2) = h_avail(i,j,1) h_frac(i,j,1) = 1.0 pres(i,j,2) = pres(i,j,1) + (GV%g_Earth*GV%H_to_RZ) * h(i,j,1) enddo ; enddo - do j=js-1,je+1 - do k=2,nz ; do i=is-1,ie+1 + ! Running sum up the column: a downward vertical recurrence, so collapse(2) over (j,i) + serial K. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif + do j=js-1,je+1 ; do i=is-1,ie+1 + do k=2,nz h_avail(i,j,k) = max(I4dt*G%areaT(i,j)*(h(i,j,k)-GV%Angstrom_H),0.0) h_avail_rsum(i,j,k+1) = h_avail_rsum(i,j,k) + h_avail(i,j,k) h_frac(i,j,k) = 0.0 ; if (h_avail(i,j,k) > 0.0) & h_frac(i,j,k) = h_avail(i,j,k) / h_avail_rsum(i,j,k+1) pres(i,j,K+1) = pres(i,j,K) + (GV%g_Earth*GV%H_to_RZ) * h(i,j,k) - enddo ; enddo - enddo - !$OMP do + enddo + enddo ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js,je ; do I=is-1,ie uhtot(I,j) = 0.0 ; Work_u(I,j) = 0.0 enddo ; enddo - !$OMP do +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do J=js-1,je ; do i=is,ie vhtot(i,J) = 0.0 ; Work_v(i,J) = 0.0 enddo ; enddo - !$OMP end parallel if (CS%id_sfn_x > 0) then ; diag_sfn_x(:,:,1) = 0.0 ; diag_sfn_x(:,:,nz+1) = 0.0 ; endif if (CS%id_sfn_y > 0) then ; diag_sfn_y(:,:,1) = 0.0 ; diag_sfn_y(:,:,nz+1) = 0.0 ; endif @@ -898,41 +983,59 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV EOSdom_v(:) = EOS_domain(G%HI) EOSdom_h1(:) = EOS_domain(G%HI, halo=1) + ! Resolve the EOS form and its unit-scaling once (host v-table lookup), so the density derivatives + ! can be evaluated through the device-callable elementwise dispatcher instead of the polymorphic + ! array interface. The scaling is always applied; for the unscaled case the factors are exactly + ! 1.0, reproducing the former array calculate_density_derivs calls bit-for-bit. + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (use_EOS) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "thickness_diffuse_full GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") +#endif + endif + + ! Zero the FGNV vertical-extent-times-N2 boundary interface values. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif + do j=js,je ; do I=is-1,ie + dzN2_u(I,j,1) = 0. ; dzN2_u(I,j,nz+1) = 0. + enddo ; enddo + + ! Calculate the zonal density gradients, slopes and unlimited streamfunction. The 3-D promoted + ! scratch (drdi_u, drdkDe_u, dzN2_u, slope2_Ratio_u, Sfn_unlim_u, Slope_x_PE, hN2_x_PE) is shared + ! so it survives into the separate flux j-loop below; each j slice is written disjointly. This + ! loop has no vertical recurrence, so on the GPU it collapses over (j,K,I). +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(3) & + !$omp private(drdiA,drdiB,drdkL,drdkR,pres_us,T_us,S_us,dT_us,dS_us, & + !$omp hg2A,hg2B,hg2L,hg2R,haA, & + !$omp N2_unlim,haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$omp dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$omp drdx,mag_grad2,Slope,calc_derivatives) +#else !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,G,GV,US,pres,T,S, & !$OMP nk_linear,IsdB,tv,h,h_neglect,e,dz,dz_neglect,dz_neglect2, & - !$OMP h_neglect2,hn_2,I_slope_max2,int_slope_u,KH_u,uhtot, & - !$OMP h_frac,h_avail_rsum,uhD,h_avail,Work_u,CS,slope_x,cg1, & - !$OMP diag_sfn_x,diag_sfn_unlim_x,N2_floor,EOSdom_u,EOSdom_h1, & - !$OMP use_stanley,present_slope_x,G_rho0,Slope_x_PE,hN2_x_PE) & - !$OMP private(drdiA,drdiB,drdkL,drdkR,pres_u,T_u,S_u,G_scale, & - !$OMP drho_dT_u,drho_dS_u,hg2A,hg2B,hg2L,hg2R,haA, & - !$OMP drho_dT_dT_h,scrap,pres_h,T_h,S_h,N2_unlim, & - !$OMP haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & - !$OMP dzg2A,dzg2B,dzaA,dzaB,dz_harm,Z_to_H, & - !$OMP drdx,mag_grad2,Slope,slope2_Ratio_u,dzN2_u, & - !$OMP Sfn_unlim_u,Rho_avg,drdi_u,drdkDe_u,c2_dz_u, & - !$OMP Sfn_safe,Sfn_est,Sfn_in_h,calc_derivatives) + !$OMP h_neglect2,hn_2,I_slope_max2,int_slope_u,KH_u, & + !$OMP h_frac,h_avail_rsum,uhD,h_avail,CS,slope_x,cg1, & + !$OMP diag_sfn_unlim_x,N2_floor,EOSdom_u,EOSdom_h1, & + !$OMP use_stanley,present_slope_x,G_rho0,Slope_x_PE,hN2_x_PE, & + !$OMP drdi_u,drdkDe_u,dzN2_u,slope2_Ratio_u,Sfn_unlim_u, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & + !$OMP private(drdiA,drdiB,drdkL,drdkR,pres_us,T_us,S_us,dT_us,dS_us, & + !$OMP hg2A,hg2B,hg2L,hg2R,haA,drho_dT_dT_h,scrap,pres_h,T_h,S_h, & + !$OMP N2_unlim,haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$OMP dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$OMP drdx,mag_grad2,Slope,calc_derivatives) +#endif do j=js,je - do I=is-1,ie ; dzN2_u(I,1) = 0. ; dzN2_u(I,nz+1) = 0. ; enddo do K=nz,2,-1 - if (find_work .and. .not.(use_EOS)) then - drdiA = 0.0 ; drdiB = 0.0 - drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL - endif - - calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & - (find_work .or. .not. present_slope_x .or. CS%use_FGNV_streamfn .or. use_stanley) - - ! Calculate the zonal fluxes and gradients. - if (calc_derivatives) then - do I=is-1,ie - pres_u(I) = 0.5*(pres(i,j,K) + pres(i+1,j,K)) - T_u(I) = 0.25*((T(i,j,k) + T(i+1,j,k)) + (T(i,j,k-1) + T(i+1,j,k-1))) - S_u(I) = 0.25*((S(i,j,k) + S(i+1,j,k)) + (S(i,j,k-1) + S(i+1,j,k-1))) - enddo - call calculate_density_derivs(T_u, S_u, pres_u, drho_dT_u, drho_dS_u, & - tv%eqn_of_state, EOSdom_u) - endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then do i=is-1,ie+1 pres_h(i) = pres(i,j,K) @@ -946,24 +1049,45 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV scrap, scrap, drho_dT_dT_h, scrap, scrap, & tv%eqn_of_state, EOSdom_h1) endif +#endif do I=is-1,ie + calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & + (find_work .or. .not. present_slope_x .or. CS%use_FGNV_streamfn .or. use_stanley) + + if (find_work .and. .not.(use_EOS)) then + drdiA = 0.0 ; drdiB = 0.0 + drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL + endif + + ! Fill the interface T, S and pressure at the u-point and evaluate the density derivatives + ! through the device-callable elementwise EOS dispatcher (fused into this loop from the + ! former separate do-I EOS-fill loop, which read the same I in the same iteration). if (calc_derivatives) then + pres_us = 0.5*(pres(i,j,K) + pres(i+1,j,K)) + T_us = 0.25*((T(i,j,k) + T(i+1,j,k)) + (T(i,j,k-1) + T(i+1,j,k-1))) + S_us = 0.25*((S(i,j,k) + S(i+1,j,k)) + (S(i,j,k-1) + S(i+1,j,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_us, eos_S_to_ppt*S_us, & + eos_RL2_T2_to_Pa*pres_us, dT_us, dS_us) + dT_us = (eos_kg_m3_to_R*eos_C_to_degC) * dT_us + dS_us = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_us + ! Estimate the horizontal density gradients along layers. - drdiA = drho_dT_u(I) * (T(i+1,j,k-1)-T(i,j,k-1)) + & - drho_dS_u(I) * (S(i+1,j,k-1)-S(i,j,k-1)) - drdiB = drho_dT_u(I) * (T(i+1,j,k)-T(i,j,k)) + & - drho_dS_u(I) * (S(i+1,j,k)-S(i,j,k)) + drdiA = dT_us * (T(i+1,j,k-1)-T(i,j,k-1)) + & + dS_us * (S(i+1,j,k-1)-S(i,j,k-1)) + drdiB = dT_us * (T(i+1,j,k)-T(i,j,k)) + & + dS_us * (S(i+1,j,k)-S(i,j,k)) ! Estimate the vertical density gradients times the grid spacing. - drdkL = (drho_dT_u(I) * (T(i,j,k)-T(i,j,k-1)) + & - drho_dS_u(I) * (S(i,j,k)-S(i,j,k-1))) - drdkR = (drho_dT_u(I) * (T(i+1,j,k)-T(i+1,j,k-1)) + & - drho_dS_u(I) * (S(i+1,j,k)-S(i+1,j,k-1))) - drdkDe_u(I,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) + drdkL = (dT_us * (T(i,j,k)-T(i,j,k-1)) + & + dS_us * (S(i,j,k)-S(i,j,k-1))) + drdkR = (dT_us * (T(i+1,j,k)-T(i+1,j,k-1)) + & + dS_us * (S(i+1,j,k)-S(i+1,j,k-1))) + drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode - drdkDe_u(I,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) + drdkDe_u(I,j,K) = (drdkR * e(i+1,j,K)) - (drdkL * e(i,j,K)) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in ! the EOS rectifying SGS temperature anomalies @@ -972,7 +1096,8 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drdiB = drdiB + 0.5 * ((drho_dT_dT_h(i+1) * tv%varT(i+1,j,k)) - & (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif - if (find_work) drdi_u(I,k) = drdiB +#endif + if (find_work) drdi_u(I,j,k) = drdiB if (k > nk_linear) then if (use_EOS) then @@ -1015,14 +1140,14 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzaA = 0.5*(dz(i,j,k-1) + dz(i+1,j,k-1)) + dz_neglect dzaB = 0.5*(dz(i,j,k) + dz(i+1,j,k)) + dz_neglect ! dzN2_u is used with the FGNV streamfunction formulation - dzN2_u(I,K) = (0.5 * ( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) + dzN2_u(I,j,K) = (0.5 * ( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) if (find_work .and. CS%GM_src_alt) & hN2_x_PE(I,j,k) = (0.5 * ( hg2A / haA + hg2B / haB )) * max(N2_unlim, N2_floor) endif if (present_slope_x) then Slope = slope_x(I,j,k) - slope2_Ratio_u(I,K) = Slope**2 * I_slope_max2 + slope2_Ratio_u(I,j,K) = Slope**2 * I_slope_max2 else ! Use the harmonic mean thicknesses to weight the horizontal gradients. ! These unnormalized weights have been rearranged to minimize divisions. @@ -1036,10 +1161,10 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV mag_grad2 = (US%Z_to_L*drdx)**2 + drdz**2 if (mag_grad2 > 0.0) then Slope = drdx / sqrt(mag_grad2) - slope2_Ratio_u(I,K) = Slope**2 * I_slope_max2 + slope2_Ratio_u(I,j,K) = Slope**2 * I_slope_max2 else ! Just in case mag_grad2 = 0 ever. Slope = 0.0 - slope2_Ratio_u(I,K) = 1.0e20 ! Force the use of the safe streamfunction. + slope2_Ratio_u(I,j,K) = 1.0e20 ! Force the use of the safe streamfunction. endif endif @@ -1047,7 +1172,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! that ignore density gradients along layers. Slope = (1.0 - int_slope_u(I,j,K)) * Slope + & int_slope_u(I,j,K) * ((e(i+1,j,K)-e(i,j,K)) * G%IdxCu(I,j)) - slope2_Ratio_u(I,K) = (1.0 - int_slope_u(I,j,K)) * slope2_Ratio_u(I,K) + slope2_Ratio_u(I,j,K) = (1.0 - int_slope_u(I,j,K)) * slope2_Ratio_u(I,j,K) if (CS%MEKE_src_slope_bug) then Slope_x_PE(I,j,k) = MIN(Slope, CS%slope_max) @@ -1056,27 +1181,29 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (Slope > CS%slope_max) Slope_x_PE(I,j,k) = CS%slope_max if (Slope < -CS%slope_max) Slope_x_PE(I,j,k) = -CS%slope_max endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope +#endif ! Estimate the streamfunction at each interface [H L2 T-1 ~> m3 s-1 or kg s-1]. - Sfn_unlim_u(I,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope + Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope ! Avoid moving dense water upslope from below the level of ! the bottom on the receiving side. - if (Sfn_unlim_u(I,K) > 0.0) then ! The flow below this interface is positive. + if (Sfn_unlim_u(I,j,K) > 0.0) then ! The flow below this interface is positive. if (e(i,j,K) < e(i+1,j,nz+1)) then - Sfn_unlim_u(I,K) = 0.0 ! This is not uhtot, because it may compensate for + Sfn_unlim_u(I,j,K) = 0.0 ! This is not uhtot, because it may compensate for ! deeper flow in very unusual cases. elseif (e(i+1,j,nz+1) > e(i,j,K+1)) then ! Scale the transport with the fraction of the donor layer above ! the bottom on the receiving side. - Sfn_unlim_u(I,K) = Sfn_unlim_u(I,K) * ((e(i,j,K) - e(i+1,j,nz+1)) / & + Sfn_unlim_u(I,j,K) = Sfn_unlim_u(I,j,K) * ((e(i,j,K) - e(i+1,j,nz+1)) / & ((e(i,j,K) - e(i,j,K+1)) + dz_neglect)) endif else - if (e(i+1,j,K) < e(i,j,nz+1)) then ; Sfn_unlim_u(I,K) = 0.0 + if (e(i+1,j,K) < e(i,j,nz+1)) then ; Sfn_unlim_u(I,j,K) = 0.0 elseif (e(i,j,nz+1) > e(i+1,j,K+1)) then - Sfn_unlim_u(I,K) = Sfn_unlim_u(I,K) * ((e(i+1,j,K) - e(i,j,nz+1)) / & + Sfn_unlim_u(I,j,K) = Sfn_unlim_u(I,j,K) * ((e(i+1,j,K) - e(i,j,nz+1)) / & ((e(i+1,j,K) - e(i+1,j,K+1)) + dz_neglect)) endif endif @@ -1087,18 +1214,38 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Slope = (e(i+1,j,K)-e(i,j,K)) * G%IdxCu_OBCmask(I,j) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_x > 0) CS%diagSlopeX(I,j,k) = Slope - Sfn_unlim_u(I,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope - dzN2_u(I,K) = GV%g_prime(K) +#endif + Sfn_unlim_u(I,j,K) = -(KH_u(I,j,K)*G%dy_Cu(I,j))*Slope + dzN2_u(I,j,K) = GV%g_prime(K) endif ! if (use_EOS) else ! if (k > nk_linear) - dzN2_u(I,K) = N2_floor * dz_neglect - Sfn_unlim_u(I,K) = 0. + dzN2_u(I,j,K) = N2_floor * dz_neglect + Sfn_unlim_u(I,j,K) = 0. endif ! if (k > nk_linear) - if (CS%id_sfn_unlim_x>0) diag_sfn_unlim_x(I,j,K) = Sfn_unlim_u(I,K) +#ifndef __NVCOMPILER_OPENMP_GPU + if (CS%id_sfn_unlim_x>0) diag_sfn_unlim_x(I,j,K) = Sfn_unlim_u(I,j,K) +#endif enddo ! i-loop enddo ! k-loop + enddo ! end of j-loop + ! Vertically accumulate the zonal transports (limited by the available mass) and, if requested, the + ! energy conversion. uhtot has a downward vertical recurrence, so this loop is collapse(2) over + ! (j,I) with a serial K inside on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) & + !$omp private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale) +#else + !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,GV,G,US,tv,h,e,hn_2, & + !$OMP nk_linear,h_frac,h_avail_rsum,h_avail,uhtot,uhD,Work_u, & + !$OMP slope2_Ratio_u,Sfn_unlim_u,drdkDe_u,drdi_u,dzN2_u,cg1, & + !$OMP dz,dz_neglect,CS,diag_sfn_x) & + !$OMP private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale,dz_harm,c2_dz_u) +#endif + do j=js,je +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%use_FGNV_streamfn) then do k=1,nz ; do I=is-1,ie ; if (G%OBCmaskCu(I,j)>0.) then dz_harm = max( dz_neglect, & @@ -1110,20 +1257,24 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV do I=is-1,ie if (G%OBCmaskCu(I,j)>0.) then do K=2,nz - Sfn_unlim_u(I,K) = (1. + CS%FGNV_scale) * Sfn_unlim_u(I,K) + Sfn_unlim_u(I,j,K) = (1. + CS%FGNV_scale) * Sfn_unlim_u(I,j,K) enddo - call streamfn_solver(nz, c2_dz_u(I,:), dzN2_u(I,:), Sfn_unlim_u(I,:)) + call streamfn_solver(nz, c2_dz_u(I,:), dzN2_u(I,j,:), Sfn_unlim_u(I,j,:)) else do K=2,nz - Sfn_unlim_u(I,K) = 0. + Sfn_unlim_u(I,j,K) = 0. enddo endif enddo endif +#endif - do K=nz,2,-1 - do I=is-1,ie + do I=is-1,ie + do K=nz,2,-1 +#ifdef __NVCOMPILER_OPENMP_GPU + Z_to_H = GV%Z_to_H +#else if (allocated(tv%SpV_avg) .and. (find_work .or. (k > nk_linear)) ) then Rho_avg = ( ((h(i,j,k) + h(i,j,k-1)) + (h(i+1,j,k) + h(i+1,j,k-1))) + 4.0*hn_2 ) / & ( (((h(i,j,k)+hn_2) * tv%SpV_avg(i,j,k)) + ((h(i,j,k-1)+hn_2) * tv%SpV_avg(i,j,k-1))) + & @@ -1133,6 +1284,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Z_to_H = GV%Z_to_H endif +#endif if (k > nk_linear) then if (use_EOS) then @@ -1145,9 +1297,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif ! Determine the actual streamfunction at each interface. - Sfn_est = (Z_to_H*Sfn_unlim_u(I,K) + slope2_Ratio_u(I,K)*Sfn_safe) / (1.0 + slope2_Ratio_u(I,K)) + Sfn_est = (Z_to_H*Sfn_unlim_u(I,j,K) + slope2_Ratio_u(I,j,K)*Sfn_safe) / (1.0 + slope2_Ratio_u(I,j,K)) else ! When use_EOS is false, the layers are constant density. - Sfn_est = Z_to_H*Sfn_unlim_u(I,K) + Sfn_est = Z_to_H*Sfn_unlim_u(I,j,K) endif ! Make sure that there is enough mass above to allow the streamfunction @@ -1159,7 +1311,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV uhD(I,j,k) = max(min((Sfn_in_H - uhtot(I,j)), h_avail(i,j,k)), & -h_avail(i+1,j,k)) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_x>0) diag_sfn_x(I,j,K) = diag_sfn_x(I,j,K+1) + uhD(I,j,k) +#endif ! sfn_x(I,j,K) = max(min(Sfn_in_h, uhtot(I,j)+h_avail(i,j,k)), & ! uhtot(I,j)-h_avail(i+1,j,K)) ! sfn_slope_x(I,j,K) = max(uhtot(I,j)-h_avail(i+1,j,k), & @@ -1177,7 +1331,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV uhD(I,j,k) = -uhtot(I,j) * h_frac(i+1,j,k) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_x>0) diag_sfn_x(I,j,K) = diag_sfn_x(I,j,K+1) + uhD(I,j,k) +#endif ! if (sfn_slope_x(I,j,K+1) <= 0.0) then ! sfn_slope_x(I,j,K) = sfn_slope_x(I,j,K+1) * (1.0 - h_frac(i,j,k)) ! else @@ -1195,57 +1351,57 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! A second order centered estimate is used for the density transferred ! between water columns. +#ifdef __NVCOMPILER_OPENMP_GPU + G_scale = GV%g_Earth * GV%H_to_Z +#else if (allocated(tv%SpV_avg)) then G_scale = GV%H_to_RZ * GV%g_Earth / Rho_avg else G_scale = GV%g_Earth * GV%H_to_Z endif +#endif Work_u(I,j) = Work_u(I,j) + G_scale * & - ( uhtot(I,j) * drdkDe_u(I,K) - & - (uhD(I,j,k) * drdi_u(I,k)) * 0.25 * & + ( uhtot(I,j) * drdkDe_u(I,j,K) - & + (uhD(I,j,k) * drdi_u(I,j,k)) * 0.25 * & ((e(i,j,K) + e(i,j,K+1)) + (e(i+1,j,K) + e(i+1,j,K+1))) ) endif - enddo - enddo ! end of k-loop + enddo ! end of k-loop + enddo ! end of i-loop enddo ! end of j-loop ! Calculate the meridional fluxes and gradients. + ! As for the u-points: 3-D scratch (drdj_v, drdkDe_v, dzN2_v, slope2_Ratio_v, Sfn_unlim_v, + ! Slope_y_PE, hN2_y_PE) is shared so it survives into the flux j-loop; no vertical recurrence, so + ! this loop collapses over (J,K,i) on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(3) & + !$omp private(drdjA,drdjB,drdkL,drdkR,pres_vs,T_vs,S_vs,dT_vs,dS_vs, & + !$omp hg2A,hg2B,hg2L,hg2R,haA, & + !$omp haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz, & + !$omp dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$omp drdy,mag_grad2,Slope,N2_unlim,calc_derivatives) +#else !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,G,GV,US,pres,T,S,dz, & !$OMP nk_linear,IsdB,tv,h,h_neglect,e,dz_neglect,dz_neglect2, & - !$OMP h_neglect2,int_slope_v,KH_v,vhtot,h_frac,h_avail_rsum, & - !$OMP I_slope_max2,vhD,h_avail,Work_v,CS,slope_y,cg1,hn_2,& - !$OMP diag_sfn_y,diag_sfn_unlim_y,N2_floor,EOSdom_v,use_stanley,& - !$OMP present_slope_y,G_rho0,Slope_y_PE,hN2_y_PE) & - !$OMP private(drdjA,drdjB,drdkL,drdkR,pres_v,T_v,S_v,S_h,S_hr, & - !$OMP drho_dT_v,drho_dS_v,hg2A,hg2B,hg2L,hg2R,haA,G_scale, & + !$OMP h_neglect2,int_slope_v,KH_v, & + !$OMP I_slope_max2,CS,slope_y, & + !$OMP diag_sfn_unlim_y,N2_floor,EOSdom_v,use_stanley,& + !$OMP present_slope_y,G_rho0,Slope_y_PE,hN2_y_PE, & + !$OMP drdj_v,drdkDe_v,dzN2_v,slope2_Ratio_v,Sfn_unlim_v, & + !$OMP eos_form,eos_kg_m3_to_R,eos_C_to_degC,eos_S_to_ppt,eos_RL2_T2_to_Pa) & + !$OMP private(drdjA,drdjB,drdkL,drdkR,pres_vs,T_vs,S_vs,dT_vs,dS_vs, & + !$OMP S_h,S_hr,hg2A,hg2B,hg2L,hg2R,haA, & !$OMP drho_dT_dT_h,drho_dT_dT_hr,scrap,pres_h,T_h,T_hr, & !$OMP haB,haL,haR,dzaL,dzaR,wtA,wtB,wtL,wtR,drdz,pres_hr, & - !$OMP dzg2A,dzg2B,dzaA,dzaB,dz_harm,Z_to_H, & - !$OMP drdy,mag_grad2,Slope,slope2_Ratio_v,dzN2_v,N2_unlim, & - !$OMP Sfn_unlim_v,Rho_avg,drdj_v,drdkDe_v,c2_dz_v, & - !$OMP Sfn_safe,Sfn_est,Sfn_in_h,calc_derivatives) + !$OMP dzg2A,dzg2B,dzaA,dzaB,Z_to_H, & + !$OMP drdy,mag_grad2,Slope,N2_unlim,calc_derivatives) +#endif do J=js-1,je do K=nz,2,-1 - if (find_work .and. .not.(use_EOS)) then - drdjA = 0.0 ; drdjB = 0.0 - drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL - endif - - calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & - (find_work .or. .not. present_slope_y .or. CS%use_FGNV_streamfn .or. use_stanley) - - if (calc_derivatives) then - do i=is,ie - pres_v(i) = 0.5*(pres(i,j,K) + pres(i,j+1,K)) - T_v(i) = 0.25*((T(i,j,k) + T(i,j+1,k)) + (T(i,j,k-1) + T(i,j+1,k-1))) - S_v(i) = 0.25*((S(i,j,k) + S(i,j+1,k)) + (S(i,j,k-1) + S(i,j+1,k-1))) - enddo - call calculate_density_derivs(T_v, S_v, pres_v, drho_dT_v, drho_dS_v, & - tv%eqn_of_state, EOSdom_v) - endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then do i=is,ie pres_h(i) = pres(i,j,K) @@ -1266,23 +1422,43 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV scrap, scrap, drho_dT_dT_hr, scrap, scrap, & tv%eqn_of_state, EOSdom_v) endif +#endif do i=is,ie + calc_derivatives = use_EOS .and. (k >= nk_linear) .and. & + (find_work .or. .not. present_slope_y .or. CS%use_FGNV_streamfn .or. use_stanley) + + if (find_work .and. .not.(use_EOS)) then + drdjA = 0.0 ; drdjB = 0.0 + drdkL = GV%Rlay(k) - GV%Rlay(k-1) ; drdkR = drdkL + endif + + ! Fill the interface T, S and pressure at the v-point and evaluate the density derivatives + ! through the device-callable elementwise EOS dispatcher (fused from the former EOS-fill loop). if (calc_derivatives) then + pres_vs = 0.5*(pres(i,j,K) + pres(i,j+1,K)) + T_vs = 0.25*((T(i,j,k) + T(i,j+1,k)) + (T(i,j,k-1) + T(i,j+1,k-1))) + S_vs = 0.25*((S(i,j,k) + S(i,j+1,k)) + (S(i,j,k-1) + S(i,j+1,k-1))) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_vs, eos_S_to_ppt*S_vs, & + eos_RL2_T2_to_Pa*pres_vs, dT_vs, dS_vs) + dT_vs = (eos_kg_m3_to_R*eos_C_to_degC) * dT_vs + dS_vs = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_vs + ! Estimate the horizontal density gradients along layers. - drdjA = drho_dT_v(i) * (T(i,j+1,k-1)-T(i,j,k-1)) + & - drho_dS_v(i) * (S(i,j+1,k-1)-S(i,j,k-1)) - drdjB = drho_dT_v(i) * (T(i,j+1,k)-T(i,j,k)) + & - drho_dS_v(i) * (S(i,j+1,k)-S(i,j,k)) + drdjA = dT_vs * (T(i,j+1,k-1)-T(i,j,k-1)) + & + dS_vs * (S(i,j+1,k-1)-S(i,j,k-1)) + drdjB = dT_vs * (T(i,j+1,k)-T(i,j,k)) + & + dS_vs * (S(i,j+1,k)-S(i,j,k)) ! Estimate the vertical density gradients times the grid spacing. - drdkL = (drho_dT_v(i) * (T(i,j,k)-T(i,j,k-1)) + & - drho_dS_v(i) * (S(i,j,k)-S(i,j,k-1))) - drdkR = (drho_dT_v(i) * (T(i,j+1,k)-T(i,j+1,k-1)) + & - drho_dS_v(i) * (S(i,j+1,k)-S(i,j+1,k-1))) - drdkDe_v(i,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) + drdkL = (dT_vs * (T(i,j,k)-T(i,j,k-1)) + & + dS_vs * (S(i,j,k)-S(i,j,k-1))) + drdkR = (dT_vs * (T(i,j+1,k)-T(i,j+1,k-1)) + & + dS_vs * (S(i,j+1,k)-S(i,j+1,k-1))) + drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) elseif (find_work) then ! This is used in pure stacked SW mode - drdkDe_v(i,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) + drdkDe_v(i,J,K) = (drdkR * e(i,j+1,K)) - (drdkL * e(i,j,K)) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (use_stanley) then ! Correction to the horizontal density gradient due to nonlinearity in ! the EOS rectifying SGS temperature anomalies @@ -1291,8 +1467,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV drdjB = drdjB + 0.5 * ((drho_dT_dT_hr(i) * tv%varT(i,j+1,k)) - & (drho_dT_dT_h(i) * tv%varT(i,j,k)) ) endif +#endif - if (find_work) drdj_v(i,k) = drdjB + if (find_work) drdj_v(i,J,k) = drdjB if (k > nk_linear) then if (use_EOS) then @@ -1337,13 +1514,13 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV dzaB = 0.5*(dz(i,j,k) + dz(i,j+1,k)) + dz_neglect ! dzN2_v is used with the FGNV streamfunction formulation - dzN2_v(i,K) = (0.5*( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) + dzN2_v(i,J,K) = (0.5*( dzg2A / dzaA + dzg2B / dzaB )) * max(N2_unlim, N2_floor) if (find_work .and. CS%GM_src_alt) & hN2_y_PE(i,J,k) = (0.5*( hg2A / haA + hg2B / haB )) * max(N2_unlim, N2_floor) endif if (present_slope_y) then Slope = slope_y(i,J,k) - slope2_Ratio_v(i,K) = Slope**2 * I_slope_max2 + slope2_Ratio_v(i,J,K) = Slope**2 * I_slope_max2 else ! Use the harmonic mean thicknesses to weight the horizontal gradients. ! These unnormalized weights have been rearranged to minimize divisions. @@ -1357,10 +1534,10 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV mag_grad2 = (US%Z_to_L*drdy)**2 + drdz**2 if (mag_grad2 > 0.0) then Slope = drdy / sqrt(mag_grad2) - slope2_Ratio_v(i,K) = Slope**2 * I_slope_max2 + slope2_Ratio_v(i,J,K) = Slope**2 * I_slope_max2 else ! Just in case mag_grad2 = 0 ever. Slope = 0.0 - slope2_Ratio_v(i,K) = 1.0e20 ! Force the use of the safe streamfunction. + slope2_Ratio_v(i,J,K) = 1.0e20 ! Force the use of the safe streamfunction. endif endif @@ -1368,7 +1545,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! that ignore density gradients along layers. Slope = (1.0 - int_slope_v(i,J,K)) * Slope + & int_slope_v(i,J,K) * ((e(i,j+1,K)-e(i,j,K)) * G%IdyCv(i,J)) - slope2_Ratio_v(i,K) = (1.0 - int_slope_v(i,J,K)) * slope2_Ratio_v(i,K) + slope2_Ratio_v(i,J,K) = (1.0 - int_slope_v(i,J,K)) * slope2_Ratio_v(i,J,K) if (CS%MEKE_src_slope_bug) then Slope_y_PE(i,J,k) = MIN(Slope, CS%slope_max) @@ -1377,26 +1554,28 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV if (Slope > CS%slope_max) Slope_y_PE(i,J,k) = CS%slope_max if (Slope < -CS%slope_max) Slope_y_PE(i,J,k) = -CS%slope_max endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope +#endif - Sfn_unlim_v(i,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) + Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) ! Avoid moving dense water upslope from below the level of ! the bottom on the receiving side. - if (Sfn_unlim_v(i,K) > 0.0) then ! The flow below this interface is positive. + if (Sfn_unlim_v(i,J,K) > 0.0) then ! The flow below this interface is positive. if (e(i,j,K) < e(i,j+1,nz+1)) then - Sfn_unlim_v(i,K) = 0.0 ! This is not vhtot, because it may compensate for + Sfn_unlim_v(i,J,K) = 0.0 ! This is not vhtot, because it may compensate for ! deeper flow in very unusual cases. elseif (e(i,j+1,nz+1) > e(i,j,K+1)) then ! Scale the transport with the fraction of the donor layer above ! the bottom on the receiving side. - Sfn_unlim_v(i,K) = Sfn_unlim_v(i,K) * ((e(i,j,K) - e(i,j+1,nz+1)) / & + Sfn_unlim_v(i,J,K) = Sfn_unlim_v(i,J,K) * ((e(i,j,K) - e(i,j+1,nz+1)) / & ((e(i,j,K) - e(i,j,K+1)) + dz_neglect)) endif else - if (e(i,j+1,K) < e(i,j,nz+1)) then ; Sfn_unlim_v(i,K) = 0.0 + if (e(i,j+1,K) < e(i,j,nz+1)) then ; Sfn_unlim_v(i,J,K) = 0.0 elseif (e(i,j,nz+1) > e(i,j+1,K+1)) then - Sfn_unlim_v(i,K) = Sfn_unlim_v(i,K) * ((e(i,j+1,K) - e(i,j,nz+1)) / & + Sfn_unlim_v(i,J,K) = Sfn_unlim_v(i,J,K) * ((e(i,j+1,K) - e(i,j,nz+1)) / & ((e(i,j+1,K) - e(i,j+1,K+1)) + dz_neglect)) endif endif @@ -1407,18 +1586,37 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Slope = (e(i,j+1,K)-e(i,j,K)) * G%IdyCv_OBCmask(i,J) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_slope_y > 0) CS%diagSlopeY(I,j,k) = Slope - Sfn_unlim_v(i,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) - dzN2_v(i,K) = GV%g_prime(K) +#endif + Sfn_unlim_v(i,J,K) = -((KH_v(i,J,K)*G%dx_Cv(i,J))*Slope) + dzN2_v(i,J,K) = GV%g_prime(K) endif ! if (use_EOS) else ! if (k > nk_linear) - dzN2_v(i,K) = N2_floor * dz_neglect - Sfn_unlim_v(i,K) = 0. + dzN2_v(i,J,K) = N2_floor * dz_neglect + Sfn_unlim_v(i,J,K) = 0. endif ! if (k > nk_linear) - if (CS%id_sfn_unlim_y>0) diag_sfn_unlim_y(i,J,K) = Sfn_unlim_v(i,K) +#ifndef __NVCOMPILER_OPENMP_GPU + if (CS%id_sfn_unlim_y>0) diag_sfn_unlim_y(i,J,K) = Sfn_unlim_v(i,J,K) +#endif enddo ! i-loop enddo ! k-loop + enddo ! end of j-loop + ! Vertically accumulate the meridional transports (limited by available mass) and, if requested, + ! the energy conversion. vhtot recurs downward, so collapse(2) over (J,i) + serial K on the GPU. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) & + !$omp private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale) +#else + !$OMP parallel do default(none) shared(nz,is,ie,js,je,find_work,use_EOS,GV,G,US,tv,h,e,hn_2, & + !$OMP nk_linear,h_frac,h_avail_rsum,h_avail,vhtot,vhD,Work_v, & + !$OMP slope2_Ratio_v,Sfn_unlim_v,drdkDe_v,drdj_v,dzN2_v,cg1, & + !$OMP dz,dz_neglect,CS,diag_sfn_y) & + !$OMP private(Rho_avg,Z_to_H,Sfn_safe,Sfn_est,Sfn_in_H,G_scale,dz_harm,c2_dz_v) +#endif + do J=js-1,je +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%use_FGNV_streamfn) then do k=1,nz ; do i=is,ie ; if (G%OBCmaskCv(i,J)>0.) then dz_harm = max( dz_neglect, & @@ -1430,19 +1628,23 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV do i=is,ie if (G%OBCmaskCv(i,J)>0.) then do K=2,nz - Sfn_unlim_v(i,K) = (1. + CS%FGNV_scale) * Sfn_unlim_v(i,K) + Sfn_unlim_v(i,J,K) = (1. + CS%FGNV_scale) * Sfn_unlim_v(i,J,K) enddo - call streamfn_solver(nz, c2_dz_v(i,:), dzN2_v(i,:), Sfn_unlim_v(i,:)) + call streamfn_solver(nz, c2_dz_v(i,:), dzN2_v(i,J,:), Sfn_unlim_v(i,J,:)) else do K=2,nz - Sfn_unlim_v(i,K) = 0. + Sfn_unlim_v(i,J,K) = 0. enddo endif enddo endif +#endif - do K=nz,2,-1 - do i=is,ie + do i=is,ie + do K=nz,2,-1 +#ifdef __NVCOMPILER_OPENMP_GPU + Z_to_H = GV%Z_to_H +#else if (allocated(tv%SpV_avg) .and. (find_work .or. (k > nk_linear)) ) then Rho_avg = ( ((h(i,j,k) + h(i,j,k-1)) + (h(i,j+1,k) + h(i,j+1,k-1))) + 4.0*hn_2 ) / & ( (((h(i,j,k)+hn_2) * tv%SpV_avg(i,j,k)) + ((h(i,j,k-1)+hn_2) * tv%SpV_avg(i,j,k-1))) + & @@ -1452,6 +1654,7 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV else Z_to_H = GV%Z_to_H endif +#endif if (k > nk_linear) then if (use_EOS) then @@ -1464,9 +1667,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV endif ! Find the actual streamfunction at each interface. - Sfn_est = (Z_to_H*Sfn_unlim_v(i,K) + slope2_Ratio_v(i,K)*Sfn_safe) / (1.0 + slope2_Ratio_v(i,K)) + Sfn_est = (Z_to_H*Sfn_unlim_v(i,J,K) + slope2_Ratio_v(i,J,K)*Sfn_safe) / (1.0 + slope2_Ratio_v(i,J,K)) else ! When use_EOS is false, the layers are constant density. - Sfn_est = Z_to_H*Sfn_unlim_v(i,K) + Sfn_est = Z_to_H*Sfn_unlim_v(i,J,K) endif ! Make sure that there is enough mass above to allow the streamfunction @@ -1477,7 +1680,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! neighboring grid cells. vhD(i,J,k) = max(min((Sfn_in_H - vhtot(i,J)), h_avail(i,j,k)), -h_avail(i,j+1,k)) +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_y>0) diag_sfn_y(i,J,K) = diag_sfn_y(i,J,K+1) + vhD(i,J,k) +#endif ! sfn_y(i,J,K) = max(min(Sfn_in_h, vhtot(i,J)+h_avail(i,j,k)), & ! vhtot(i,J)-h_avail(i,j+1,k)) ! sfn_slope_y(i,J,K) = max(vhtot(i,J)-h_avail(i,j+1,k), & @@ -1495,7 +1700,9 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV vhD(i,J,k) = -vhtot(i,J) * h_frac(i,j+1,k) endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%id_sfn_y>0) diag_sfn_y(i,J,K) = diag_sfn_y(i,J,K+1) + vhD(i,J,k) +#endif ! if (sfn_slope_y(i,J,K+1) <= 0.0) then ! sfn_slope_y(i,J,K) = sfn_slope_y(i,J,K+1) * (1.0 - h_frac(i,j,k)) ! else @@ -1512,114 +1719,134 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV ! A second order centered estimate is used for the density transferred ! between water columns. +#ifdef __NVCOMPILER_OPENMP_GPU + G_scale = GV%g_Earth * GV%H_to_Z +#else if (allocated(tv%SpV_avg)) then G_scale = GV%H_to_RZ * GV%g_Earth / Rho_avg else G_scale = GV%g_Earth * GV%H_to_Z endif +#endif Work_v(i,J) = Work_v(i,J) + G_scale * & - ( vhtot(i,J) * drdkDe_v(i,K) - & - (vhD(i,J,k) * drdj_v(i,k)) * 0.25 * & + ( vhtot(i,J) * drdkDe_v(i,J,K) - & + (vhD(i,J,k) * drdj_v(i,J,k)) * 0.25 * & ((e(i,j,K) + e(i,j,K+1)) + (e(i,j+1,K) + e(i,j+1,K+1))) ) endif - enddo - enddo ! end of k-loop + enddo ! end of k-loop + enddo ! end of i-loop enddo ! end of j-loop ! In layer 1, enforce the boundary conditions that Sfn(z=0) = 0.0 if (.not.find_work .or. .not.(use_EOS)) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do j=js,je ; do I=is-1,ie ; uhD(I,j,1) = -uhtot(I,j) ; enddo ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) +#endif do J=js-1,je ; do i=is,ie ; vhD(i,J,1) = -vhtot(i,J) ; enddo ; enddo else - EOSdom_u(1) = (is-1) - (G%IsdB-1) ; EOSdom_u(2) = ie - (G%IsdB-1) - !$OMP parallel do default(shared) private(pres_u,T_u,S_u,drho_dT_u,drho_dS_u,drdiB,G_scale) - do j=js,je - if (use_EOS) then - do I=is-1,ie - pres_u(I) = 0.5*(pres(i,j,1) + pres(i+1,j,1)) - T_u(I) = 0.5*(T(i,j,1) + T(i+1,j,1)) - S_u(I) = 0.5*(S(i,j,1) + S(i+1,j,1)) - enddo - call calculate_density_derivs(T_u, S_u, pres_u, drho_dT_u, drho_dS_u, & - tv%eqn_of_state, EOSdom_u ) + ! In this branch use_EOS is necessarily true. The density derivatives use the device-callable + ! elementwise dispatcher (v-points too), fused into the work loop and scalarized. +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(pres_us,T_us,S_us,dT_us,dS_us,drdiB,G_scale) +#else + !$OMP parallel do default(shared) private(pres_us,T_us,S_us,dT_us,dS_us,drdiB,G_scale) +#endif + do j=js,je ; do I=is-1,ie + pres_us = 0.5*(pres(i,j,1) + pres(i+1,j,1)) + T_us = 0.5*(T(i,j,1) + T(i+1,j,1)) + S_us = 0.5*(S(i,j,1) + S(i+1,j,1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_us, eos_S_to_ppt*S_us, & + eos_RL2_T2_to_Pa*pres_us, dT_us, dS_us) + dT_us = (eos_kg_m3_to_R*eos_C_to_degC) * dT_us + dS_us = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_us + + uhD(I,j,1) = -uhtot(I,j) + + G_scale = GV%g_Earth * GV%H_to_Z + drdiB = dT_us * (T(i+1,j,1)-T(i,j,1)) + dS_us * (S(i+1,j,1)-S(i,j,1)) +#ifndef __NVCOMPILER_OPENMP_GPU + if (allocated(tv%SpV_avg)) then + G_scale = GV%H_to_RZ * GV%g_Earth * & + ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i+1,j,1)+hn_2) * tv%SpV_avg(i+1,j,1)) ) / & + ( (h(i,j,1) + h(i+1,j,1)) + 2.0*hn_2 ) ) endif - do I=is-1,ie - uhD(I,j,1) = -uhtot(I,j) +#endif + if (CS%use_GM_work_bug) then + Work_u(I,j) = Work_u(I,j) + G_scale * & + ( (uhD(I,j,1) * drdiB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) + else + Work_u(I,j) = Work_u(I,j) - G_scale * & + ( (uhD(I,j,1) * drdiB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) + endif + enddo ; enddo - G_scale = GV%g_Earth * GV%H_to_Z - if (use_EOS) then - drdiB = drho_dT_u(I) * (T(i+1,j,1)-T(i,j,1)) + & - drho_dS_u(I) * (S(i+1,j,1)-S(i,j,1)) - if (allocated(tv%SpV_avg)) then - G_scale = GV%H_to_RZ * GV%g_Earth * & - ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i+1,j,1)+hn_2) * tv%SpV_avg(i+1,j,1)) ) / & - ( (h(i,j,1) + h(i+1,j,1)) + 2.0*hn_2 ) ) - endif - endif - if (CS%use_GM_work_bug) then - Work_u(I,j) = Work_u(I,j) + G_scale * & - ( (uhD(I,j,1) * drdiB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) - else - Work_u(I,j) = Work_u(I,j) - G_scale * & - ( (uhD(I,j,1) * drdiB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i+1,j,1) + e(i+1,j,2))) ) - endif - enddo - enddo +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(pres_vs,T_vs,S_vs,dT_vs,dS_vs,drdjB,G_scale) +#else + !$OMP parallel do default(shared) private(pres_vs,T_vs,S_vs,dT_vs,dS_vs,drdjB,G_scale) +#endif + do J=js-1,je ; do i=is,ie + pres_vs = 0.5*(pres(i,j,1) + pres(i,j+1,1)) + T_vs = 0.5*(T(i,j,1) + T(i,j+1,1)) + S_vs = 0.5*(S(i,j,1) + S(i,j+1,1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_vs, eos_S_to_ppt*S_vs, & + eos_RL2_T2_to_Pa*pres_vs, dT_vs, dS_vs) + dT_vs = (eos_kg_m3_to_R*eos_C_to_degC) * dT_vs + dS_vs = (eos_kg_m3_to_R*eos_S_to_ppt) * dS_vs + + vhD(i,J,1) = -vhtot(i,J) + + G_scale = GV%g_Earth * GV%H_to_Z + drdjB = dT_vs * (T(i,j+1,1)-T(i,j,1)) + dS_vs * (S(i,j+1,1)-S(i,j,1)) +#ifndef __NVCOMPILER_OPENMP_GPU + if (allocated(tv%SpV_avg)) then + G_scale = GV%H_to_RZ * GV%g_Earth * & + ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i,j+1,1)+hn_2) * tv%SpV_avg(i,j+1,1)) ) / & + ( (h(i,j,1) + h(i,j+1,1)) + 2.0*hn_2 ) ) + endif +#endif + Work_v(i,J) = Work_v(i,J) - G_scale * & + ( (vhD(i,J,1) * drdjB) * 0.25 * & + ((e(i,j,1) + e(i,j,2)) + (e(i,j+1,1) + e(i,j+1,2))) ) + enddo ; enddo + endif - EOSdom_v(:) = EOS_domain(G%HI) - !$OMP parallel do default(shared) private(pres_v,T_v,S_v,drho_dT_v,drho_dS_v,drdjB,G_scale) - do J=js-1,je - if (use_EOS) then - do i=is,ie - pres_v(i) = 0.5*(pres(i,j,1) + pres(i,j+1,1)) - T_v(i) = 0.5*(T(i,j,1) + T(i,j+1,1)) - S_v(i) = 0.5*(S(i,j,1) + S(i,j+1,1)) + ! Assemble the GM work / energy source (writing the hoisted local arrays, not the pointer/ + ! allocatable components). + if (find_work) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(Work_h) +#endif + do j=js,je ; do i=is,ie + ! Note that the units of Work_v and Work_u are [R Z L4 T-3 ~> W], while Work_h is in [R Z L2 T-3 ~> W m-2]. + Work_h = 0.5 * G%IareaT(i,j) * & + ((Work_u(I-1,j) + Work_u(I,j)) + (Work_v(i,J-1) + Work_v(i,J))) + if (have_GMwork) GMwork_loc(i,j) = Work_h + if (.not. CS%GM_src_alt) then ; if (have_GM_src) then + GM_src_loc(i,j) = GM_src_loc(i,j) + Work_h + endif ; endif +#ifndef __NVCOMPILER_OPENMP_GPU + if (skeb_use_gm) then + skeb_gm_work(i,j) = STOCH%skeb_gm_coef * Work_h + skeb_ebt_norm2(i,j) = 0.0 + do k=1,nz + skeb_ebt_norm2(i,j) = skeb_ebt_norm2(i,j) + h(i,j,k) * VarMix%ebt_struct(i,j,k)**2 enddo - call calculate_density_derivs(T_v, S_v, pres_v, drho_dT_v, drho_dS_v, & - tv%eqn_of_state, EOSdom_v) + skeb_ebt_norm2(i,j) = GV%H_to_RZ * (skeb_ebt_norm2(i,j) + h_neglect) endif - do i=is,ie - vhD(i,J,1) = -vhtot(i,J) - - G_scale = GV%g_Earth * GV%H_to_Z - if (use_EOS) then - drdjB = drho_dT_v(i) * (T(i,j+1,1)-T(i,j,1)) + & - drho_dS_v(i) * (S(i,j+1,1)-S(i,j,1)) - if (allocated(tv%SpV_avg)) then - G_scale = GV%H_to_RZ * GV%g_Earth * & - ( ( ((h(i,j,1)+hn_2) * tv%SpV_avg(i,j,1)) + ((h(i,j+1,1)+hn_2) * tv%SpV_avg(i,j+1,1)) ) / & - ( (h(i,j,1) + h(i,j+1,1)) + 2.0*hn_2 ) ) - endif - endif - Work_v(i,J) = Work_v(i,J) - G_scale * & - ( (vhD(i,J,1) * drdjB) * 0.25 * & - ((e(i,j,1) + e(i,j,2)) + (e(i,j+1,1) + e(i,j+1,2))) ) - enddo - enddo +#endif + enddo ; enddo endif - if (find_work) then ; do j=js,je ; do i=is,ie - ! Note that the units of Work_v and Work_u are [R Z L4 T-3 ~> W], while Work_h is in [R Z L2 T-3 ~> W m-2]. - Work_h = 0.5 * G%IareaT(i,j) * & - ((Work_u(I-1,j) + Work_u(I,j)) + (Work_v(i,J-1) + Work_v(i,J))) - if (allocated(CS%GMwork)) CS%GMwork(i,j) = Work_h - if (.not. CS%GM_src_alt) then ; if (allocated(MEKE%GM_src)) then - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + Work_h - endif ; endif - if (skeb_use_gm) then - skeb_gm_work(i,j) = STOCH%skeb_gm_coef * Work_h - skeb_ebt_norm2(i,j) = 0.0 - do k=1,nz - skeb_ebt_norm2(i,j) = skeb_ebt_norm2(i,j) + h(i,j,k) * VarMix%ebt_struct(i,j,k)**2 - enddo - skeb_ebt_norm2(i,j) = GV%H_to_RZ * (skeb_ebt_norm2(i,j) + h_neglect) - endif - enddo ; enddo ; endif - +#ifndef __NVCOMPILER_OPENMP_GPU if (skeb_use_gm) then ! This block spreads the GM work down through the column using the ebt vertical structure, squared. ! Note the sign convention. @@ -1628,38 +1855,71 @@ subroutine thickness_diffuse_full(h, e, Kh_u, Kh_v, tv, uhD, vhD, cg1, dt, G, GV VarMix%ebt_struct(i,j,k)**2 / skeb_ebt_norm2(i,j) enddo ; enddo ; enddo endif +#endif - if (find_work .and. CS%GM_src_alt) then ; if (allocated(MEKE%GM_src)) then + if (find_work .and. CS%GM_src_alt) then ; if (have_GM_src) then + ! Accumulate the S^2 N^2 kappa energy release down each column (a serial-K sum into GM_src_loc). if (CS%MEKE_src_answer_date >= 20240601) then +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(PE_release_h) +#endif do j=js,je ; do i=is,ie ; do k=nz,1,-1 PE_release_h = -0.25 * GV%H_to_RZ * & ( ((KH_u(I,j,k)*(Slope_x_PE(I,j,k)**2) * hN2_x_PE(I,j,k)) + & (Kh_u(I-1,j,k)*(Slope_x_PE(I-1,j,k)**2) * hN2_x_PE(I-1,j,k))) + & ((Kh_v(i,J,k)*(Slope_y_PE(i,J,k)**2) * hN2_y_PE(i,J,k)) + & (Kh_v(i,J-1,k)*(Slope_y_PE(i,J-1,k)**2) * hN2_y_PE(i,J-1,k))) ) - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + PE_release_h + GM_src_loc(i,j) = GM_src_loc(i,j) + PE_release_h enddo ; enddo ; enddo else +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target teams loop collapse(2) private(PE_release_h) +#endif do j=js,je ; do i=is,ie ; do k=nz,1,-1 PE_release_h = -0.25 * GV%H_to_RZ * & ((KH_u(I,j,k)*(Slope_x_PE(I,j,k)**2) * hN2_x_PE(I,j,k)) + & (Kh_u(I-1,j,k)*(Slope_x_PE(I-1,j,k)**2) * hN2_x_PE(I-1,j,k)) + & (Kh_v(i,J,k)*(Slope_y_PE(i,J,k)**2) * hN2_y_PE(i,J,k)) + & (Kh_v(i,J-1,k)*(Slope_y_PE(i,J-1,k)**2) * hN2_y_PE(i,J-1,k))) - MEKE%GM_src(i,j) = MEKE%GM_src(i,j) + PE_release_h + GM_src_loc(i,j) = GM_src_loc(i,j) + PE_release_h enddo ; enddo ; enddo endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%debug) then - call hchksum(MEKE%GM_src, 'MEKE%GM_src', G%HI, unscale=US%RZ3_T3_to_W_m2*US%L_to_Z**2) + call hchksum(GM_src_loc, 'MEKE%GM_src', G%HI, unscale=US%RZ3_T3_to_W_m2*US%L_to_Z**2) call uvchksum("KH_[uv]", Kh_u, Kh_v, G%HI, unscale=US%L_to_m**2*US%s_to_T, & scalar_pair=.true.) call uvchksum("Slope_[xy]_PE", Slope_x_PE, Slope_y_PE, G%HI, unscale=US%Z_to_L) call uvchksum("hN2_[xy]_PE", hN2_x_PE, hN2_y_PE, G%HI, unscale=GV%H_to_mks*US%L_to_Z**2*US%s_to_T**2, & scalar_pair=.true.) endif +#endif endif ; endif + ! ------------------------------------------------------------------------------------------ + ! Close the device data region: copy the outputs back and release every mapping (release, not + ! delete — a per-call delete forces the refcount to zero and would destroy any outer/persistent + ! mapping of the same storage; KNOWLEDGE §8), then copy the hoisted GM work back to the host. + ! ------------------------------------------------------------------------------------------ +#ifdef __NVCOMPILER_OPENMP_GPU + !$omp target exit data map(from: uhD, vhD, GM_src_loc, GMwork_loc) & + !$omp map(release: e, dz, T, S, Kh_u, Kh_v, int_slope_u, int_slope_v, CS, p_surf_loc, & + !$omp Slope_x_PE, Slope_y_PE, hN2_x_PE, hN2_y_PE, pres, h_avail, h_frac, & + !$omp h_avail_rsum, uhtot, vhtot, Work_u, Work_v, drdi_u, drdj_v, & + !$omp drdkDe_u, drdkDe_v, dzN2_u, dzN2_v, Sfn_unlim_u, Sfn_unlim_v, & + !$omp slope2_Ratio_u, slope2_Ratio_v) & + !$omp map(release: h) + if (present_slope_x) then + !$omp target exit data map(release: slope_x) + endif + if (present_slope_y) then + !$omp target exit data map(release: slope_y) + endif +#endif + if (have_GM_src) then ; do j=js,je ; do i=is,ie ; MEKE%GM_src(i,j) = GM_src_loc(i,j) ; enddo ; enddo ; endif + if (have_GMwork) then ; do j=js,je ; do i=is,ie ; CS%GMwork(i,j) = GMwork_loc(i,j) ; enddo ; enddo ; endif + if (CS%id_slope_x > 0) call post_data(CS%id_slope_x, CS%diagSlopeX, CS%diag) if (CS%id_slope_y > 0) call post_data(CS%id_slope_y, CS%diagSlopeY, CS%diag) if (CS%id_sfn_x > 0) call post_data(CS%id_sfn_x, diag_sfn_x, CS%diag) diff --git a/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 b/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 index 5b41cef038..24b4f3aadb 100644 --- a/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 +++ b/src/parameterizations/vertical/MOM_bkgnd_mixing.F90 @@ -309,21 +309,23 @@ subroutine bkgnd_mixing_init(Time, G, GV, US, param_file, diag, CS, physical_OBL end subroutine bkgnd_mixing_init !> Calculates the vertical background diffusivities/viscosities -subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, GV, US, CS) +subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, jstart, jend, nj, G, GV, US, CS) type(ocean_grid_type), intent(in) :: G !< Grid structure. type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), intent(in) :: h !< Layer thickness [H ~> m or kg m-2]. type(thermo_var_ptrs), intent(in) :: tv !< Thermodynamics structure. - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: N2_lay !< squared buoyancy frequency associated + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: N2_lay !< squared buoyancy frequency associated !! with layers [T-2 ~> s-2] - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: Kd_lay !< The background diapycnal diffusivity of each + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: Kd_lay !< The background diapycnal diffusivity of each !! layer [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(out) :: Kd_int !< The background diapycnal diffusivity of each + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(out) :: Kd_int !< The background diapycnal diffusivity of each !! interface [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(out) :: Kv_bkgnd !< The background vertical viscosity at + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(out) :: Kv_bkgnd !< The background vertical viscosity at !! each interface [H Z T-1 ~> m2 s-1 or Pa s] - integer, intent(in) :: j !< Meridional grid index + integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block + integer, intent(in) :: nj !< Number of j-rows in this block type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type type(bkgnd_mixing_cs), pointer :: CS !< The control structure returned by !! a previous call to bkgnd_mixing_init. @@ -332,9 +334,9 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, real, dimension(SZK_(GV)+1) :: depth_int !< Distance from surface of the interfaces [m] real, dimension(SZK_(GV)+1) :: Kd_col !< Diffusivities at the interfaces [m2 s-1] real, dimension(SZK_(GV)+1) :: Kv_col !< Viscosities at the interfaces [m2 s-1] - real, dimension(SZI_(G)) :: Kd_sfc !< Surface value of the diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G)) :: depth !< Distance from surface of an interface [H ~> m or kg m-2] - real, dimension(SZI_(G),SZK_(GV)) :: dz !< Height change across layers [Z ~> m] + real, dimension(SZI_(G),nj) :: Kd_sfc !< Surface value of the diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),nj) :: depth !< Distance from surface of an interface [H ~> m or kg m-2] + real, dimension(SZI_(G),SZK_(GV),nj) :: dz !< Height change across layers [Z ~> m] real :: depth_c !< depth of the center of a layer [H ~> m or kg m-2] real :: I_Hmix !< inverse of fixed mixed layer thickness [H-1 ~> m-1 or m2 kg-1] real :: I_x30 !< 2/acos(2) = 1/(sin(30 deg) * acosh(1/sin(30 deg))) [nondim] @@ -343,7 +345,7 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, real :: min_sinlat ! The minimum value of the sine of latitude [nondim] real :: bckgrnd_vdc_psin !< PSI diffusivity in northern hemisphere [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real :: bckgrnd_vdc_psis !< PSI diffusivity in southern hemisphere [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - integer :: i, k, is, ie, js, je, nz + integer :: i, j, k, is, ie, js, je, nz, jj is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke @@ -352,108 +354,122 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, min_sinlat = 1.e-10 ! Start with a constant value that may be replaced below. - Kd_lay(:,:) = CS%Kd - Kv_bkgnd(:,:) = 0.0 + do j=jstart,jend ; jj = j - jstart + 1 + Kd_lay(:,:,jj) = CS%Kd + Kv_bkgnd(:,:,jj) = 0.0 + enddo ! Set up the background diffusivity. if (CS%Bryan_Lewis_diffusivity) then - call thickness_to_dz(h, tv, dz, j, G, GV) - - do i=is,ie - depth_int(1) = 0.0 - do k=2,nz+1 - depth_int(k) = depth_int(k-1) + US%Z_to_m*dz(i,k-1) - enddo + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo - call CVMix_init_bkgnd(max_nlev=nz, & - zw = depth_int(:), & !< interface depths relative to the surface in m, must be positive. - bl1 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c1, & - bl2 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c2, & - bl3 = US%m_to_Z*CS%Bryan_Lewis_c3, & - bl4 = US%Z_to_m*CS%Bryan_Lewis_c4, & - prandtl = CS%prandtl_bkgnd) - - Kd_col(:) = 0.0 ; Kv_col(:) = 0.0 ! Is this line necessary? - call CVMix_coeffs_bkgnd(Mdiff_out=Kv_col, Tdiff_out=Kd_col, nlev=nz, max_nlev=nz) - - ! Update Kd and Kv. - do K=1,nz+1 - Kv_bkgnd(i,K) = GV%m2_s_to_HZ_T * Kv_col(K) - Kd_int(i,K) = GV%m2_s_to_HZ_T*Kd_col(K) - enddo - do k=1,nz - Kd_lay(i,k) = Kd_lay(i,k) + 0.5 * GV%m2_s_to_HZ_T * (Kd_col(K) + Kd_col(K+1)) - enddo - enddo ! i loop + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + depth_int(1) = 0.0 + do k=2,nz+1 + depth_int(k) = depth_int(k-1) + US%Z_to_m*dz(i,k-1,jj) + enddo + + call CVMix_init_bkgnd(max_nlev=nz, & + zw = depth_int(:), & !< interface depths relative to the surface in m, must be positive. + bl1 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c1, & + bl2 = US%Z2_T_to_m2_s*CS%Bryan_Lewis_c2, & + bl3 = US%m_to_Z*CS%Bryan_Lewis_c3, & + bl4 = US%Z_to_m*CS%Bryan_Lewis_c4, & + prandtl = CS%prandtl_bkgnd) + + Kd_col(:) = 0.0 ; Kv_col(:) = 0.0 ! Is this line necessary? + call CVMix_coeffs_bkgnd(Mdiff_out=Kv_col, Tdiff_out=Kd_col, nlev=nz, max_nlev=nz) + + ! Update Kd and Kv. + do K=1,nz+1 + Kv_bkgnd(i,K,jj) = GV%m2_s_to_HZ_T * Kv_col(K) + Kd_int(i,K,jj) = GV%m2_s_to_HZ_T*Kd_col(K) + enddo + do k=1,nz + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + 0.5 * GV%m2_s_to_HZ_T * (Kd_col(K) + Kd_col(K+1)) + enddo + enddo ! i loop + enddo elseif (CS%horiz_varying_background) then !### Note that there are lots of hard-coded parameters (mostly latitudes and longitudes) here. - do i=is,ie - bckgrnd_vdc_psis = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)+28.9))**2) - bckgrnd_vdc_psin = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)-28.9))**2) - Kd_int(i,1) = (CS%bckgrnd_vdc_eq + bckgrnd_vdc_psin) + bckgrnd_vdc_psis - - if (G%geoLatT(i,j) < -10.0) then - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 - elseif (G%geoLatT(i,j) <= 10.0) then - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 * (G%geoLatT(i,j)/10.0)**2 - else - Kd_int(i,1) = Kd_int(i,1) + CS%bckgrnd_vdc1 - endif - - ! North Banda Sea - if ( (G%geoLatT(i,j) < -1.0) .and. (G%geoLatT(i,j) > -4.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 103.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 134.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif - - ! Middle Banda Sea - if ( (G%geoLatT(i,j) <= -4.0) .and. (G%geoLatT(i,j) > -7.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 106.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 140.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif - - ! South Banda Sea - if ( (G%geoLatT(i,j) <= -7.0) .and. (G%geoLatT(i,j) > -8.3) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) > 111.0) .and. & - ( mod(G%geoLonT(i,j)+360.0,360.0) < 142.0) ) then - Kd_int(i,1) = CS%bckgrnd_vdc_Banda - endif + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + bckgrnd_vdc_psis = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)+28.9))**2) + bckgrnd_vdc_psin = CS%bckgrnd_vdc_psim * exp(-(0.4*(G%geoLatT(i,j)-28.9))**2) + Kd_int(i,1,jj) = (CS%bckgrnd_vdc_eq + bckgrnd_vdc_psin) + bckgrnd_vdc_psis + + if (G%geoLatT(i,j) < -10.0) then + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 + elseif (G%geoLatT(i,j) <= 10.0) then + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 * (G%geoLatT(i,j)/10.0)**2 + else + Kd_int(i,1,jj) = Kd_int(i,1,jj) + CS%bckgrnd_vdc1 + endif + + ! North Banda Sea + if ( (G%geoLatT(i,j) < -1.0) .and. (G%geoLatT(i,j) > -4.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 103.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 134.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + ! Middle Banda Sea + if ( (G%geoLatT(i,j) <= -4.0) .and. (G%geoLatT(i,j) > -7.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 106.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 140.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + + ! South Banda Sea + if ( (G%geoLatT(i,j) <= -7.0) .and. (G%geoLatT(i,j) > -8.3) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) > 111.0) .and. & + ( mod(G%geoLonT(i,j)+360.0,360.0) < 142.0) ) then + Kd_int(i,1,jj) = CS%bckgrnd_vdc_Banda + endif + + enddo enddo ! Update interior values of Kd and Kv (uniform profile; no interpolation needed) - do K=1,nz+1 ; do i=is,ie - Kd_int(i,K) = Kd_int(i,1) - Kv_bkgnd(i,K) = Kd_int(i,1) * CS%prandtl_bkgnd - enddo ; enddo - do k=1,nz ; do i=is,ie - Kd_lay(i,k) = Kd_int(i,1) - enddo ; enddo + do K=1,nz+1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int(i,K,jj) = Kd_int(i,1,jj) + Kv_bkgnd(i,K,jj) = Kd_int(i,1,jj) * CS%prandtl_bkgnd + enddo ; enddo ; enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_lay(i,k,jj) = Kd_int(i,1,jj) + enddo ; enddo ; enddo else ! Set a potentially spatially varying surface value of diffusivity. if (CS%Henyey_IGW_background) then I_x30 = 2.0 / invcosh(CS%N0_2Omega*2.0) ! This is evaluated at 30 deg. - do i=is,ie - abs_sinlat = abs(sin(G%geoLatT(i,j)*deg_to_rad)) - if (abs(G%geoLatT(i,j))>CS%Henyey_max_lat) abs_sinlat = min_sinlat - Kd_sfc(i) = max(CS%Kd_min, CS%Kd * & - ((abs_sinlat * invcosh(CS%N0_2Omega / max(min_sinlat, abs_sinlat))) * I_x30) ) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + abs_sinlat = abs(sin(G%geoLatT(i,j)*deg_to_rad)) + if (abs(G%geoLatT(i,j))>CS%Henyey_max_lat) abs_sinlat = min_sinlat + Kd_sfc(i,jj) = max(CS%Kd_min, CS%Kd * & + ((abs_sinlat * invcosh(CS%N0_2Omega / max(min_sinlat, abs_sinlat))) * I_x30) ) + enddo enddo elseif (CS%Kd_tanh_lat_fn) then - do i=is,ie - ! The transition latitude and latitude range are hard-scaled here, since - ! this is not really intended for wide-spread use, but rather for - ! comparison with CM2M / CM2.1 settings. - Kd_sfc(i) = max(CS%Kd_min, CS%Kd * (1.0 + & - CS%Kd_tanh_lat_scale * 0.5*tanh((abs(G%geoLatT(i,j)) - 35.0)/5.0) )) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + ! The transition latitude and latitude range are hard-scaled here, since + ! this is not really intended for wide-spread use, but rather for + ! comparison with CM2M / CM2.1 settings. + Kd_sfc(i,jj) = max(CS%Kd_min, CS%Kd * (1.0 + & + CS%Kd_tanh_lat_scale * 0.5*tanh((abs(G%geoLatT(i,j)) - 35.0)/5.0) )) + enddo enddo else ! Use a spatially constant surface value. - do i=is,ie - Kd_sfc(i) = CS%Kd + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + Kd_sfc(i,jj) = CS%Kd + enddo enddo endif @@ -462,33 +478,37 @@ subroutine calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay, Kd_int, Kv_bkgnd, j, G, ! This is a crude way to put in a diffusive boundary layer without an explicit boundary ! layer turbulence scheme. It should not be used for any realistic ocean models. I_Hmix = 1.0 / (CS%Hmix + GV%H_subroundoff) - do i=is,ie ; depth(i) = 0.0 ; enddo - do k=1,nz ; do i=is,ie - depth_c = depth(i) + 0.5*h(i,j,k) - if (depth_c <= CS%Hmix) then ; Kd_lay(i,k) = CS%Kd_tot_ml - elseif (depth_c >= 2.0*CS%Hmix) then ; Kd_lay(i,k) = Kd_sfc(i) + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; depth(i,jj) = 0.0 ; enddo + enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + depth_c = depth(i,jj) + 0.5*h(i,j,k) + if (depth_c <= CS%Hmix) then ; Kd_lay(i,k,jj) = CS%Kd_tot_ml + elseif (depth_c >= 2.0*CS%Hmix) then ; Kd_lay(i,k,jj) = Kd_sfc(i,jj) else - Kd_lay(i,k) = ((Kd_sfc(i) - CS%Kd_tot_ml) * I_Hmix) * depth_c + (2.0*CS%Kd_tot_ml - Kd_sfc(i)) + Kd_lay(i,k,jj) = ((Kd_sfc(i,jj) - CS%Kd_tot_ml) * I_Hmix) * depth_c + (2.0*CS%Kd_tot_ml - Kd_sfc(i,jj)) endif - depth(i) = depth(i) + h(i,j,k) - enddo ; enddo + depth(i,jj) = depth(i,jj) + h(i,j,k) + enddo ; enddo ; enddo else ! There is no vertical structure to the background diffusivity. - do k=1,nz ; do i=is,ie - Kd_lay(i,k) = Kd_sfc(i) - enddo ; enddo + do k=1,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_lay(i,k,jj) = Kd_sfc(i,jj) + enddo ; enddo ; enddo endif ! Update Kd_int and Kv_bkgnd, based on Kd_lay. These might be just used for diagnostic purposes. - do i=is,ie - Kd_int(i,1) = 0.0 ; Kv_bkgnd(i,1) = 0.0 - Kd_int(i,nz+1) = 0.0 ; Kv_bkgnd(i,nz+1) = 0.0 + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + Kd_int(i,1,jj) = 0.0 ; Kv_bkgnd(i,1,jj) = 0.0 + Kd_int(i,nz+1,jj) = 0.0 ; Kv_bkgnd(i,nz+1,jj) = 0.0 + enddo enddo - do K=2,nz ; do i=is,ie - Kd_int(i,K) = 0.5*(Kd_lay(i,k-1) + Kd_lay(i,k)) - Kv_bkgnd(i,K) = Kd_int(i,K) * CS%prandtl_bkgnd - enddo ; enddo + do K=2,nz ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int(i,K,jj) = 0.5*(Kd_lay(i,k-1,jj) + Kd_lay(i,k,jj)) + Kv_bkgnd(i,K,jj) = Kd_int(i,K,jj) * CS%prandtl_bkgnd + enddo ; enddo ; enddo endif end subroutine calculate_bkgnd_mixing diff --git a/src/parameterizations/vertical/MOM_energetic_PBL.F90 b/src/parameterizations/vertical/MOM_energetic_PBL.F90 index 6930007bd1..3aa1e83e62 100644 --- a/src/parameterizations/vertical/MOM_energetic_PBL.F90 +++ b/src/parameterizations/vertical/MOM_energetic_PBL.F90 @@ -17,12 +17,13 @@ module MOM_energetic_PBL use MOM_forcing_type, only : forcing use MOM_grid, only : ocean_grid_type use MOM_interface_heights, only : thickness_to_dz -use MOM_intrinsic_functions, only : cuberoot +use MOM_intrinsic_functions, only : cuberoot, exp_reprod, log_reprod use MOM_string_functions, only : uppercase use MOM_unit_scaling, only : unit_scale_type use MOM_variables, only : thermo_var_ptrs, vertvisc_type use MOM_verticalGrid, only : verticalGrid_type -use MOM_wave_interface, only : wave_parameters_CS, Get_Langmuir_Number +use MOM_wave_interface, only : wave_parameters_CS, Get_Langmuir_Number, Get_Langmuir_Number_LF17 +use MOM_wave_interface, only : wave_LF17_params, set_wave_LF17_params use MOM_stochastics, only : stochastic_CS implicit none ; private @@ -277,6 +278,10 @@ module MOM_energetic_PBL ! The next options are used when passively diagnosing sensitivities from parameter choices integer :: id_opt_diff_Kd_ePBL = -1, id_opt_maxdiff_Kd_ePBL = -1, id_opt_diff_hML_depth = -1 !>@} + type(wave_LF17_params) :: wave_lf17 !< A device-mappable bundle of the loop-invariant LF17 wave + !! scalars, populated from the Waves CS each call so the on-device + !! Langmuir path reads these (allocatable-free, maps cleanly) instead + !! of dereferencing the Waves pointer. end type energetic_PBL_CS !>@{ Enumeration values for mstar_scheme @@ -321,6 +326,39 @@ module MOM_energetic_PBL integer :: BBL_its !< The number of iterations used to find a self-consistent bottom boundary layer depth end type ePBL_column_diags +!> GPU port: make the per-column ePBL helper kernels device-callable so ePBL_column (which runs on the +!! device once the driver loop is offloaded) can call them. They are pure scalar computations, so the +!! same source serves the host and device paths. +!$omp declare target(exp_decay_TKE_adjust, find_PE_chg, find_PE_chg_orig, find_Kd_from_PE_chg) +!$omp declare target(find_mstar, mstar_Langmuir) +!$omp declare target(ePBL_column) + +! RESOLVED(gpu-bitwise-repro) 2026-07-20: ePBL is now GPU/CPU bit-for-bit -- the on-device +! transcendentals below were routed through the reproducible kernels exp_reprod/log_reprod/ +! erfc_reprod (MOM_intrinsic_functions), and a same-source nvfortran-CPU build (ocean_only/cpu_build) +! reproduces the GPU benchmark_ALE ocean.stats bit-for-bit (empty diff). These prototype reprod +! kernels change answers slightly vs the intrinsics (like cuberoot vs x**(1/3)); coordinate with +! Marshall's transcendental-repro pass, which may supersede them. The sites that were wired (kept as +! a record; live in benchmark_ALE) were the device-vs-host divergence before the fix: +! - find_PE_chg / find_PE_chg_orig : exp() in the PE-change integrals (core, per interface x +! iteration -- the likely dominant contributor). +! - exp_decay_TKE_adjust (~:1533) : exp(-h*Idecay_len_TKE) TKE decay. +! - find_mstar mstar_N (~:3753) : log(...) Ekman-limit term (EPBL_MSTAR_SCHEME=OM4). +! - mstar_Langmuir (~:3875,:3878) : Convect_Langmuir_Number**LT_enhance_exp (LT_ENHANCE_EXP=-1.33, +! an ARBITRARY real power -- no cuberoot/nth_root form; = exp(y*log x) under the hood). +! - MixLen_shape (~:1492) : (...)**MixLenExponent with MixLenExponent=1.0 == pow(x,1.0); +! a trivial local stopgap would be an `==1.0` identity fast-path next to the existing `==2.0` +! one, but it alone does NOT restore bitwise (the exp/log above dominate), so left for Marshall. +! (MKE_src exp() sites ~:1755+ are multiplied by MKE_TO_TKE_EFFIC=0.0 in benchmark_ALE -> harmless.) +! See the [[gpu-transcendental-bitwise-plan]] auto-memory. Until this lands, ePBL offload increments +! are gated on "no NEW diff vs the ePBL-device baseline", not bit-for-bit vs the golden. + +!> GPU port: fixes the per-column private scratch in ePBL_column to a compile-time-constant size on +!! GPU builds, so each device thread gets stack (local memory) arrays rather than runtime-sized +!! device-heap automatics (NVFORTRAN-W-0155 / KNOWLEDGE row 21). Checked against GV%ke in +!! energetic_PBL_init. Unused in CPU builds, where the local declarations keep their exact sizes. +integer, parameter :: GPU_nk_max = 128 + contains !> This subroutine determines the diffusivities from the integrated energetics @@ -406,7 +444,16 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, v_2d ! A 2-d slice of the meridional velocity [L T-1 ~> m s-1]. real, dimension(SZI_(G),SZK_(GV)+1) :: & Kd_2d ! A 2-d version of the diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: & + dz_3d ! The vertical distance across layers over the whole domain [Z ~> m], computed + ! once per call to replace the former per-j thickness_to_dz slice. +#ifdef __NVCOMPILER_OPENMP_GPU + ! On GPU these per-column column scratch arrays are teams-loop private() and so must have a + ! compile-time-constant size (NVFORTRAN-W-0155); GPU_nk_max is guarded by a FATAL on GV%ke. + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif h, & ! The layer thickness [H ~> m or kg m-2]. dz, & ! The vertical distance across layers [Z ~> m]. T0, & ! The initial layer temperatures [C ~> degC]. @@ -416,13 +463,14 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, TKE_forcing, & ! Forcing of the TKE in the layer coming from TKE_forced [R Z3 T-2 ~> J m-2]. u, & ! The zonal velocity [L T-1 ~> m s-1]. v ! The meridional velocity [L T-1 ~> m s-1]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif Kd, & ! The diapycnal diffusivity due to ePBL [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. mixvel, & ! A turbulent mixing velocity [Z T-1 ~> m s-1]. mixlen, & ! A turbulent mixing length [Z ~> m]. - mixvel_BBL, & ! A bottom boundary layer turbulent mixing velocity [Z T-1 ~> m s-1]. - mixlen_BBL, & ! A bottom boundary layer turbulent mixing length [Z ~> m]. - Kd_BBL, & ! The bottom boundary layer diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. SpV_dt, & ! Specific volume interpolated to interfaces divided by dt or 1.0 / (dt * Rho0), ! in [R-1 T-1 ~> m3 kg-1 s-1], used to convert local TKE into a turbulence velocity cubed. SpV_dt_cf ! Specific volume interpolated to interfaces divided by dt or 1.0 / (dt * Rho0) @@ -430,6 +478,12 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! [m3 Z-3 R-1 T2 s-3 ~> m3 kg-1 s-1] or without the conversion factors for ! answer dates of 20240101 and later in [R-1 T-1 ~> m3 kg-1 s-1], used to ! convert local TKE into a turbulence velocity cubed. + ! These bottom-boundary-layer column arrays are only used on the host (BBL_mixing / options_diff + ! paths, both FATAL-guarded off the GPU port), so they keep the runtime size. + real, dimension(SZK_(GV)+1) :: & + mixvel_BBL, & ! A bottom boundary layer turbulent mixing velocity [Z T-1 ~> m s-1]. + mixlen_BBL, & ! A bottom boundary layer turbulent mixing length [Z ~> m]. + Kd_BBL ! The bottom boundary layer diapycnal diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1]. real :: h_neglect ! A thickness that is so small it is usually lost ! in roundoff and can be neglected [H ~> m or kg m-2]. @@ -507,6 +561,16 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! can be modified to test for sensitivities logical :: BBL_mixing ! If true, there is bottom boundary layer mixing. integer :: i, j, k, is, ie, js, je, nz +#ifdef __NVCOMPILER_OPENMP_GPU + ! Plain-array hoists for the GPU port: pointer-member forcing and the CS%ML_depth first guess are + ! copied to these on the host (where associated() is evaluated) so the device loop touches neither + ! fluxes% pointer members nor the CS allocatable components. + real, dimension(SZI_(G),SZJ_(G)) :: & + ustar_2d, & ! Surface friction velocity [Z T-1 ~> m s-1] + ustar_gustless_2d, & ! Gustless surface friction velocity [Z T-1 ~> m s-1] + ML_depth_2d, & ! Mixed layer depth guess (in) / result (out) [H ~> m or kg m-2] + BBL_depth_2d ! Bottom boundary layer depth result [H ~> m or kg m-2] +#endif is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke @@ -587,50 +651,114 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, if (CS%id_opt_diff_hML_depth > 0) diff_hML_depth(:,:) = 0.0 endif - !!OMP parallel do default(private) shared(js,je,nz,is,ie,h_3d,u_3d,v_3d,tv,dt,I_dt,BBL_mixing, & - !!OMP CS,G,GV,US,fluxes,TKE_forced,dSV_dT,dSV_dS,Kd_int) + ! Compute the vertical layer extent over the whole domain once, replacing the per-j + ! thickness_to_dz slice that used to sit inside the j-loop (identical elementwise result). + call thickness_to_dz(h_3d, tv, dz_3d, G, GV, US) + +#ifdef __NVCOMPILER_OPENMP_GPU + ! ---- GPU port: FATAL-guard the features this port does not implement on device ---- + if (.not.GV%Boussinesq) call MOM_error(FATAL, & + "energetic_PBL: the GPU port only supports Boussinesq mode.") + if (BBL_mixing) call MOM_error(FATAL, & + "energetic_PBL: bottom boundary layer mixing is not supported in the GPU port.") + if (CS%options_diff > 0) call MOM_error(FATAL, & + "energetic_PBL: options_diff>0 (sensitivity diagnostics) is not supported in the GPU port.") + if (stoch_CS%pert_epbl) call MOM_error(FATAL, & + "energetic_PBL: stochastic ePBL perturbations are not supported in the GPU port.") + if (CS%TKE_diagnostics) call MOM_error(FATAL, & + "energetic_PBL: the ePBL TKE-budget diagnostics are not supported in the GPU port.") + if (report_avg_its) call MOM_error(FATAL, & + "energetic_PBL: report_avg_its is not supported in the GPU port.") + if (CS%debug) call MOM_error(FATAL, & + "energetic_PBL: CS%debug is not supported in the GPU port.") + if ((CS%id_Mixing_Length>0) .or. (CS%id_Velocity_Scale>0) .or. (CS%id_ustar_ePBL>0) .or. & + (CS%id_Kd_ePBL_col_by_col>0) .or. (CS%id_mstar_sfc>0) .or. (CS%id_mstar_bbl>0) .or. & + (CS%id_mstar_LT>0) .or. (CS%id_LA>0) .or. (CS%id_LA_mod>0)) call MOM_error(FATAL, & + "energetic_PBL: the per-column ePBL diagnostics are not supported in the GPU port.") + + ! Hoist pointer-member forcing and the CS%ML_depth first guess onto plain arrays (associated() + ! tests happen here on the host); the device loop then reads only these plain arrays and never + ! dereferences a fluxes% pointer member or a CS allocatable component. + if (.not.(associated(fluxes%ustar) .and. (GV%Boussinesq .or. .not.associated(fluxes%tau_mag)))) & + call MOM_error(FATAL, "energetic_PBL: the GPU port requires the fluxes%ustar forcing path.") + if (associated(fluxes%ustar_shelf) .and. associated(fluxes%frac_shelf_h)) call MOM_error(FATAL, & + "energetic_PBL: ice-shelf ustar blending is not supported in the GPU port.") + do j=js,je ; do i=is,ie + ustar_2d(i,j) = fluxes%ustar(i,j) ; ustar_gustless_2d(i,j) = fluxes%ustar_gustless(i,j) + ML_depth_2d(i,j) = CS%ML_depth(i,j) + enddo ; enddo + ! Bundle the loop-invariant LF17 wave scalars into CS%wave_lf17 (allocatable-free, so it rides the + ! CS map onto the device) — the on-device Langmuir path reads these instead of the Waves pointer. + if (associated(Waves)) call set_wave_LF17_params(Waves, CS%wave_lf17) + + ! One driver-level data region (find_N2 lesson: map full-domain arrays ONCE, never per-column). + !$omp target enter data map(to: h_3d, u_3d, v_3d, dz_3d, TKE_forced, dSV_dT, dSV_dS, buoy_flux, & + !$omp ustar_2d, ustar_gustless_2d, ML_depth_2d, CS, tv, tv%T, tv%S) & + !$omp map(alloc: BBL_depth_2d, Kd_int) + if (associated(Waves)) then + !$omp target enter data map(to: Waves) + endif + ! tv%T/tv%S are persistently resident via CS%tv and were mutated on the host earlier in diabatic, + ! so refresh them here (the map-to-present-no-copy trap); h_3d may likewise be persistently mapped. + !$omp target update to(tv%T, tv%S, h_3d) + + !$omp target teams loop collapse(2) & + !$omp private(eCD, u_star, u_star_mean, mech_TKE, absf, B_flux, MLD_io, BBLD_io, MLD_in, & + !$omp h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, TKE_forcing, Kd, mixvel, mixlen, & + !$omp SpV_dt, SpV_dt_cf) & + !$omp firstprivate(is, ie, nz, dt, I_dt, I_rho, I_rho0dt) +#endif do j=js,je - ! Copy the thicknesses and other fields to 2-d arrays. - do k=1,nz ; do i=is,ie - h_2d(i,k) = h_3d(i,j,k) ; u_2d(i,k) = u_3d(i,j,k) ; v_2d(i,k) = v_3d(i,j,k) - T_2d(i,k) = tv%T(i,j,k) ; S_2d(i,k) = tv%S(i,j,k) - TKE_forced_2d(i,k) = TKE_forced(i,j,k) - dSV_dT_2d(i,k) = dSV_dT(i,j,k) ; dSV_dS_2d(i,k) = dSV_dS(i,j,k) - enddo ; enddo - call thickness_to_dz(h_3d, tv, dz_2d, j, G, GV) - - ! Set the inverse density used to translating local TKE into a turbulence velocity - SpV_dt(:) = 0.0 - if ((dt > 0.0) .and. GV%Boussinesq .or. .not.allocated(tv%SpV_avg)) then - if (CS%answer_date < 20240101) then - do K=1,nz+1 - SpV_dt(K) = 1.0 / (dt*GV%Rho0) - enddo - else - do K=1,nz+1 - SpV_dt(K) = I_rho0dt - enddo - endif - endif - - ! Determine the initial mech_TKE and conv_PErel, including the energy required - ! to mix surface heating through the topmost cell, the energy released by mixing - ! surface cooling & brine rejection down through the topmost cell, and - ! homogenizing the shortwave heating within that cell. This sets the energy - ! and ustar and wstar available to drive mixing at the first interior - ! interface. + ! For each ocean column determine the initial mech_TKE and conv_PErel, including the energy + ! required to mix surface heating through the topmost cell, the energy released by mixing surface + ! cooling & brine rejection down through the topmost cell, and homogenizing the shortwave heating + ! within that cell. This sets the energy and ustar and wstar available to drive mixing at the + ! first interior interface. do i=is,ie ; if (G%mask2dT(i,j) > 0.0) then - ! Copy the thicknesses and other fields to 1-d arrays. + ! Copy the thicknesses and other fields to 1-d column arrays, read directly from the 3-d + ! inputs (the former per-j 2-d staging slabs are no longer needed). do k=1,nz - h(k) = h_2d(i,k) + GV%H_subroundoff ; dz(k) = dz_2d(i,k) + GV%dZ_subroundoff - u(k) = u_2d(i,k) ; v(k) = v_2d(i,k) - T0(k) = T_2d(i,k) ; S0(k) = S_2d(i,k) ; TKE_forcing(k) = TKE_forced_2d(i,k) - dSV_dT_1d(k) = dSV_dT_2d(i,k) ; dSV_dS_1d(k) = dSV_dS_2d(i,k) + h(k) = h_3d(i,j,k) + GV%H_subroundoff ; dz(k) = dz_3d(i,j,k) + GV%dZ_subroundoff + u(k) = u_3d(i,j,k) ; v(k) = v_3d(i,j,k) + T0(k) = tv%T(i,j,k) ; S0(k) = tv%S(i,j,k) ; TKE_forcing(k) = TKE_forced(i,j,k) + dSV_dT_1d(k) = dSV_dT(i,j,k) ; dSV_dS_1d(k) = dSV_dS(i,j,k) enddo do K=1,nz+1 ; Kd(K) = 0.0 ; enddo + ! Set the inverse density used to translate local TKE into a turbulence velocity. +#ifdef __NVCOMPILER_OPENMP_GPU + do K=1,nz+1 ; SpV_dt(K) = 0.0 ; enddo + if (dt > 0.0) then + if (CS%answer_date < 20240101) then + do K=1,nz+1 ; SpV_dt(K) = 1.0 / (dt*GV%Rho0) ; enddo + else + do K=1,nz+1 ; SpV_dt(K) = I_rho0dt ; enddo + endif + endif +#else + SpV_dt(:) = 0.0 + if ((dt > 0.0) .and. GV%Boussinesq .or. .not.allocated(tv%SpV_avg)) then + if (CS%answer_date < 20240101) then + do K=1,nz+1 + SpV_dt(K) = 1.0 / (dt*GV%Rho0) + enddo + else + do K=1,nz+1 + SpV_dt(K) = I_rho0dt + enddo + endif + endif +#endif + ! Make local copies of surface forcing and process them. +#ifdef __NVCOMPILER_OPENMP_GPU + ! GPU port: Boussinesq fluxes%ustar path only (guarded above); read the hoisted plain arrays. + u_star = ustar_2d(i,j) + u_star_Mean = ustar_gustless_2d(i,j) + mech_TKE = dt * GV%Rho0 * u_star**3 + B_flux = buoy_flux(i,j) +#else if (associated(fluxes%ustar) .and. (GV%Boussinesq .or. .not.associated(fluxes%tau_mag))) then u_star = fluxes%ustar(i,j) u_star_Mean = fluxes%ustar_gustless(i,j) @@ -661,6 +789,7 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, u_star = (1.0 - fluxes%frac_shelf_h(i,j)) * u_star + & fluxes%frac_shelf_h(i,j) * fluxes%ustar_shelf(i,j) endif +#endif if (u_star < CS%ustar_min) u_star = CS%ustar_min if (CS%omega_frac >= 1.0) then absf = 2.0*CS%omega @@ -673,7 +802,11 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, ! Perhaps provide a first guess for MLD based on a stored previous value. MLD_io = -1.0 +#ifdef __NVCOMPILER_OPENMP_GPU + if (CS%MLD_iteration_guess .and. (ML_depth_2d(i,j) > 0.0)) MLD_io = ML_depth_2d(i,j) +#else if (CS%MLD_iteration_guess .and. (CS%ML_depth(i,j) > 0.0)) MLD_io = CS%ML_depth(i,j) +#endif BBLD_io = 0.0 ! Store the initial guesses at the boundary layer depths for testing sensitivities. @@ -684,19 +817,24 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, else do K=1,nz+1 ; SpV_dt_cf(K) = SpV_dt(K) ; enddo endif +#ifndef __NVCOMPILER_OPENMP_GPU if (stoch_CS%pert_epbl) then ! stochastics are active call ePBL_column(h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, SpV_dt_cf, TKE_forcing, B_flux, absf, & u_star, u_star_mean, mech_TKE, dt, MLD_io, Kd, mixvel, mixlen, GV, & US, CS, eCD, Waves, G, i, j, & TKE_gen_stoch=stoch_CS%epbl1_wts(i,j), TKE_diss_stoch=stoch_CS%epbl2_wts(i,j)) else +#endif call ePBL_column(h, dz, u, v, T0, S0, dSV_dT_1d, dSV_dS_1d, SpV_dt_cf, TKE_forcing, B_flux, absf, & u_star, u_star_mean, mech_TKE, dt, MLD_io, Kd, mixvel, mixlen, GV, & US, CS, eCD, Waves, G, i, j) +#ifndef __NVCOMPILER_OPENMP_GPU endif if (CS%id_Kd_ePBL_col_by_col > 0) & call post_data_3d_by_column(CS%id_Kd_ePBL_col_by_col, Kd, CS%diag, i, j) +#endif +#ifndef __NVCOMPILER_OPENMP_GPU ! Add the diffusivity due to bottom boundary layer mixing, if there is energy to drive this mixing. if (BBL_mixing) then if (CS%MLD_iteration_guess .and. (CS%BBL_depth(i,j) > 0.0)) BBLD_io = CS%BBL_depth(i,j) @@ -734,14 +872,21 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, if ((CS%id_BBL_decay_scale > 0) .and. (CS%TKE_decay * absf > 0)) & diag_BBL_decay_scale(i,j) = u_star_BBL / (CS%TKE_decay * absf) endif +#endif - ! Copy the diffusivities to a 2-d array. + ! Copy the diffusivities to the output interface diffusivity array. do K=1,nz+1 - Kd_2d(i,K) = Kd(K) + Kd_int(i,j,K) = Kd(K) enddo +#ifdef __NVCOMPILER_OPENMP_GPU + ML_depth_2d(i,j) = MLD_io + BBL_depth_2d(i,j) = BBLD_io +#else CS%ML_depth(i,j) = MLD_io CS%BBL_depth(i,j) = BBLD_io +#endif +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%TKE_diagnostics) then diag_TKE_MKE(i,j) = diag_TKE_MKE(i,j) + eCD%dTKE_MKE diag_TKE_conv(i,j) = diag_TKE_conv(i,j) + eCD%dTKE_conv @@ -820,17 +965,34 @@ subroutine energetic_PBL(h_3d, u_3d, v_3d, tv, fluxes, visc, dt, Kd_int, G, GV, endif if (CS%id_opt_diff_hML_depth > 0) diff_hML_depth(i,j) = BLD_1 - BLD_2 endif +#endif else ! End of the ocean-point part of the i-loop ! For masked points, Kd_int must still be set (to 0) because it has intent out. - do K=1,nz+1 ; Kd_2d(i,K) = 0. ; enddo + do K=1,nz+1 ; Kd_int(i,j,K) = 0. ; enddo +#ifdef __NVCOMPILER_OPENMP_GPU + ML_depth_2d(i,j) = 0.0 + BBL_depth_2d(i,j) = 0.0 +#else CS%ML_depth(i,j) = 0.0 CS%BBL_depth(i,j) = 0.0 +#endif endif ; enddo ! Close of i-loop - Note the unusual loop order, with k-loops inside i-loops. - do K=1,nz+1 ; do i=is,ie ; Kd_int(i,j,K) = Kd_2d(i,K) ; enddo ; enddo - enddo ! j-loop + +#ifdef __NVCOMPILER_OPENMP_GPU + ! Copy the results back to the host and tear down the driver-level data region. + !$omp target exit data map(from: Kd_int, ML_depth_2d, BBL_depth_2d) & + !$omp map(release: h_3d, u_3d, v_3d, dz_3d, TKE_forced, dSV_dT, dSV_dS, buoy_flux, & + !$omp ustar_2d, ustar_gustless_2d, CS, tv%T, tv%S, tv) + if (associated(Waves)) then + !$omp target exit data map(release: Waves) + endif + do j=js,je ; do i=is,ie + CS%ML_depth(i,j) = ML_depth_2d(i,j) ; CS%BBL_depth(i,j) = BBL_depth_2d(i,j) + enddo ; enddo +#endif if (CS%id_Kd_ePBL_col_by_col > 0) call post_data_3d_final(CS%id_Kd_ePBL_col_by_col, CS%diag) if (CS%debug .and. BBL_mixing) then @@ -966,7 +1128,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, ! mixing. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif pres_Z, & ! Interface pressures with a rescaling factor to convert interface height ! movements into changes in column potential energy [R Z2 T-2 ~> kg m-1 s-2]. hb_hs ! The distance from the bottom over the thickness of the @@ -983,7 +1149,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real :: Idecay_len_TKE ! The inverse of a turbulence decay length scale [H-1 ~> m-1 or m2 kg-1]. real :: dz_sum ! The total thickness of the water column [Z ~> m]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif dT_to_dColHt, & ! Partial derivative of the total column height with the temperature changes ! within a layer [Z C-1 ~> m degC-1]. dS_to_dColHt, & ! Partial derivative of the total column height with the salinity changes @@ -1020,7 +1190,11 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, ! mixing effects with other yet lower layers [C H ~> degC m or degC kg m-2]. Sh_b ! An effective salinity times a thickness in the layer below, including implicit ! mixing effects with other yet lower layers [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif MixLen_shape, & ! A nondimensional shape factor for the mixing length that ! gives it an appropriate asymptotic value at the bottom of ! the boundary layer [nondim]. @@ -1153,6 +1327,10 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real, dimension(20) :: Kddt_h_itt ! The value of Kddt_h_guess after each iteration [H ~> m or kg m-2] real, dimension(20) :: dPEa_dKd_itt ! The value of dPEc_dKd after each iteration [R Z3 T-2 H-1 ~> J m-3 or J kg-1] real, dimension(20) :: MKE_src_itt ! The value of MKE_src after each iteration [R Z3 T-2 ~> J m-2] +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: mech_TKE_k, conv_PErel_k, nstar_k, dT_expect, dS_expect + integer, dimension(GPU_nk_max) :: num_itts +#else real, dimension(SZK_(GV)) :: mech_TKE_k ! The mechanically generated turbulent kinetic energy ! available for mixing over a time step for each layer [R Z3 T-2 ~> J m-2]. real, dimension(SZK_(GV)) :: conv_PErel_k ! The potential energy that has been convectively released @@ -1162,6 +1340,7 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, real, dimension(SZK_(GV)) :: dT_expect ! Expected temperature changes [C ~> degC] real, dimension(SZK_(GV)) :: dS_expect ! Expected salinity changes [S ~> ppt] integer, dimension(SZK_(GV)) :: num_itts +#endif integer :: k, nz, itt, max_itt @@ -1251,8 +1430,12 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, !/ Here we get mstar, which is the ratio of convective TKE driven mixing to UStar**3 if (CS%Use_LT) then +#ifdef __NVCOMPILER_OPENMP_GPU + call get_Langmuir_Number_LF17(LA, GV, US, abs(MLD_guess), u_star_mean, CS%wave_lf17) +#else call get_Langmuir_Number(LA, G, GV, US, abs(MLD_guess), u_star_mean, i, j, dz, Waves, & U_H=u, V_H=v) +#endif call find_mstar(CS, US, B_flux, u_star, MLD_guess, absf, .false., & mstar_total, Langmuir_Number=La, Convect_Langmuir_Number=LAmod,& mstar_LT=mstar_LT) @@ -1315,13 +1498,18 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, dz_rsum = 0.0 MixLen_shape(1) = 1.0 if (CS%eqdisc) then ! update Kd as per Machine Learning equation discovery +#ifndef __NVCOMPILER_OPENMP_GPU call kappa_eqdisc(MixLen_shape, CS, GV, h, absf, B_flux, u_star, MLD_guess) +#endif else do K=2,nz+1 dz_rsum = dz_rsum + dz(k-1) if (CS%MixLenExponent==2.0) then MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**2 ! CS%MixLenExponent + elseif (CS%MixLenExponent==1.0) then + MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & + max(0.0, (MLD_guess - dz_rsum)*I_MLD) ! CS%MixLenExponent==1.0 identity else MixLen_shape(K) = CS%transLay_scale + (1.0 - CS%transLay_scale) * & (max(0.0, (MLD_guess - dz_rsum)*I_MLD) )**CS%MixLenExponent @@ -1331,11 +1519,13 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, endif v0_ML_turb_vel_scale = 0.0 ! a variable that gets passed on to get_eqdisc_v0 & get_eqdisc_v0h +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%eqdisc_v0) then call get_eqdisc_v0(CS,absf,B_flux,u_star,v0_ML_turb_vel_scale) elseif (CS%eqdisc_v0h) then call get_eqdisc_v0h(CS,B_flux,u_star,MLD_guess,v0_ML_turb_vel_scale) endif +#endif Kd(1) = 0.0 ; Kddt_h(1) = 0.0 hp_a(1) = h(1) @@ -1363,7 +1553,7 @@ subroutine ePBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, TKE_forcing, Idecay_len_TKE = (CS%TKE_decay * absf) / (h_dz_int(K) * u_star) endif exp_kh = 1.0 - if (Idecay_len_TKE > 0.0) exp_kh = exp(-h(k-1)*Idecay_len_TKE) + if (Idecay_len_TKE > 0.0) exp_kh = exp_reprod(-h(k-1)*Idecay_len_TKE) if (CS%TKE_diagnostics) & eCD%dTKE_mech_decay = eCD%dTKE_mech_decay + (exp_kh-1.0) * mech_TKE * I_dtdiag if (present(TKE_diss_stoch)) then ! perturb the TKE destruction @@ -2021,7 +2211,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & ! energy that is supplied as an argument to this routine. ! Local variables +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif pres_Z, & ! Interface pressures with a rescaling factor to convert interface height ! movements into changes in column potential energy [R Z2 T-2 ~> kg m-1 s-2]. dztop_dztot ! The distance from the surface divided by the thickness of the @@ -2039,7 +2233,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & real :: Idecay_len_TKE ! The inverse of a turbulence decay length scale [H-1 ~> m-1 or m2 kg-1]. real :: dz_sum ! The total thickness of the water column [Z ~> m]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif dT_to_dColHt, & ! Partial derivative of the total column height with the temperature changes ! within a layer [Z C-1 ~> m degC-1]. dS_to_dColHt, & ! Partial derivative of the total column height with the salinity changes @@ -2089,7 +2287,11 @@ subroutine ePBL_BBL_column(h, dz, u, v, T0, S0, dSV_dT, dSV_dS, SpV_dt, absf, & ! mixing effects with other yet lower layers [C H ~> degC m or degC kg m-2]. Sh_b ! An effective salinity times a thickness in the layer below, including implicit ! mixing effects with other yet lower layers [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif MixLen_shape, & ! A nondimensional shape factor for the mixing length that ! gives it an appropriate asymptotic value at the bottom of ! the boundary layer [nondim]. @@ -3042,17 +3244,17 @@ function exp_decay_TKE_adjust(hb, ha, Idecay) result(TKE_to_PE_scale) ! TKE_to_PE_scale = (0.5 * (khb + kha)) / & ! ((exp(-khb) - (1.0 - khb)) / khb + (exp(kha) - (1.0 + kha)) / kha) TKE_to_PE_scale = (0.5 * (khb + kha) * (kha * khb)) / & - (kha * (exp(-khb) - (1.0 - khb)) + khb * (exp(kha) - (1.0 + kha))) + (kha * (exp_reprod(-khb) - (1.0 - khb)) + khb * (exp_reprod(kha) - (1.0 + kha))) elseif (khb > 2.2e-4) then ! For small values of kha, approximate (exp(kha) - (1.0 + hha)) by the first two ! terms of its Taylor series: 0.5*kha**2 + C1_6*kha**3 + ... + kha**n/n! + ... ! which is more accurate when kha**4/24. < 1e-16 or kha < ~ 2.21e-4. TKE_to_PE_scale = (0.5 * (khb + kha) * khb) / & - ((exp(-khb) - (1.0 - khb)) + 0.5*(khb * kha) * (1.0 + C1_3*kha)) + ((exp_reprod(-khb) - (1.0 - khb)) + 0.5*(khb * kha) * (1.0 + C1_3*kha)) elseif (kha > 2.2e-4) then ! Use a Taylor series expansion for small values of khb TKE_to_PE_scale = (0.5 * (khb + kha) * kha) / & - (0.5 * (kha * khb) * (1.0 - C1_3*Khb) + (exp(kha) - (1.0 + kha))) + (0.5 * (kha * khb) * (1.0 - C1_3*Khb) + (exp_reprod(kha) - (1.0 + kha))) else ! (kha < 2.2e-4) .and. (khb < 2.2e-4) - use Taylor series approximations for both TKE_to_PE_scale = 1.0 / (1.0 + C1_3*(kha - khb)) endif @@ -3565,13 +3767,13 @@ subroutine find_mstar(CS, US, Buoyancy_Flux, UStar, & mstar_S = CS%mstar_coef*sqrt(max(0.0,Buoyancy_Flux) / UStar**2 / & (Abs_Coriolis + 1.e-10*US%T_to_s) ) ! The limit for rotation (Ekman length) limited mixing - mstar_N = CS%C_Ek * log( max( 1., UStar / (Abs_Coriolis + 1.e-10*US%T_to_s) / BLD ) ) + mstar_N = CS%C_Ek * log_reprod( max( 1., UStar / (Abs_Coriolis + 1.e-10*US%T_to_s) / BLD ) ) else ! The limit for the balance of rotation and stabilizing is f(L_Ekman,L_Obukhov) mstar_S = CS%mstar_coef*sqrt(max(0.0, Buoyancy_Flux) / (UStar**2 * max(Abs_Coriolis, 1.e-20*US%T_to_s))) ! The limit for rotation (Ekman length) limited mixing mstar_N = 0.0 - if (UStar > Abs_Coriolis * BLD) mstar_N = CS%C_Ek * log(UStar / (Abs_Coriolis * BLD)) + if (UStar > Abs_Coriolis * BLD) mstar_N = CS%C_Ek * log_reprod(UStar / (Abs_Coriolis * BLD)) endif ! Here 1.25 is about .5/von Karman, which gives the Obukhov limit. @@ -3693,10 +3895,10 @@ subroutine mstar_Langmuir(CS, US, Abs_Coriolis, Buoyancy_Flux, UStar, BLD, Langm if (CS%LT_enhance_form == Langmuir_rescale) then ! Enhancement is multiplied (added mst_lt set to 0) Enhance_mstar = min(CS%Max_Enhance_M, & - (1. + CS%LT_enhance_coef * Convect_Langmuir_Number**CS%LT_enhance_exp) ) + (1. + CS%LT_enhance_coef * exp_reprod(CS%LT_enhance_exp*log_reprod(Convect_Langmuir_Number))) ) elseif (CS%LT_enhance_form == Langmuir_add) then ! or Enhancement is additive (multiplied enhance_m set to 1) - mstar_LT_add = CS%LT_enhance_coef * Convect_Langmuir_Number**CS%LT_enhance_exp + mstar_LT_add = CS%LT_enhance_coef * exp_reprod(CS%LT_enhance_exp*log_reprod(Convect_Langmuir_Number)) endif endif @@ -4431,6 +4633,15 @@ subroutine energetic_PBL_init(Time, G, GV, US, param_file, diag, CS) call safe_alloc_alloc(CS%ML_depth, isd, ied, jsd, jed) call safe_alloc_alloc(CS%BBL_depth, isd, ied, jsd, jed) +#ifdef __NVCOMPILER_OPENMP_GPU + ! GPU port guards: ePBL_column runs on the device with fixed-size (GPU_nk_max) per-column scratch, + ! and the equation-discovery (eqdisc) mixing-length/velocity paths are not device-callable. + if (GV%ke > GPU_nk_max) call MOM_error(FATAL, & + "energetic_PBL GPU build: GV%ke exceeds GPU_nk_max; increase GPU_nk_max in MOM_energetic_PBL.F90.") + if (CS%eqdisc .or. CS%eqdisc_v0 .or. CS%eqdisc_v0h) call MOM_error(FATAL, & + "energetic_PBL GPU build: the EPBL_EQD_DIFFUSIVITY (equation-discovery) paths are not device-callable.") +#endif + end subroutine energetic_PBL_init !> Clean up and deallocate memory associated with the energetic_PBL module. diff --git a/src/parameterizations/vertical/MOM_kappa_shear.F90 b/src/parameterizations/vertical/MOM_kappa_shear.F90 index 1cff115fe5..b5306d6b88 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 @@ -129,6 +131,17 @@ 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) + +!> 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 @@ -170,15 +183,24 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & diag_S2_init, & ! Diagnostic of S2 as provided to this routine [T-2 ~> s-2] diag_N2_mean, & ! Diagnostic of N2 averaged over the timestep applied [T-2 ~> s-2] diag_S2_mean ! Diagnostic of S2 averaged over the timestep applied [T-2 ~> s-2] - real, dimension(SZI_(G),SZK_(GV)) :: & - h_2d, & ! A 2-D version of h [H ~> m or kg m-2]. - dz_2d, & ! Vertical distance between interface heights [Z ~> m]. - u_2d, v_2d, & ! 2-D versions of u_in and v_in, converted to [L T-1 ~> m s-1]. - T_2d, S_2d, rho_2d ! 2-D versions of T [C ~> degC], S [S ~> ppt], and rho [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)+1) :: & - kappa_2d, & ! 2-D version of kappa_io [H Z T-1 ~> m2 s-1 or Pa s] - tke_2d ! 2-D version tke_io [Z2 T-2 ~> m2 s-2]. + ! GPU port: the per-J 2-D slabs (h_2d/dz_2d/u_2d/v_2d/T_2d/S_2d/rho_2d) and the per-J kappa_2d/ + ! tke_2d are gone; the device column loop reads the mapped 3-D inputs (h, u_in, v_in, tv%T, tv%S, + ! dz_3d) directly and stages its results in kappa_3d/tke_3d (was the per-J kappa_2d/tke_2d). + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)) :: & + dz_3d ! Vertical distance between interface heights [Z ~> m]. + real, dimension(SZI_(G),SZJ_(G),SZK_(GV)+1) :: & + kappa_3d, & ! Device staging array for the columns' averaged kappa [H Z T-1 ~> m2 s-1 or Pa s] + tke_3d ! Device staging array for the columns' TKE [Z2 T-2 ~> m2 s-2]. + real, dimension(SZI_(G),SZJ_(G)) :: & + surface_pres_2d ! The surface pressure at tracer points [R L2 T-2 ~> Pa]. + ! In GPU builds the per-column private scratch has compile-time-constant sizes so each device + ! thread gets stack ("local memory") arrays; runtime-sized privates are device-heap allocated + ! per column, which exhausts the default heap and serializes on the allocator. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif Idz, & ! The inverse of the thickness of the merged layers [H-1 ~> m2 kg-1]. h_lay, & ! The layer thickness [H ~> m or kg m-2] dz_lay, & ! The geometric layer thickness in height units [Z ~> m] @@ -187,7 +209,11 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & T0xdz, & ! The initial temperature times thickness [C H ~> degC m or degC kg m-2] or if ! temperature is not a state variable, the density times thickness [R H ~> kg m-2 or kg2 m-5] S0xdz ! The initial salinity times dz [S H ~> ppt m or ppt kg m-2]. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: & +#else real, dimension(SZK_(GV)+1) :: & +#endif kappa, & ! The shear-driven diapycnal diffusivity at an interface [H Z T-1 ~> m2 s-1 or Pa s] tke, & ! The Turbulent Kinetic Energy per unit mass at an interface [Z2 T-2 ~> m2 s-2]. kappa_avg, & ! The time-weighted average of kappa [H Z T-1 ~> m2 s-1 or Pa s] @@ -206,17 +232,43 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & logical :: use_temperature ! If true, temperature and salinity have been ! allocated and are being used as state variables. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max+1) :: kc ! The index map between the original +#else integer, dimension(SZK_(GV)+1) :: kc ! The index map between the original +#endif ! interfaces and the interfaces with massless layers ! merged into nearby massive layers. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max+1) :: kf ! The fractional weight of interface kc+1 for +#else real, dimension(SZK_(GV)+1) :: kf ! The fractional weight of interface kc+1 for +#endif ! interpolating back to the original index space [nondim]. integer :: is, ie, js, je, i, j, k, nz, nzc + integer :: eos_form ! The equation-of-state form id, resolved host-side for the GPU EOS path. + 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) @@ -225,29 +277,62 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & if (CS%id_N2_mean>0) diag_N2_mean(:,:,:) = 0.0 if (CS%id_S2_mean>0) diag_S2_mean(:,:,:) = 0.0 - !$OMP parallel do default(private) shared(js,je,is,ie,nz,h,u_in,v_in,use_temperature,tv,G,GV,US, & - !$OMP CS,kappa_io,dz_massless,k0dt,p_surf,dt,tke_io,kv_io, & - !$OMP diag_N2_init,diag_S2_init,diag_N2_mean,diag_S2_mean) - do j=js,je - - ! Convert layer thicknesses into geometric thickness in height units. - call thickness_to_dz(h, tv, dz_2d, j, G, GV) - - do k=1,nz ; do i=is,ie - h_2d(i,k) = h(i,j,k) - u_2d(i,k) = u_in(i,j,k) ; v_2d(i,k) = v_in(i,j,k) + ! GPU port: interpolate p_surf to tracer points on the host, ahead of the device column region + ! (p_surf is a possibly-unassociated pointer, simpler kept off the device). Verbatim from the + ! former in-column expression, evaluated for every tracer point; land values are never read. + surface_pres_2d(:,:) = 0.0 + if (associated(p_surf)) then + do j=js,je ; do i=is,ie + surface_pres_2d(i,j) = p_surf(i,j) enddo ; enddo - if (use_temperature) then ; do k=1,nz ; do i=is,ie - T_2d(i,k) = tv%T(i,j,k) ; S_2d(i,k) = tv%S(i,j,k) - enddo ; enddo ; else ; do k=1,nz ; do i=is,ie - rho_2d(i,k) = GV%Rlay(k) ! Could be tv%Rho(i,j,k) ? - enddo ; enddo ; endif + endif + + ! Convert layer thicknesses into geometric thickness in height units, over the whole compute + ! domain (was a per-J call inside the loop; hoisted so dz_3d can be mapped once). + call thickness_to_dz(h, tv, dz_3d, G, GV, US, halo_size=0) + + ! --- GPU port: run the per-column solver on the device. The columns iterate as a target teams + ! loop collapsed over (j,i); every piece of per-column scratch is private (the declare-target + ! solver's own locals are automatically private per device thread). h is host-authoritative in + ! the (host) diabatic stack -> refresh. tv%T/tv%S are persistently mapped (MOM.F90) but this + ! non-full_convection path reads the diabatic-mutated fields directly, so refresh them too + ! (a map(to:) on an already-present object copies nothing). u_in/v_in (set_diffusivity's + ! u_h/v_h) and dz_3d are fresh host locals. CS is all scalars plus a diag pointer that is never + ! dereferenced in device code, so a per-call shallow map(to:) suffices. kappa_io/tke_io/kv_io + ! are mapped to: (not alloc) and refreshed so the host values (halos, and kappa_io's previous + ! contents outside the compute domain) survive the full-array update from below. diag_* are + ! mapped after their (conditional) host zeroing. kappa_3d/tke_3d are device-only staging for + ! what was the per-J kappa_2d/tke_2d. + !$omp target enter data map(to: h) + !$omp target update to(h) + ! tv%T/tv%S are read on the device in the column setup below. The T/S *data* is persistently + ! mapped via MOM's CS%tv (MOM.F90), but that attaches it to a different descriptor; this routine's + ! tv dummy must be mapped here so the device can resolve tv%T/tv%S (refcount bump + attach, no + ! copy), then update to refreshes the diabatic-mutated values. (set_viscosity uses this idiom.) + if (use_temperature) then + !$omp target enter data map(to: tv, tv%T, tv%S) + !$omp target update to(tv%T, tv%S) + endif + !$omp target enter data map(to: u_in, v_in, dz_3d) + !$omp target update to(u_in, v_in) + !$omp target enter data map(to: CS) + !$omp target enter data map(to: surface_pres_2d) + !$omp target enter data map(to: kappa_io, tke_io, kv_io) + !$omp target update to(kappa_io, tke_io, kv_io) + !$omp target enter data map(to: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target enter data map(alloc: kappa_3d, tke_3d) !--------------------------------------- ! Work on each column. !--------------------------------------- + !$omp target teams loop collapse(2) & + !$omp private(nzc, kc, kf, Idz, h_lay, dz_lay, u0xdz, v0xdz, T0xdz, S0xdz, dz_in_lay, & + !$omp f2, surface_pres, kappa, tke, kappa_avg, tke_avg, N2_init, S2_init, & + !$omp N2_mean, S2_mean, k) & + !$omp firstprivate(nz, dt, k0dt, dz_massless, use_temperature, eos_form, & + !$omp eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa) + do j=js,je do i=is,ie ; if (G%mask2dT(i,j) > 0.0) then - ! call cpu_clock_begin(id_clock_setup) ! Store a transposed version of the initial arrays. ! Any elimination of massless layers would occur here. @@ -259,25 +344,20 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & T0xdz(k) = 0.0 ; S0xdz(k) = 0.0 ! Add a new layer if this one has mass. -! if ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless)) nzc = nzc+1 if ((k>CS%nkml) .and. (h_lay(nzc) > 0.0) .and. & - (h_2d(i,k) > dz_massless)) nzc = nzc+1 - - ! Only merge clusters of massless layers. -! if ((h_lay(nzc) > dz_massless) .or. & -! ((h_lay(nzc) > 0.0) .and. (h_2d(i,k) > dz_massless))) nzc = nzc+1 + (h(i,j,k) > dz_massless)) nzc = nzc+1 kc(k) = nzc - h_lay(nzc) = h_lay(nzc) + h_2d(i,k) - dz_lay(nzc) = dz_lay(nzc) + dz_2d(i,k) - u0xdz(nzc) = u0xdz(nzc) + u_2d(i,k)*h_2d(i,k) - v0xdz(nzc) = v0xdz(nzc) + v_2d(i,k)*h_2d(i,k) + h_lay(nzc) = h_lay(nzc) + h(i,j,k) + dz_lay(nzc) = dz_lay(nzc) + dz_3d(i,j,k) + u0xdz(nzc) = u0xdz(nzc) + u_in(i,j,k)*h(i,j,k) + v0xdz(nzc) = v0xdz(nzc) + v_in(i,j,k)*h(i,j,k) if (use_temperature) then - T0xdz(nzc) = T0xdz(nzc) + T_2d(i,k)*h_2d(i,k) - S0xdz(nzc) = S0xdz(nzc) + S_2d(i,k)*h_2d(i,k) + T0xdz(nzc) = T0xdz(nzc) + tv%T(i,j,k)*h(i,j,k) + S0xdz(nzc) = S0xdz(nzc) + tv%S(i,j,k)*h(i,j,k) else - T0xdz(nzc) = T0xdz(nzc) + rho_2d(i,k)*h_2d(i,k) - S0xdz(nzc) = S0xdz(nzc) + rho_2d(i,k)*h_2d(i,k) + T0xdz(nzc) = T0xdz(nzc) + GV%Rlay(k)*h(i,j,k) + S0xdz(nzc) = S0xdz(nzc) + GV%Rlay(k)*h(i,j,k) endif enddo kc(nz+1) = nzc+1 @@ -287,28 +367,28 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & ! Now determine kf, the fractional weight of interface kc when ! interpolating between interfaces kc and kc+1. - kf(1) = 0.0 ; dz_in_lay = h_2d(i,1) + kf(1) = 0.0 ; dz_in_lay = h(i,j,1) do k=2,nz if (kc(k) > kc(k-1)) then - kf(k) = 0.0 ; dz_in_lay = h_2d(i,k) + kf(k) = 0.0 ; dz_in_lay = h(i,j,k) else - kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h_2d(i,k) + kf(k) = dz_in_lay*Idz(kc(k)) ; dz_in_lay = dz_in_lay + h(i,j,k) endif enddo kf(nz+1) = 0.0 else do k=1,nz - h_lay(k) = h_2d(i,k) - dz_lay(k) = dz_2d(i,k) - u0xdz(k) = u_2d(i,k)*h_lay(k) ; v0xdz(k) = v_2d(i,k)*h_lay(k) + h_lay(k) = h(i,j,k) + dz_lay(k) = dz_3d(i,j,k) + u0xdz(k) = u_in(i,j,k)*h_lay(k) ; v0xdz(k) = v_in(i,j,k)*h_lay(k) enddo if (use_temperature) then do k=1,nz - T0xdz(k) = T_2d(i,k)*h_lay(k) ; S0xdz(k) = S_2d(i,k)*h_lay(k) + T0xdz(k) = tv%T(i,j,k)*h_lay(k) ; S0xdz(k) = tv%S(i,j,k)*h_lay(k) enddo else do k=1,nz - T0xdz(k) = rho_2d(i,k)*h_lay(k) ; S0xdz(k) = rho_2d(i,k)*h_lay(k) + T0xdz(k) = GV%Rlay(k)*h_lay(k) ; S0xdz(k) = GV%Rlay(k)*h_lay(k) enddo endif nzc = nz @@ -316,28 +396,26 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & endif f2 = 0.25 * ((G%Coriolis2Bu(I,J) + G%Coriolis2Bu(I-1,J-1)) + & (G%Coriolis2Bu(I,J-1) + G%Coriolis2Bu(I-1,J))) - surface_pres = 0.0 ; if (associated(p_surf)) surface_pres = p_surf(i,j) - - ! ---------------------------------------------------- I_Ld2_1d, dz_Int_1d + surface_pres = surface_pres_2d(i,j) ! Set the initial guess for kappa, here defined at interfaces. - ! ---------------------------------------------------- do K=1,nzc+1 ; kappa(K) = CS%kappa_seed ; enddo call kappa_shear_column(kappa, tke, dt, nzc, f2, surface_pres, & 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, & + use_temperature) - ! call cpu_clock_begin(id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_2d(i,K) = kappa_avg(K) + kappa_3d(i,j,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_2d(i,K) = tke(K) + tke_3d(i,j,K) = tke(K) else - tke_2d(i,K) = tke_avg(K) + tke_3d(i,j,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -355,11 +433,11 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & else do K=1,nz+1 if (kf(K) == 0.0) then - kappa_2d(i,K) = kappa_avg(kc(K)) - tke_2d(i,K) = tke_avg(kc(K)) + kappa_3d(i,j,K) = kappa_avg(kc(K)) + tke_3d(i,j,K) = tke_avg(kc(K)) else - kappa_2d(i,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) - tke_2d(i,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) + kappa_3d(i,j,K) = (1.0-kf(K)) * kappa_avg(kc(K)) + kf(K) * kappa_avg(kc(K)+1) + tke_3d(i,j,K) = (1.0-kf(K)) * tke_avg(kc(K)) + kf(K) * tke_avg(kc(K)+1) endif enddo do K=1,nz+1 @@ -380,21 +458,38 @@ subroutine Calculate_kappa_shear(u_in, v_in, h, tv, p_surf, kappa_io, tke_io, & endif enddo endif - ! call cpu_clock_end(id_clock_setup) else ! Land points, still inside the i-loop. do K=1,nz+1 - kappa_2d(i,K) = 0.0 ; tke_2d(i,K) = 0.0 + kappa_3d(i,j,K) = 0.0 ; tke_3d(i,j,K) = 0.0 enddo endif ; enddo ! i-loop + enddo ! end of j-loop - do K=1,nz+1 ; do i=is,ie - kappa_io(i,j,K) = G%mask2dT(i,j) * kappa_2d(i,K) - tke_io(i,j,K) = G%mask2dT(i,j) * tke_2d(i,K) - kv_io(i,j,K) = ( G%mask2dT(i,j) * kappa_2d(i,K) ) * CS%Prandtl_turb + ! Store the columns' results back in the output arrays (masked), on the device. + do concurrent (K=1:nz+1, j=js:je, i=is:ie) + kappa_io(i,j,K) = G%mask2dT(i,j) * kappa_3d(i,j,K) + tke_io(i,j,K) = G%mask2dT(i,j) * tke_3d(i,j,K) + kv_io(i,j,K) = ( G%mask2dT(i,j) * kappa_3d(i,j,K) ) * CS%Prandtl_turb + enddo - enddo ; enddo + ! The checksums and post_data below are still on the host; copy back the outputs (and, guarded + ! by the same conditions as their consumers, the diag_* arrays), then mirror the enter data. + !$omp target update from(kappa_io, tke_io, kv_io) + if ((CS%id_N2_init>0) .or. (CS%id_S2_init>0) .or. (CS%id_N2_mean>0) .or. (CS%id_S2_mean>0) & + .or. CS%debug) then + !$omp target update from(diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + endif - enddo ! end of j-loop + !$omp target exit data map(release: kappa_3d, tke_3d) + !$omp target exit data map(release: diag_N2_init, diag_S2_init, diag_N2_mean, diag_S2_mean) + !$omp target exit data map(release: kappa_io, tke_io, kv_io) + !$omp target exit data map(release: surface_pres_2d) + !$omp target exit data map(release: CS) + !$omp target exit data map(release: u_in, v_in, dz_3d) + if (use_temperature) then + !$omp target exit data map(release: tv, tv%T, tv%S) + endif + !$omp target exit data map(release: h) if (CS%debug) then call hchksum(diag_N2_init, "kappa_shear N2_init", G%HI, unscale=US%s_to_T**2) @@ -467,16 +562,24 @@ 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),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)) :: & + 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) :: & + 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]. + ! 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] @@ -484,7 +587,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] @@ -500,16 +607,23 @@ 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 ! 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 @@ -518,6 +632,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 @@ -530,87 +646,170 @@ 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 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) + ! --- 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) + 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 + ! --- 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) + + ! 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 - !$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 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 + ! --- 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) !--------------------------------------- ! 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) @@ -624,25 +823,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 @@ -652,28 +851,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 @@ -681,19 +880,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. @@ -702,16 +889,18 @@ 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, & + use_temperature) ! call cpu_clock_begin(Id_clock_setup) ! Extrapolate from the vertically reduced grid back to the original layers. if (nz == nzc) then do K=1,nz+1 - kappa_2d(I,K) = kappa_avg(K) + kappa_3d(I,J,K) = kappa_avg(K) if (CS%all_layer_TKE_bug) then - tke_2d(I,K) = tke(K) + tke_3d(I,J,K) = tke(K) else - tke_2d(I,K) = tke_avg(K) + tke_3d(I,J,K) = tke_avg(K) endif enddo if (CS%id_N2_mean>0) then ; do K=1,nz+1 @@ -729,11 +918,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 @@ -757,30 +946,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 - do K=1,nz+1 ; do I=IsB,IeB - h_vert(I,J,k) = h_2d(I,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 + ! 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 - 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 + h_vert(I,J,k) = h_slab(I,J,k) 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. @@ -860,13 +1063,26 @@ 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) + !$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) + !$omp target exit data map(release: h) + 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, & + 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] @@ -905,12 +1121,29 @@ 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. + 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] + 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 + ! 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]. @@ -919,7 +1152,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] @@ -970,6 +1207,10 @@ 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 (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]. 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. @@ -993,8 +1234,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., @@ -1021,7 +1260,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. @@ -1133,9 +1371,36 @@ 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 (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), & + 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 +#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, & @@ -1145,6 +1410,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 @@ -1410,7 +1680,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] @@ -1524,7 +1798,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 @@ -1542,10 +1817,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 @@ -1623,7 +1906,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]. @@ -2330,6 +2617,15 @@ 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 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 !> This function indicates to other modules whether the Jackson et al shear mixing diff --git a/src/parameterizations/vertical/MOM_set_diffusivity.F90 b/src/parameterizations/vertical/MOM_set_diffusivity.F90 index 971e5f6226..84a48cee3d 100644 --- a/src/parameterizations/vertical/MOM_set_diffusivity.F90 +++ b/src/parameterizations/vertical/MOM_set_diffusivity.F90 @@ -18,6 +18,8 @@ module MOM_set_diffusivity use MOM_diagnose_kdwork, only : vbf_CS use MOM_debugging, only : hchksum, uvchksum, Bchksum, hchksum_pair use MOM_EOS, only : calculate_density, calculate_density_derivs, EOS_domain +use MOM_EOS, only : calculate_density_derivs_elem_loc, get_EOS_form_and_scaling +use MOM_EOS, only : EOS_ROQUET_RHO, EOS_WRIGHT use MOM_error_handler, only : MOM_error, is_root_pe, FATAL, WARNING, NOTE use MOM_error_handler, only : callTree_showQuery use MOM_error_handler, only : callTree_enter, callTree_leave, callTree_waypoint @@ -239,6 +241,8 @@ module MOM_set_diffusivity integer :: id_clock_kappaShear, id_clock_CVMix_ddiff !>@} +integer, parameter :: njblock = 1 !< Number of j-rows per block in the tiled j-loop + contains subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_int, & @@ -285,10 +289,10 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] ! local variables - real :: N2_bot(SZI_(G)) ! Bottom squared buoyancy frequency [T-2 ~> s-2] - real :: rho_bot(SZI_(G)) ! In situ near-bottom density [T-2 ~> s-2] - real :: h_bot(SZI_(G)) ! Bottom boundary layer thickness [H ~> m or kg m-2] - integer :: k_bot(SZI_(G)) ! Bottom boundary layer thickness top layer index + real :: N2_bot(SZI_(G),njblock) ! Bottom squared buoyancy frequency [T-2 ~> s-2] + real :: rho_bot(SZI_(G),njblock) ! In situ near-bottom density [T-2 ~> s-2] + real :: h_bot(SZI_(G),njblock) ! Bottom boundary layer thickness [H ~> m or kg m-2] + integer :: k_bot(SZI_(G),njblock) ! Bottom boundary layer thickness top layer index type(diffusivity_diags) :: dd ! structure with arrays of available diags @@ -297,29 +301,33 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i T_f, S_f ! Temperature and salinity [C ~> degC] and [S ~> ppt] with properties in massless layers ! filled vertically by diffusion or the properties after full convective adjustment. - real, dimension(SZI_(G),SZK_(GV)) :: & - N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] + real, dimension(SZI_(G),SZK_(GV),njblock) :: & + N2_lay, & !< Squared buoyancy frequency associated with layers [T-2 ~> s-2] Kd_lay_2d, & !< The layer diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - dz, & !< Height change across layers [Z ~> m] maxTKE, & !< Energy required to entrain to h_max [H Z2 T-3 ~> m3 s-3 or W m-2] + TKE_to_Kd, & !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between + !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] + dz !< Height change across layers [Z ~> m] + + real, dimension(SZI_(G),SZK_(GV)) :: & prof_leak_2d, & !< vertical profile for leakage [Z-1 ~> m-1] prof_quad_2d, & !< vertical profile for bottom drag [Z-1 ~> m-1] prof_itidal_2d, & !< vertical profile for wave drag [Z-1 ~> m-1] prof_Froude_2d, & !< vertical profile for Froude drag [Z-1 ~> m-1] - prof_slope_2d, & !< vertical profile for critical slopes [Z-1 ~> m-1] - TKE_to_Kd !< Conversion rate (~1.0 / (G_Earth + dRho_lay)) between - !< TKE dissipated within a layer and Kd in that layer [T2 Z-1 ~> s2 m-1] + prof_slope_2d !< vertical profile for critical slopes [Z-1 ~> m-1] - real, dimension(SZI_(G),SZK_(GV)+1) :: & + real, dimension(SZI_(G),SZK_(GV)+1,njblock) :: & N2_int, & !< squared buoyancy frequency associated at interfaces [T-2 ~> s-2] + dRho_int, & !< Locally referenced potential density difference across interfaces [R ~> kg m-3] Kd_int_2d, & !< The interface diffusivities [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - Kv_bkgnd, & !< The background diffusion related interface viscosities [H Z T-1 ~> m2 s-1 or Pa s] + Kv_bkgnd !< The background diffusion related interface viscosities [H Z T-1 ~> m2 s-1 or Pa s] + + real, dimension(SZI_(G),SZK_(GV)+1) :: & Kd_leak_2d, & !< internal tides leakage diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_quad_2d, & !< internal tides bottom drag diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_itidal_2d, & !< internal tides wave drag diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_Froude_2d, & !< internal tides high Froude diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] Kd_slope_2d, & !< internal tides critical slopes diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - dRho_int, & !< Locally referenced potential density difference across interfaces [R ~> kg m-3] KT_extra, & !< Double diffusion diffusivity of temperature [H Z T-1 ~> m2 s-1 or kg m-1 s-1] KS_extra !< Double diffusion diffusivity of salinity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] @@ -328,11 +336,12 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i logical :: use_EOS ! If true, compute density from T/S using equation of state. logical :: TKE_to_Kd_used ! If true, TKE_to_Kd and maxTKE need to be calculated. - integer :: kb(SZI_(G)) ! The index of the lightest layer denser than the - ! buffer layer, or -1 without a bulk mixed layer. + integer :: kb(SZI_(G),njblock) ! The index of the lightest layer denser than the + ! buffer layer, or -1 without a bulk mixed layer. logical :: showCallTree ! If true, show the call tree. integer :: i, j, k, is, ie, js, je, nz, isd, ied, jsd, jed + integer :: jstart, jend, jj real :: kappa_dt_fill ! diffusivity times a timestep used to fill massless layers [H Z ~> m2 or kg m-1] @@ -472,211 +481,256 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! be an appropriate place to add a depth-dependent parameterization or another explicit ! parameterization of Kd. - !$OMP parallel do default(shared) private(dRho_int,N2_lay,Kd_lay_2d,Kd_int_2d,Kv_bkgnd,N2_int,dz, & - !$OMP N2_bot,rho_bot,h_bot,k_bot,KT_extra,KS_extra,TKE_to_Kd,maxTKE,dissip,kb) & - !$OMP if(.not. CS%use_CVMix_ddiff) - do j=js,je + ! GPU port: map the full-domain EOS inputs (T_f/S_f and the tv descriptor) once for the whole + ! j-loop and refresh the diabatic-mutated tv%T/tv%S here. find_N2 offloads per block (njblock=1), + ! so mapping/refreshing these inside it would re-copy full-domain arrays on every j-row and dominate + ! the runtime. h is already persistently device-resident. + !$omp target enter data map(to: T_f, S_f) + if (associated(tv%eqn_of_state)) then + !$omp target enter data map(to: tv, tv%T, tv%S) + !$omp target update to(tv%T, tv%S) + endif + + do jstart=js,je,njblock + jend = min(jstart+njblock-1, je) + + do j=jstart,jend ; jj = j - jstart + 1 + call thickness_to_dz(h, tv, dz(:,:,jj), j, G, GV) + enddo ! Set up variables related to the stratification. - call find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot) + call find_N2(h, tv, T_f, S_f, fluxes, njblock, jstart, jend, G, GV, US, CS, dRho_int, N2_lay, N2_int, N2_bot, rho_bot, h_bot, k_bot, dz) + + ! Add background mixing + call calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay_2d, Kd_int_2d, Kv_bkgnd, jstart, jend, njblock, G, GV, US, CS%bkgnd_mixing_csp) if (associated(dd%N2_3d)) then - do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K) ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie ; dd%N2_3d(i,j,K) = N2_int(i,K,jj) ; enddo ; enddo + enddo endif - - ! Add background mixing - call calculate_bkgnd_mixing(h, tv, N2_lay, Kd_lay_2d, Kd_int_2d, Kv_bkgnd, j, G, GV, US, CS%bkgnd_mixing_csp) ! Update Kv and 3-d diffusivity diagnostics. - if (associated(visc%Kv_slow)) then ; do K=1,nz+1 ; do i=is,ie - visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K) - enddo ; enddo ; endif - if (CS%id_Kv_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_bkgnd > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K) - enddo ; enddo ; endif - if (associated(VBF%Kd_bkgnd)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K) - enddo ; enddo ; endif - - ! Double-diffusion (old method) - if (CS%double_diffusion) then - call double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, KT_extra, KS_extra) - ! One of Kd_extra_T and Kd_extra_S is always 0. Kd_extra_S is positive for salt fingering. - ! Kd_extra_T is positive for double diffusive convection. - do K=2,nz ; do i=is,ie - if (KS_extra(i,K) > KT_extra(i,K)) then ! salt fingering - Kd_lay_2d(i,k-1) = Kd_lay_2d(i,k-1) + 0.5 * KT_extra(i,K) - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * KT_extra(i,K) - Kd_extra_S(i,j,K) = KS_extra(i,K) - KT_extra(i,K) - Kd_extra_T(i,j,K) = 0.0 - elseif (KT_extra(i,K) > 0.0) then ! double-diffusive convection - Kd_lay_2d(i,k-1) = Kd_lay_2d(i,k-1) + 0.5 * KS_extra(i,K) - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * KS_extra(i,K) - Kd_extra_T(i,j,K) = KT_extra(i,K) - KS_extra(i,K) - Kd_extra_S(i,j,K) = 0.0 - else ! There is no double diffusion at this interface. - Kd_extra_T(i,j,K) = 0.0 - Kd_extra_S(i,j,K) = 0.0 - endif - enddo ; enddo - if (associated(dd%KT_extra)) then ; do K=1,nz+1 ; do i=is,ie - dd%KT_extra(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - - if (associated(dd%KS_extra)) then ; do K=1,nz+1 ; do i=is,ie - dd%KS_extra(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif - - if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif + if (associated(visc%Kv_slow)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + visc%Kv_slow(i,j,K) = visc%Kv_slow(i,j,K) + Kv_bkgnd(i,K,jj) + enddo ; enddo + enddo + endif + if (CS%id_Kv_bkgnd > 0) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + dd%Kv_bkgnd(i,j,K) = Kv_bkgnd(i,K,jj) + enddo ; enddo + enddo + endif + if (CS%id_Kd_bkgnd > 0) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + dd%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo + endif + if (associated(VBF%Kd_bkgnd)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + VBF%Kd_bkgnd(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo endif - ! Apply double diffusion via CVMix - ! GMM, we need to pass HBL to compute_ddiff_coeffs, but it is not yet available. - if (CS%use_CVMix_ddiff) then - call cpu_clock_begin(id_clock_CVMix_ddiff) - if (associated(dd%drho_rat)) then - call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, & - CS%CVMix_ddiff_csp, dd%drho_rat) - else - call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, CS%CVMix_ddiff_csp) + do j=jstart,jend ; jj = j - jstart + 1 + ! Double-diffusion (old method) + if (CS%double_diffusion) then + call double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, KT_extra, KS_extra) + ! One of Kd_extra_T and Kd_extra_S is always 0. Kd_extra_S is positive for salt fingering. + ! Kd_extra_T is positive for double diffusive convection. + do K=2,nz ; do i=is,ie + if (KS_extra(i,K) > KT_extra(i,K)) then ! salt fingering + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KT_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KT_extra(i,K) + Kd_extra_S(i,j,K) = KS_extra(i,K) - KT_extra(i,K) + Kd_extra_T(i,j,K) = 0.0 + elseif (KT_extra(i,K) > 0.0) then ! double-diffusive convection + Kd_lay_2d(i,k-1,jj) = Kd_lay_2d(i,k-1,jj) + 0.5 * KS_extra(i,K) + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * KS_extra(i,K) + Kd_extra_T(i,j,K) = KT_extra(i,K) - KS_extra(i,K) + Kd_extra_S(i,j,K) = 0.0 + else ! There is no double diffusion at this interface. + Kd_extra_T(i,j,K) = 0.0 + Kd_extra_S(i,j,K) = 0.0 + endif + enddo ; enddo + if (associated(dd%KT_extra)) then ; do K=1,nz+1 ; do i=is,ie + dd%KT_extra(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + + if (associated(dd%KS_extra)) then ; do K=1,nz+1 ; do i=is,ie + dd%KS_extra(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif + + if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif endif - if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) - enddo ; enddo ; endif - if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) - enddo ; enddo ; endif - call cpu_clock_end(id_clock_CVMix_ddiff) - endif + + ! Apply double diffusion via CVMix + ! GMM, we need to pass HBL to compute_ddiff_coeffs, but it is not yet available. + if (CS%use_CVMix_ddiff) then + call cpu_clock_begin(id_clock_CVMix_ddiff) + if (associated(dd%drho_rat)) then + call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, & + CS%CVMix_ddiff_csp, dd%drho_rat) + else + call compute_ddiff_coeffs(h, tv, G, GV, US, j, Kd_extra_T, Kd_extra_S, CS%CVMix_ddiff_csp) + endif + if (associated(VBF%Kd_ddiff_T)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_T(i,j,K) = KT_extra(i,K) + enddo ; enddo ; endif + if (associated(VBF%Kd_ddiff_S)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_ddiff_S(i,j,K) = KS_extra(i,K) + enddo ; enddo ; endif + call cpu_clock_end(id_clock_CVMix_ddiff) + endif + enddo ! Calculate conversion ratios from TKE to layer diffusivities. if (TKE_to_Kd_used) then - call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb) - if (associated(dd%maxTKE)) then ; do k=1,nz ; do i=is,ie - dd%maxTKE(i,j,k) = maxTKE(i,k) - enddo ; enddo ; endif - if (associated(dd%TKE_to_Kd)) then ; do k=1,nz ; do i=is,ie - dd%TKE_to_Kd(i,j,k) = TKE_to_Kd(i,k) - enddo ; enddo ; endif + call find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, njblock, dt, G, GV, US, CS, TKE_to_Kd, maxTKE, kb, dz) + if (associated(dd%maxTKE)) then + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + dd%maxTKE(i,j,k) = maxTKE(i,k,jj) + enddo ; enddo ; enddo + endif + if (associated(dd%TKE_to_Kd)) then + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + dd%TKE_to_Kd(i,j,k) = TKE_to_Kd(i,k,jj) + enddo ; enddo ; enddo + endif endif ! Add the input turbulent diffusivity. if (CS%useKappaShear .or. CS%use_CVMix_shear) then - do K=2,nz ; do i=is,ie - Kd_int_2d(i,K) = visc%Kd_shear(i,j,K) + 0.5 * (Kd_lay_2d(i,k-1) + Kd_lay_2d(i,k)) - enddo ; enddo - do i=is,ie - Kd_int_2d(i,1) = visc%Kd_shear(i,j,1) ! This isn't actually used. It could be 0. - Kd_int_2d(i,nz+1) = 0.0 - enddo - do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) + do j=jstart,jend ; jj = j - jstart + 1 ; do K=2,nz ; do i=is,ie + Kd_int_2d(i,K,jj) = visc%Kd_shear(i,j,K) + 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int_2d(i,1,jj) = visc%Kd_shear(i,j,1) ! This isn't actually used. It could be 0. + Kd_int_2d(i,nz+1,jj) = 0.0 enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + 0.5 * (visc%Kd_shear(i,j,K) + visc%Kd_shear(i,j,K+1)) + enddo ; enddo ; enddo else - do i=is,ie - Kd_int_2d(i,1) = Kd_lay_2d(i,1) ; Kd_int_2d(i,nz+1) = 0.0 - enddo - do K=2,nz ; do i=is,ie - Kd_int_2d(i,K) = 0.5 * (Kd_lay_2d(i,k-1) + Kd_lay_2d(i,k)) + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + Kd_int_2d(i,1,jj) = Kd_lay_2d(i,1,jj) ; Kd_int_2d(i,nz+1,jj) = 0.0 enddo ; enddo - endif - - if (CS%ML_radiation .or. CS%use_tidal_mixing .or. associated(dd%Kd_Work)) then - call thickness_to_dz(h, tv, dz, j, G, GV) + do j=jstart,jend ; jj = j - jstart + 1 ; do K=2,nz ; do i=is,ie + Kd_int_2d(i,K,jj) = 0.5 * (Kd_lay_2d(i,k-1,jj) + Kd_lay_2d(i,k,jj)) + enddo ; enddo ; enddo endif ! Add the ML_Rad diffusivity. if (CS%ML_radiation) then - call add_MLrad_diffusivity(dz, fluxes, tv, j, Kd_int_2d, G, GV, US, CS, TKE_to_Kd, Kd_lay_2d) + do j=jstart,jend ; jj = j - jstart + 1 + call add_MLrad_diffusivity(dz(:,:,jj), fluxes, tv, j, Kd_int_2d(:,:,jj), G, GV, US, CS, TKE_to_Kd(:,:,jj), Kd_lay_2d(:,:,jj)) + enddo endif ! Add the Nikurashin and / or tidal bottom-driven mixing - if (CS%use_tidal_mixing) & - call calculate_tidal_mixing(dz, j, N2_bot, rho_bot, N2_lay, N2_int, TKE_to_Kd, & - maxTKE, G, GV, US, CS%tidal_mixing, & - CS%Kd_max, visc%Kv_slow, Kd_lay_2d, Kd_int_2d, VBF) + if (CS%use_tidal_mixing) then + do j=jstart,jend ; jj = j - jstart + 1 + call calculate_tidal_mixing(dz(:,:,jj), j, N2_bot(:,jj), rho_bot(:,jj), N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), & + maxTKE(:,:,jj), G, GV, US, CS%tidal_mixing, & + CS%Kd_max, visc%Kv_slow, Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), VBF) + enddo + endif ! Add diffusivity from internal tides ray tracing if (CS%use_int_tides) then - - call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot, k_bot, j, N2_lay, N2_int, TKE_to_Kd, CS%Kd_max, & - CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & - Kd_lay_2d, Kd_int_2d, prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & - prof_slope_2d) - - if (CS%id_kbbl > 0) then ; do i=is,ie - dd%kbbl(i,j) = k_bot(i) - enddo ; endif - if (CS%id_bbl_thick > 0) then ; do i=is,ie - dd%bbl_thick(i,j) = h_bot(i) - enddo ; endif - if (CS%id_Kd_leak > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_leak(i,j,K) = Kd_leak_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_quad > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_quad(i,j,K) = Kd_quad_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_itidal > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_itidal(i,j,K) = Kd_itidal_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_Froude > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_Froude(i,j,K) = Kd_Froude_2d(i,K) - enddo ; enddo ; endif - if (CS%id_Kd_slope > 0) then ; do K=1,nz+1 ; do i=is,ie - dd%Kd_slope(i,j,K) = Kd_slope_2d(i,K) - enddo ; enddo ; endif - if (associated (VBF%Kd_leak)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_leak(i,j,K) = min(Kd_leak_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_quad)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_quad(i,j,K) = min(Kd_quad_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_itidal)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_itidal(i,j,K) = min(Kd_itidal_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_Froude)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_Froude(i,j,K) = min(Kd_Froude_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - if (associated (VBF%Kd_slope)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_slope(i,j,K) = min(Kd_slope_2d(i,K), CS%Kd_max) - enddo ; enddo ; endif - - if (CS%id_prof_leak > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_leak(i,j,k) = prof_leak_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_quad > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_quad(i,j,k) = prof_quad_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_itidal > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_itidal(i,j,k) = prof_itidal_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_Froude > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_Froude(i,j,k) = prof_Froude_2d(i,k) - enddo ; enddo ; endif - if (CS%id_prof_slope > 0) then ; do k=1,nz ; do i=is,ie - dd%prof_slope(i,j,k) = prof_slope_2d(i,k) - enddo ; enddo ; endif + do j=jstart,jend ; jj = j - jstart + 1 + call get_lowmode_diffusivity(G, GV, h, tv, US, h_bot(:,jj), k_bot(:,jj), j, N2_lay(:,:,jj), N2_int(:,:,jj), TKE_to_Kd(:,:,jj), CS%Kd_max, & + CS%int_tide_CSp, Kd_leak_2d, Kd_quad_2d, Kd_itidal_2d, Kd_Froude_2d, Kd_slope_2d, & + Kd_lay_2d(:,:,jj), Kd_int_2d(:,:,jj), prof_leak_2d, prof_quad_2d, prof_itidal_2d, prof_froude_2d, & + prof_slope_2d) + + if (CS%id_kbbl > 0) then ; do i=is,ie + dd%kbbl(i,j) = k_bot(i,jj) + enddo ; endif + if (CS%id_bbl_thick > 0) then ; do i=is,ie + dd%bbl_thick(i,j) = h_bot(i,jj) + enddo ; endif + if (CS%id_Kd_leak > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_leak(i,j,K) = Kd_leak_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_quad > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_quad(i,j,K) = Kd_quad_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_itidal > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_itidal(i,j,K) = Kd_itidal_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_Froude > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_Froude(i,j,K) = Kd_Froude_2d(i,K) + enddo ; enddo ; endif + if (CS%id_Kd_slope > 0) then ; do K=1,nz+1 ; do i=is,ie + dd%Kd_slope(i,j,K) = Kd_slope_2d(i,K) + enddo ; enddo ; endif + if (associated (VBF%Kd_leak)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_leak(i,j,K) = min(Kd_leak_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_quad)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_quad(i,j,K) = min(Kd_quad_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_itidal)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_itidal(i,j,K) = min(Kd_itidal_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_Froude)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_Froude(i,j,K) = min(Kd_Froude_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + if (associated (VBF%Kd_slope)) then ; do K=1,nz+1 ; do i=is,ie + VBF%Kd_slope(i,j,K) = min(Kd_slope_2d(i,K), CS%Kd_max) + enddo ; enddo ; endif + + if (CS%id_prof_leak > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_leak(i,j,k) = prof_leak_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_quad > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_quad(i,j,k) = prof_quad_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_itidal > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_itidal(i,j,k) = prof_itidal_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_Froude > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_Froude(i,j,k) = prof_Froude_2d(i,k) + enddo ; enddo ; endif + if (CS%id_prof_slope > 0) then ; do k=1,nz ; do i=is,ie + dd%prof_slope(i,j,k) = prof_slope_2d(i,k) + enddo ; enddo ; endif + enddo endif - ! This adds the diffusion sustained by the energy extracted from the flow by the bottom drag. + ! This adds the diffusion sustained by drag TKE. + if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0) .and. .not. CS%use_LOTW_BBL_diffusivity) & + call add_drag_diffusivity(h, u, v, tv, fluxes, visc, jstart, jend, njblock, TKE_to_Kd, & + maxTKE, kb, rho_bot, G, GV, US, CS, Kd_lay_2d, Kd_int_2d, dd%Kd_BBL) + if (CS%bottomdraglaw .and. (CS%BBL_effic > 0.0)) then if (CS%use_LOTW_BBL_diffusivity) then - call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bot, Kd_int_2d, & - G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d) - else - call add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, & - maxTKE, kb, rho_bot, G, GV, US, CS, Kd_lay_2d, Kd_int_2d, dd%Kd_BBL) + do j=jstart,jend ; jj = j - jstart + 1 + call add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int(:,:,jj), Rho_bot(:,jj), & + Kd_int_2d(:,:,jj), G, GV, US, CS, dd%Kd_BBL, Kd_lay_2d(:,:,jj), dz(:,:,jj)) + enddo + endif + if (associated(VBF%Kd_BBL)) then + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) + enddo ; enddo + enddo endif - if (associated(VBF%Kd_BBL)) then ; do K=1,nz+1 ; do i=is,ie - VBF%Kd_BBL(i,j,K) = dd%Kd_BBL(i,j,K) - enddo ; enddo ; endif endif if (CS%limit_dissipation) then @@ -685,27 +739,33 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 1) a global constant, ! 2) a dissipation proportional to N (aka Gargett) and ! 3) dissipation corresponding to a (nearly) constant diffusivity. - do K=2,nz ; do i=is,ie - dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K)), & ! Floor aka Gargett - CS%dissip_N2 * N2_int(i,K)) ! Floor of Kd_min*rho0/F_Ri - Kd_int_2d(i,K) = max(Kd_int_2d(i,K) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K) + Omega2)))) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=2,nz ; do i=is,ie + dissip = max( CS%dissip_min, & ! Const. floor on dissip. + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_int(i,K,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_int(i,K,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_int_2d(i,K,jj) = max(Kd_int_2d(i,K,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_int(i,K,jj) + Omega2)))) + enddo ; enddo + enddo endif ! Optionally add a uniform diffusivity at the interfaces. if (CS%Kd_add > 0.0) then - do K=1,nz+1 ; do i=is,ie - Kd_int_2d(i,K) = Kd_int_2d(i,K) + CS%Kd_add - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + Kd_int_2d(i,K,jj) = Kd_int_2d(i,K,jj) + CS%Kd_add + enddo ; enddo + enddo VBF%Kd_add = CS%Kd_add endif ! Copy the 2-d slices into the 3-d array that is exported. - do K=1,nz+1 ; do i=is,ie - Kd_int(i,j,K) = Kd_int_2d(i,K) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do K=1,nz+1 ; do i=is,ie + Kd_int(i,j,K) = Kd_int_2d(i,K,jj) + enddo ; enddo + enddo if (CS%limit_dissipation) then ! This calculates the layer dissipation ONLY from Kd calculated in this routine @@ -713,39 +773,56 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i ! 1) a global constant, ! 2) a dissipation proportional to N (aka Gargett) and ! 3) dissipation corresponding to a (nearly) constant diffusivity. - do k=2,nz-1 ; do i=is,ie - dissip = max( CS%dissip_min, & ! Const. floor on dissip. - CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k)), & ! Floor aka Gargett - CS%dissip_N2 * N2_lay(i,k)) ! Floor of Kd_min*rho0/F_Ri - Kd_lay_2d(i,k) = max(Kd_lay_2d(i,k) , & ! Apply floor to Kd - dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k) + Omega2)))) - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=2,nz-1 ; do i=is,ie + dissip = max( CS%dissip_min, & ! Const. floor on dissip. + CS%dissip_N0 + CS%dissip_N1 * sqrt(N2_lay(i,k,jj)), & ! Floor aka Gargett + CS%dissip_N2 * N2_lay(i,k,jj)) ! Floor of Kd_min*rho0/F_Ri + Kd_lay_2d(i,k,jj) = max(Kd_lay_2d(i,k,jj) , & ! Apply floor to Kd + dissip * (CS%FluxRi_max / (GV%H_to_RZ * (N2_lay(i,k,jj) + Omega2)))) + enddo ; enddo + enddo endif if (associated(dd%Kd_Work)) then - do k=1,nz ; do i=is,ie - dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k) * N2_lay(i,k) * dz(i,k) ! Watt m-2 = kg s-3 - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + dd%Kd_Work(i,j,k) = GV%H_to_RZ * Kd_lay_2d(i,k,jj) * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 + enddo ; enddo + enddo endif ! Optionally add a uniform diffusivity to the layers. if ((CS%Kd_add > 0.0) .and. (present(Kd_lay))) then - do k=1,nz ; do i=is,ie - Kd_lay_2d(i,k) = Kd_lay_2d(i,k) + CS%Kd_add - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + Kd_lay_2d(i,k,jj) = Kd_lay_2d(i,k,jj) + CS%Kd_add + enddo ; enddo + enddo endif if (associated(dd%Kd_Work_added)) then - do k=1,nz ; do i=is,ie - dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k) * dz(i,k) ! Watt m-2 = kg s-3 - enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + dd%Kd_Work_added(i,j,k) = GV%H_to_RZ * CS%Kd_add * N2_lay(i,k,jj) * dz(i,k,jj) ! Watt m-2 = kg s-3 + enddo ; enddo + enddo endif ! Copy the 2-d slices into the 3-d array that is exported; this was done above for Kd_int. - if (present(Kd_lay)) then ; do k=1,nz ; do i=is,ie - Kd_lay(i,j,k) = Kd_lay_2d(i,k) - enddo ; enddo ; endif - enddo ! j-loop + if (present(Kd_lay)) then + do j=jstart,jend ; jj = j - jstart + 1 + do k=1,nz ; do i=is,ie + Kd_lay(i,j,k) = Kd_lay_2d(i,k,jj) + enddo ; enddo + enddo + endif + enddo ! jstart-loop + + !$omp target exit data map(release: T_f, S_f) + if (associated(tv%eqn_of_state)) then + !$omp target exit data map(release: tv, tv%T, tv%S) + endif if (CS%user_change_diff) then call user_change_diff(h, tv, G, GV, US, CS%user_change_diff_CSp, Kd_lay, Kd_int, & @@ -868,8 +945,8 @@ subroutine set_diffusivity(u, v, h, u_h, v_h, tv, fluxes, optics, visc, dt, Kd_i end subroutine set_diffusivity !> Convert turbulent kinetic energy to diffusivity -subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & - TKE_to_Kd, maxTKE, kb) +subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, jstart, jend, nj, dt, G, GV, US, CS, & + TKE_to_Kd, maxTKE, kb, dz) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -877,24 +954,27 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & intent(in) :: h !< Layer thicknesses [H ~> m or kg m-2] type(thermo_var_ptrs), intent(in) :: tv !< Structure containing pointers to any available !! thermodynamic fields. - real, dimension(SZI_(G),SZK_(GV)+1), intent(in) :: dRho_int !< Change in locally referenced potential density + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(in) :: dRho_int !< Change in locally referenced potential density !! across each interface [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: N2_lay !< The squared buoyancy frequency of the + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: N2_lay !< The squared buoyancy frequency of the !! layers [T-2 ~> s-2]. - integer, intent(in) :: j !< j-index of row to work on + integer, intent(in) :: jstart !< Start j-index of the j-block + integer, intent(in) :: jend !< End j-index of the j-block + integer, intent(in) :: nj !< Number of j-rows in the block real, intent(in) :: dt !< Time increment [T ~> s]. type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: TKE_to_Kd !< The conversion rate between the + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: TKE_to_Kd !< The conversion rate between the !! TKE dissipated within a layer and the !! diapycnal diffusivity within that layer, !! usually (~Rho_0 / (G_Earth * dRho_lay)) !! [T2 Z-1 ~> s2 m-1] - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: maxTKE !< The energy required to for a layer to entrain to its + real, dimension(SZI_(G),SZK_(GV),nj), intent(out) :: maxTKE !< The energy required to for a layer to entrain to its !! maximum realizable thickness [H Z2 T-3 ~> m3 s-3 or W m-2] - integer, dimension(SZI_(G)), intent(out) :: kb !< Index of lightest layer denser than the buffer + integer, dimension(SZI_(G),nj), intent(out) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer. + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real, dimension(SZI_(G),SZK_(GV)) :: & + real, dimension(SZI_(G),nj,SZK_(GV)) :: & ds_dsp1, & ! coordinate variable (sigma-2) difference across an ! interface divided by the difference across the interface ! below it [nondim] @@ -902,12 +982,11 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & ! across an interface times the difference across the ! interface above it [nondim] rho_0, & ! Layer potential densities relative to surface pressure [R ~> kg m-3] - dz, & ! Height change across layers [Z ~> m] maxEnt ! maxEnt is the maximum value of entrainment from below (with ! compensating entrainment from above to keep the layer ! density from changing) that will not deplete all of the ! layers above or below a layer within a timestep [H ~> m or kg m-2]. - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & htot, & ! total thickness above or below a layer, or the ! integrated thickness in the BBL [H ~> m or kg m-2]. mFkb, & ! total thickness in the mixed and buffer layers times ds_dsp1 [H ~> m or kg m-2] @@ -927,10 +1006,10 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & real :: I_dt ! 1/dt [T-1 ~> s-1] real :: dz_neglect ! A negligibly small height change [Z ~> m] real :: hN2pO2 ! h (N^2 + Omega^2), in [Z T-2 ~> m s-2]. - logical :: do_i(SZI_(G)) + logical :: do_i(SZI_(G),nj) integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, is, ie, nz, i_rem, kmb, kb_min + integer :: i, j, k, is, ie, nz, i_rem, kmb, kb_min, jj is = G%isc ; ie = G%iec ; nz = GV%ke I_dt = 1.0 / dt @@ -944,156 +1023,163 @@ subroutine find_TKE_to_Kd(h, tv, dRho_int, N2_lay, j, dt, G, GV, US, CS, & G_IRho0 = GV%H_to_Z*G_Rho0 endif - ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) - ! Simple but coordinate-independent estimate of Kd/TKE if (CS%simple_TKE_to_Kd) then - do k=1,nz ; do i=is,ie - hN2pO2 = dz(i,k) * (N2_lay(i,k) + Omega2) ! Units of Z T-2. + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz ; do i=is,ie + hN2pO2 = dz(i,k,jj) * (N2_lay(i,k,jj) + Omega2) ! Units of Z T-2. if (hN2pO2 > 0.) then - TKE_to_Kd(i,k) = 1.0 / hN2pO2 ! Units of T2 H-1. - else ; TKE_to_Kd(i,k) = 0. ; endif + TKE_to_Kd(i,k,jj) = 1.0 / hN2pO2 ! Units of T2 H-1. + else ; TKE_to_Kd(i,k,jj) = 0. ; endif ! The maximum TKE conversion we allow is really a statement ! about the upper diffusivity we allow. Kd_max must be set. - maxTKE(i,k) = hN2pO2 * CS%Kd_max ! Units of H Z2 T-3. - enddo ; enddo - kb(is:ie) = -1 ! kb should not be used by any code in non-layered mode -AJA + maxTKE(i,k,jj) = hN2pO2 * CS%Kd_max ! Units of H Z2 T-3. + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + kb(is:ie,jj) = -1 ! kb should not be used by any code in non-layered mode -AJA + enddo return endif ! Determine kb - the index of the shallowest active interior layer. if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies - do i=is,ie ; p_0(i) = 0.0 ; p_ref(i) = tv%P_Ref ; enddo + kb_min = kmb+1 EOSdom(:) = EOS_domain(G%HI) - do k=1,nz - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0, rho_0(:,k), tv%eqn_of_state, EOSdom) + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + p_0(i,jj) = 0.0 ; p_ref(i,jj) = tv%P_Ref + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do k=1,nz + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_0(:,jj), rho_0(:,jj,k), tv%eqn_of_state, EOSdom) + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 + call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref(:,jj), Rcv_kmb(:,jj), tv%eqn_of_state, EOSdom) enddo - call calculate_density(tv%T(:,j,kmb), tv%S(:,j,kmb), p_ref, Rcv_kmb, tv%eqn_of_state, EOSdom) - - kb_min = kmb+1 - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie ! Determine the next denser layer than the buffer layer in the ! coordinate density (sigma-2). - do k=kmb+1,nz-1 ; if (Rcv_kmb(i) <= GV%Rlay(k)) exit ; enddo - kb(i) = k + do k=kmb+1,nz-1 ; if (Rcv_kmb(i,jj) <= GV%Rlay(k)) exit ; enddo + kb(i,jj) = k ! Backtrack, in case there are massive layers above that are stable ! in sigma-0. - do k=kb(i)-1,kmb+1,-1 - if (rho_0(i,kmb) > rho_0(i,k)) exit - if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i) = k + do k=kb(i,jj)-1,kmb+1,-1 + if (rho_0(i,jj,kmb) > rho_0(i,jj,k)) exit + if (h(i,j,k)>2.0*GV%Angstrom_H) kb(i,jj) = k enddo - enddo - - call set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) + enddo ; enddo + call set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp1, rho_0) else ! not bulkmixedlayer kb_min = 2 ; kmb = 0 - do i=is,ie ; kb(i) = 1 ; enddo - call set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1) + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + kb(i,jj) = 1 + enddo ; enddo + call set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp1) endif ! Determine maxEnt - the maximum permitted entrainment from below by each ! interior layer. - do k=2,nz-1 ; do i=is,ie - dsp1_ds(i,k) = 1.0 / ds_dsp1(i,k) + do k=2,nz-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + dsp1_ds(i,jj,k) = 1.0 / ds_dsp1(i,jj,k) + enddo ; enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + dsp1_ds(i,jj,nz) = 0.0 enddo ; enddo - do i=is,ie ; dsp1_ds(i,nz) = 0.0 ; enddo if (CS%bulkmixedlayer) then kmb = GV%nk_rho_varies - do i=is,ie - htot(i) = h(i,j,kmb) - mFkb(i) = 0.0 - if (kb(i) < nz) mFkb(i) = ds_dsp1(i,kb(i)) * (h(i,j,kmb) - GV%Angstrom_H) - enddo - do k=1,kmb-1 ; do i=is,ie - htot(i) = htot(i) + h(i,j,k) - mFkb(i) = mFkb(i) + ds_dsp1(i,k+1)*(h(i,j,k) - GV%Angstrom_H) + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + htot(i,jj) = h(i,j,kmb) + mFkb(i,jj) = 0.0 + if (kb(i,jj) < nz) mFkb(i,jj) = ds_dsp1(i,jj,kb(i,jj)) * (h(i,j,kmb) - GV%Angstrom_H) enddo ; enddo + do k=1,kmb-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + htot(i,jj) = htot(i,jj) + h(i,j,k) + mFkb(i,jj) = mFkb(i,jj) + ds_dsp1(i,jj,k+1)*(h(i,j,k) - GV%Angstrom_H) + enddo ; enddo ; enddo else - do i=is,i - maxEnt(i,1) = 0.0 ; htot(i) = h(i,j,1) - GV%Angstrom_H - enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + maxEnt(i,jj,1) = 0.0 ; htot(i,jj) = h(i,j,1) - GV%Angstrom_H + enddo ; enddo endif - do k=kb_min,nz-1 ; do i=is,ie - if (k == kb(i)) then - maxEnt(i,kb(i)) = mFkb(i) - elseif (k > kb(i)) then + do k=kb_min,nz-1 ; do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + if (k == kb(i,jj)) then + maxEnt(i,jj,kb(i,jj)) = mFkb(i,jj) + elseif (k > kb(i,jj)) then if (CS%answer_date < 20190101) then - maxEnt(i,k) = (1.0/dsp1_ds(i,k))*(maxEnt(i,k-1) + htot(i)) + maxEnt(i,jj,k) = (1.0/dsp1_ds(i,jj,k))*(maxEnt(i,jj,k-1) + htot(i,jj)) else - maxEnt(i,k) = ds_dsp1(i,k)*(maxEnt(i,k-1) + htot(i)) + maxEnt(i,jj,k) = ds_dsp1(i,jj,k)*(maxEnt(i,jj,k-1) + htot(i,jj)) endif - htot(i) = htot(i) + (h(i,j,k) - GV%Angstrom_H) + htot(i,jj) = htot(i,jj) + (h(i,j,k) - GV%Angstrom_H) endif - enddo ; enddo + enddo ; enddo ; enddo - do i=is,ie - htot(i) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,nz) = 0.0 - do_i(i) = (G%mask2dT(i,j) > 0.0) - enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + htot(i,jj) = h(i,j,nz) - GV%Angstrom_H ; maxEnt(i,jj,nz) = 0.0 + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + enddo ; enddo do k=nz-1,kb_min,-1 - i_rem = 0 - do i=is,ie ; if (do_i(i)) then - if (k Calculate Brunt-Vaisala frequency, N^2. -subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & - N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot) +subroutine find_N2(h, tv, T_f, S_f, fluxes, nj, jstart, jend, G, GV, US, CS, dRho_int, & + N2_lay, N2_int, N2_bot, Rho_bot, h_bot, k_bot, dz) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -1108,29 +1194,30 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & intent(in) :: S_f !< Layer salinities with values in massless !! layers filled vertically by diffusion [S ~> ppt]. type(forcing), intent(in) :: fluxes !< A structure of thermodynamic surface fluxes - integer, intent(in) :: j !< j-index of row to work on + integer, intent(in) :: nj !< Number of j-rows in this block + integer, intent(in) :: jstart !< Starting j-index of this block + integer, intent(in) :: jend !< Ending j-index of this block type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)+1), & + real, dimension(SZI_(G),SZK_(GV)+1,nj), & intent(out) :: dRho_int !< Change in locally referenced potential density !! across each interface [R ~> kg m-3]. - real, dimension(SZI_(G),SZK_(GV)+1), & + real, dimension(SZI_(G),SZK_(GV)+1,nj), & intent(out) :: N2_int !< The squared buoyancy frequency at the interfaces [T-2 ~> s-2]. - real, dimension(SZI_(G),SZK_(GV)), & + real, dimension(SZI_(G),SZK_(GV),nj), & intent(out) :: N2_lay !< The squared buoyancy frequency of the layers [T-2 ~> s-2]. - real, dimension(SZI_(G)), intent(out) :: N2_bot !< The near-bottom squared buoyancy frequency [T-2 ~> s-2]. - real, dimension(SZI_(G)), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. - real, dimension(SZI_(G)), optional, intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. - integer, dimension(SZI_(G)), optional, intent(out) :: k_bot !< Bottom boundary layer top layer index. + real, dimension(SZI_(G),nj), intent(out) :: N2_bot !< The near-bottom squared buoyancy frequency [T-2 ~> s-2]. + real, dimension(SZI_(G),nj), intent(out) :: Rho_bot !< Near-bottom density [R ~> kg m-3]. + real, dimension(SZI_(G),nj), intent(out) :: h_bot !< Bottom boundary layer thickness [H ~> m or kg m-2]. + integer, dimension(SZI_(G),nj), intent(out) :: k_bot !< Bottom boundary layer top layer index. + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real, dimension(SZI_(G),SZK_(GV)+1) :: & - pres, & ! pressure at each interface [R L2 T-2 ~> Pa] + real, dimension(SZI_(G),nj,SZK_(GV)+1) :: & + pres, & ! pressure at each interface [R L2 T-2 ~> Pa] dRho_int_unfilt, & ! unfiltered density differences across interfaces [R ~> kg m-3] dRho_dT, & ! partial derivative of density wrt temp [R C-1 ~> kg m-3 degC-1] dRho_dS ! partial derivative of density wrt saln [R S-1 ~> kg m-3 ppt-1] - real, dimension(SZI_(G),SZK_(GV)) :: & - dz ! Height change across layers [Z ~> m] - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & Temp_int, & ! temperature at each interface [C ~> degC] Salin_int, & ! salinity at each interface [S ~> ppt] drho_bot, & ! A density difference [R ~> kg m-3] @@ -1144,128 +1231,191 @@ subroutine find_N2(h, tv, T_f, S_f, fluxes, j, G, GV, US, CS, dRho_int, & ! times some unit conversion factors [H T-2 R-1 ~> m4 s-2 kg-1 or m s-2]. real :: H_neglect ! A negligibly small thickness [H ~> m or kg m-2] - logical :: do_i(SZI_(G)), do_any + logical :: do_i(SZI_(G),nj), do_any integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, is, ie, nz + integer :: i, j, k, is, ie, nz, jj + ! GPU port: EOS form + unit scaling resolved once on the host so the density derivatives can be + ! evaluated on the device through the declare-target dispatcher; per-interface scalar scratch. + integer :: eos_form + real :: eos_kg_m3_to_R, eos_C_to_degC, eos_S_to_ppt, eos_RL2_T2_to_Pa + real :: T_int_s, S_int_s, dRdT_s, dRdS_s + real :: dz_BBL_min_l ! A host copy of CS%dz_BBL_avg_min, for use in a device region [Z ~> m] + logical :: have_p_surf is = G%isc ; ie = G%iec ; nz = GV%ke G_Rho0 = GV%g_Earth_Z_T2 / GV%H_to_RZ H_neglect = GV%H_subroundoff + ! GPU port: resolve the EOS form + unit scaling on the host (the accessor is not device-callable). + eos_form = -1 + eos_kg_m3_to_R = 1.0 ; eos_C_to_degC = 1.0 ; eos_S_to_ppt = 1.0 ; eos_RL2_T2_to_Pa = 1.0 + if (associated(tv%eqn_of_state)) then + call get_EOS_form_and_scaling(tv%eqn_of_state, eos_form, eos_kg_m3_to_R, eos_C_to_degC, & + eos_S_to_ppt, eos_RL2_T2_to_Pa) +#ifdef __NVCOMPILER_OPENMP_GPU + if ((eos_form /= EOS_ROQUET_RHO) .and. (eos_form /= EOS_WRIGHT)) call MOM_error(FATAL, & + "find_N2 GPU build: EQN_OF_STATE has no device-callable density-derivs kernel "// & + "(only ROQUET_RHO and WRIGHT are supported); use a CPU build or add a _loc kernel.") +#endif + endif + have_p_surf = associated(fluxes%p_surf) + dz_BBL_min_l = CS%dz_BBL_avg_min + + ! GPU port: the surface pressure (K=1, from the possibly-null p_surf pointer) and the topographic + ! roughness h_amp (set by the host tidal_mixing_h_amp when active) are computed on the host and + ! mapped in; the per-column pressure recurrence + density derivatives + buoyancy frequencies then + ! run on the device. One device data region brackets all of find_N2's per-column loops. + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie + if (have_p_surf) then ; pres(i,jj,1) = fluxes%p_surf(i,j) ; else ; pres(i,jj,1) = 0.0 ; endif + h_amp(i,jj) = 0.0 + enddo + enddo + if (CS%use_tidal_mixing) then + do j=jstart,jend ; jj = j - jstart + 1 + call tidal_mixing_h_amp(h_amp(:,jj), G, j, CS%tidal_mixing) + enddo + endif + + ! h, T_f, S_f and tv/tv%T/tv%S are mapped/refreshed once by the caller (set_diffusivity) around the + ! whole jstart loop -- doing it here would re-copy full-domain arrays on every block (njblock=1 => a + ! per-j-row copy), which dominates the runtime. Only the per-block scratch is mapped here. + !$omp target enter data map(to: dz, pres, h_amp) & + !$omp map(alloc: dRho_int, dRho_int_unfilt, N2_int, N2_lay, N2_bot, hb, drho_bot, z_from_bot, do_i, dz_BBL_avg) + ! Find the (limited) density jump across each interface. - do i=is,ie - dRho_int(i,1) = 0.0 ; dRho_int(i,nz+1) = 0.0 - dRho_int_unfilt(i,1) = 0.0 ; dRho_int_unfilt(i,nz+1) = 0.0 + do concurrent (jj=1:nj, i=is:ie) + dRho_int(i,1,jj) = 0.0 ; dRho_int(i,nz+1,jj) = 0.0 + dRho_int_unfilt(i,jj,1) = 0.0 ; dRho_int_unfilt(i,jj,nz+1) = 0.0 enddo if (associated(tv%eqn_of_state)) then - if (associated(fluxes%p_surf)) then - do i=is,ie ; pres(i,1) = fluxes%p_surf(i,j) ; enddo - else - do i=is,ie ; pres(i,1) = 0.0 ; enddo - endif - EOSdom(:) = EOS_domain(G%HI) - do K=2,nz - do i=is,ie - pres(i,K) = pres(i,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,k-1) - Temp_Int(i) = 0.5 * (T_f(i,j,k) + T_f(i,j,k-1)) - Salin_Int(i) = 0.5 * (S_f(i,j,k) + S_f(i,j,k-1)) - enddo - call calculate_density_derivs(Temp_int, Salin_int, pres(:,K), dRho_dT(:,K), dRho_dS(:,K), & - tv%eqn_of_state, EOSdom) + ! Per-column pressure recurrence (serial in K) + density derivatives via the device-callable + ! EOS dispatcher. The unit-scaling is always applied; for the unscaled (unity) case the factors + ! are exactly 1.0, so this reproduces the former array calculate_density_derivs call bit-for-bit. + !$omp target teams loop collapse(2) private(jj, K, T_int_s, S_int_s, dRdT_s, dRdS_s) + do j=jstart,jend do i=is,ie - dRho_int(i,K) = max(dRho_dT(i,K)*(T_f(i,j,k) - T_f(i,j,k-1)) + & - dRho_dS(i,K)*(S_f(i,j,k) - S_f(i,j,k-1)), 0.0) - dRho_int_unfilt(i,K) = max(dRho_dT(i,K)*(tv%T(i,j,k) - tv%T(i,j,k-1)) + & - dRho_dS(i,K)*(tv%S(i,j,k) - tv%S(i,j,k-1)), 0.0) + jj = j - jstart + 1 + do K=2,nz + pres(i,jj,K) = pres(i,jj,K-1) + (GV%g_Earth*GV%H_to_RZ)*h(i,j,K-1) + T_int_s = 0.5 * (T_f(i,j,K) + T_f(i,j,K-1)) + S_int_s = 0.5 * (S_f(i,j,K) + S_f(i,j,K-1)) + call calculate_density_derivs_elem_loc(eos_form, eos_C_to_degC*T_int_s, & + eos_S_to_ppt*S_int_s, eos_RL2_T2_to_Pa*pres(i,jj,K), dRdT_s, dRdS_s) + dRdT_s = (eos_kg_m3_to_R*eos_C_to_degC) * dRdT_s + dRdS_s = (eos_kg_m3_to_R*eos_S_to_ppt) * dRdS_s + dRho_int(i,K,jj) = max(dRdT_s*(T_f(i,j,K) - T_f(i,j,K-1)) + & + dRdS_s*(S_f(i,j,K) - S_f(i,j,K-1)), 0.0) + dRho_int_unfilt(i,jj,K) = max(dRdT_s*(tv%T(i,j,K) - tv%T(i,j,K-1)) + & + dRdS_s*(tv%S(i,j,K) - tv%S(i,j,K-1)), 0.0) + enddo enddo enddo else - do K=2,nz ; do i=is,ie - dRho_int(i,K) = GV%Rlay(k) - GV%Rlay(k-1) - enddo ; enddo + do concurrent (jj=1:nj, K=2:nz, i=is:ie) + dRho_int(i,K,jj) = GV%Rlay(K) - GV%Rlay(K-1) + enddo endif - ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) - ! Set the buoyancy frequencies. - do k=1,nz ; do i=is,ie - N2_lay(i,k) = G_Rho0 * 0.5*(dRho_int(i,K) + dRho_int(i,K+1)) / & - (h(i,j,k) + H_neglect) - enddo ; enddo - do i=is,ie ; N2_int(i,1) = 0.0 ; N2_int(i,nz+1) = 0.0 ; enddo - do K=2,nz ; do i=is,ie - N2_int(i,K) = G_Rho0 * dRho_int(i,K) / & - (0.5*(h(i,j,k-1) + h(i,j,k) + H_neglect)) - enddo ; enddo + do concurrent (k=1:nz, jj=1:nj, i=is:ie) + N2_lay(i,k,jj) = G_Rho0 * 0.5*(dRho_int(i,k,jj) + dRho_int(i,k+1,jj)) / & + (h(i,jstart+jj-1,k) + H_neglect) + enddo + do concurrent (jj=1:nj, i=is:ie) + N2_int(i,1,jj) = 0.0 ; N2_int(i,nz+1,jj) = 0.0 + enddo + do concurrent (K=2:nz, jj=1:nj, i=is:ie) + N2_int(i,K,jj) = G_Rho0 * dRho_int(i,K,jj) / & + (0.5*(h(i,jstart+jj-1,K-1) + h(i,jstart+jj-1,K) + H_neglect)) + enddo ! Find the bottom boundary layer stratification, and use this in the deepest layers. - do i=is,ie - hb(i) = 0.0 ; dRho_bot(i) = 0.0 ; h_amp(i) = 0.0 - z_from_bot(i) = 0.5*dz(i,nz) - do_i(i) = (G%mask2dT(i,j) > 0.0) + ! (h_amp was set on the host above, including the tidal_mixing_h_amp path.) + do concurrent (jj=1:nj, i=is:ie) + hb(i,jj) = 0.0 ; dRho_bot(i,jj) = 0.0 + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,jstart+jj-1) > 0.0) enddo - if (CS%use_tidal_mixing) call tidal_mixing_h_amp(h_amp, G, j, CS%tidal_mixing) - do k=nz,2,-1 - do_any = .false. - do i=is,ie ; if (do_i(i)) then - dz_int = 0.5*(dz(i,k) + dz(i,k-1)) - z_from_bot(i) = z_from_bot(i) + dz_int ! middle of the layer above - - hb(i) = hb(i) + 0.5*(h(i,j,k) + h(i,j,k-1)) - drho_bot(i) = drho_bot(i) + dRho_int(i,K) - - if (z_from_bot(i) > h_amp(i)) then - if (k>2) then - ! Always include at least one full layer. - hb(i) = hb(i) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) - drho_bot(i) = drho_bot(i) + dRho_int(i,K-1) + ! Accumulate the near-bottom thickness/density difference per column (serial down-K); the former + ! cross-column do_any/exit early-out is dropped -- each device thread runs its own K loop and stops + ! via its do_i flag, which is bit-for-bit equivalent to the original. + !$omp target teams loop collapse(2) private(jj, k, dz_int) + do j=jstart,jend + do i=is,ie + jj = j - jstart + 1 + do k=nz,2,-1 + if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k) + h(i,j,k-1)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,k,jj) + + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) then + ! Always include at least one full layer. + hb(i,jj) = hb(i,jj) + 0.5*(h(i,j,k-1) + h(i,j,k-2)) + drho_bot(i,jj) = drho_bot(i,jj) + dRho_int(i,k-1,jj) + endif + do_i(i,jj) = .false. + endif endif - do_i(i) = .false. - else - do_any = .true. - endif - endif ; enddo - if (.not.do_any) exit + enddo + enddo enddo - do i=is,ie - if (hb(i) > 0.0) then - N2_bot(i) = (G_Rho0 * drho_bot(i)) / hb(i) - else ; N2_bot(i) = 0.0 ; endif - z_from_bot(i) = 0.5*dz(i,nz) - do_i(i) = (G%mask2dT(i,j) > 0.0) + do concurrent (jj=1:nj, i=is:ie) + if (hb(i,jj) > 0.0) then + N2_bot(i,jj) = (G_Rho0 * drho_bot(i,jj)) / hb(i,jj) + else ; N2_bot(i,jj) = 0.0 ; endif + z_from_bot(i,jj) = 0.5*dz(i,nz,jj) + do_i(i,jj) = (G%mask2dT(i,jstart+jj-1) > 0.0) enddo - do k=nz,2,-1 - do_any = .false. - do i=is,ie ; if (do_i(i)) then - dz_int = 0.5*(dz(i,k) + dz(i,k-1)) - z_from_bot(i) = z_from_bot(i) + dz_int ! middle of the layer above - - N2_int(i,K) = N2_bot(i) - if (k>2) N2_lay(i,k-1) = N2_bot(i) - - if (z_from_bot(i) > h_amp(i)) then - if (k>2) N2_int(i,K-1) = N2_bot(i) - do_i(i) = .false. - else - do_any = .true. - endif - endif ; enddo - if (.not.do_any) exit + ! Overwrite the near-bottom interface/layer N2 with the bottom-BL value (serial down-K per column; + ! cross-column do_any/exit dropped, per-column do_i flag preserves bit-for-bit behaviour). + !$omp target teams loop collapse(2) private(jj, k, dz_int) + do j=jstart,jend + do i=is,ie + jj = j - jstart + 1 + do k=nz,2,-1 + if (do_i(i,jj)) then + dz_int = 0.5*(dz(i,k,jj) + dz(i,k-1,jj)) + z_from_bot(i,jj) = z_from_bot(i,jj) + dz_int ! middle of the layer above + + N2_int(i,k,jj) = N2_bot(i,jj) + if (k>2) N2_lay(i,k-1,jj) = N2_bot(i,jj) + + if (z_from_bot(i,jj) > h_amp(i,jj)) then + if (k>2) N2_int(i,k-1,jj) = N2_bot(i,jj) + do_i(i,jj) = .false. + endif + endif + enddo + enddo enddo if (associated(tv%eqn_of_state)) then - do K=1,nz+1 ; do i=is,ie - dRho_int(i,K) = dRho_int_unfilt(i,K) - enddo ; enddo + do concurrent (K=1:nz+1, jj=1:nj, i=is:ie) + dRho_int(i,K,jj) = dRho_int_unfilt(i,jj,K) + enddo endif ! Average over the larger of the envelope of the topography or a minimal distance. - do i=is,ie ; dz_BBL_avg(i) = max(h_amp(i), CS%dz_BBL_avg_min) ; enddo - call find_rho_bottom(G, GV, US, tv, h, dz, pres, dz_BBL_avg, j, Rho_bot, h_bot, k_bot) + do concurrent (jj=1:nj, i=is:ie) + dz_BBL_avg(i,jj) = max(h_amp(i,jj), dz_BBL_min_l) + enddo + + ! GPU port: copy the device results back for the host find_rho_bottom and the (host) set_diffusivity + ! consumers (dRho_int/N2_int/N2_lay/N2_bot are find_N2's outputs; pres/dz_BBL_avg feed find_rho_bottom), + ! then close the device data region. + !$omp target update from(dRho_int, N2_int, N2_lay, N2_bot, pres, dz_BBL_avg) + !$omp target exit data map(release: dz, pres, h_amp, hb, drho_bot, z_from_bot, do_i, & + !$omp dRho_int, dRho_int_unfilt, N2_int, N2_lay, N2_bot, dz_BBL_avg) + + call find_rho_bottom(G, GV, US, tv, h, dz, pres, dz_BBL_avg, jstart, jend, nj, Rho_bot, h_bot, k_bot) end subroutine find_N2 @@ -1361,7 +1511,7 @@ subroutine double_diffusion(tv, h, T_f, S_f, j, G, GV, US, CS, Kd_T_dd, Kd_S_dd) end subroutine double_diffusion !> This routine adds diffusion sustained by flow energy extracted by bottom drag. -subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, & +subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, jstart, jend, nj, TKE_to_Kd, maxTKE, & kb, rho_bot, G, GV, US, CS, Kd_lay, Kd_int, Kd_BBL) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure @@ -1377,22 +1527,24 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, type(forcing), intent(in) :: fluxes !< A structure of thermodynamic surface fluxes type(vertvisc_type), intent(in) :: visc !< Structure containing vertical viscosities, bottom !! boundary layer properties and related fields - integer, intent(in) :: j !< j-index of row to work on - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: TKE_to_Kd !< The conversion rate between the TKE + integer, intent(in) :: jstart !< j-index of first row to work on + integer, intent(in) :: jend !< j-index of last row to work on + integer, intent(in) :: nj !< Number of j-rows in a j-block + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: TKE_to_Kd !< The conversion rate between the TKE !! TKE dissipated within a layer and the !! diapycnal diffusivity within that layer, !! usually (~Rho_0 / (G_Earth * dRho_lay)) !! [T2 Z-1 ~> s2 m-1] - real, dimension(SZI_(G),SZK_(GV)), intent(in) :: maxTKE !< The energy required to for a layer to entrain to its + real, dimension(SZI_(G),SZK_(GV),nj), intent(in) :: maxTKE !< The energy required to for a layer to entrain to its !! maximum-realizable thickness [H Z2 T-3 ~> m3 s-3 or W m-2] - integer, dimension(SZI_(G)), intent(in) :: kb !< Index of lightest layer denser than the buffer + integer, dimension(SZI_(G),nj), intent(in) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer - real, dimension(SZI_(G)), intent(in) :: rho_bot !< In situ density averaged over a near-bottom + real, dimension(SZI_(G),nj), intent(in) :: rho_bot !< In situ density averaged over a near-bottom !! region [R ~> kg m-3] type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure - real, dimension(SZI_(G),SZK_(GV)), intent(inout) :: Kd_lay !< The diapycnal diffusivity in layers, + real, dimension(SZI_(G),SZK_(GV),nj), intent(inout) :: Kd_lay !< The diapycnal diffusivity in layers, !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] - real, dimension(SZI_(G),SZK_(GV)+1), intent(inout) :: Kd_int !< The diapycnal diffusivity at interfaces, + real, dimension(SZI_(G),SZK_(GV)+1,nj), intent(inout) :: Kd_int !< The diapycnal diffusivity at interfaces, !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real, dimension(:,:,:), pointer :: Kd_BBL !< Interface BBL diffusivity !! [H Z T-1 ~> m2 s-1 or kg m-1 s-1] @@ -1401,7 +1553,7 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, real, dimension(SZK_(GV)+1) :: & Rint ! coordinate density of an interface [R ~> kg m-3] - real, dimension(SZI_(G)) :: & + real, dimension(SZI_(G),nj) :: & htot, & ! total thickness above or below a layer, or the ! integrated thickness in the BBL [H ~> m or kg m-2]. rho_htot, & ! running integral with depth of density [R H ~> kg m-2 or kg2 m-5] @@ -1426,10 +1578,10 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, ! defined in visc, on the assumption that this ! extracted energy also drives diapycnal mixing. - logical :: domore, do_i(SZI_(G)) + logical :: domore, do_i(SZI_(G),nj) logical :: do_diag_Kd_BBL - integer :: i, k, is, ie, nz, i_rem, kb_min + integer :: i, j, k, is, ie, nz, i_rem, kb_min, jj is = G%isc ; ie = G%iec ; nz = GV%ke do_diag_Kd_BBL = associated(Kd_BBL) @@ -1449,11 +1601,11 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, ! The turbulence decay scale is 0.5*ustar/f from K&E & MOM_vertvisc.F90 ! Any turbulence that makes it into the mixed layers is assumed ! to be relatively small and is discarded. - do i=is,ie + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie ustar_h = visc%ustar_BBL(i,j) if (associated(fluxes%ustar_tidal)) then if (allocated(tv%SpV_avg)) then - ustar_h = ustar_h + GV%RZ_to_H*rho_bot(i) * fluxes%ustar_tidal(i,j) + ustar_h = ustar_h + GV%RZ_to_H*rho_bot(i,jj) * fluxes%ustar_tidal(i,j) else ustar_h = ustar_h + GV%Z_to_H * fluxes%ustar_tidal(i,j) endif @@ -1461,142 +1613,146 @@ subroutine add_drag_diffusivity(h, u, v, tv, fluxes, visc, j, TKE_to_Kd, maxTKE, absf = 0.25 * ((abs(G%CoriolisBu(I-1,J-1)) + abs(G%CoriolisBu(I,J))) + & (abs(G%CoriolisBu(I-1,J)) + abs(G%CoriolisBu(I,J-1)))) if ((ustar_h > 0.0) .and. (absf > 0.5*CS%IMax_decay*ustar_h)) then - I2decay(i) = absf / ustar_h + I2decay(i,jj) = absf / ustar_h else ! The maximum decay scale should be something of order 200 m. ! If ustar_h = 0, this is land so this value doesn't matter. - I2decay(i) = 0.5*CS%IMax_decay + I2decay(i,jj) = 0.5*CS%IMax_decay endif if (CS%drag_diff_answer_date <= 20250301) then - TKE(i) = ((CS%BBL_effic * cdrag_sqrt) * exp(-I2decay(i)*h(i,j,nz)) ) * visc%BBL_meanKE_loss_sqrtCd(i,j) + TKE(i,jj) = ((CS%BBL_effic * cdrag_sqrt) * exp(-I2decay(i,jj)*h(i,j,nz)) ) * visc%BBL_meanKE_loss_sqrtCd(i,j) else - TKE(i) = (CS%BBL_effic * exp(-I2decay(i)*h(i,j,nz)) ) * visc%BBL_meanKE_loss(i,j) + TKE(i,jj) = (CS%BBL_effic * exp(-I2decay(i,jj)*h(i,j,nz)) ) * visc%BBL_meanKE_loss(i,j) endif if (associated(fluxes%BBL_tidal_dis)) & - TKE(i) = TKE(i) + fluxes%BBL_tidal_dis(i,j) * GV%RZ_to_H * & - (CS%BBL_effic * exp(-I2decay(i)*h(i,j,nz))) + TKE(i,jj) = TKE(i,jj) + fluxes%BBL_tidal_dis(i,j) * GV%RZ_to_H * & + (CS%BBL_effic * exp(-I2decay(i,jj)*h(i,j,nz))) ! Distribute the work over a BBL of depth 20^2 ustar^2 / g' following ! Killworth & Edwards (1999) and Zilitikevich & Mironov (1996). ! Rho_top is determined by finding the density where ! integral(bottom, Z) (rho(z') - rho(Z)) dz' = rho_0 400 ustar^2 / g - gh_sum_top(i) = R0_g * 400.0 * ustar_h**2 + gh_sum_top(i,jj) = R0_g * 400.0 * ustar_h**2 - do_i(i) = (G%mask2dT(i,j) > 0.0) - htot(i) = h(i,j,nz) - rho_htot(i) = GV%Rlay(nz)*(h(i,j,nz)) - Rho_top(i) = GV%Rlay(1) - if (CS%bulkmixedlayer .and. do_i(i)) Rho_top(i) = GV%Rlay(kb(i)-1) - enddo - - do k=nz-1,2,-1 ; domore = .false. - do i=is,ie ; if (do_i(i)) then - htot(i) = htot(i) + h(i,j,k) - rho_htot(i) = rho_htot(i) + GV%Rlay(k)*(h(i,j,k)) - if (htot(i)*GV%Rlay(k-1) <= (rho_htot(i) - gh_sum_top(i))) then - ! The top of the mixing is in the interface atop the current layer. - Rho_top(i) = (rho_htot(i) - gh_sum_top(i)) / htot(i) - do_i(i) = .false. - elseif (k <= kb(i)) then ; do_i(i) = .false. - else ; domore = .true. ; endif - endif ; enddo - if (.not.domore) exit - enddo ! k-loop - - do i=is,ie ; do_i(i) = (G%mask2dT(i,j) > 0.0) ; enddo - do k=nz-1,kb_min,-1 - i_rem = 0 - do i=is,ie ; if (do_i(i)) then - if (k 0.0) then - if (Rint(K) <= Rho_top(i)) then - TKE_to_layer = TKE(i) - else - dRl = Rint(K+1) - Rint(K) ; dRbot = Rint(K+1) - Rho_top(i) - TKE_to_layer = TKE(i) * dRl * & - (3.0*dRbot*(Rint(K) - Rho_top(i)) + dRl**2) / (dRbot**3) - endif - else ; TKE_to_layer = 0.0 ; endif - - ! TKE_Ray has been initialized to 0 above. - if (Rayleigh_drag) TKE_Ray = 0.5*CS%BBL_effic * G%IareaT(i,j) * & - (((G%areaCu(I-1,j) * visc%Ray_u(I-1,j,k) * u(I-1,j,k)**2) + & - (G%areaCu(I,j) * visc%Ray_u(I,j,k) * u(I,j,k)**2)) + & - ((G%areaCv(i,J-1) * visc%Ray_v(i,J-1,k) * v(i,J-1,k)**2) + & - (G%areaCv(i,J) * visc%Ray_v(i,J,k) * v(i,J,k)**2))) - - if (TKE_to_layer + TKE_Ray > 0.0) then - if (CS%BBL_mixing_as_max) then - if (TKE_to_layer + TKE_Ray > maxTKE(i,k)) & - TKE_to_layer = maxTKE(i,k) - TKE_Ray + do_i(i,jj) = (G%mask2dT(i,j) > 0.0) + htot(i,jj) = h(i,j,nz) + rho_htot(i,jj) = GV%Rlay(nz)*(h(i,j,nz)) + Rho_top(i,jj) = GV%Rlay(1) + if (CS%bulkmixedlayer .and. do_i(i,jj)) Rho_top(i,jj) = GV%Rlay(kb(i,jj)-1) + enddo ; enddo - TKE(i) = TKE(i) - TKE_to_layer + do j=jstart,jend ; jj = j - jstart + 1 + do k=nz-1,2,-1 ; domore = .false. + do i=is,ie ; if (do_i(i,jj)) then + htot(i,jj) = htot(i,jj) + h(i,j,k) + rho_htot(i,jj) = rho_htot(i,jj) + GV%Rlay(k)*(h(i,j,k)) + if (htot(i,jj)*GV%Rlay(k-1) <= (rho_htot(i,jj) - gh_sum_top(i,jj))) then + ! The top of the mixing is in the interface atop the current layer. + Rho_top(i,jj) = (rho_htot(i,jj) - gh_sum_top(i,jj)) / htot(i,jj) + do_i(i,jj) = .false. + elseif (k <= kb(i,jj)) then ; do_i(i,jj) = .false. + else ; domore = .true. ; endif + endif ; enddo + if (.not.domore) exit + enddo ! k-loop + enddo - if (Kd_lay(i,k) < (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k)) then - delta_Kd = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k) - Kd_lay(i,k) - if ((CS%Kd_max >= 0.0) .and. (delta_Kd > CS%Kd_max)) then - delta_Kd = CS%Kd_max - Kd_lay(i,k) = Kd_lay(i,k) + delta_Kd - else - Kd_lay(i,k) = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k) - endif - Kd_int(i,K) = Kd_int(i,K) + 0.5 * delta_Kd - Kd_int(i,K+1) = Kd_int(i,K+1) + 0.5 * delta_Kd - if (do_diag_Kd_BBL) then - Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd - Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd - endif - endif - else - if (Kd_lay(i,k) >= maxTKE(i,k) * TKE_to_Kd(i,k)) then - TKE_here = 0.0 - TKE(i) = TKE(i) + TKE_Ray - elseif (Kd_lay(i,k) + (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k) > & - maxTKE(i,k) * TKE_to_Kd(i,k)) then - TKE_here = ((TKE_to_layer + TKE_Ray) + Kd_lay(i,k) / TKE_to_Kd(i,k)) - maxTKE(i,k) - TKE(i) = (TKE(i) - TKE_here) + TKE_Ray + do j=jstart,jend ; jj = j - jstart + 1 + do i=is,ie ; do_i(i,jj) = (G%mask2dT(i,j) > 0.0) ; enddo + do k=nz-1,kb_min,-1 + i_rem = 0 + do i=is,ie ; if (do_i(i,jj)) then + if (k 0.0) then + if (Rint(K) <= Rho_top(i,jj)) then + TKE_to_layer = TKE(i,jj) else - TKE_here = TKE_to_layer + TKE_Ray - TKE(i) = TKE(i) - TKE_to_layer + dRl = Rint(K+1) - Rint(K) ; dRbot = Rint(K+1) - Rho_top(i,jj) + TKE_to_layer = TKE(i,jj) * dRl * & + (3.0*dRbot*(Rint(K) - Rho_top(i,jj)) + dRl**2) / (dRbot**3) endif - if (TKE(i) < 0.0) TKE(i) = 0.0 ! This should be unnecessary? - - if (TKE_here > 0.0) then - delta_Kd = TKE_here * TKE_to_Kd(i,k) - if (CS%Kd_max >= 0.0) delta_Kd = min(delta_Kd, CS%Kd_max) - Kd_lay(i,k) = Kd_lay(i,k) + delta_Kd - Kd_int(i,K) = Kd_int(i,K) + 0.5 * delta_Kd - Kd_int(i,K+1) = Kd_int(i,K+1) + 0.5 * delta_Kd - if (do_diag_Kd_BBL) then - Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd - Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + else ; TKE_to_layer = 0.0 ; endif + + ! TKE_Ray has been initialized to 0 above. + if (Rayleigh_drag) TKE_Ray = 0.5*CS%BBL_effic * G%IareaT(i,j) * & + (((G%areaCu(I-1,j) * visc%Ray_u(I-1,j,k) * u(I-1,j,k)**2) + & + (G%areaCu(I,j) * visc%Ray_u(I,j,k) * u(I,j,k)**2)) + & + ((G%areaCv(i,J-1) * visc%Ray_v(i,J-1,k) * v(i,J-1,k)**2) + & + (G%areaCv(i,J) * visc%Ray_v(i,J,k) * v(i,J,k)**2))) + + if (TKE_to_layer + TKE_Ray > 0.0) then + if (CS%BBL_mixing_as_max) then + if (TKE_to_layer + TKE_Ray > maxTKE(i,k,jj)) & + TKE_to_layer = maxTKE(i,k,jj) - TKE_Ray + + TKE(i,jj) = TKE(i,jj) - TKE_to_layer + + if (Kd_lay(i,k,jj) < (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj)) then + delta_Kd = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) - Kd_lay(i,k,jj) + if ((CS%Kd_max >= 0.0) .and. (delta_Kd > CS%Kd_max)) then + delta_Kd = CS%Kd_max + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + delta_Kd + else + Kd_lay(i,k,jj) = (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) + endif + Kd_int(i,K,jj) = Kd_int(i,K,jj) + 0.5 * delta_Kd + Kd_int(i,K+1,jj) = Kd_int(i,K+1,jj) + 0.5 * delta_Kd + if (do_diag_Kd_BBL) then + Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd + Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + endif + endif + else + if (Kd_lay(i,k,jj) >= maxTKE(i,k,jj) * TKE_to_Kd(i,k,jj)) then + TKE_here = 0.0 + TKE(i,jj) = TKE(i,jj) + TKE_Ray + elseif (Kd_lay(i,k,jj) + (TKE_to_layer + TKE_Ray) * TKE_to_Kd(i,k,jj) > & + maxTKE(i,k,jj) * TKE_to_Kd(i,k,jj)) then + TKE_here = ((TKE_to_layer + TKE_Ray) + Kd_lay(i,k,jj) / TKE_to_Kd(i,k,jj)) - maxTKE(i,k,jj) + TKE(i,jj) = (TKE(i,jj) - TKE_here) + TKE_Ray + else + TKE_here = TKE_to_layer + TKE_Ray + TKE(i,jj) = TKE(i,jj) - TKE_to_layer + endif + if (TKE(i,jj) < 0.0) TKE(i,jj) = 0.0 ! This should be unnecessary? + + if (TKE_here > 0.0) then + delta_Kd = TKE_here * TKE_to_Kd(i,k,jj) + if (CS%Kd_max >= 0.0) delta_Kd = min(delta_Kd, CS%Kd_max) + Kd_lay(i,k,jj) = Kd_lay(i,k,jj) + delta_Kd + Kd_int(i,K,jj) = Kd_int(i,K,jj) + 0.5 * delta_Kd + Kd_int(i,K+1,jj) = Kd_int(i,K+1,jj) + 0.5 * delta_Kd + if (do_diag_Kd_BBL) then + Kd_BBL(i,j,K) = Kd_BBL(i,j,K) + 0.5 * delta_Kd + Kd_BBL(i,j,K+1) = Kd_BBL(i,j,K+1) + 0.5 * delta_Kd + endif endif endif endif - endif - ! This may be risky - in the case that there are exactly zero - ! velocities at 4 neighboring points, but nonzero velocities - ! above the iterations would stop too soon. I don't see how this - ! could happen in practice. RWH - if ((TKE(i)<= 0.0) .and. (TKE_Ray == 0.0)) then - do_i(i) = .false. ; i_rem = i_rem - 1 - endif + ! This may be risky - in the case that there are exactly zero + ! velocities at 4 neighboring points, but nonzero velocities + ! above the iterations would stop too soon. I don't see how this + ! could happen in practice. RWH + if ((TKE(i,jj)<= 0.0) .and. (TKE_Ray == 0.0)) then + do_i(i,jj) = .false. ; i_rem = i_rem - 1 + endif - endif ; enddo - if (i_rem == 0) exit - enddo ! k-loop + endif ; enddo + if (i_rem == 0) exit + enddo ! k-loop + enddo end subroutine add_drag_diffusivity @@ -1604,7 +1760,7 @@ end subroutine add_drag_diffusivity !! wall turbulent viscosity, up to a BBL height where the energy used for mixing has !! consumed the mechanical TKE input. subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bot, Kd_int, & - G, GV, US, CS, Kd_BBL, Kd_lay) + G, GV, US, CS, Kd_BBL, Kd_lay, dz) type(ocean_grid_type), intent(in) :: G !< Grid structure type(verticalGrid_type), intent(in) :: GV !< Vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type @@ -1629,10 +1785,10 @@ subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bo type(set_diffusivity_CS), pointer :: CS !< Diffusivity control structure real, dimension(:,:,:), pointer :: Kd_BBL !< Interface BBL diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] real, dimension(SZI_(G),SZK_(GV)), & - optional, intent(inout) :: Kd_lay !< Layer net diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + optional, intent(inout) :: Kd_lay !< Layer net diffusivity [H Z T-1 ~> m2 s-1 or kg m-1 s-1] + real, dimension(SZI_(G),SZK_(GV)), intent(in) :: dz !< Height change across layers [Z ~> m] ! Local variables - real :: dz(SZI_(G),SZK_(GV)) ! Height change across layers [Z ~> m] real :: dz_above(SZK_(GV)+1) ! Distance from each interface to the surface [Z ~> m] real :: TKE_column ! net TKE input into the column [H Z2 T-3 ~> m3 s-3 or W m-2] real :: BBL_meanKE_dis ! Sum of tidal and mean kinetic energy dissipation in the bottom boundary layer, which @@ -1671,9 +1827,6 @@ subroutine add_LOTW_BBL_diffusivity(h, u, v, tv, fluxes, visc, j, N2_int, Rho_bo if (allocated(visc%Ray_u) .and. allocated(visc%Ray_v)) Rayleigh_drag = .true. cdrag_sqrt = sqrt(CS%cdrag) - ! Find the vertical distances across layers. - call thickness_to_dz(h, tv, dz, j, G, GV) - do i=G%isc,G%iec ! Developed in single-column mode ! Column-wise parameters. @@ -2151,7 +2304,7 @@ subroutine set_BBL_TKE(u, v, h, tv, fluxes, visc, G, GV, US, CS, OBC) end subroutine set_BBL_TKE -subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) +subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, jstart, jend, nj, ds_dsp1, rho_0) type(ocean_grid_type), intent(in) :: G !< The ocean's grid structure. type(verticalGrid_type), intent(in) :: GV !< The ocean's vertical grid structure. real, dimension(SZI_(G),SZJ_(G),SZK_(GV)), & @@ -2159,17 +2312,19 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) type(thermo_var_ptrs), intent(in) :: tv !< Structure containing pointers to any !! available thermodynamic fields; absent !! fields have NULL ptrs. - integer, dimension(SZI_(G)), intent(in) :: kb !< Index of lightest layer denser than the buffer + integer, dimension(SZI_(G),nj), intent(in) :: kb !< Index of lightest layer denser than the buffer !! layer, or -1 without a bulk mixed layer. type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type type(set_diffusivity_CS), pointer :: CS !< Control structure returned by previous !! call to diabatic_entrain_init. - integer, intent(in) :: j !< Meridional index upon which to work. - real, dimension(SZI_(G),SZK_(GV)), intent(out) :: ds_dsp1 !< Coordinate variable (sigma-2) + integer, intent(in) :: jstart !< Start j-index of the j-block + integer, intent(in) :: jend !< End j-index of the j-block + integer, intent(in) :: nj !< Number of j-rows in the block + real, dimension(SZI_(G),nj,SZK_(GV)), intent(out) :: ds_dsp1 !< Coordinate variable (sigma-2) !! difference across an interface divided by !! the difference across the interface below !! it [nondim] - real, dimension(SZI_(G),SZK_(GV)), & + real, dimension(SZI_(G),nj,SZK_(GV)), & optional, intent(in) :: rho_0 !< Layer potential densities relative to !! surface press [R ~> kg m-3]. @@ -2177,47 +2332,43 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) real :: g_R0 ! g_R0 is a rescaled version of g/Rho [L2 Z-1 R-1 T-2 ~> m4 kg-1 s-2] real :: eps, tmp ! nondimensional temporary variables [nondim] real :: a(SZK_(GV)), a_0(SZK_(GV)) ! nondimensional temporary variables [nondim] - real :: p_ref(SZI_(G)) ! an array of tv%P_Ref pressures [R L2 T-2 ~> Pa] - real :: Rcv(SZI_(G),SZK_(GV)) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] + real :: p_ref(SZI_(G),nj) ! an array of tv%P_Ref pressures [R L2 T-2 ~> Pa] + real :: Rcv(SZI_(G),nj,SZK_(GV)) ! coordinate density in the mixed and buffer layers [R ~> kg m-3] real :: I_Drho ! The inverse of the coordinate density difference between ! layers [R-1 ~> m3 kg-1] integer, dimension(2) :: EOSdom ! The i-computational domain for the equation of state - integer :: i, k, k3, is, ie, nz, kmb + integer :: i, j, k, k3, is, ie, nz, kmb, jj is = G%isc ; ie = G%iec ; nz = GV%ke - do k=2,nz-1 + do j=jstart,jend ; jj = j - jstart + 1 ; do k=2,nz-1 ; do i=is,ie if (GV%g_prime(k+1) /= 0.0) then if (GV%Boussinesq .or. GV%Semi_Boussinesq) then - do i=is,ie - ds_dsp1(i,k) = GV%g_prime(k) / GV%g_prime(k+1) - enddo + ds_dsp1(i,jj,k) = GV%g_prime(k) / GV%g_prime(k+1) else ! Use a mathematically equivalent form that avoids any dependency on RHO_0. - do i=is,ie - ds_dsp1(i,k) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) - enddo + ds_dsp1(i,jj,k) = (GV%Rlay(k) - GV%Rlay(k-1)) / (GV%Rlay(k+1) - GV%Rlay(k)) endif else - do i=is,ie - ds_dsp1(i,k) = 1. - enddo + ds_dsp1(i,jj,k) = 1. endif - enddo + enddo ; enddo ; enddo if (CS%bulkmixedlayer) then g_R0 = GV%g_Earth / (GV%Rho0) kmb = GV%nk_rho_varies eps = 0.1 - do i=is,ie ; p_ref(i) = tv%P_Ref ; enddo EOSdom(:) = EOS_domain(G%HI) - do k=1,kmb - call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref, Rcv(:,k), tv%eqn_of_state, EOSdom) - enddo - do i=is,ie - if (kb(i) <= nz-1) then + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + p_ref(i,jj) = tv%P_Ref + enddo ; enddo + do k=1,kmb ; do j=jstart,jend ; jj = j - jstart + 1 + call calculate_density(tv%T(:,j,k), tv%S(:,j,k), p_ref(:,jj), Rcv(:,jj,k), tv%eqn_of_state, EOSdom) + enddo ; enddo + do j=jstart,jend ; jj = j - jstart + 1 ; do i=is,ie + if (kb(i,jj) <= nz-1) then ! Set up appropriately limited ratios of the reduced gravities of the ! interfaces above and below the buffer layer and the next denser layer. - k = kb(i) + k = kb(i,jj) if (GV%Boussinesq .or. GV%Semi_Boussinesq) then I_Drho = g_R0 / GV%g_prime(k+1) @@ -2226,44 +2377,44 @@ subroutine set_density_ratios(h, tv, kb, G, GV, US, CS, j, ds_dsp1, rho_0) endif ! The indexing convention for a is appropriate for the interfaces. do k3=1,kmb - a(k3+1) = (GV%Rlay(k) - Rcv(i,k3)) * I_Drho + a(k3+1) = (GV%Rlay(k) - Rcv(i,k3,jj)) * I_Drho enddo - if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,k))) then + if ((present(rho_0)) .and. (a(kmb+1) < 2.0*eps*ds_dsp1(i,jj,k))) then ! If the buffer layer nearly matches the density of the layer below in the ! coordinate variable (sigma-2), use the sigma-0-based density ratio if it is ! greater (and stable). - if ((rho_0(i,k) > rho_0(i,kmb)) .and. & - (rho_0(i,k+1) > rho_0(i,k))) then - I_Drho = 1.0 / (rho_0(i,k+1)-rho_0(i,k)) - a_0(kmb+1) = min((rho_0(i,k)-rho_0(i,kmb)) * I_Drho, ds_dsp1(i,k)) + if ((rho_0(i,jj,k) > rho_0(i,jj,kmb)) .and. & + (rho_0(i,jj,k+1) > rho_0(i,jj,k))) then + I_Drho = 1.0 / (rho_0(i,jj,k+1)-rho_0(i,jj,k)) + a_0(kmb+1) = min((rho_0(i,jj,k)-rho_0(i,jj,kmb)) * I_Drho, ds_dsp1(i,jj,k)) if (a_0(kmb+1) > a(kmb+1)) then do k3=2,kmb - a_0(k3) = a_0(kmb+1) + (rho_0(i,kmb)-rho_0(i,k3-1)) * I_Drho + a_0(k3) = a_0(kmb+1) + (rho_0(i,jj,kmb)-rho_0(i,jj,k3-1)) * I_Drho enddo - if (a(kmb+1) <= eps*ds_dsp1(i,k)) then + if (a(kmb+1) <= eps*ds_dsp1(i,jj,k)) then do k3=2,kmb+1 ; a(k3) = a_0(k3) ; enddo else -! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,k)) - 1.0)) ) - tmp = a(kmb+1)/(eps*ds_dsp1(i,k)) - 1.0 +! Alternative... tmp = 0.5*(1.0 - cos(PI*(a(K2+1)/(eps*ds_dsp1(i,jj,k)) - 1.0)) ) + tmp = a(kmb+1)/(eps*ds_dsp1(i,jj,k)) - 1.0 do k3=2,kmb+1 ; a(k3) = tmp*a(k3) + (1.0-tmp)*a_0(k3) ; enddo endif endif endif endif - ds_dsp1(i,k) = MAX(a(kmb+1),1e-5) + ds_dsp1(i,jj,k) = MAX(a(kmb+1),1e-5) do k3=2,kmb -! ds_dsp1(i,k3) = MAX(a(k3),1e-5) +! ds_dsp1(i,jj,k3) = MAX(a(k3),1e-5) ! Deliberately treat convective instabilities of the upper mixed ! and buffer layers with respect to the deepest buffer layer as ! though they don't exist. They will be eliminated by the upcoming ! call to the mixedlayer code anyway. ! The indexing convention is appropriate for the interfaces. - ds_dsp1(i,k3) = MAX(a(k3),ds_dsp1(i,k)) + ds_dsp1(i,jj,k3) = MAX(a(k3),ds_dsp1(i,jj,k)) enddo - endif ! (kb(i) <= nz-1) - enddo ! I-loop. + endif ! (kb(i,jj) <= nz-1) + enddo ; enddo ! i-loop / j-loop endif ! bulkmixedlayer end subroutine set_density_ratios diff --git a/src/tracer/MOM_tracer_hor_diff.F90 b/src/tracer/MOM_tracer_hor_diff.F90 index 93bcd8549f..ac46b4d621 100644 --- a/src/tracer/MOM_tracer_hor_diff.F90 +++ b/src/tracer/MOM_tracer_hor_diff.F90 @@ -39,6 +39,13 @@ module MOM_tracer_hor_diff public tracer_hordiff, tracer_hor_diff_init, tracer_hor_diff_end +!> On GPU builds this makes the per-column private scratch in tracer_epipycnal_ML_diff's device +!! regions fixed-size (stack) arrays instead of runtime-sized privates, which nvfortran cannot +!! place in local memory (NVFORTRAN-W-0155) and instead spills to slow device global memory. +!! Checked against GV%ke in tracer_epipycnal_ML_diff. Unused in CPU builds, where the +!! declarations keep their exact SZK_(GV) sizes. +integer, parameter :: GPU_nk_max = 128 + !> The control structure for along-layer and epineutral tracer diffusion type, public :: tracer_hor_diff_CS ; private real :: KhTr !< The along-isopycnal tracer diffusivity [L2 T-1 ~> m2 s-1]. @@ -800,7 +807,11 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & integer, dimension(SZI_(G),SZK_(GV), SZJ_(G)) :: & k0_srt ! The original k-index that each layer of the sorted column corresponds to. +#ifdef __NVCOMPILER_OPENMP_GPU + real, dimension(GPU_nk_max) :: & +#else real, dimension(SZK_(GV)) :: & +#endif h_demand_L, & ! The thickness in the left column that is demanded to match the thickness ! in the counterpart [H ~> m or kg m-2]. h_demand_R, & ! The thickness in the right column that is demanded to match the thickness @@ -843,10 +854,18 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & ! The total number of pairings is usually much less than twice the number of layers, but ! the memory in these 1-d columns of pairings can be allocated generously for safety. +#ifdef __NVCOMPILER_OPENMP_GPU + integer, dimension(GPU_nk_max*2) :: & +#else integer, dimension(SZK_(GV)*2) :: & +#endif kbs_Lp, & ! The sorted indices of the Left and Right columns for kbs_Rp ! each pairing. +#ifdef __NVCOMPILER_OPENMP_GPU + logical, dimension(GPU_nk_max*2) :: & +#else logical, dimension(SZK_(GV)*2) :: & +#endif left_set, & ! If true, the left or right point determines the density of right_set ! of the trio. If densities are exactly equal, both are true. @@ -863,6 +882,11 @@ subroutine tracer_epipycnal_ML_diff(h, dt, Tr, ntr, khdt_epi_x, khdt_epi_y, G, & is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = GV%ke isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed IsdB = G%IsdB ; IedB = G%IedB ; JsdB = G%JsdB ; JedB = G%JedB +#ifdef __NVCOMPILER_OPENMP_GPU + if (nz > GPU_nk_max) call MOM_error(FATAL, & + "tracer_epipycnal_ML_diff: GPU builds require GV%ke <= GPU_nk_max because the device column "//& + "regions use fixed-size private scratch arrays; increase GPU_nk_max in MOM_tracer_hor_diff.F90.") +#endif Idt = 1.0 / dt nkmb = GV%nk_rho_varies diff --git a/src/user/MOM_wave_interface.F90 b/src/user/MOM_wave_interface.F90 index 23f4c8cb7d..35eb6ae33e 100644 --- a/src/user/MOM_wave_interface.F90 +++ b/src/user/MOM_wave_interface.F90 @@ -15,6 +15,7 @@ module MOM_wave_interface use MOM_forcing_type, only : mech_forcing use MOM_grid, only : ocean_grid_type use MOM_hor_index, only : hor_index_type +use MOM_intrinsic_functions, only : exp_reprod, erfc_reprod, log_reprod use MOM_io, only : file_exists, get_var_sizes, read_variable use MOM_io, only : vardesc, var_desc use MOM_safe_alloc, only : safe_alloc_ptr @@ -37,6 +38,8 @@ module MOM_wave_interface ! called in step_mom. public get_Langmuir_Number ! Public interface to compute Langmuir number called from ! ePBL or KPP routines. +public get_Langmuir_Number_LF17 ! Device-callable LF17-only Langmuir-number kernel (GPU offload). +public set_wave_LF17_params ! Populate a device-mappable wave_LF17_params bundle from a Waves CS. public Stokes_PGF ! Public interface to compute Stokes-shear induced pressure gradient force anomaly public StokesMixing ! NOT READY - Public interface to add down-Stokes gradient ! momentum mixing (e.g. the approach of Harcourt 2013/2015) @@ -260,6 +263,26 @@ module MOM_wave_interface end type wave_parameters_CS +!> A small, allocatable-free bundle of the loop-invariant wave_parameters_CS scalars needed by the +!! LF17 statistical-wave Langmuir path. Unlike the wave_parameters_CS pointer (whose allocatable +!! components do not attach under nvfortran `-gpu=mem:separate`), a plain scalar derived type maps to +!! the device cleanly, so the on-device LF17 chain reads these instead of dereferencing Waves. +type, public :: wave_LF17_params + real :: LA_FracHBL !< Fraction of OSBL for averaging Langmuir number [nondim] + real :: rho_air !< A typical density of air at sea level [R ~> kg m-3] + real :: nu_air !< The viscosity of air [Z2 T-1 ~> m2 s-1] + real :: rho_ocn !< A typical surface density of seawater [R ~> kg m-3] + real :: SWH_from_u10sq !< Factor converting the square of the 10 m wind speed to + !! significant wave height [Z T2 L-2 ~> s2 m-1] + real :: vonKar !< The von Karman coefficient as used in the wave code [nondim] + real :: Charnock_slope_U10 !< Slope of the Charnock coefficient/U10 relationship [T L-1 ~> s m-1] + real :: Charnock_min !< Minimum value of the Charnock coefficient [nondim] + real :: Charnock_intercept !< Intercept of the Charnock coefficient/U10 relationship [nondim] + real :: I_g_Earth !< The inverse of the gravitational acceleration [T2 Z L-2 ~> s2 m-1] + integer :: answer_date !< The vintage of the order of arithmetic and expressions in the + !! wave calculations. +end type wave_LF17_params + ! Switches needed in import_stokes_drift !>@{ Enumeration values for the wave method integer, parameter :: TESTPROF = 0, SURFBANDS = 1, DHH85 = 2, LF17 = 3, EFACTOR = 4, NULL_WaveMethod = -99 @@ -276,6 +299,18 @@ module MOM_wave_interface character*(4), parameter :: LF17_STRING = "LF17" !< LF17 wave method string character*(7), parameter :: EFACTOR_STRING = "EFACTOR" !< EFACTOR (based on vr12-ma) wave method string +!> GPU port: the Li & Fox-Kemper statistical-wave Langmuir-number kernels are device-callable so the +!! ePBL column solver can evaluate the Langmuir number on the device; on GPU builds only the LF17 +!! (USE_LA_LI2016) statistical-wave path is compiled (the wave-model branches are host-only). +!$omp declare target(get_StokesSL_LiFoxKemper, ust_2_u10_coare3p5) + +! RESOLVED(gpu-bitwise-repro) 2026-07-20: this LF17 Langmuir chain runs on-device (called from +! ePBL_column); its transcendentals are now the reproducible kernels -- one_minus_exp_x exp() -> +! exp_reprod; get_StokesSL_LiFoxKemper r5 erfc() -> erfc_reprod; ust_2_u10_coare3p5 COARE-Cd log() -> +! log_reprod; PI = 4*atan(1) -> a literal. (sqrt() is IEEE-exact.) With these + the ePBL side, the GPU +! and a same-source nvfortran-CPU build reproduce benchmark_ALE bit-for-bit. Prototype kernels; +! coordinate with Marshall's transcendental-repro pass. See MOM_energetic_PBL.F90's resolved audit. + contains !> Initializes parameters related to MOM_wave_interface @@ -1048,7 +1083,7 @@ real function one_minus_exp_x(x) ! The Taylor series expression for exp(-x) gives a more accurate expression for 64-bit reals. one_minus_exp_x = 1.0 - x * (0.5 - C1_6*x) else - one_minus_exp_x = (1.0 - exp(-x)) / x + one_minus_exp_x = (1.0 - exp_reprod(-x)) / x endif end function one_minus_exp_x @@ -1216,11 +1251,13 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & !Local Variables + real :: LA_STK ! Surface-layer averaged Stokes drift magnitude [L T-1 ~> m s-1] + type(wave_LF17_params) :: p_lf17 ! Device-mappable bundle of LF17 wave scalars (for the LF17 branch) real :: Top, Bottom, MidPoint ! Positions within each layer [Z ~> m] real :: Dpt_LASL ! Averaging depth for Stokes drift [Z ~> m] real :: ShearDirection ! Shear angular direction from atan2 [radians] real :: WaveDirection ! Wave angular direction from atan2 [radians] - real :: LA_STKx, LA_STKy, LA_STK ! Stokes velocities in [L T-1 ~> m s-1] + real :: LA_STKx, LA_STKy ! Stokes velocities in [L T-1 ~> m s-1] logical :: ContinueLoop, USE_MA real, dimension(SZK_(GV)) :: US_H, VS_H ! Profiles of Stokes velocities [L T-1 ~> m s-1] real, allocatable :: StkBand_X(:), StkBand_Y(:) ! Stokes drifts by band [L T-1 ~> m s-1] @@ -1286,7 +1323,8 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & call Get_SL_Average_Prof( GV, Dpt_LASL, dz, VS_H, LA_STKy) LA_STK = sqrt((LA_STKX**2) + (LA_STKY**2)) elseif (Waves%WaveMethod==LF17) then - call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, Waves, LA_STK, LA) + call set_wave_LF17_params(Waves, p_lf17) + call get_StokesSL_LiFoxKemper(ustar, HBL*Waves%LA_FracHBL, GV, US, p_lf17, LA_STK, LA) elseif (Waves%WaveMethod==Null_WaveMethod) then call MOM_error(FATAL, "Get_Langmuir_number called without defining a WaveMethod. "//& "Suggest to make sure USE_LT is set/overridden to False or choose "//& @@ -1307,6 +1345,43 @@ subroutine get_Langmuir_Number( LA, G, GV, US, HBL, ustar, i, j, dz, Waves, & end subroutine get_Langmuir_Number +!> Device-callable Langmuir-number kernel for the LF17 statistical-wave method (USE_LA_LI2016). +!! This is the GPU-offload entry point (e.g. from ePBL_column): it has no optional arguments — which +!! nvfortran cannot emit `declare target` device code for — and covers only the LF17 path, which is +!! the only WaveMethod supported on the GPU port. It is bit-for-bit the LF17 branch of +!! get_Langmuir_Number (get_StokesSL_LiFoxKemper sets LA directly; no misalignment is applied). +subroutine get_Langmuir_Number_LF17(LA, GV, US, HBL, ustar, p) + real, intent(out) :: LA !< Langmuir number [nondim] + type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure + type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type + real, intent(in) :: HBL !< (Positive) thickness of boundary layer [Z ~> m] + real, intent(in) :: ustar !< Friction velocity [Z T-1 ~> m s-1] + type(wave_LF17_params), intent(in) :: p !< Device-mappable bundle of LF17 wave scalars. +!$omp declare target + + real :: LA_STK ! Surface-layer averaged Stokes drift magnitude [L T-1 ~> m s-1] + + call get_StokesSL_LiFoxKemper(ustar, HBL*p%LA_FracHBL, GV, US, p, LA_STK, LA) +end subroutine get_Langmuir_Number_LF17 + +!> Populate a device-mappable wave_LF17_params bundle from a Waves control structure (host-side). +subroutine set_wave_LF17_params(Waves, p) + type(Wave_parameters_CS), pointer, intent(in) :: Waves !< Surface wave control structure. + type(wave_LF17_params), intent(out) :: p !< The populated scalar bundle. + + p%LA_FracHBL = Waves%LA_FracHBL + p%rho_air = Waves%rho_air + p%nu_air = Waves%nu_air + p%rho_ocn = Waves%rho_ocn + p%SWH_from_u10sq = Waves%SWH_from_u10sq + p%vonKar = Waves%vonKar + p%Charnock_slope_U10 = Waves%Charnock_slope_U10 + p%Charnock_min = Waves%Charnock_min + p%Charnock_intercept = Waves%Charnock_intercept + p%I_g_Earth = Waves%I_g_Earth + p%answer_date = Waves%answer_date +end subroutine set_wave_LF17_params + !> function to return the wave method string set in the param file function get_wave_method(CS) character(:), allocatable :: get_wave_method @@ -1353,7 +1428,7 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) real, intent(in) :: hbl !< boundary layer depth [Z ~> m]. type(verticalGrid_type), intent(in) :: GV !< Ocean vertical grid structure type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type - type(wave_parameters_CS), pointer :: CS !< Wave parameter Control structure + type(wave_LF17_params), intent(in) :: CS !< Device-mappable bundle of LF17 wave scalars real, intent(out) :: UStokes_SL !< Surface layer averaged Stokes drift [L T-1 ~> m s-1] real, intent(out) :: LA !< Langmuir number [nondim] ! Local variables @@ -1377,9 +1452,8 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) real :: root_2kz ! The square root of twice the peak wavenumber times the ! boundary layer depth [nondim] real :: u10 ! The 10 m wind speed [L T-1 ~> m s-1] - real :: PI ! 3.1415926535... [nondim] + real, parameter :: PI = 3.14159265358979323846 ! 3.1415926535... [nondim] - PI = 4.0*atan(1.0) UStokes_sl = 0.0 LA = 1.e8 if (ustar > 0.0) then @@ -1453,8 +1527,8 @@ subroutine get_StokesSL_LiFoxKemper(ustar, hbl, GV, US, CS, UStokes_SL, LA) ! It has been verified that these two expressions for r5 are the same to 6 decimal places for ! root_2kz between 1e-10 and 1e-3, but that the first one degrades for smaller values. if (root_2kz > 1e-3) then - r5 = sqrt(PI) * (root_2kz * (-0.84 * erfc(root_2kz) + 0.2 * erfc(1.6*root_2kz)) + & - 0.1182 * (erfc(1.6*root_2kz) - erfc(root_2kz)) / root_2kz) + r5 = sqrt(PI) * (root_2kz * (-0.84 * erfc_reprod(root_2kz) + 0.2 * erfc_reprod(1.6*root_2kz)) + & + 0.1182 * (erfc_reprod(1.6*root_2kz) - erfc_reprod(root_2kz)) / root_2kz) else ! It is more accurate to replace erf with the first two terms of its Taylor series ! erf(z) = (2/sqrt(pi)) * z * (1. - (1/3)*z**2 + (1/10)*z**4 - (1/42)*z**6 + ...) @@ -2064,7 +2138,7 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) real, intent(out) :: U10 !< 10-m neutral wind speed [L T-1 ~> m s-1] type(verticalGrid_type), intent(in) :: GV !< vertical grid type type(unit_scale_type), intent(in) :: US !< A dimensional unit scaling type - type(wave_parameters_CS), pointer :: CS !< Wave parameter Control structure + type(wave_LF17_params), intent(in) :: CS !< Device-mappable bundle of LF17 wave scalars ! Local variables real :: z0sm, z0, z0rough ! Roughness lengths [Z ~> m] @@ -2084,8 +2158,10 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) ! Note in Edson et al. 2013, eq. 13 m is given as 0.017. However, ! m=0.0017 reproduces the curve in their figure 6. +#ifndef __NVCOMPILER_OPENMP_GPU if (CS%vonKar < 0.0) call MOM_error(FATAL, & "ust_2_u10_coare3p5 called with a negative value of Waves%vonKar") +#endif z0sm = 0.11 * CS%nu_air / USTair ! Compute z0smooth from ustar guess u10a = 1000.0*US%m_s_to_L_T ! An insanely large upper bound for u10. @@ -2126,7 +2202,7 @@ subroutine ust_2_u10_coare3p5(USTair, U10, GV, US, CS) alpha = min(CS%Charnock_min, CS%Charnock_slope_U10 * u10 + CS%Charnock_intercept) z0rough = alpha * (CS%I_g_Earth * USTair**2) ! Compute z0rough from ustar guess z0 = z0sm + z0rough - I_sqrtCd = abs(log(z0 * I_ten_m_scale)) * I_vonKar ! Compute Cd from derived roughness + I_sqrtCd = abs(log_reprod(z0 * I_ten_m_scale)) * I_vonKar ! Compute Cd from derived roughness u10 = US%Z_to_L*USTair * I_sqrtCd ! Compute new u10 from the derived Cd. enddo