diff --git a/src/ALE/MOM_hybgen_remap.F90 b/src/ALE/MOM_hybgen_remap.F90 index 68a193cebe..29a5cf2b02 100644 --- a/src/ALE/MOM_hybgen_remap.F90 +++ b/src/ALE/MOM_hybgen_remap.F90 @@ -13,7 +13,8 @@ module MOM_hybgen_remap contains !> Set up the coefficients for PLM remapping of a set of scalars -subroutine hybgen_plm_coefs(si, dpi, slope, nk, ns, thin, PCM_lay) +subroutine hybgen_plm_coefs(si, dpi, slope, nk, ns, thin, PCM_lay) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: nk !< The number of input layers integer, intent(in) :: ns !< The number of scalar fields to work on real, intent(in) :: si(nk,ns) !< The cell-averaged input scalar fields [A] @@ -90,7 +91,8 @@ end subroutine hybgen_plm_coefs !> Set up the coefficients for PPM remapping of a set of scalars -subroutine hybgen_ppm_coefs(s, h_src, edges, nk, ns, thin, PCM_lay) +subroutine hybgen_ppm_coefs(s, h_src, edges, nk, ns, thin, PCM_lay) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: nk !< The number of input layers integer, intent(in) :: ns !< The scalar fields to work on real, intent(in) :: s(nk,ns) !< The input scalar fields [A] @@ -225,7 +227,8 @@ end subroutine hybgen_ppm_coefs !> Set up the coefficients for PPM remapping of a set of scalars -subroutine hybgen_weno_coefs(s, h_src, edges, nk, ns, thin, PCM_lay) +subroutine hybgen_weno_coefs(s, h_src, edges, nk, ns, thin, PCM_lay) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: nk !< The number of input layers integer, intent(in) :: ns !< The number of scalar fields to work on real, intent(in) :: s(nk,ns) !< The input scalar fields [A] diff --git a/src/ALE/MOM_remapping.F90 b/src/ALE/MOM_remapping.F90 index 9c0634310f..34fb2f40e6 100644 --- a/src/ALE/MOM_remapping.F90 +++ b/src/ALE/MOM_remapping.F90 @@ -41,6 +41,11 @@ module MOM_remapping implicit none ; private +!> Maximum number of vertical levels supported for GPU-resident local arrays. +!! Variable-length locals inside declare-target routines force nvfortran to use +!! GPU heap (NVCOMPILER_ACC_CUDA_HEAPSIZE). Fixed-size arrays use stack instead. +integer, parameter, public :: NK_GPU_MAX = 500 + !> Container for remapping parameters type, public :: remapping_CS ; private !> Determines which reconstruction to use @@ -234,7 +239,8 @@ end subroutine extract_member_remapping_CS !! !! \todo Remove h_neglect argument by moving into remapping_CS !! \todo Remove PCM_cell argument by adding new method in Recon1D class -subroutine remapping_core_h(CS, n0, h0, u0, n1, h1, u1, net_err, PCM_cell) +subroutine remapping_core_h(CS, n0, h0, u0, n1, h1, u1, net_err, PCM_cell) ! GPU PORT DIAGNOSTICS + !$omp declare target type(remapping_CS), intent(in) :: CS !< Remapping control structure integer, intent(in) :: n0 !< Number of cells on source grid real, dimension(n0), intent(in) :: h0 !< Cell widths on source grid [H] @@ -269,39 +275,27 @@ subroutine remapping_core_h(CS, n0, h0, u0, n1, h1, u1, net_err, PCM_cell) call intersect_src_tgt_grids(n0, h0, n1, h1, h_sub, h0_eff, & isrc_start, isrc_end, isrc_max, itgt_start, itgt_end, isub_src) - if (CS%remapping_scheme == REMAPPING_VIA_CLASS) then - -! if (CS%debug) call CS%reconstruction%set_debug() ! Sets an internal flag - - call CS%reconstruction%reconstruct(h0, u0) - - ! Adjust h_sub so that the Hallberg conservation trick works properly -! call adjust_h_sub( n0, h0, n1, isrc_start, isrc_end, isrc_max, h_sub ) - - ! Loop over each sub-cell to calculate average/integral values within each sub-cell. - ! Uses: h_sub, isrc_start, isrc_end, isrc_max, isub_src - ! Sets: u_sub, uh_sub - call CS%reconstruction%remap_to_sub_grid(h0, u0, n1, h_sub, & - isrc_start, isrc_end, isrc_max, isub_src, & - u_sub, uh_sub, u02_err) - - ! Loop over each target cell summing the integrals from sub-cells within the target cell. - ! Uses: itgt_start, itgt_end, h1, h_sub, uh_sub, u_sub - ! Sets: u1, uh_err - call remap_sub_to_tgt_grid(n0, n1, h1, h_sub, u_sub, uh_sub, itgt_start, itgt_end, & - CS%force_bounds_in_target, CS%offset_tgt_summation, & - CS%better_force_bounds_in_target, u1, uh_err) - - ! Include the error remapping from source to sub-cells in the estimate of total remapping error - uh_err = uh_err + u02_err - - else ! Uses the OM4-era reconstruction functions + ! GPU: The class-based (polymorphic) remapping path has been removed from this + ! routine for GPU compatibility. nvfortran cannot compile class dispatch for device. + ! The OM4-era polynomial path handles all reconstruction schemes. + ! GPU: original code had class-based dispatch here: + ! if (CS%remapping_scheme == REMAPPING_VIA_CLASS) then + ! call CS%reconstruction%reconstruct(h0, u0) + ! call CS%reconstruction%remap_to_sub_grid(h0, u0, n1, h_sub, & + ! isrc_start, isrc_end, isrc_max, isub_src, & + ! u_sub, uh_sub, u02_err) + ! call remap_sub_to_tgt_grid(n0, n1, h1, h_sub, u_sub, uh_sub, itgt_start, itgt_end, & + ! CS%force_bounds_in_target, CS%offset_tgt_summation, & + ! CS%better_force_bounds_in_target, u1, uh_err) + ! uh_err = uh_err + u02_err + ! else ! Uses the OM4-era reconstruction functions call build_reconstructions_1d(CS, n0, h0, u0, ppoly_r_coefs, ppoly_r_E, ppoly_r_S, iMethod, & CS%h_neglect, CS%h_neglect_edge, PCM_cell, debug=CS%debug) - if (CS%check_reconstruction) call check_reconstructions_1d(n0, h0, u0, CS%degree, & - CS%boundary_extrapolation, ppoly_r_coefs, ppoly_r_E) + ! GPU: formatted I/O not GPU-safe, skip on device + ! if (CS%check_reconstruction) call check_reconstructions_1d(n0, h0, u0, CS%degree, & + ! CS%boundary_extrapolation, ppoly_r_coefs, ppoly_r_E) ! Loop over each sub-cell to calculate average/integral values within each sub-cell. ! Uses: h_sub, h0_eff, isub_src @@ -328,12 +322,15 @@ subroutine remapping_core_h(CS, n0, h0, u0, n1, h1, u1, net_err, PCM_cell) ! Include the error remapping from source to sub-cells in the estimate of total remapping error uh_err = uh_err + u02_err - if (CS%check_remapping) call check_remapped_values(n0, h0, u0, ppoly_r_E, CS%degree, ppoly_r_coefs, & - n1, h1, u1, iMethod, uh_err, "remapping_core_h") + ! GPU: formatted I/O not GPU-safe, skip on device + ! if (CS%check_remapping) call check_remapped_values(n0, h0, u0, ppoly_r_E, CS%degree, ppoly_r_coefs, & + ! n1, h1, u1, iMethod, uh_err, "remapping_core_h") - endif + ! GPU: original code (class path ended with endif here): + ! endif - if (present(net_err)) net_err = uh_err + ! GPU: present() not supported on GPU — net_err is never passed from diag remap callers + ! if (present(net_err)) net_err = uh_err end subroutine remapping_core_h @@ -410,9 +407,10 @@ subroutine remapping_core_w( CS, n0, h0, u0, n1, dx, u1) end subroutine remapping_core_w !> Creates polynomial reconstructions of u0 on the source grid h0. -subroutine build_reconstructions_1d( CS, n0, h0, u0, ppoly_r_coefs, & +subroutine build_reconstructions_1d( CS, n0, h0, u0, ppoly_r_coefs, & ! GPU PORT DIAGNOSTICS ppoly_r_E, ppoly_r_S, iMethod, h_neglect, & h_neglect_edge, PCM_cell, debug ) + !$omp declare target type(remapping_CS), intent(in) :: CS !< Remapping control structure integer, intent(in) :: n0 !< Number of cells on source grid real, dimension(n0), intent(in) :: h0 !< Cell widths on source grid [H] @@ -439,9 +437,12 @@ subroutine build_reconstructions_1d( CS, n0, h0, u0, ppoly_r_coefs, & integer :: k, n logical :: deb ! Do debugging - deb=.false.; if (present(debug)) deb=debug - - h_neg_edge = h_neglect ; if (present(h_neglect_edge)) h_neg_edge = h_neglect_edge + ! GPU: present() not supported on GPU — use safe defaults + ! GPU: original code: + ! deb=.false.; if (present(debug)) deb=debug + ! h_neg_edge = h_neglect ; if (present(h_neglect_edge)) h_neg_edge = h_neglect_edge + deb=.false. + h_neg_edge = h_neglect ! Reset polynomial ppoly_r_E(:,:) = 0.0 @@ -533,22 +534,28 @@ subroutine build_reconstructions_1d( CS, n0, h0, u0, ppoly_r_coefs, & endif iMethod = INTEGRATION_PQM case ( REMAPPING_VIA_CLASS ) - call MOM_error( FATAL, 'MOM_remapping, build_reconstructions_1d: '//& - 'Should not reach this point if using Recon1d class for remapping' ) + stop 1 ! Should not reach this point if using Recon1d class + ! GPU: original code: + ! call MOM_error( FATAL, 'MOM_remapping, build_reconstructions_1d: '//& + ! 'Should not reach this point if using Recon1d class for remapping' ) case default - call MOM_error( FATAL, 'MOM_remapping, build_reconstructions_1d: '//& - 'The selected remapping method is invalid' ) + stop 1 ! The selected remapping method is invalid + ! GPU: original code: + ! call MOM_error( FATAL, 'MOM_remapping, build_reconstructions_1d: '//& + ! 'The selected remapping method is invalid' ) end select - if (present(PCM_cell)) then - ! Change the coefficients to those for the piecewise constant method in indicated cells. - do k=1,n0 ; if (PCM_cell(k)) then - ppoly_r_coefs(k,1) = u0(k) - ppoly_r_E(k,1:2) = u0(k) - ppoly_r_S(k,1:2) = 0.0 - do n=2,CS%degree+1 ; ppoly_r_coefs(k,n) = 0.0 ; enddo - endif ; enddo - endif + ! GPU: PCM_cell present() check removed — not supported on GPU. + ! PCM_cell is not passed from the diagnostic remapping callers. + ! GPU: original code: + ! if (present(PCM_cell)) then + ! do k=1,n0 ; if (PCM_cell(k)) then + ! ppoly_r_coefs(k,1) = u0(k) + ! ppoly_r_E(k,1:2) = u0(k) + ! ppoly_r_S(k,1:2) = 0.0 + ! do n=2,CS%degree+1 ; ppoly_r_coefs(k,n) = 0.0 ; enddo + ! endif ; enddo + ! endif end subroutine build_reconstructions_1d @@ -642,8 +649,9 @@ end subroutine check_reconstructions_1d !! at the left (top). !! If the total column thicknesses are the same, then the right (bottom) interfaces are also aligned and !! so the last subcell will also be vanished. -subroutine intersect_src_tgt_grids( n0, h0, n1, h1, h_sub, h0_eff, & +subroutine intersect_src_tgt_grids( n0, h0, n1, h1, h_sub, h0_eff, & ! GPU PORT DIAGNOSTICS isrc_start, isrc_end, isrc_max, itgt_start, itgt_end, isub_src ) + !$omp declare target integer, intent(in) :: n0 !< Number of cells in source grid real, intent(in) :: h0(n0) !< Source grid widths (size n0) [H] integer, intent(in) :: n1 !< Number of cells in target grid @@ -845,9 +853,10 @@ end subroutine intersect_src_tgt_grids !! !! This includes an error for the scenario where the source grid is much thicker than !! the target grid and extrapolation is needed. -subroutine remap_src_to_sub_grid_om4(n0, h0, u0, ppoly0_E, ppoly0_coefs, n1, h_sub, & +subroutine remap_src_to_sub_grid_om4(n0, h0, u0, ppoly0_E, ppoly0_coefs, n1, h_sub, & ! GPU PORT DIAGNOSTICS h0_eff, isrc_start, isrc_end, isrc_max, isub_src, & method, force_bounds_in_subcell, u_sub, uh_sub, u02_err) + !$omp declare target integer, intent(in) :: n0 !< Number of cells in source grid real, intent(in) :: h0(n0) !< Source grid widths (size n0) [H] real, intent(in) :: u0(n0) !< Source grid widths (size n0) [H] @@ -962,9 +971,10 @@ subroutine remap_src_to_sub_grid_om4(n0, h0, u0, ppoly0_E, ppoly0_coefs, n1, h_s end subroutine remap_src_to_sub_grid_om4 !> Remaps column of n0 values u0 on grid h0 to subgrid h_sub -subroutine remap_src_to_sub_grid(n0, h0, u0, ppoly0_E, ppoly0_coefs, n1, h_sub, & +subroutine remap_src_to_sub_grid(n0, h0, u0, ppoly0_E, ppoly0_coefs, n1, h_sub, & ! GPU PORT DIAGNOSTICS isrc_start, isrc_end, isrc_max, isub_src, & method, force_bounds_in_subcell, u_sub, uh_sub, u02_err) + !$omp declare target integer, intent(in) :: n0 !< Number of cells in source grid real, intent(in) :: h0(n0) !< Source grid widths (size n0) [H] real, intent(in) :: u0(n0) !< Source grid widths (size n0) [H] @@ -1103,8 +1113,9 @@ end subroutine remap_src_to_sub_grid !> Remaps column of n0+n1+1 values usub on sub-grid h_sub to targets on grid h1 !! using the OM4-era algorithm -subroutine remap_sub_to_tgt_grid_om4(n0, n1, h1, h_sub, u_sub, uh_sub, & +subroutine remap_sub_to_tgt_grid_om4(n0, n1, h1, h_sub, u_sub, uh_sub, & ! GPU PORT DIAGNOSTICS itgt_start, itgt_end, force_bounds_in_target, u1, uh_err) + !$omp declare target integer, intent(in) :: n0 !< Number of cells in source grid integer, intent(in) :: n1 !< Number of cells in target grid real, intent(in) :: h1(n1) !< Target grid widths (size n1) [H] @@ -1247,7 +1258,8 @@ subroutine remap_sub_to_tgt_grid(n0, n1, h1, h_sub, u_sub, uh_sub, & end subroutine remap_sub_to_tgt_grid !> Linearly interpolate interface data, u_src, from grid h_src to a grid h_dest -subroutine interpolate_column(nsrc, h_src, u_src, ndest, h_dest, u_dest, mask_edges) +subroutine interpolate_column(nsrc, h_src, u_src, ndest, h_dest, u_dest, mask_edges) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: nsrc !< Number of source cells real, dimension(nsrc), intent(in) :: h_src !< Thickness of source cells [H] real, dimension(nsrc+1), intent(in) :: u_src !< Values at source cell interfaces [A] @@ -1260,9 +1272,9 @@ subroutine interpolate_column(nsrc, h_src, u_src, ndest, h_dest, u_dest, mask_ed ! Local variables real :: x_dest ! Relative position of target interface [H] real :: dh ! Source cell thickness [H] - real :: frac_pos(ndest+1) ! Fractional position of the destination interface + real :: frac_pos(NK_GPU_MAX+1) ! Fractional position of the destination interface ! within the source layer [nondim], 0 <= frac_pos <= 1. - integer :: k_src(ndest+1) ! Source grid layer index of destination interface, 1 <= k_src <= ndest. + integer :: k_src(NK_GPU_MAX+1) ! Source grid layer index of destination interface, 1 <= k_src <= ndest. integer :: ks, k_dest ! Index of cell in src and dest columns ! The following forces the "do while" loop to do one cycle that will set u1, u2, dh. @@ -1391,7 +1403,8 @@ end subroutine reintegrate_column !> Returns the average value of a reconstruction within a single source cell, i0, !! between the non-dimensional positions xa and xb (xa<=xb) with dimensional !! separation dh. -real function average_value_ppoly( n0, u0, ppoly0_E, ppoly0_coefs, method, i0, xa, xb) +real function average_value_ppoly( n0, u0, ppoly0_E, ppoly0_coefs, method, i0, xa, xb) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: n0 !< Number of cells in source grid real, intent(in) :: u0(n0) !< Cell means [A] real, intent(in) :: ppoly0_E(n0,2) !< Edge value of polynomial [A] @@ -1450,7 +1463,9 @@ real function average_value_ppoly( n0, u0, ppoly0_E, ppoly0_coefs, method, i0, x + ( ppoly0_coefs(i0,4) * 0.25* ( xa2pxb2 * xapxb ) & + ppoly0_coefs(i0,5) * 0.2 * ( ( xb*xb_2 + xa*xa_2 ) * xapxb + xa_2*xb_2 ) ) ) ) ) case default - call MOM_error( FATAL,'The selected integration method is invalid' ) + stop 1 ! The selected integration method is invalid + ! GPU: original code: + ! call MOM_error( FATAL,'The selected integration method is invalid' ) end select else ! dh == 0. select case ( method ) @@ -1489,7 +1504,9 @@ real function average_value_ppoly( n0, u0, ppoly0_E, ppoly0_coefs, method, i0, x + xa * ppoly0_coefs(i0,5) ) ) ) case default u_ave = 0. - call MOM_error( FATAL,'The selected integration method is invalid' ) + stop 1 ! The selected integration method is invalid + ! GPU: original code: + ! call MOM_error( FATAL,'The selected integration method is invalid' ) end select endif average_value_ppoly = u_ave diff --git a/src/ALE/PCM_functions.F90 b/src/ALE/PCM_functions.F90 index dff25e5fc6..b9a222c2d8 100644 --- a/src/ALE/PCM_functions.F90 +++ b/src/ALE/PCM_functions.F90 @@ -17,7 +17,8 @@ module PCM_functions !! !! It is assumed that the dimension of 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed. -subroutine PCM_reconstruction( N, u, edge_values, ppoly_coef ) +subroutine PCM_reconstruction( N, u, edge_values, ppoly_coef ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: u !< cell averages in arbitrary units [A] real, dimension(:,:), intent(inout) :: edge_values !< Edge value of polynomial, diff --git a/src/ALE/PLM_functions.F90 b/src/ALE/PLM_functions.F90 index ab70541747..34a9a7755a 100644 --- a/src/ALE/PLM_functions.F90 +++ b/src/ALE/PLM_functions.F90 @@ -187,7 +187,8 @@ end function PLM_extrapolate_slope !! !! It is assumed that the size of the array 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed here. -subroutine PLM_reconstruction( N, h, u, edge_values, ppoly_coef, h_neglect ) +subroutine PLM_reconstruction( N, h, u, edge_values, ppoly_coef, h_neglect ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell averages (size N) in arbitrary units [A] @@ -266,7 +267,8 @@ end subroutine PLM_reconstruction !! !! It is assumed that the size of the array 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed here. -subroutine PLM_boundary_extrapolation( N, h, u, edge_values, ppoly_coef, h_neglect ) +subroutine PLM_boundary_extrapolation( N, h, u, edge_values, ppoly_coef, h_neglect ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell averages (size N) in arbitrary units [A] diff --git a/src/ALE/PPM_functions.F90 b/src/ALE/PPM_functions.F90 index ad8fe2adb6..e9d53c2a63 100644 --- a/src/ALE/PPM_functions.F90 +++ b/src/ALE/PPM_functions.F90 @@ -20,7 +20,8 @@ module PPM_functions contains !> Builds quadratic polynomials coefficients from cell mean and edge values. -subroutine PPM_reconstruction( N, h, u, edge_values, ppoly_coef, h_neglect, answer_date) +subroutine PPM_reconstruction( N, h, u, edge_values, ppoly_coef, h_neglect, answer_date) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< Cell widths [H] real, dimension(N), intent(in) :: u !< Cell averages in arbitrary coordinates [A] @@ -54,7 +55,8 @@ end subroutine PPM_reconstruction !> Adjusts edge values using the standard PPM limiter (Colella & Woodward, JCP 1984) !! after first checking that the edge values are bounded by neighbors cell averages !! and that the edge values are monotonic between cell averages. -subroutine PPM_limiter_standard( N, h, u, edge_values, h_neglect, answer_date ) +subroutine PPM_limiter_standard( N, h, u, edge_values, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell average properties (size N) [A] @@ -124,7 +126,8 @@ end subroutine PPM_limiter_standard !> Adjusts edge values using the original monotonicity constraint (Colella & Woodward, JCP 1984) !! Based on hybgen_ppm_coefs -subroutine PPM_monotonicity( N, u, edge_values ) +subroutine PPM_monotonicity( N, u, edge_values ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: u !< cell average properties (size N) [A] real, dimension(:,:), intent(inout) :: edge_values !< Potentially modified edge values [A] @@ -154,7 +157,7 @@ end subroutine PPM_monotonicity !------------------------------------------------------------------------------ !> Reconstruction by parabolas within boundary cells -subroutine PPM_boundary_extrapolation( N, h, u, edge_values, ppoly_coef, h_neglect) +subroutine PPM_boundary_extrapolation( N, h, u, edge_values, ppoly_coef, h_neglect) ! GPU PORT DIAGNOSTICS !------------------------------------------------------------------------------ ! Reconstruction by parabolas within boundary cells. ! @@ -179,6 +182,7 @@ subroutine PPM_boundary_extrapolation( N, h, u, edge_values, ppoly_coef, h_negle ! defining 'grid' and 'ppoly'. No consistency check is performed here. !------------------------------------------------------------------------------ + !$omp declare target ! Arguments integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] diff --git a/src/ALE/PQM_functions.F90 b/src/ALE/PQM_functions.F90 index d0bd58a9fe..0fbdc73759 100644 --- a/src/ALE/PQM_functions.F90 +++ b/src/ALE/PQM_functions.F90 @@ -17,7 +17,8 @@ module PQM_functions !! !! It is assumed that the dimension of 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed. -subroutine PQM_reconstruction( N, h, u, edge_values, edge_slopes, ppoly_coef, h_neglect, answer_date ) +subroutine PQM_reconstruction( N, h, u, edge_values, edge_slopes, ppoly_coef, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell averages (size N) [A] @@ -72,7 +73,8 @@ end subroutine PQM_reconstruction !! !! It is assumed that the dimension of 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed. -subroutine PQM_limiter( N, h, u, edge_values, edge_slopes, h_neglect, answer_date ) +subroutine PQM_limiter( N, h, u, edge_values, edge_slopes, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell average properties (size N) [A] @@ -496,7 +498,8 @@ end subroutine PQM_boundary_extrapolation !! !! It is assumed that the size of the array 'u' is equal to the number of cells !! defining 'grid' and 'ppoly'. No consistency check is performed here. -subroutine PQM_boundary_extrapolation_v1( N, h, u, edge_values, edge_slopes, ppoly_coef, h_neglect ) +subroutine PQM_boundary_extrapolation_v1( N, h, u, edge_values, edge_slopes, ppoly_coef, h_neglect ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(:), intent(in) :: h !< cell widths (size N) [H] real, dimension(:), intent(in) :: u !< cell averages (size N) [A] diff --git a/src/ALE/polynomial_functions.F90 b/src/ALE/polynomial_functions.F90 index 0b232dc359..d3112583c0 100644 --- a/src/ALE/polynomial_functions.F90 +++ b/src/ALE/polynomial_functions.F90 @@ -18,7 +18,8 @@ module polynomial_functions !! where C refers to the array 'coeff'. !! The number of coefficients is given by ncoef and x !! is the coordinate where the polynomial is to be evaluated. -real function evaluation_polynomial( coeff, ncoef, x ) +real function evaluation_polynomial( coeff, ncoef, x ) ! GPU PORT DIAGNOSTICS + !$omp declare target real, dimension(:), intent(in) :: coeff !< The coefficients of the polynomial, in units that !! vary with the index k as [A H^(k-1)] integer, intent(in) :: ncoef !< The number of polynomial coefficients @@ -45,7 +46,8 @@ end function evaluation_polynomial !! where C refers to the array 'coeff'. !! The number of coefficients is given by ncoef and x !! is the coordinate where the polynomial's derivative is to be evaluated. -real function first_derivative_polynomial( coeff, ncoef, x ) +real function first_derivative_polynomial( coeff, ncoef, x ) ! GPU PORT DIAGNOSTICS + !$omp declare target real, dimension(:), intent(in) :: coeff !< The coefficients of the polynomial, in units that !! vary with the index k as [A H^(k-1)] integer, intent(in) :: ncoef !< The number of polynomial coefficients @@ -67,7 +69,8 @@ end function first_derivative_polynomial !> Exact integration of polynomial of degree npoly in arbitrary units of [A H] !! !! The array of coefficients (Coeff) must be of size npoly+1. -real function integration_polynomial( xi0, xi1, Coeff, npoly ) +real function integration_polynomial( xi0, xi1, Coeff, npoly ) ! GPU PORT DIAGNOSTICS + !$omp declare target real, intent(in) :: xi0 !< The lower bound of the integral in arbitrary !! thickness units [H] real, intent(in) :: xi1 !< The upper bound of the integral in arbitrary diff --git a/src/ALE/regrid_edge_values.F90 b/src/ALE/regrid_edge_values.F90 index f9a0bacd25..d15e556cf6 100644 --- a/src/ALE/regrid_edge_values.F90 +++ b/src/ALE/regrid_edge_values.F90 @@ -38,7 +38,8 @@ module regrid_edge_values !! Both boundary edge values are set equal to the boundary cell averages. !! Any extrapolation scheme is applied after this routine has been called. !! Therefore, boundary cells are treated as if they were local extrema. -subroutine bound_edge_values( N, h, u, edge_val, h_neglect, answer_date ) +subroutine bound_edge_values( N, h, u, edge_val, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -131,7 +132,8 @@ end subroutine average_discontinuous_edge_values !! !! For each interior edge, check whether the edge values are discontinuous. !! If so and if they are not monotonic, replace each edge value by their average. -subroutine check_discontinuous_edge_values( N, u, edge_val ) +subroutine check_discontinuous_edge_values( N, u, edge_val ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: u !< cell averages in arbitrary units [A] real, dimension(N,2), intent(inout) :: edge_val !< Cell edge values [A]; the @@ -212,7 +214,8 @@ end subroutine edge_values_explicit_h2 !! available interpolant. !! !! For this fourth-order scheme, at least four cells must exist. -subroutine edge_values_explicit_h4( N, h, u, edge_val, h_neglect, answer_date ) +subroutine edge_values_explicit_h4( N, h, u, edge_val, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -365,7 +368,8 @@ end subroutine edge_values_explicit_h4 !! i-1/2 !! !! For this fourth-order scheme, at least four cells must exist. -subroutine edge_values_explicit_h4cw( N, h, u, edge_val, h_neglect ) +subroutine edge_values_explicit_h4cw( N, h, u, edge_val, h_neglect ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -472,7 +476,8 @@ end subroutine edge_values_explicit_h4cw !! !! There are N+1 unknowns and we are able to write N-1 equations. The !! boundary conditions close the system. -subroutine edge_values_implicit_h4( N, h, u, edge_val, h_neglect, answer_date ) +subroutine edge_values_implicit_h4( N, h, u, edge_val, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -510,7 +515,10 @@ subroutine edge_values_implicit_h4( N, h, u, edge_val, h_neglect, answer_date ) tri_x ! tridiagonal system (solution vector) [A] logical :: use_2018_answers ! If true use older, less accurate expressions. - use_2018_answers = .true. ; if (present(answer_date)) use_2018_answers = (answer_date < 20190101) + ! present() not supported on GPU — answer_date always passed from callers + ! GPU: original code: + ! use_2018_answers = .true. ; if (present(answer_date)) use_2018_answers = (answer_date < 20190101) + use_2018_answers = (answer_date < 20190101) ! Loop on cells (except last one) do i = 1,N-1 @@ -633,7 +641,8 @@ end subroutine edge_values_implicit_h4 !> Determine a one-sided 4th order polynomial fit of u to the data points for the purposes of specifying !! edge values, as described in the appendix of White and Adcroft JCP 2008. -subroutine end_value_h4(dz, u, Csys) +subroutine end_value_h4(dz, u, Csys) ! GPU PORT DIAGNOSTICS + !$omp declare target real, dimension(4), intent(in) :: dz !< The thicknesses of 4 layers, starting at the edge [H]. !! The values of dz must be positive. real, dimension(4), intent(in) :: u !< The average properties of 4 layers, starting at the edge [A] @@ -778,7 +787,8 @@ end subroutine end_value_h4 !! !! There are N+1 unknowns and we are able to write N-1 equations. The !! boundary conditions close the system. -subroutine edge_slopes_implicit_h3( N, h, u, edge_slopes, h_neglect, answer_date ) +subroutine edge_slopes_implicit_h3( N, h, u, edge_slopes, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -950,7 +960,8 @@ end subroutine edge_slopes_implicit_h3 !------------------------------------------------------------------------------ !> Compute ih5 edge slopes (implicit fifth order accurate) -subroutine edge_slopes_implicit_h5( N, h, u, edge_slopes, h_neglect, answer_date ) +subroutine edge_slopes_implicit_h5( N, h, u, edge_slopes, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties in arbitrary units [A] @@ -1222,7 +1233,8 @@ end subroutine edge_slopes_implicit_h5 !! become computationally expensive if regridding is carried out !! often. Figuring out closed-form expressions for these coefficients !! on nonuniform meshes turned out to be intractable. -subroutine edge_values_implicit_h6( N, h, u, edge_val, h_neglect, answer_date ) +subroutine edge_values_implicit_h6( N, h, u, edge_val, h_neglect, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< Number of cells real, dimension(N), intent(in) :: h !< cell widths [H] real, dimension(N), intent(in) :: u !< cell average properties (size N) in arbitrary units [A] diff --git a/src/ALE/regrid_solvers.F90 b/src/ALE/regrid_solvers.F90 index 2ac2230ec4..f7a3f30110 100644 --- a/src/ALE/regrid_solvers.F90 +++ b/src/ALE/regrid_solvers.F90 @@ -18,7 +18,8 @@ module regrid_solvers !! This routine uses Gauss's algorithm to transform the system's original !! matrix into an upper triangular matrix. Back substitution yields the answer. !! The matrix A must be square, with the first index varing down the column. -subroutine solve_linear_system( A, R, X, N, answer_date ) +subroutine solve_linear_system( A, R, X, N, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< The size of the system real, dimension(N,N), intent(inout) :: A !< The matrix being inverted in arbitrary units [A] on !! input, but internally modified to become nondimensional @@ -58,8 +59,10 @@ subroutine solve_linear_system( A, R, X, N, answer_date ) ! If no pivot could be found, the system is singular. if ( .NOT. found_pivot ) then - write(0,*) ' A=',A - call MOM_error( FATAL, 'The linear system is singular !' ) + stop 1 ! Singular system + ! GPU: original code used formatted I/O + MOM_error (not GPU-safe): + ! write(0,*) ' A=',A + ! call MOM_error( FATAL, 'The linear system is singular !' ) endif ! If the pivot is in a row that is different than row i, that is if @@ -114,7 +117,8 @@ end subroutine solve_linear_system !! This routine uses Gauss's algorithm to transform the system's original !! matrix into an upper triangular matrix. Back substitution then yields the answer. !! The matrix A must be square, with the first index varing along the row. -subroutine linear_solver( N, A, R, X ) +subroutine linear_solver( N, A, R, X ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< The size of the system real, dimension(N,N), intent(inout) :: A !< The matrix being inverted in arbitrary units [A] on !! input, but internally modified to become nondimensional @@ -136,8 +140,10 @@ subroutine linear_solver( N, A, R, X ) ! pivot is in a row other than i, swap them. If no valid pivot is found, i = N+1 after this loop. do k=i,N ; if ( abs( A(i,k) ) > eps ) exit ; enddo ! end loop to find pivot if ( k > N ) then ! No pivot could be found and the system is singular. - write(0,*) ' A=',A - call MOM_error( FATAL, 'The linear system is singular !' ) + stop 1 ! Singular system + ! GPU: original code used formatted I/O + MOM_error (not GPU-safe): + ! write(0,*) ' A=',A + ! call MOM_error( FATAL, 'The linear system is singular !' ) endif ! If the pivot is in a row that is different than row i, swap those two rows, noting that both @@ -164,8 +170,10 @@ subroutine linear_solver( N, A, R, X ) enddo ! end loop on i if (A(N,N) == 0.0) then - ! no pivot could be found, and the sytem is singular - call MOM_error(FATAL, 'The final pivot in linear_solver is zero.') + ! no pivot could be found, and the system is singular + stop 1 ! Singular system + ! GPU: original code: + ! call MOM_error(FATAL, 'The final pivot in linear_solver is zero.') end if ! Solve the system by back substituting into what is now an upper-right matrix. @@ -182,7 +190,8 @@ end subroutine linear_solver !! !! This routine uses Thomas's algorithm to solve the tridiagonal system AX = R. !! (A is made up of lower, middle and upper diagonals) -subroutine solve_tridiagonal_system( Al, Ad, Au, R, X, N, answer_date ) +subroutine solve_tridiagonal_system( Al, Ad, Au, R, X, N, answer_date ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< The size of the system real, dimension(N), intent(in) :: Ad !< Matrix center diagonal in arbitrary units [A] real, dimension(N), intent(in) :: Al !< Matrix lower diagonal [A] @@ -245,7 +254,8 @@ end subroutine solve_tridiagonal_system !! lower (Al) and upper diagonals (Au) and a central diagonal Ad = Ac+Al+Au, where !! Al, Au, and Ac are all positive (or negative) definite. However when Ac is smaller than !! roundoff compared with (Al+Au), the answers are prone to inaccuracy. -subroutine solve_diag_dominant_tridiag( Al, Ac, Au, R, X, N ) +subroutine solve_diag_dominant_tridiag( Al, Ac, Au, R, X, N ) ! GPU PORT DIAGNOSTICS + !$omp declare target integer, intent(in) :: N !< The size of the system real, dimension(N), intent(in) :: Ac !< Matrix center diagonal offset from Al + Au in arbitrary units [A] real, dimension(N), intent(in) :: Al !< Matrix lower diagonal [A] diff --git a/src/framework/MOM_diag_mediator.F90 b/src/framework/MOM_diag_mediator.F90 index 71e22f0926..e77fbf8aed 100644 --- a/src/framework/MOM_diag_mediator.F90 +++ b/src/framework/MOM_diag_mediator.F90 @@ -774,13 +774,18 @@ end subroutine set_axes_info_dsamp !> set_masks_for_axes sets up the 2d and 3d masks for diagnostics using the current grid !! recorded after calling diag_update_remap_grids() -subroutine set_masks_for_axes(G, diag_cs) +subroutine set_masks_for_axes(G, diag_cs) ! GPU PORT DIAGNOSTICS type(ocean_grid_type), target, intent(in) :: G !< The ocean grid type. type(diag_ctrl), pointer :: diag_cs !< A pointer to a type with many variables !! used for diagnostics ! Local variables integer :: c, nk, i, j, k type(axes_grp), pointer :: axes => NULL(), h_axes => NULL() ! Current axes, for convenience + ! Local pointer aliases to avoid derived-type components in OpenMP map clauses + real, pointer :: mTL(:,:,:), mCuL(:,:,:), mCvL(:,:,:), mBL(:,:,:) + real, pointer :: mTi(:,:,:), mCui(:,:,:), mCvi(:,:,:), mBi(:,:,:) + real, pointer :: h_mask(:,:,:) + real, pointer :: h_remap(:,:,:) ! Local alias for remap grid thicknesses do c=1, diag_cs%num_diag_coords ! This vertical coordinate has been configured so can be used. @@ -790,17 +795,26 @@ subroutine set_masks_for_axes(G, diag_cs) axes => diag_cs%remap_axesTL(c) nk = axes%nz allocate( axes%mask3d(G%isd:G%ied,G%jsd:G%jed,nk), source=0. ) - call diag_remap_calc_hmask(diag_cs%diag_remap_cs(c), G, axes%mask3d) + mTL => axes%mask3d + !$omp target enter data map(alloc: mTL) + h_remap => diag_cs%diag_remap_cs(c)%h + !$omp target enter data map(to: h_remap) + call diag_remap_calc_hmask(diag_cs%diag_remap_cs(c), G, axes%mask3d, h_remap) + !$omp target exit data map(release: h_remap) h_axes => diag_cs%remap_axesTL(c) ! Use the h-point masks to generate the u-, v- and q- masks + h_mask => h_axes%mask3d ! Level/layer u-points in diagnostic coordinate axes => diag_cs%remap_axesCuL(c) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at u-layers') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%IsdB:G%IedB,G%jsd:G%jed,nk), source=0. ) + mCuL => axes%mask3d + !$omp target enter data map(alloc: mCuL) + !$omp target teams distribute parallel do collapse(3) do k = 1, nk ; do j=G%jsc,G%jec ; do I=G%isc-1,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i+1,j,k) > 0.) axes%mask3d(I,j,k) = 1. + if (h_mask(i,j,k) + h_mask(i+1,j,k) > 0.) mCuL(I,j,k) = 1. enddo ; enddo ; enddo ! Level/layer v-points in diagnostic coordinate @@ -808,8 +822,11 @@ subroutine set_masks_for_axes(G, diag_cs) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at v-layers') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%isd:G%ied,G%JsdB:G%JedB,nk), source=0. ) + mCvL => axes%mask3d + !$omp target enter data map(alloc: mCvL) + !$omp target teams distribute parallel do collapse(3) do k = 1, nk ; do J=G%jsc-1,G%jec ; do i=G%isc,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i,j+1,k) > 0.) axes%mask3d(i,J,k) = 1. + if (h_mask(i,j,k) + h_mask(i,j+1,k) > 0.) mCvL(i,J,k) = 1. enddo ; enddo ; enddo ! Level/layer q-points in diagnostic coordinate @@ -817,9 +834,12 @@ subroutine set_masks_for_axes(G, diag_cs) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at q-layers') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%IsdB:G%IedB,G%JsdB:G%JedB,nk), source=0. ) + mBL => axes%mask3d + !$omp target enter data map(alloc: mBL) + !$omp target teams distribute parallel do collapse(3) do k = 1, nk ; do J=G%jsc-1,G%jec ; do I=G%isc-1,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i+1,j+1,k) + & - h_axes%mask3d(i+1,j,k) + h_axes%mask3d(i,j+1,k) > 0.) axes%mask3d(I,J,k) = 1. + if (h_mask(i,j,k) + h_mask(i+1,j+1,k) + & + h_mask(i+1,j,k) + h_mask(i,j+1,k) > 0.) mBL(I,J,k) = 1. enddo ; enddo ; enddo ! Interface h-points in diagnostic coordinate (w-point) @@ -827,23 +847,30 @@ subroutine set_masks_for_axes(G, diag_cs) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at h-interfaces') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%isd:G%ied,G%jsd:G%jed,nk+1), source=0. ) - do J=G%jsc-1,G%jec+1 ; do i=G%isc-1,G%iec+1 - if (h_axes%mask3d(i,j,1) > 0.) axes%mask3d(i,J,1) = 1. + mTi => axes%mask3d + !$omp target enter data map(alloc: mTi) + !$omp target teams distribute parallel do collapse(2) private(K) + do j=G%jsc-1,G%jec+1 ; do i=G%isc-1,G%iec+1 + if (h_mask(i,j,1) > 0.) mTi(i,j,1) = 1. do K = 2, nk - if (h_axes%mask3d(i,j,k-1) + h_axes%mask3d(i,j,k) > 0.) axes%mask3d(i,J,k) = 1. + if (h_mask(i,j,k-1) + h_mask(i,j,k) > 0.) mTi(i,j,K) = 1. enddo - if (h_axes%mask3d(i,j,nk) > 0.) axes%mask3d(i,J,nk+1) = 1. + if (h_mask(i,j,nk) > 0.) mTi(i,j,nk+1) = 1. enddo ; enddo h_axes => diag_cs%remap_axesTi(c) ! Use the w-point masks to generate the u-, v- and q- masks + h_mask => h_axes%mask3d ! Interface u-points in diagnostic coordinate axes => diag_cs%remap_axesCui(c) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at u-interfaces') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%IsdB:G%IedB,G%jsd:G%jed,nk+1), source=0. ) + mCui => axes%mask3d + !$omp target enter data map(alloc: mCui) + !$omp target teams distribute parallel do collapse(3) do k = 1, nk+1 ; do j=G%jsc,G%jec ; do I=G%isc-1,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i+1,j,k) > 0.) axes%mask3d(I,j,k) = 1. + if (h_mask(i,j,k) + h_mask(i+1,j,k) > 0.) mCui(I,j,k) = 1. enddo ; enddo ; enddo ! Interface v-points in diagnostic coordinate @@ -851,8 +878,11 @@ subroutine set_masks_for_axes(G, diag_cs) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at v-interfaces') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%isd:G%ied,G%JsdB:G%JedB,nk+1), source=0. ) + mCvi => axes%mask3d + !$omp target enter data map(alloc: mCvi) + !$omp target teams distribute parallel do collapse(3) do k = 1, nk+1 ; do J=G%jsc-1,G%jec ; do i=G%isc,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i,j+1,k) > 0.) axes%mask3d(i,J,k) = 1. + if (h_mask(i,j,k) + h_mask(i,j+1,k) > 0.) mCvi(i,J,k) = 1. enddo ; enddo ; enddo ! Interface q-points in diagnostic coordinate @@ -860,10 +890,16 @@ subroutine set_masks_for_axes(G, diag_cs) call assert(axes%nz == nk, 'set_masks_for_axes: vertical size mismatch at q-interfaces') call assert(.not. associated(axes%mask3d), 'set_masks_for_axes: already associated') allocate( axes%mask3d(G%IsdB:G%IedB,G%JsdB:G%JedB,nk+1), source=0. ) - do k = 1, nk ; do J=G%jsc-1,G%jec ; do I=G%isc-1,G%iec - if (h_axes%mask3d(i,j,k) + h_axes%mask3d(i+1,j+1,k) + & - h_axes%mask3d(i+1,j,k) + h_axes%mask3d(i,j+1,k) > 0.) axes%mask3d(I,J,k) = 1. + mBi => axes%mask3d + !$omp target enter data map(alloc: mBi) + !$omp target teams distribute parallel do collapse(3) + do k = 1, nk+1 ; do J=G%jsc-1,G%jec ; do I=G%isc-1,G%iec + if (h_mask(i,j,k) + h_mask(i+1,j+1,k) + & + h_mask(i+1,j,k) + h_mask(i,j+1,k) > 0.) mBi(I,J,k) = 1. enddo ; enddo ; enddo + + ! Copy all 8 masks back to host for use by set_masks_for_axes_dsamp (maybe someone else?) + !$omp target exit data map(from: mTL, mCuL, mCvL, mBL, mTi, mCui, mCvi, mBi) endif enddo @@ -872,13 +908,16 @@ subroutine set_masks_for_axes(G, diag_cs) end subroutine set_masks_for_axes -subroutine set_masks_for_axes_dsamp(G, diag_cs) +subroutine set_masks_for_axes_dsamp(G, diag_cs) ! GPU PORT DIAGNOSTICS type(ocean_grid_type), target, intent(in) :: G !< The ocean grid type. type(diag_ctrl), pointer :: diag_cs !< A pointer to a type with many variables !! used for diagnostics ! Local variables integer :: c, dl type(axes_grp), pointer :: axes => NULL() ! Current axes, for convenience + ! Local aliases for downsample outputs to avoid derived-type references in map clauses + real, pointer :: ds_mTL(:,:,:), ds_mCuL(:,:,:), ds_mCvL(:,:,:), ds_mBL(:,:,:) + real, pointer :: ds_mTi(:,:,:), ds_mCui(:,:,:), ds_mCvi(:,:,:), ds_mBi(:,:,:) !Each downsampled axis needs both downsampled and non-downsampled mask !The downsampled mask is needed for sending out the diagnostics output via diag_manager @@ -934,6 +973,17 @@ subroutine set_masks_for_axes_dsamp(G, diag_cs) dl, G%IscB, G%JscB, G%IsdB, G%JsdB, & G%HId2%IscB, G%HId2%IecB, G%HId2%JscB, G%HId2%JecB, G%HId2%IsdB, G%HId2%IedB, G%HId2%JsdB, G%HId2%JedB) diag_cs%dsamp(dl)%remap_axesBi(c)%mask3d => axes%mask3d !set non-downsampled mask + + ! Bind local aliases to the downsampled outputs for bulk map(from:) + ds_mTL => diag_cs%dsamp(dl)%remap_axesTL(c)%dsamp(dl)%mask3d + ds_mCuL => diag_cs%dsamp(dl)%remap_axesCuL(c)%dsamp(dl)%mask3d + ds_mCvL => diag_cs%dsamp(dl)%remap_axesCvL(c)%dsamp(dl)%mask3d + ds_mBL => diag_cs%dsamp(dl)%remap_axesBL(c)%dsamp(dl)%mask3d + ds_mTi => diag_cs%dsamp(dl)%remap_axesTi(c)%dsamp(dl)%mask3d + ds_mCui => diag_cs%dsamp(dl)%remap_axesCui(c)%dsamp(dl)%mask3d + ds_mCvi => diag_cs%dsamp(dl)%remap_axesCvi(c)%dsamp(dl)%mask3d + ds_mBi => diag_cs%dsamp(dl)%remap_axesBi(c)%dsamp(dl)%mask3d + !$omp target exit data map(from: ds_mTL, ds_mCuL, ds_mCvL, ds_mBL, ds_mTi, ds_mCui, ds_mCvi, ds_mBi) enddo enddo end subroutine set_masks_for_axes_dsamp @@ -1417,7 +1467,7 @@ end subroutine post_data_2d !> Make a real 2-d array diagnostic available for averaging or output !! using a diag_type instead of an integer id. -subroutine post_data_2d_low(diag, field, diag_cs, is_static, mask) +subroutine post_data_2d_low(diag, field, diag_cs, is_static, mask) ! GPU PORT DIAGNOSTICS type(diag_type), intent(in) :: diag !< A structure describing the diagnostic to post real, target, intent(in) :: field(:,:) !< 2-d array being offered for output or averaging !! in internally scaled arbitrary units [A ~> a] @@ -1491,6 +1541,7 @@ subroutine post_data_2d_low(diag, field, diag_cs, is_static, mask) endif enddo ; enddo locfield(isv:iev,jsv:jev) = field(isv:iev,jsv:jev) * diag%conversion_factor + !$omp target enter data map(to: locfield) else locfield => field endif @@ -1507,7 +1558,10 @@ subroutine post_data_2d_low(diag, field, diag_cs, is_static, mask) if (dl > 1) then isv_o = isv ; jsv_o = jsv call downsample_diag_field(locfield, locfield_dsamp, dl, diag_cs, diag,isv,iev,jsv,jev, mask) - if ((diag%conversion_factor /= 0.) .and. (diag%conversion_factor /= 1.)) deallocate( locfield ) + if ((diag%conversion_factor /= 0.) .and. (diag%conversion_factor /= 1.)) then + !$omp target exit data map(delete: locfield) + deallocate( locfield ) + endif locfield => locfield_dsamp if (present(mask)) then call downsample_field_2d(locmask, locmask_dsamp, dl, MSK, locmask, diag_cs,diag,isv_o,jsv_o,isv,iev,jsv,jev) @@ -1571,7 +1625,7 @@ subroutine post_data_2d_low(diag, field, diag_cs, is_static, mask) end subroutine post_data_2d_low !> Make a real 3-d array diagnostic available for averaging or output. -subroutine post_data_3d(diag_field_id, field, diag_cs, is_static, mask, alt_h) +subroutine post_data_3d(diag_field_id, field, diag_cs, is_static, mask, alt_h) ! GPU PORT DIAGNOSTICS integer, intent(in) :: diag_field_id !< The id for an output variable returned by a !! previous call to register_diag_field. @@ -1643,6 +1697,7 @@ subroutine post_data_3d(diag_field_id, field, diag_cs, is_static, mask, alt_h) if (diag%v_extensive .and. .not.diag%axes%is_native) then ! The field is vertically integrated and needs to be re-gridded + ! print *, "[DIAG_REMAP_PATH] vertically_reintegrate: ", trim(diag%debug_str) if (present(mask)) then call MOM_error(FATAL,"post_data_3d: no mask for regridded field.") endif @@ -1674,6 +1729,7 @@ subroutine post_data_3d(diag_field_id, field, diag_cs, is_static, mask, alt_h) if (id_clock_diag_remap>0) call cpu_clock_end(id_clock_diag_remap) elseif (diag%axes%needs_remapping) then ! Remap this field to another vertical coordinate. + ! print *, "[DIAG_REMAP_PATH] do_remap (remapping_core_h): ", trim(diag%debug_str) if (present(mask)) then call MOM_error(FATAL,"post_data_3d: no mask for regridded field.") endif @@ -1703,6 +1759,7 @@ subroutine post_data_3d(diag_field_id, field, diag_cs, is_static, mask, alt_h) if (id_clock_diag_remap>0) call cpu_clock_end(id_clock_diag_remap) elseif (diag%axes%needs_interpolating) then ! Interpolate this field to another vertical coordinate. + ! print *, "[DIAG_REMAP_PATH] interpolate_column: ", trim(diag%debug_str) if (present(mask)) then call MOM_error(FATAL,"post_data_3d: no mask for regridded field.") endif @@ -1744,7 +1801,7 @@ end subroutine post_data_3d !> Make a real 3-d array diagnostic available for averaging or output !! using a diag_type instead of an integer id. -subroutine post_data_3d_low(diag, field, diag_cs, is_static, mask) +subroutine post_data_3d_low(diag, field, diag_cs, is_static, mask) ! GPU PORT DIAGNOSTICS type(diag_type), intent(in) :: diag !< A structure describing the diagnostic to post real, target, intent(in) :: field(:,:,:) !< 3-d array being offered for output or averaging !! in internally scaled arbitrary units [A ~> a] @@ -1837,6 +1894,7 @@ subroutine post_data_3d_low(diag, field, diag_cs, is_static, mask) locfield(i,j,k) = field(i,j,k) * diag%conversion_factor endif enddo ; enddo ; enddo + !$omp target enter data map(to: locfield) else locfield => field endif @@ -1853,7 +1911,10 @@ subroutine post_data_3d_low(diag, field, diag_cs, is_static, mask) if (dl > 1) then isv_o = isv ; jsv_o = jsv call downsample_diag_field(locfield, locfield_dsamp, dl, diag_cs, diag,isv,iev,jsv,jev, mask) - if ((diag%conversion_factor /= 0.) .and. (diag%conversion_factor /= 1.)) deallocate( locfield ) + if ((diag%conversion_factor /= 0.) .and. (diag%conversion_factor /= 1.)) then + !$omp target exit data map(delete: locfield) + deallocate( locfield ) + endif locfield => locfield_dsamp if (present(mask)) then call downsample_field_3d(locmask, locmask_dsamp, dl, MSK, locmask, diag_cs,diag,isv_o,jsv_o,isv,iev,jsv,jev) @@ -4122,13 +4183,11 @@ end subroutine diag_grid_storage_end !< Allocate and initialize the masks for downsampled diagostics in diag_cs !! The downsampled masks in the axes would later "point" to these. -subroutine downsample_diag_masks_set(G, nz, diag_cs) +subroutine downsample_diag_masks_set(G, nz, diag_cs) ! GPU PORT DIAGNOSTICS type(ocean_grid_type), target, intent(in) :: G !< The ocean grid type. integer, intent(in) :: nz !< The number of layers in the model's native grid. type(diag_ctrl), pointer :: diag_cs !< A pointer to a type with many variables !! used for diagnostics - ! Local variables - integer :: k, dl !print*,'original c extents ',G%isc,G%iec,G%jsc,G%jec !print*,'original c extents ',G%iscb,G%iecb,G%jscb,G%jecb @@ -4143,9 +4202,24 @@ subroutine downsample_diag_masks_set(G, nz, diag_cs) ! original dB-nonsym extents 1 56 1 56 ! original dB-sym extents 0 56 0 56 ! coarse d extents 1 28 1 28 + ! Local variables + integer :: i, j, k, dl + ! Local pointer aliases used in OpenMP target map clauses to avoid derived-type + ! deep-copy issues. Each is pointer-associated to the corresponding dsamp component + ! after allocation so the compiler sees a simple array pointer rather than a + ! multi-level derived-type component reference. + real, pointer :: m2dT(:,:), m2dBu(:,:), m2dCu(:,:), m2dCv(:,:) + real, pointer :: m3dTL(:,:,:), m3dBL(:,:,:), m3dCuL(:,:,:), m3dCvL(:,:,:) + real, pointer :: m3dTi(:,:,:), m3dBi(:,:,:), m3dCui(:,:,:), m3dCvi(:,:,:) + + ! Map the four 2D input masks to the device once, before the dl loop. + ! downsample_mask_2d/3d expect field_in to already be present on device. + !$omp target enter data map(to: G%mask2dT, G%mask2dBu, G%mask2dCu, G%mask2dCv) do dl=2,MAX_DSAMP_LEV - ! 2d mask + ! --- 2D masks --- + ! downsample_mask allocates field_out on the host and maps it to the device + ! with map(alloc:) but does NOT copy it back — we do that below. call downsample_mask(G%mask2dT, diag_cs%dsamp(dl)%mask2dT, dl, G%isc, G%jsc, G%isd, G%jsd, & G%HId2%isc, G%HId2%iec, G%HId2%jsc, G%HId2%jec, G%HId2%isd, G%HId2%ied, G%HId2%jsd, G%HId2%jed) call downsample_mask(G%mask2dBu, diag_cs%dsamp(dl)%mask2dBu, dl,G%IscB, G%JscB, G%IsdB, G%JsdB, & @@ -4154,29 +4228,56 @@ subroutine downsample_diag_masks_set(G, nz, diag_cs) G%HId2%IscB,G%HId2%IecB, G%HId2%jsc, G%HId2%jec,G%HId2%IsdB,G%HId2%IedB,G%HId2%jsd, G%HId2%jed) call downsample_mask(G%mask2dCv, diag_cs%dsamp(dl)%mask2dCv, dl,G %isc ,G%JscB, G%isd, G%JsdB, & G%HId2%isc ,G%HId2%iec, G%HId2%JscB,G%HId2%JecB,G%HId2%isd ,G%HId2%ied, G%HId2%JsdB,G%HId2%JedB) + + ! Bind local aliases to the (now device-mapped) 2D mask arrays so we can + ! reference them cleanly in map clauses below. + m2dT => diag_cs%dsamp(dl)%mask2dT + m2dBu => diag_cs%dsamp(dl)%mask2dBu + m2dCu => diag_cs%dsamp(dl)%mask2dCu + m2dCv => diag_cs%dsamp(dl)%mask2dCv + + ! --- 3D layer masks (nz levels) --- ! 3d native masks are needed by diag_manager but the native variables - ! can only be masked 2d - for ocean points, all layers exists. + ! can only be masked 2d - for ocean points, all layers exist. allocate(diag_cs%dsamp(dl)%mask3dTL(G%HId2%isd:G%HId2%ied,G%HId2%jsd:G%HId2%jed,1:nz)) allocate(diag_cs%dsamp(dl)%mask3dBL(G%HId2%IsdB:G%HId2%IedB,G%HId2%JsdB:G%HId2%JedB,1:nz)) allocate(diag_cs%dsamp(dl)%mask3dCuL(G%HId2%IsdB:G%HId2%IedB,G%HId2%jsd:G%HId2%jed,1:nz)) allocate(diag_cs%dsamp(dl)%mask3dCvL(G%HId2%isd:G%HId2%ied,G%HId2%JsdB:G%HId2%JedB,1:nz)) - do k=1,nz - diag_cs%dsamp(dl)%mask3dTL(:,:,k) = diag_cs%dsamp(dl)%mask2dT(:,:) - diag_cs%dsamp(dl)%mask3dBL(:,:,k) = diag_cs%dsamp(dl)%mask2dBu(:,:) - diag_cs%dsamp(dl)%mask3dCuL(:,:,k) = diag_cs%dsamp(dl)%mask2dCu(:,:) - diag_cs%dsamp(dl)%mask3dCvL(:,:,k) = diag_cs%dsamp(dl)%mask2dCv(:,:) - enddo + m3dTL => diag_cs%dsamp(dl)%mask3dTL + m3dBL => diag_cs%dsamp(dl)%mask3dBL + m3dCuL => diag_cs%dsamp(dl)%mask3dCuL + m3dCvL => diag_cs%dsamp(dl)%mask3dCvL + !$omp target enter data map(alloc: m3dTL, m3dBL, m3dCuL, m3dCvL) + call copy_2d_into_3d(m2dT, m3dTL) + call copy_2d_into_3d(m2dBu, m3dBL) + call copy_2d_into_3d(m2dCu, m3dCuL) + call copy_2d_into_3d(m2dCv, m3dCvL) + + ! --- 3D interface masks (nz+1 levels) --- allocate(diag_cs%dsamp(dl)%mask3dTi(G%HId2%isd:G%HId2%ied,G%HId2%jsd:G%HId2%jed,1:nz+1)) allocate(diag_cs%dsamp(dl)%mask3dBi(G%HId2%IsdB:G%HId2%IedB,G%HId2%JsdB:G%HId2%JedB,1:nz+1)) allocate(diag_cs%dsamp(dl)%mask3dCui(G%HId2%IsdB:G%HId2%IedB,G%HId2%jsd:G%HId2%jed,1:nz+1)) allocate(diag_cs%dsamp(dl)%mask3dCvi(G%HId2%isd:G%HId2%ied,G%HId2%JsdB:G%HId2%JedB,1:nz+1)) - do k=1,nz+1 - diag_cs%dsamp(dl)%mask3dTi(:,:,k) = diag_cs%dsamp(dl)%mask2dT(:,:) - diag_cs%dsamp(dl)%mask3dBi(:,:,k) = diag_cs%dsamp(dl)%mask2dBu(:,:) - diag_cs%dsamp(dl)%mask3dCui(:,:,k) = diag_cs%dsamp(dl)%mask2dCu(:,:) - diag_cs%dsamp(dl)%mask3dCvi(:,:,k) = diag_cs%dsamp(dl)%mask2dCv(:,:) - enddo + m3dTi => diag_cs%dsamp(dl)%mask3dTi + m3dBi => diag_cs%dsamp(dl)%mask3dBi + m3dCui => diag_cs%dsamp(dl)%mask3dCui + m3dCvi => diag_cs%dsamp(dl)%mask3dCvi + !$omp target enter data map(alloc: m3dTi, m3dBi, m3dCui, m3dCvi) + call copy_2d_into_3d(m2dT, m3dTi) + call copy_2d_into_3d(m2dBu, m3dBi) + call copy_2d_into_3d(m2dCu, m3dCui) + call copy_2d_into_3d(m2dCv, m3dCvi) + + ! Copy all output masks back to the host. The masks are consumed by host-side + ! routines (post_data_*_low via send_data_infra), so we need host copies. + !$omp target exit data map(from: m2dT, m2dBu, m2dCu, m2dCv) + !$omp target exit data map(from: m3dTL, m3dBL, m3dCuL, m3dCvL) + !$omp target exit data map(from: m3dTi, m3dBi, m3dCui, m3dCvi) enddo + + ! Release the input masks from the device now that all dl levels are done. + !$omp target exit data map(release: G%mask2dT, G%mask2dBu, G%mask2dCu, G%mask2dCv) + end subroutine downsample_diag_masks_set !> Get the diagnostics-compute indices (to be passed to send_data) based on the shape of @@ -4370,7 +4471,7 @@ end subroutine downsample_diag_field_2d !> This subroutine allocates and computes a down sampled 3d array given an input array !! The down sample method is based on the "cell_methods" for the diagnostics as explained !! in the above table -subroutine downsample_field_3d(field_in, field_out, dl, method, mask, diag_cs, diag,isv_o,jsv_o,isv_d,iev_d,jsv_d,jev_d) +subroutine downsample_field_3d(field_in, field_out, dl, method, mask, diag_cs, diag,isv_o,jsv_o,isv_d,iev_d,jsv_d,jev_d) ! GPU PORT DIAGNOSTICS real, dimension(:,:,:), pointer :: field_in !< Original field to be downsampled in arbitrary units [A ~> a] real, dimension(:,:,:), allocatable :: field_out !< Downsampled field in the same arbtrary units [A ~> a] integer, intent(in) :: dl !< Level of down sampling @@ -4397,15 +4498,14 @@ subroutine downsample_field_3d(field_in, field_out, dl, method, mask, diag_cs, d real :: eps_vol ! A negligibly small volume or mass [H L2 ~> m3 or kg] real :: eps_area ! A negligibly small area [L2 ~> m2] real :: eps_face ! A negligibly small face area [H L ~> m2 or kg m-1] + ! Local pointer aliases for device-resident grid arrays + real, pointer :: areaT(:,:), dyCu(:,:), dxCv(:,:), h_ptr(:,:,:) ks = 1 ; ke = size(field_in,3) eps_face = 1.0e-20 * diag_cs%G%US%m_to_L * diag_cs%GV%m_to_H eps_area = 1.0e-20 * diag_cs%G%US%m_to_L**2 eps_vol = 1.0e-20 * diag_cs%G%US%m_to_L**2 * diag_cs%GV%m_to_H - ! Allocate the down sampled field on the down sampled data domain -! allocate(field_out(diag_cs%dsamp(dl)%isd:diag_cs%dsamp(dl)%ied,diag_cs%dsamp(dl)%jsd:diag_cs%dsamp(dl)%jed,ks:ke)) -! allocate(field_out(1:size(field_in,1)/dl,1:size(field_in,2)/dl,ks:ke)) f_in1 = size(field_in,1) f_in2 = size(field_in,2) f1 = f_in1/dl @@ -4417,121 +4517,145 @@ subroutine downsample_field_3d(field_in, field_out, dl, method, mask, diag_cs, d endif allocate(field_out(1:f1,1:f2,ks:ke)) - ! Fill the down sampled field on the down sampled diagnostics (almost always compuate) domain + ! Bind local aliases to device-resident grid arrays (mapped in MOM.F90 init). + areaT => diag_cs%G%areaT + dyCu => diag_cs%G%dyCu + dxCv => diag_cs%G%dxCv + h_ptr => diag_cs%h + + !$omp target enter data map(alloc: field_out) + + ! Fill the down sampled field on the down sampled diagnostics (almost always compute) domain !### The averaging used here is not rotationally invariant. if (method == MMM) then + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 -! do ii=i0,i0+dl-1 ; do jj=j0,j0+dl-1 !This seems to be faster!!!! - weight = mask(ii,jj,k) * diag_cs%G%areaT(ii,jj) * diag_cs%h(ii,jj,k) + weight = mask(ii,jj,k) * areaT(ii,jj) * h_ptr(ii,jj,k) total_weight = total_weight + weight - ave = ave+field_in(ii,jj,k) * weight + ave = ave + field_in(ii,jj,k) * weight enddo ; enddo - field_out(i,j,k) = ave/(total_weight + eps_vol) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j,k) = ave / (total_weight + eps_vol) enddo ; enddo ; enddo elseif (method == SSS) then !e.g., volcello + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 - weight = mask(ii,jj,k) - ave = ave+field_in(ii,jj,k)*weight + ave = ave + field_in(ii,jj,k) * mask(ii,jj,k) enddo ; enddo - field_out(i,j,k) = ave !Masked Sum (total_weight=1) + field_out(i,j,k) = ave !Masked Sum (total_weight=1) enddo ; enddo ; enddo elseif (method == MMP .or. method == MMS) then !e.g., T_advection_xy + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 -! do ii=i0,i0+dl-1 ; do jj=j0,j0+dl-1 - weight = mask(ii,jj,k) * diag_cs%G%areaT(ii,jj) + weight = mask(ii,jj,k) * areaT(ii,jj) total_weight = total_weight + weight - ave = ave+field_in(ii,jj,k)*weight + ave = ave + field_in(ii,jj,k) * weight enddo ; enddo - field_out(i,j,k) = ave / (total_weight+eps_area) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j,k) = ave / (total_weight + eps_area) enddo ; enddo ; enddo elseif (method == PMM) then + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 - ii=i0 + ii = i0 do jj=j0,j0+dl-1 - weight = mask(ii,jj,k) * diag_cs%G%dyCu(ii,jj) * diag_cs%h(ii,jj,k) - total_weight = total_weight +weight - ave = ave+field_in(ii,jj,k)*weight + weight = mask(ii,jj,k) * dyCu(ii,jj) * h_ptr(ii,jj,k) + total_weight = total_weight + weight + ave = ave + field_in(ii,jj,k) * weight enddo - field_out(i,j,k) = ave/(total_weight+eps_face) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j,k) = ave / (total_weight + eps_face) enddo ; enddo ; enddo elseif (method == PSS) then !e.g. umo + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 - ii=i0 + ii = i0 do jj=j0,j0+dl-1 - weight = mask(ii,jj,k) - ave = ave+field_in(ii,jj,k)*weight + ave = ave + field_in(ii,jj,k) * mask(ii,jj,k) enddo - field_out(i,j,k) = ave !Masked Sum (total_weight=1) + field_out(i,j,k) = ave !Masked Sum (total_weight=1) enddo ; enddo ; enddo elseif (method == SPS) then !e.g. vmo + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 - jj=j0 + jj = j0 do ii=i0,i0+dl-1 - weight = mask(ii,jj,k) - ave = ave+field_in(ii,jj,k)*weight + ave = ave + field_in(ii,jj,k) * mask(ii,jj,k) enddo - field_out(i,j,k) = ave !Masked Sum (total_weight=1) + field_out(i,j,k) = ave !Masked Sum (total_weight=1) enddo ; enddo ; enddo elseif (method == MPM) then + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 - jj=j0 + jj = j0 do ii=i0,i0+dl-1 - weight = mask(ii,jj,k) * diag_cs%G%dxCv(ii,jj) * diag_cs%h(ii,jj,k) + weight = mask(ii,jj,k) * dxCv(ii,jj) * h_ptr(ii,jj,k) total_weight = total_weight + weight - ave = ave+field_in(ii,jj,k)*weight + ave = ave + field_in(ii,jj,k) * weight enddo - field_out(i,j,k) = ave/(total_weight+eps_face) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j,k) = ave / (total_weight + eps_face) enddo ; enddo ; enddo elseif (method == MSK) then !The input field is a mask, subsample - field_out(:,:,:) = 0.0 + !$omp target teams distribute parallel do collapse(3) + do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d + field_out(i,j,k) = 0.0 + enddo ; enddo ; enddo + !$omp target teams distribute parallel do collapse(3) & + !$omp private(i0,j0,ave,ii,jj) do k=ks,ke ; do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 - ave = ave+field_in(ii,jj,k) + ave = ave + field_in(ii,jj,k) enddo ; enddo - if (ave > 0.0) field_out(i,j,k)=1.0 + if (ave > 0.0) field_out(i,j,k) = 1.0 enddo ; enddo ; enddo else write (mesg,*) " unknown sampling method: ",method call MOM_error(FATAL, "downsample_field_3d: "//trim(mesg)//" "//trim(diag%debug_str)) endif + !$omp target exit data map(from: field_out) + end subroutine downsample_field_3d !> This subroutine allocates and computes a down sampled 2d array given an input array !! The down sample method is based on the "cell_methods" for the diagnostics as explained !! in the above table -subroutine downsample_field_2d(field_in, field_out, dl, method, mask, diag_cs, diag, & +subroutine downsample_field_2d(field_in, field_out, dl, method, mask, diag_cs, diag, & ! GPU PORT DIAGNOSTICS isv_o, jsv_o, isv_d, iev_d, jsv_d, jev_d) real, dimension(:,:), pointer :: field_in !< Original field to be downsampled in arbitrary units [A ~> a] real, dimension(:,:), allocatable :: field_out !< Downsampled field in the same arbtrary units [A ~> a] @@ -4554,14 +4678,12 @@ subroutine downsample_field_2d(field_in, field_out, dl, method, mask, diag_cs, d real :: total_weight ! The sum of weights contributing to a point [nondim] or [L2 ~> m2] real :: eps_area ! A negligibly small area [L2 ~> m2] real :: eps_len ! A negligibly small horizontal length [L ~> m] + ! Local pointer aliases for device-resident grid arrays + real, pointer :: areaT(:,:), dyCu(:,:), dxCv(:,:) eps_len = 1.0e-20 * diag_cs%G%US%m_to_L eps_area = 1.0e-20 * diag_cs%G%US%m_to_L**2 - ! Allocate the down sampled field on the down sampled data domain -! allocate(field_out(diag_cs%dsamp(dl)%isd:diag_cs%dsamp(dl)%ied,diag_cs%dsamp(dl)%jsd:diag_cs%dsamp(dl)%jed)) -! allocate(field_out(1:size(field_in,1)/dl,1:size(field_in,2)/dl)) - ! Fill the down sampled field on the down sampled diagnostics (almost always compuate) domain f_in1 = size(field_in,1) f_in2 = size(field_in,2) f1 = f_in1/dl @@ -4573,106 +4695,127 @@ subroutine downsample_field_2d(field_in, field_out, dl, method, mask, diag_cs, d endif allocate(field_out(1:f1,1:f2)) + ! Bind local aliases to device-resident grid arrays (mapped in MOM.F90 init). + areaT => diag_cs%G%areaT + dyCu => diag_cs%G%dyCu + dxCv => diag_cs%G%dxCv + + !$omp target enter data map(alloc: field_out) + if (method == MMP) then + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 -! do ii=i0,i0+dl-1 ; do jj=j0,j0+dl-1 - weight = mask(ii,jj)*diag_cs%G%areaT(ii,jj) + weight = mask(ii,jj) * areaT(ii,jj) total_weight = total_weight + weight - ave = ave+field_in(ii,jj)*weight + ave = ave + field_in(ii,jj) * weight enddo ; enddo - field_out(i,j) = ave/(total_weight + eps_area) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j) = ave / (total_weight + eps_area) enddo ; enddo elseif (method == SSP) then ! e.g., T_dfxy_cont_tendency_2d + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,ii,jj) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 -! do ii=i0,i0+dl-1 ; do jj=j0,j0+dl-1 - weight = mask(ii,jj) - ave = ave+field_in(ii,jj)*weight + ave = ave + field_in(ii,jj) * mask(ii,jj) enddo ; enddo field_out(i,j) = ave !Masked Sum (total_weight=1) enddo ; enddo elseif (method == PSP) then ! e.g., umo_2d + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,ii,jj) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 - ii=i0 + ii = i0 do jj=j0,j0+dl-1 - weight = mask(ii,jj) - ave = ave+field_in(ii,jj)*weight + ave = ave + field_in(ii,jj) * mask(ii,jj) enddo field_out(i,j) = ave !Masked Sum (total_weight=1) enddo ; enddo elseif (method == SPP) then ! e.g., vmo_2d + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,ii,jj) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 - jj=j0 + jj = j0 do ii=i0,i0+dl-1 - weight = mask(ii,jj) - ave = ave+field_in(ii,jj)*weight + ave = ave + field_in(ii,jj) * mask(ii,jj) enddo field_out(i,j) = ave !Masked Sum (total_weight=1) enddo ; enddo elseif (method == PMP) then + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 - ii=i0 + ii = i0 do jj=j0,j0+dl-1 - weight = mask(ii,jj) * diag_cs%G%dyCu(ii,jj)!*diag_cs%h(ii,jj,1) !Niki? - total_weight = total_weight +weight - ave = ave+field_in(ii,jj)*weight + weight = mask(ii,jj) * dyCu(ii,jj) + total_weight = total_weight + weight + ave = ave + field_in(ii,jj) * weight enddo - field_out(i,j) = ave/(total_weight+eps_len) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j) = ave / (total_weight + eps_len) enddo ; enddo elseif (method == MPP) then + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,total_weight,ii,jj,weight) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 total_weight = 0.0 - jj=j0 + jj = j0 do ii=i0,i0+dl-1 - weight = mask(ii,jj)* diag_cs%G%dxCv(ii,jj)!*diag_cs%h(ii,jj,1) !Niki? - total_weight = total_weight +weight - ave = ave+field_in(ii,jj)*weight + weight = mask(ii,jj) * dxCv(ii,jj) + total_weight = total_weight + weight + ave = ave + field_in(ii,jj) * weight enddo - field_out(i,j) = ave/(total_weight+eps_len) !Avoid zero mask at all aggregating cells where ave=0.0 + field_out(i,j) = ave / (total_weight + eps_len) enddo ; enddo elseif (method == MSK) then !The input field is a mask, subsample - field_out(:,:) = 0.0 + !$omp target teams distribute parallel do collapse(2) + do j=jsv_d,jev_d ; do i=isv_d,iev_d + field_out(i,j) = 0.0 + enddo ; enddo + !$omp target teams distribute parallel do collapse(2) & + !$omp private(i0,j0,ave,ii,jj) do j=jsv_d,jev_d ; do i=isv_d,iev_d i0 = isv_o+dl*(i-isv_d) j0 = jsv_o+dl*(j-jsv_d) ave = 0.0 do jj=j0,j0+dl-1 ; do ii=i0,i0+dl-1 - ave = ave+field_in(ii,jj) + ave = ave + field_in(ii,jj) enddo ; enddo - if (ave > 0.0) field_out(i,j)=1.0 + if (ave > 0.0) field_out(i,j) = 1.0 enddo ; enddo else write (mesg,*) " unknown sampling method: ",method call MOM_error(FATAL, "downsample_field_2d: "//trim(mesg)//" "//trim(diag%debug_str)) endif + !$omp target exit data map(from: field_out) + end subroutine downsample_field_2d !> Allocate and compute the 2d down sampled mask !! The masks are down sampled based on a minority rule, i.e., a coarse cell is open (1) !! if at least one of the sub-cells are open, otherwise it's closed (0) -subroutine downsample_mask_2d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_o, & +subroutine downsample_mask_2d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_o, & ! GPU PORT DIAGNOSTICS isc_d, iec_d, jsc_d, jec_d, isd_d, ied_d, jsd_d, jed_d) integer, intent(in) :: isd_o !< Original data domain i-start index integer, intent(in) :: jsd_o !< Original data domain j-start index @@ -4693,8 +4836,17 @@ subroutine downsample_mask_2d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_ integer :: i,j,ii,jj,i0,j0 real :: tot_non_zero ! The sum of values in the down-scaled cell [A] ! down sampled mask = 0 unless the mask value of one of the down sampling cells is 1 + + ! field_in should already be on the device allocate(field_out(isd_d:ied_d,jsd_d:jed_d)) - field_out(:,:) = 0.0 + !$omp target enter data map(alloc:field_out) + + !$omp target teams distribute parallel do collapse(2) + do j = jsd_d,jed_d ; do i = isd_d,ied_d + field_out(i,j) = 0.0 + enddo ; enddo + + !$omp target teams distribute parallel do collapse(2) private(i0,j0,tot_non_zero,ii,jj) do j=jsc_d,jec_d ; do i=isc_d,iec_d i0 = isc_o+dl*(i-isc_d) j0 = jsc_o+dl*(j-jsc_d) @@ -4704,12 +4856,13 @@ subroutine downsample_mask_2d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_ enddo ; enddo if (tot_non_zero > 0.0) field_out(i,j)=1.0 enddo ; enddo + end subroutine downsample_mask_2d !> Allocate and compute the 3d down sampled mask !! The masks are down sampled based on a minority rule, i.e., a coarse cell is open (1) !! if at least one of the sub-cells are open, otherwise it's closed (0) -subroutine downsample_mask_3d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_o, & +subroutine downsample_mask_3d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_o, & ! GPU PORT DIAGNOSTICS isc_d, iec_d, jsc_d, jec_d, isd_d, ied_d, jsd_d, jed_d) integer, intent(in) :: isd_o !< Original data domain i-start index integer, intent(in) :: jsd_o !< Original data domain j-start index @@ -4730,9 +4883,19 @@ subroutine downsample_mask_3d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_ integer :: i,j,ii,jj,i0,j0,k,ks,ke real :: tot_non_zero ! The sum of values in the down-scaled cell [A] ! down sampled mask = 0 unless the mask value of one of the down sampling cells is 1 + ! field_in should be on the device already ks = lbound(field_in,3) ; ke = ubound(field_in,3) + allocate(field_out(isd_d:ied_d,jsd_d:jed_d,ks:ke)) - field_out(:,:,:) = 0.0 + + !$omp target enter data map(alloc:field_out) + + !$omp target teams distribute parallel do collapse(3) + do k=ks,ke ; do j = jsd_d,jed_d ; do i = isd_d,ied_d + field_out(i,j,k) = 0.0 + enddo ; enddo ; enddo + + !$omp target teams distribute parallel do collapse(3) private(i0,j0,tot_non_zero,ii,jj) do k=ks,ke ; do j=jsc_d,jec_d ; do i=isc_d,iec_d i0 = isc_o+dl*(i-isc_d) j0 = jsc_o+dl*(j-jsc_d) @@ -4742,6 +4905,9 @@ subroutine downsample_mask_3d(field_in, field_out, dl, isc_o, jsc_o, isd_o, jsd_ enddo ; enddo if (tot_non_zero > 0.0) field_out(i,j,k)=1.0 enddo ; enddo ; enddo + + !$omp target exit data map(delete:field_out) + end subroutine downsample_mask_3d !> Fakes a register of a diagnostic to find out if an obsolete @@ -4760,4 +4926,23 @@ logical function found_in_diagtable(diag, varName) end function found_in_diagtable +!> Broadcast a 2D mask into every k-level of a pre-allocated, device-mapped 3D array. +!! Both arrays must already be present on the device before this routine is called +!! (the caller is responsible for map(alloc:) of field_3d and the 2D source must +!! already be device-resident via a prior enter data map). +subroutine copy_2d_into_3d(field_2d, field_3d) + real, intent(in) :: field_2d(:,:) !< 2D source mask [nondim] + real, intent(inout) :: field_3d(:,:,:) !< 3D destination mask [nondim] + ! Locals + integer :: i, j, k + !$omp target teams distribute parallel do collapse(3) + do k = 1, size(field_3d, 3) + do j = 1, size(field_3d, 2) + do i = 1, size(field_3d, 1) + field_3d(i,j,k) = field_2d(i,j) + enddo + enddo + enddo +end subroutine copy_2d_into_3d + end module MOM_diag_mediator diff --git a/src/framework/MOM_diag_remap.F90 b/src/framework/MOM_diag_remap.F90 index a15b815214..4abd8bdbb8 100644 --- a/src/framework/MOM_diag_remap.F90 +++ b/src/framework/MOM_diag_remap.F90 @@ -44,6 +44,7 @@ module MOM_diag_remap use MOM_EOS, only : EOS_type use MOM_remapping, only : remapping_CS, initialize_remapping, remapping_core_h use MOM_remapping, only : interpolate_column, reintegrate_column +use MOM_remapping, only : NK_GPU_MAX use MOM_regridding, only : regridding_CS, initialize_regridding, end_regridding use MOM_regridding, only : set_regrid_params, get_regrid_size use MOM_regridding, only : getCoordinateInterfaces, set_h_neglect, set_dz_neglect @@ -95,6 +96,7 @@ module MOM_diag_remap integer :: interface_axes_id = 0 !< Vertical axes id for remapping at interfaces integer :: layer_axes_id = 0 !< Vertical axes id for remapping on layers logical :: om4_remap_via_sub_cells !< Use the OM4-era ramap_via_sub_cells + integer :: answer_date !< The vintage of the order of arithmetic and expressions !! to use for remapping. Values below 20190101 recover !! the answers from 2018, while higher values use more @@ -102,6 +104,16 @@ module MOM_diag_remap end type diag_remap_ctrl +! Module-level persistent GPU work buffers for vertically_interpolate_field. +! Must be module-level (not type members) because nvfortran can't resolve +! derived-type component addresses inside omp target regions. +real, dimension(:,:,:), allocatable :: gpu_h_work !< Source thickness work buffer [H] +real, dimension(:,:,:), allocatable :: gpu_htgt_work !< Target thickness work buffer [H] +real, dimension(:,:,:), allocatable :: gpu_fld_work !< Field work buffer [A] +real, dimension(:,:,:), allocatable :: gpu_out_work !< Output work buffer [A] +real, dimension(:,:), allocatable :: gpu_mask_work !< Mask work buffer [nondim] +logical :: gpu_bufs_ready = .false. !< Whether module-level GPU buffers are mapped + contains !> Initialize a diagnostic remapping type with the given vertical coordinate. @@ -266,6 +278,7 @@ function diag_remap_axes_configured(remap_cs) !! coordinates then technically we should also regenerate the !! target grid whenever T/S change. subroutine diag_remap_update(remap_cs, G, GV, US, h, T, S, eqn_of_state, h_target) +! GPU PORT DIAGNOSTICS type(diag_remap_ctrl), intent(inout) :: remap_cs !< Diagnostic coordinate control structure type(ocean_grid_type), pointer :: G !< The ocean's grid type type(verticalGrid_type), intent(in) :: GV !< ocean vertical grid structure @@ -313,6 +326,8 @@ subroutine diag_remap_update(remap_cs, G, GV, US, h, T, S, eqn_of_state, h_targe if (.not. remap_cs%initialized) then ! Initialize remapping and regridding on the first call + ! print *, "[DIAG_REMAP_INIT] coord=", trim(remap_cs%diag_coord_name), & + ! " scheme=PPM_IH4 om4_subcells=", remap_cs%om4_remap_via_sub_cells call initialize_remapping(remap_cs%remap_cs, 'PPM_IH4', boundary_extrapolation=.false., & om4_remap_via_sub_cells=remap_cs%om4_remap_via_sub_cells, & answer_date=remap_cs%answer_date, & @@ -412,8 +427,9 @@ subroutine diag_remap_do_remap(remap_cs, G, GV, US, h, staggered_in_x, staggered end subroutine diag_remap_do_remap !> The internal routine to remap a diagnostic field to an alternative vertical grid. -subroutine do_remap(remap_cs, G, GV, US, isdf, jsdf, h, staggered_in_x, staggered_in_y, & +subroutine do_remap(remap_cs, G, GV, US, isdf, jsdf, h, staggered_in_x, staggered_in_y, & field, remapped_field, mask) + ! GPU PORT DIAGNOSTICS type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure type(ocean_grid_type), intent(in) :: G !< Ocean grid structure type(verticalGrid_type), intent(in) :: GV !< ocean vertical grid structure @@ -433,111 +449,242 @@ subroutine do_remap(remap_cs, G, GV, US, isdf, jsdf, h, staggered_in_x, staggere optional, intent(in) :: mask !< A mask for the field [nondim] ! Local variables - real, dimension(remap_cs%nz) :: h_dest ! Destination thicknesses [H ~> m or kg m-2] or [Z ~> m] - real, dimension(size(h,3)) :: h_src ! A column of source thicknesses [H ~> m or kg m-2] or [Z ~> m] + type(remapping_CS) :: rcs ! Local copy of remapping CS for device mapping + real :: h_src(NK_GPU_MAX), h_dest(NK_GPU_MAX) ! Thread-private column thickness arrays + real :: fld_loc(NK_GPU_MAX), out_loc(NK_GPU_MAX) ! Thread-private field/output arrays integer :: nz_src, nz_dest ! The number of layers on the native and remapped grids - integer :: i, j ! Grid index + integer :: i, j, k ! Grid index + logical :: have_mask nz_src = size(field,3) nz_dest = remap_cs%nz - remapped_field(:,:,:) = 0. + have_mask = present(mask) + + ! Local copy of remapping control structure for device mapping + rcs = remap_cs%remap_cs + + ! make sure buffers are there + call ensure_gpu_buffers(remap_cs, G, nz_src) + + ! Copy source data into persistent host-side buffers, then sync to device + gpu_h_work(G%isd:G%ied, G%jsd:G%jed, :) = h(:,:,:) + gpu_htgt_work(G%isd:G%ied, G%jsd:G%jed, :) = remap_cs%h(:,:,:) + do k=1,nz_src + do j=jsdf,jsdf+size(field,2)-1 ; do i=isdf,isdf+size(field,1)-1 + gpu_fld_work(i,j,k) = field(i,j,k) + enddo ; enddo + enddo + if (have_mask) then + do j=jsdf,jsdf+size(mask,2)-1 ; do i=isdf,isdf+size(mask,1)-1 + gpu_mask_work(i,j) = mask(i,j) + enddo ; enddo + endif + + !$omp target update to(gpu_h_work, gpu_htgt_work, gpu_fld_work) + if (have_mask) then + !$omp target update to(gpu_mask_work) + endif + + ! Zero-fill the output on device + !$omp target teams distribute parallel do collapse(3) + do k=1,nz_dest + do j=lbound(gpu_out_work,2),ubound(gpu_out_work,2) + do i=lbound(gpu_out_work,1),ubound(gpu_out_work,1) + gpu_out_work(i,j,k) = 0. + enddo + enddo + enddo if (staggered_in_x .and. .not. staggered_in_y) then ! U-points - if (present(mask)) then - do j=G%jsc,G%jec ; do I=G%IscB,G%IecB ; if (mask(I,j) > 0.) then - h_src(:) = 0.5 * (h(i,j,:) + h(i+1,j,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i+1,j,:)) - call remapping_core_h(remap_cs%remap_cs, nz_src, h_src(:), field(I,j,:), & - nz_dest, h_dest(:), remapped_field(I,j,:)) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) + do j=G%jsc,G%jec ; do I=G%IscB,G%IecB + if (gpu_mask_work(I,j) > 0.0) then + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i+1,j,k)) + fld_loc(k) = gpu_fld_work(I,j,k) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i+1,j,k)) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(I,j,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) do j=G%jsc,G%jec ; do I=G%IscB,G%IecB - h_src(:) = 0.5 * (h(i,j,:) + h(i+1,j,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i+1,j,:)) - call remapping_core_h(remap_cs%remap_cs, nz_src, h_src(:), field(I,j,:), & - nz_dest, h_dest(:), remapped_field(I,j,:)) + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i+1,j,k)) + fld_loc(k) = gpu_fld_work(I,j,k) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i+1,j,k)) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(I,j,k) = out_loc(k) + enddo enddo ; enddo endif elseif (staggered_in_y .and. .not. staggered_in_x) then ! V-points - if (present(mask)) then - do J=G%jscB,G%jecB ; do i=G%isc,G%iec ; if (mask(i,j) > 0.) then - h_src(:) = 0.5 * (h(i,j,:) + h(i,j+1,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i,j+1,:)) - call remapping_core_h(remap_cs%remap_cs, nz_src, h_src(:), field(i,J,:), & - nz_dest, h_dest(:), remapped_field(i,J,:)) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) + do J=G%jscB,G%jecB ; do i=G%isc,G%iec + if (gpu_mask_work(i,J) > 0.0) then + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i,j+1,k)) + fld_loc(k) = gpu_fld_work(i,J,k) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i,j+1,k)) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(i,J,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) do J=G%jscB,G%jecB ; do i=G%isc,G%iec - h_src(:) = 0.5 * (h(i,j,:) + h(i,j+1,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i,j+1,:)) - call remapping_core_h(remap_cs%remap_cs, nz_src, h_src(:), field(i,J,:), & - nz_dest, h_dest(:), remapped_field(i,J,:)) + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i,j+1,k)) + fld_loc(k) = gpu_fld_work(i,J,k) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i,j+1,k)) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(i,J,k) = out_loc(k) + enddo enddo ; enddo endif elseif ((.not. staggered_in_x) .and. (.not. staggered_in_y)) then ! H-points - if (present(mask)) then - do j=G%jsc,G%jec ; do i=G%isc,G%iec ; if (mask(i,j) > 0.) then - call remapping_core_h(remap_cs%remap_cs, nz_src, h(i,j,:), field(i,j,:), & - nz_dest, remap_cs%h(i,j,:), remapped_field(i,j,:)) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) + do j=G%jsc,G%jec ; do i=G%isc,G%iec + if (gpu_mask_work(i,j) > 0.0) then + do k=1,nz_src + h_src(k) = gpu_h_work(i,j,k) + fld_loc(k) = gpu_fld_work(i,j,k) + enddo + do k=1,nz_dest + h_dest(k) = gpu_htgt_work(i,j,k) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(i,j,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) & + !$omp map(to: rcs) private(h_src, h_dest, fld_loc, out_loc) do j=G%jsc,G%jec ; do i=G%isc,G%iec - call remapping_core_h(remap_cs%remap_cs, nz_src, h(i,j,:), field(i,j,:), & - nz_dest, remap_cs%h(i,j,:), remapped_field(i,j,:)) + do k=1,nz_src + h_src(k) = gpu_h_work(i,j,k) + fld_loc(k) = gpu_fld_work(i,j,k) + enddo + do k=1,nz_dest + h_dest(k) = gpu_htgt_work(i,j,k) + enddo + call remapping_core_h(rcs, nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc) + do k=1,nz_dest + gpu_out_work(i,j,k) = out_loc(k) + enddo enddo ; enddo endif else call assert(.false., 'diag_remap_do_remap: Unsupported axis combination') endif + ! Copy result back from device + !$omp target update from(gpu_out_work) + do k=1,nz_dest + do j=jsdf,jsdf+size(remapped_field,2)-1 + do i=isdf,isdf+size(remapped_field,1)-1 + remapped_field(i,j,k) = gpu_out_work(i,j,k) + enddo + enddo + enddo + end subroutine do_remap -!> Calculate masks for target grid -subroutine diag_remap_calc_hmask(remap_cs, G, mask) +!> Calculate masks for target grid. +!! Both mask and h must already be present on the device (via prior enter data map). +!! G%mask2dT must also be device-resident. +subroutine diag_remap_calc_hmask(remap_cs, G, mask, h) +! GPU PORT DIAGNOSTICS type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure type(ocean_grid_type), intent(in) :: G !< Ocean grid structure real, dimension(G%isd:,G%jsd:,:), & intent(out) :: mask !< h-point mask for target grid [nondim] + real, dimension(G%isd:,G%jsd:,:), & + intent(in) :: h !< Remap grid thicknesses [H ~> m or kg m-2] or [Z ~> m] ! Local variables - real, dimension(remap_cs%nz) :: h_dest ! Destination thicknesses [H ~> m or kg m-2] or [Z ~> m] - integer :: i, j, k + integer :: i, j, k, nz logical :: mask_vanished_layers real :: h_tot ! Sum of all thicknesses [H ~> m or kg m-2] or [Z ~> m] real :: h_err ! An estimate of a negligible thickness [H ~> m or kg m-2] or [Z ~> m] + real :: h_k ! Thickness of the current layer [H ~> m or kg m-2] or [Z ~> m] call assert(remap_cs%initialized, 'diag_remap_calc_hmask: remap_cs not initialized.') ! Only z*-like diagnostic coordinates should have a 3d mask mask_vanished_layers = (remap_cs%vertical_coord == coordinateMode('ZSTAR')) - mask(:,:,:) = 0. + nz = remap_cs%nz - do j=G%jsc-1,G%jec+1 ; do i=G%isc-1,G%iec+1 - if (G%mask2dT(i,j)>0.) then - if (mask_vanished_layers) then - h_dest(:) = remap_cs%h(i,j,:) + ! Zero-fill the mask on device + !$omp target teams distribute parallel do collapse(3) + do k=1,nz ; do j=G%jsd,G%jed ; do i=G%isd,G%ied + mask(i,j,k) = 0. + enddo ; enddo ; enddo + + if (mask_vanished_layers) then + ! z*-like coordinate: mask out vanished layers. + ! The k-accumulation (h_tot, h_err) is sequential per column, + ! but columns are independent → collapse(2) over (i,j). + !$omp target teams distribute parallel do collapse(2) & + !$omp private(k, h_tot, h_err, h_k) + do j=G%jsc-1,G%jec+1 ; do i=G%isc-1,G%iec+1 + if (G%mask2dT(i,j) > 0.) then h_tot = 0. h_err = 0. - do k=1, remap_cs%nz - h_tot = h_tot + h_dest(k) + do k=1, nz + h_k = h(i,j,k) + h_tot = h_tot + h_k ! This is an overestimate of how thick a vanished layer might be, that ! appears due to round-off. h_err = h_err + epsilon(h_tot) * h_tot ! Mask out vanished layers - if (h_dest(k)<=8.*h_err) then - mask(i,j,k) = 0. - else + if (h_k > 8.*h_err) then mask(i,j,k) = 1. endif enddo - else ! all layers might contain data - mask(i,j,:) = 1. endif - endif - enddo ; enddo + enddo ; enddo + else + ! rho/sigma coordinates: all layers might contain data + !$omp target teams distribute parallel do collapse(3) + do k=1,nz ; do j=G%jsc-1,G%jec+1 ; do i=G%isc-1,G%iec+1 + if (G%mask2dT(i,j) > 0.) then + mask(i,j,k) = 1. + endif + enddo ; enddo ; enddo + endif end subroutine diag_remap_calc_hmask @@ -660,8 +807,9 @@ subroutine vertically_reintegrate_field(remap_cs, G, isdf, jsdf, h, h_target, st end subroutine vertically_reintegrate_field !> Vertically interpolate diagnostic field to alternative vertical grid. -subroutine vertically_interpolate_diag_field(remap_cs, G, h, staggered_in_x, staggered_in_y, & +subroutine vertically_interpolate_diag_field(remap_cs, G, h, staggered_in_x, staggered_in_y, & mask, field, interpolated_field) +! GPU PORT DIAGNOSTICS type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure type(ocean_grid_type), intent(in) :: G !< Ocean grid structure real, dimension(:,:,:), intent(in) :: h !< The current thicknesses [H ~> m or kg m-2] or [Z ~> m], @@ -693,87 +841,234 @@ subroutine vertically_interpolate_diag_field(remap_cs, G, h, staggered_in_x, sta end subroutine vertically_interpolate_diag_field +!> Ensure persistent GPU work buffers are allocated and device-mapped. +!! Buffers are sized to the full domain so they +!! work for h-, u-, and v-point fields without reallocation. Will it die for multi MPI? +subroutine ensure_gpu_buffers(remap_cs, G, nz_src) +! GPU PORT DIAGNOSTICS + type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure + type(ocean_grid_type), intent(in) :: G !< Ocean grid structure + integer, intent(in) :: nz_src !< Number of source vertical levels + + integer :: nz_dest, isd, ied, jsd, jed + + if (gpu_bufs_ready) return + + nz_dest = remap_cs%nz + + ! Use the widest bounds to cover h-, u-, and v-point fields + isd = min(G%isd, G%IsdB) ; ied = max(G%ied, G%IedB) + jsd = min(G%jsd, G%JsdB) ; jed = max(G%jed, G%JedB) + + allocate(gpu_h_work (isd:ied, jsd:jed, nz_src)) + allocate(gpu_htgt_work(isd:ied, jsd:jed, nz_dest)) + allocate(gpu_fld_work (isd:ied, jsd:jed, nz_src+1)) + allocate(gpu_out_work (isd:ied, jsd:jed, nz_dest+1)) + allocate(gpu_mask_work(isd:ied, jsd:jed)) + + !$omp target enter data map(alloc: gpu_h_work, gpu_htgt_work) + !$omp target enter data map(alloc: gpu_fld_work, gpu_out_work) + !$omp target enter data map(alloc: gpu_mask_work) + + gpu_bufs_ready = .true. +end subroutine ensure_gpu_buffers + !> Internal routine to vertically interpolate a diagnostic field to an alternative vertical grid. +!! Vertically interpolate a diagnostic field to an alternative vertical grid on GPU. +!! Uses persistent device-mapped work buffers to avoid per-call alloc/map overhead. subroutine vertically_interpolate_field(remap_cs, G, isdf, jsdf, h, staggered_in_x, staggered_in_y, & field, interpolated_field, mask) - type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure - type(ocean_grid_type), intent(in) :: G !< Ocean grid structure - integer, intent(in) :: isdf !< The starting i-index in memory for field - integer, intent(in) :: jsdf !< The starting j-index in memory for field + type(diag_remap_ctrl), intent(in) :: remap_cs !< Diagnostic coordinate control structure + type(ocean_grid_type), intent(in) :: G !< Ocean grid structure + integer, intent(in) :: isdf !< The starting i-index in memory for field + integer, intent(in) :: jsdf !< The starting j-index in memory for field real, dimension(G%isd:,G%jsd:,:), & - intent(in) :: h !< The current thicknesses [H ~> m or kg m-2] or [Z ~> m], - !! depending on the value of remap_cs%Z_based_coord - logical, intent(in) :: staggered_in_x !< True is the x-axis location is at u or q points - logical, intent(in) :: staggered_in_y !< True is the y-axis location is at v or q points + intent(in) :: h !< The current thicknesses [H ~> m or kg m-2] or [Z ~> m], + !! depending on the value of remap_cs%Z_based_coord + logical, intent(in) :: staggered_in_x !< True is the x-axis location is at u or q points + logical, intent(in) :: staggered_in_y !< True is the y-axis location is at v or q points real, dimension(isdf:,jsdf:,:), & - intent(in) :: field !< The diagnostic field to be remapped [A] + intent(in) :: field !< The diagnostic field to be remapped [A] real, dimension(isdf:,jsdf:,:), & - intent(out) :: interpolated_field !< Field argument remapped to alternative coordinate [A] + intent(out) :: interpolated_field !< Field argument remapped to alternative coordinate [A] real, dimension(isdf:,jsdf:), & - optional, intent(in) :: mask !< A mask for the field [nondim] + optional, intent(in) :: mask !< A mask for the field [nondim] ! Local variables - real, dimension(remap_cs%nz) :: h_dest ! Destination thicknesses [H ~> m or kg m-2] or [Z ~> m] - real, dimension(size(h,3)) :: h_src ! A column of source thicknesses [H ~> m or kg m-2] or [Z ~> m] + real :: h_src(NK_GPU_MAX), h_dest(NK_GPU_MAX) ! Column thickness work arrays for staggered averaging + real :: fld_loc(NK_GPU_MAX+1), out_loc(NK_GPU_MAX+1) ! Thread-private copies to avoid strided-slice heap temps integer :: nz_src, nz_dest ! The number of layers on the native and remapped grids - integer :: i, j !< Grid index - - interpolated_field(:,:,:) = 0. + integer :: i, j, k !< Grid index + logical :: have_mask nz_src = size(h,3) nz_dest = remap_cs%nz + have_mask = present(mask) + + call ensure_gpu_buffers(remap_cs, G, nz_src) + + ! Copy source data into persistent host-side buffers, then sync to device. + ! This avoids per-call allocate/map/delete overhead — only the data transfer remains. + gpu_h_work(G%isd:G%ied, G%jsd:G%jed, :) = h(:,:,:) + gpu_htgt_work(G%isd:G%ied, G%jsd:G%jed, :) = remap_cs%h(:,:,:) + ! Field and output use isdf:,jsdf: bounds — copy into same region of work buffer + do k=1,size(field,3) + do j=jsdf,jsdf+size(field,2)-1 ; do i=isdf,isdf+size(field,1)-1 + gpu_fld_work(i,j,k) = field(i,j,k) + enddo ; enddo + enddo + if (have_mask) then + do j=jsdf,jsdf+size(mask,2)-1 ; do i=isdf,isdf+size(mask,1)-1 + gpu_mask_work(i,j) = mask(i,j) + enddo ; enddo + endif + + !$omp target update to(gpu_h_work, gpu_htgt_work, gpu_fld_work) + if (have_mask) then + !$omp target update to(gpu_mask_work) + endif + + ! Zero-fill the output on device + !$omp target teams distribute parallel do collapse(3) + do k=1,nz_dest+1 + do j=lbound(gpu_out_work,2),ubound(gpu_out_work,2) + do i=lbound(gpu_out_work,1),ubound(gpu_out_work,1) + gpu_out_work(i,j,k) = 0. + enddo + enddo + enddo if (staggered_in_x .and. .not. staggered_in_y) then ! U-points - if (present(mask)) then - do j=G%jsc,G%jec ; do I=G%IscB,G%IecB ; if (mask(I,j) > 0.0) then - h_src(:) = 0.5 * (h(i,j,:) + h(i+1,j,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i+1,j,:)) - call interpolate_column(nz_src, h_src, field(I,j,:), & - nz_dest, h_dest, interpolated_field(I,j,:), .true.) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) + do j=G%jsc,G%jec ; do I=G%IscB,G%IecB + if (gpu_mask_work(I,j) > 0.0) then + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i+1,j,k)) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i+1,j,k)) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(I,j,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(I,j,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) do j=G%jsc,G%jec ; do I=G%IscB,G%IecB - h_src(:) = 0.5 * (h(i,j,:) + h(i+1,j,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i+1,j,:)) - call interpolate_column(nz_src, h_src, field(I,j,:), & - nz_dest, h_dest, interpolated_field(I,j,:), .true.) + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i+1,j,k)) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i+1,j,k)) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(I,j,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(I,j,k) = out_loc(k) + enddo enddo ; enddo endif elseif (staggered_in_y .and. .not. staggered_in_x) then ! V-points - if (present(mask)) then - do J=G%jscB,G%jecB ; do i=G%isc,G%iec ; if (mask(I,j) > 0.0) then - h_src(:) = 0.5 * (h(i,j,:) + h(i,j+1,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i,j+1,:)) - call interpolate_column(nz_src, h_src, field(i,J,:), & - nz_dest, h_dest, interpolated_field(i,J,:), .true.) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) + do J=G%jscB,G%jecB ; do i=G%isc,G%iec + if (gpu_mask_work(i,J) > 0.0) then + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i,j+1,k)) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i,j+1,k)) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(i,J,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(i,J,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) do J=G%jscB,G%jecB ; do i=G%isc,G%iec - h_src(:) = 0.5 * (h(i,j,:) + h(i,j+1,:)) - h_dest(:) = 0.5 * (remap_cs%h(i,j,:) + remap_cs%h(i,j+1,:)) - call interpolate_column(nz_src, h_src, field(i,J,:), & - nz_dest, h_dest, interpolated_field(i,J,:), .true.) + do k=1,nz_src + h_src(k) = 0.5 * (gpu_h_work(i,j,k) + gpu_h_work(i,j+1,k)) + enddo + do k=1,nz_dest + h_dest(k) = 0.5 * (gpu_htgt_work(i,j,k) + gpu_htgt_work(i,j+1,k)) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(i,J,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(i,J,k) = out_loc(k) + enddo enddo ; enddo endif elseif ((.not. staggered_in_x) .and. (.not. staggered_in_y)) then ! H-points - if (present(mask)) then - do j=G%jsc,G%jec ; do i=G%isc,G%iec ; if (mask(i,j) > 0.0) then - call interpolate_column(nz_src, h(i,j,:), field(i,j,:), & - nz_dest, remap_cs%h(i,j,:), interpolated_field(i,j,:), .true.) - endif ; enddo ; enddo + if (have_mask) then + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) + do j=G%jsc,G%jec ; do i=G%isc,G%iec + if (gpu_mask_work(i,j) > 0.0) then + do k=1,nz_src + h_src(k) = gpu_h_work(i,j,k) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(i,j,k) + enddo + do k=1,nz_dest + h_dest(k) = gpu_htgt_work(i,j,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(i,j,k) = out_loc(k) + enddo + endif + enddo ; enddo else + !$omp target teams distribute parallel do collapse(2) private(h_src, h_dest, fld_loc, out_loc) do j=G%jsc,G%jec ; do i=G%isc,G%iec - call interpolate_column(nz_src, h(i,j,:), field(i,j,:), & - nz_dest, remap_cs%h(i,j,:), interpolated_field(i,j,:), .true.) + do k=1,nz_src + h_src(k) = gpu_h_work(i,j,k) + enddo + do k=1,nz_src+1 + fld_loc(k) = gpu_fld_work(i,j,k) + enddo + do k=1,nz_dest + h_dest(k) = gpu_htgt_work(i,j,k) + enddo + call interpolate_column(nz_src, h_src, fld_loc, nz_dest, h_dest, out_loc, .true.) + do k=1,nz_dest+1 + gpu_out_work(i,j,k) = out_loc(k) + enddo enddo ; enddo endif else call assert(.false., 'vertically_interpolate_diag_field: Q point remapping is not coded yet.') endif + ! Copy result back from device + !$omp target update from(gpu_out_work) + do k=1,nz_dest+1 + do j=jsdf,jsdf+size(interpolated_field,2)-1 + do i=isdf,isdf+size(interpolated_field,1)-1 + interpolated_field(i,j,k) = gpu_out_work(i,j,k) + enddo + enddo + enddo + end subroutine vertically_interpolate_field !> Horizontally average a diagnostic field