Merge remote-tracking branch 'GFDL/gfdl-to-main-2025-09-25' into test/GFDL-20250925

Author: jiandewang

ac/makedep

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 from __future__ import print_function
 

config_src/drivers/timing_tests/time_MOM_ANN.F90

@@ -0,0 +1,189 @@
+program time_MOM_ANN
+
+! This file is part of MOM6. See LICENSE.md for the license.
+
+use MOM_ANN, only : ANN_CS
+use MOM_ANN, only : ANN_allocate, ANN_apply, ANN_end
+use MOM_ANN, only : ANN_apply_vector_orig, ANN_apply_vector_oi
+use MOM_ANN, only : ANN_apply_array_sio
+use MOM_ANN, only : ANN_random
+
+implicit none
+
+! Command line options
+integer :: nargs ! Number of command line arguments
+character(len=12) :: cmd_ln_arg !< Command line argument (if any)
+
+! ANN parameters
+integer :: nlayers ! Number of layers
+integer :: nin ! Number of inputs
+integer :: layer_width ! Width of hidden layers
+integer :: nout ! Number of outputs
+! Timing parameters
+integer :: nsamp ! Number of measurements
+integer :: nits ! Number of calls to time
+integer :: nxy ! Spatial dimension
+
+nlayers = 7; nin = 4; layer_width = 16; nout = 1 ! Deep network
+!nlayers = 4; nin = 4; layer_width = 48; nout = 1 ! Shallow-wide network
+!nlayers = 3; nin = 4; layer_width = 20; nout = 1 ! Small network
+
+nsamp = 100
+nits = 20000
+!nits = 300000 ! Needed for robust measurements on small networks
+nxy = 100 ! larger array
+!nxy = 10 ! small array
+
+! Optionally grab ANN and timing parameters from the command line
+nargs = command_argument_count()
+if (nargs==7) then
+  call get_command_argument(1, cmd_ln_arg)
+  read(cmd_ln_arg,*) nlayers
+  call get_command_argument(2, cmd_ln_arg)
+  read(cmd_ln_arg,*) nin
+  call get_command_argument(3, cmd_ln_arg)
+  read(cmd_ln_arg,*) layer_width
+  call get_command_argument(4, cmd_ln_arg)
+  read(cmd_ln_arg,*) nout
+  call get_command_argument(5, cmd_ln_arg)
+  read(cmd_ln_arg,*) nsamp
+  call get_command_argument(6, cmd_ln_arg)
+  read(cmd_ln_arg,*) nits
+  call get_command_argument(7, cmd_ln_arg)
+  read(cmd_ln_arg,*) nxy
+endif
+
+! Fastest variants on Intel Xeon W-2223 CPU @ 3.60GHz (gfortran-13.2 -O3)
+!                   | vector(nxy=1)  |   nxy = 10   |   nxy = 100
+! ----------------------------------------------------------------------------
+! Small ANN         |   vector_oi    |   array_soi  |   array_sio
+! Shallow-wide ANN  |   vector_oi    |   array_ois  |   array_sio
+! Deep ANN          |   vector_oi    |   array_ois  |   array_sio
+
+write(*,'(a)') "{"
+
+call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
+              0, "MOM_ANN:ANN_apply(vector)")
+write(*,"(',')")
+call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
+              1, "MOM_ANN:ANN_apply_vector_orig(array)")
+write(*,"(',')")
+call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
+              2, "MOM_ANN:ANN_apply_vector_oi(array)")
+write(*,"(',')")
+call time_ANN(nlayers, nin, layer_width, nout, nsamp, nits, nxy, &
+              12, "MOM_ANN:ANN_apply_array_sio(array)")
+write(*,"()")
+
+write(*,'(a)') "}"
+
+contains
+
+!> Time ANN inference.
+!!
+!! Times are measured over the "nits effective calls" and appropriately scaled to the
+!! time per call per single vector of input features. For array inputs, the number of
+!! actual calls is reduced by the size of the array.  The timing measurement is repeated
+!! "nsamp" times, to check the statistics of the timing measurement.
+subroutine time_ANN(nlayers, nin, width, nout, nsamp, nits, nxy, impl, label)
+  integer,          intent(in)  :: nlayers !< Number of layers
+  integer,          intent(in)  :: nin     !< Number of inputs
+  integer,          intent(in)  :: width   !< Width of hidden layers
+  integer,          intent(in)  :: nout    !< Number of outputs
+  integer,          intent(in)  :: nsamp   !< Number of measurements
+  integer,          intent(in)  :: nits    !< Number of calls to time
+  integer,          intent(in)  :: nxy     !< Spatial dimension
+  integer,          intent(in)  :: impl    !< Implementation to time
+  character(len=*), intent(in)  :: label   !< Label for YAML output
+  ! Local variables
+  type(ANN_CS) :: ANN ! ANN
+  integer :: widths(nlayers) ! Width of each layer
+  real :: x_s(nin) ! Inputs (just features) [nondim]
+  real :: y_s(nin) ! Outputs (just features) [nondim]
+  real :: x_fs(nin,nxy) ! Inputs (feature, space) [nondim]
+  real :: y_fs(nin,nxy) ! Outputs (feature, space) [nondim]
+  real :: x_sf(nin,nxy) ! Inputs (space, feature) [nondim]
+  real :: y_sf(nin,nxy) ! Outputs (space, feature) [nondim]
+  integer :: iter, samp ! Loop counters
+  integer :: ij ! Horizontal loop index
+  real :: start, finish, timing ! CPU times [s]
+  real :: tmin, tmax, tmean, tstd ! Min, max, mean, and standard deviation, of CPU times [s]
+  integer :: asamp ! Actual samples of timings
+  integer :: aits ! Actual iterations
+  real :: words_per_sec ! Operations per sec estimated from parameters [# s-1]
+
+  widths(:) = width
+  widths(1) = nin
+  widths(nlayers) = nout
+
+  call ANN_random(ANN, nlayers, widths)
+  call random_number(x_fs)
+  call random_number(x_sf)
+
+
+  tmin = 1e9
+  tmax = 0.
+  tmean = 0.
+  tstd = 0.
+  asamp = nits ! Most cases below use this
+  aits = nits / nxy ! Most cases below use this
+
+  do samp = 1, nsamp
+    select case (impl)
+      case (0)
+        aits = nits
+        call cpu_time(start)
+        do iter = 1, nits ! Make many passes to reduce sampling error
+          call ANN_apply(x_s, y_s, ANN)
+        enddo
+        call cpu_time(finish)
+      case (1)
+        call cpu_time(start)
+        do iter = 1, aits ! Make many passes to reduce sampling error
+          do ij = 1, nxy
+            call ANN_apply_vector_orig(x_fs(:,ij), y_fs(:,ij), ANN)
+          enddo
+        enddo
+        call cpu_time(finish)
+      case (2)
+        call cpu_time(start)
+        do iter = 1, aits ! Make many passes to reduce sampling error
+          do ij = 1, nxy
+            call ANN_apply_vector_oi(x_fs(:,ij), y_fs(:,ij), ANN)
+          enddo
+        enddo
+        call cpu_time(finish)
+      case (12)
+        call cpu_time(start)
+        do iter = 1, aits ! Make many passes to reduce sampling error
+          call ANN_apply_array_sio(nxy, x_sf(:,:), y_sf(:,:), ANN)
+        enddo
+        call cpu_time(finish)
+        asamp = nsamp * aits ! Account for working on whole arrays
+    end select
+
+    timing = ( finish - start ) / real(nits) ! Average time per call
+
+    tmin = min( tmin, timing )
+    tmax = max( tmax, timing )
+    tmean = tmean + timing
+    tstd = tstd + timing**2
+  enddo
+
+  tmean = tmean / real(nsamp)
+  tstd = tstd / real(nsamp) ! convert to mean of squares
+  tstd = tstd - tmean**2  ! convert to variance
+  tstd = sqrt( tstd * real(nsamp) / real(nsamp-1) ) ! convert to standard deviation
+  words_per_sec = ANN%parameters / ( tmean * 1024 * 1024 )
+
+  write(*,"(2x,3a)") '"', trim(label), '": {'
+  write(*,"(4x,a,1pe11.4,',')") '"min": ', tmin
+  write(*,"(4x,a,1pe11.4,',')") '"mean":', tmean
+  write(*,"(4x,a,1pe11.4,',')") '"std": ', tstd
+  write(*,"(4x,a,i0,',')") '"n_samples": ', asamp
+  write(*,"(4x,a,1pe11.4,',')") '"max": ', tmax
+  write(*,"(4x,a,1pe11.4,'}')", advance="no") '"MBps": ', words_per_sec
+
+end subroutine time_ANN
+
+end program time_MOM_ANN

config_src/drivers/unit_tests/test_MOM_ANN.F90

@@ -0,0 +1,10 @@
+program test_MOM_ANN
+
+use MOM_ANN, only : ANN_unit_tests
+use MOM_error_handler, only : set_skip_mpi
+
+call set_skip_mpi(.true.) ! This unit tests is not expecting MPI to be used
+
+if ( ANN_unit_tests(.true.) ) stop 1
+
+end program test_MOM_ANN

src/ALE/MOM_ALE.F90

@@ -165,8 +165,9 @@ module MOM_ALE
 !! before the main time integration loop to initialize the regridding stuff.
 !! We read the MOM_input file to register the values of different
 !! regridding/remapping parameters.
-subroutine ALE_init( param_file, GV, US, max_depth, CS)
+subroutine ALE_init( param_file, G, GV, US, max_depth, CS)
   type(param_file_type),   intent(in) :: param_file !< Parameter file
+  type(ocean_grid_type),   intent(in) :: G          !< Grid structure
   type(verticalGrid_type), intent(in) :: GV         !< Ocean vertical grid structure
   type(unit_scale_type),   intent(in) :: US         !< A dimensional unit scaling type
   real,                    intent(in) :: max_depth  !< The maximum depth of the ocean [Z ~> m].
@@ -205,8 +206,9 @@ subroutine ALE_init( param_file, GV, US, max_depth, CS)
                  default=.false.)
 
   ! Initialize and configure regridding
-  call ALE_initRegridding(GV, US, max_depth, param_file, mdl, CS%regridCS)
-  call regridding_preadjust_reqs(CS%regridCS, CS%do_conv_adj, CS%use_hybgen_unmix, hybgen_CS=hybgen_regridCS)
+  call ALE_initRegridding(G, GV, US, max_depth, param_file, mdl, CS%regridCS)
+  call regridding_preadjust_reqs(CS%regridCS, CS%do_conv_adj, CS%use_hybgen_unmix, &
+                                 hybgen_CS=hybgen_regridCS)
 
   ! Initialize and configure remapping that is orchestrated by ALE.
   call get_param(param_file, mdl, "REMAPPING_SCHEME", string, &
@@ -321,12 +323,12 @@ subroutine ALE_init( param_file, GV, US, max_depth, CS)
                  default=-0.001, units="m", scale=GV%m_to_H)
   call get_param(param_file, mdl, "REMAP_VEL_MASK_H_THIN", CS%h_vel_mask, &
                  "A thickness at velocity points below which near-bottom layers are zeroed out "//&
-                 "after remapping, following practice with Hybgen remapping, or a negative value "//&
-                 "to avoid such filtering altogether.", &
+                 "after remapping, following practice with Hybgen remapping, "//&
+                 "or a negative value to avoid such filtering altogether.", &
                  default=1.0e-6, units="m", scale=GV%m_to_H, do_not_log=(CS%BBL_h_vel_mask<=0.0))
 
   if (CS%use_hybgen_unmix) &
-    call init_hybgen_unmix(CS%hybgen_unmixCS, GV, US, param_file, hybgen_regridCS)
+      call init_hybgen_unmix(CS%hybgen_unmixCS, GV, US, param_file, hybgen_regridCS)
 
   call get_param(param_file, mdl, "REMAP_VEL_CONSERVE_KE", CS%conserve_ke, &
                  "If true, a correction is applied to the baroclinic component of velocity "//&
@@ -640,7 +642,8 @@ end subroutine ALE_offline_inputs
 
 !> For a state-based coordinate, accelerate the process of regridding by
 !! repeatedly applying the grid calculation algorithm
-subroutine ALE_regrid_accelerated(CS, G, GV, US, h, tv, n_itt, u, v, OBC, Reg, dt, dzRegrid, initial)
+subroutine ALE_regrid_accelerated(CS, G, GV, US, h, tv, n_itt, u, v, OBC, Reg, dt, &
+                                  dzRegrid, initial)
   type(ALE_CS),            pointer       :: CS     !< ALE control structure
   type(ocean_grid_type),   intent(inout) :: G      !< Ocean grid
   type(verticalGrid_type), intent(in)    :: GV     !< Vertical grid
@@ -689,7 +692,7 @@ subroutine ALE_regrid_accelerated(CS, G, GV, US, h, tv, n_itt, u, v, OBC, Reg, d
 
   ! initial total interface displacement due to successive regridding
   if (CS%remap_uv_using_old_alg) &
-    dzIntTotal(:,:,:) = 0.
+      dzIntTotal(:,:,:) = 0.
 
   call create_group_pass(pass_T_S_h, T, G%domain)
   call create_group_pass(pass_T_S_h, S, G%domain)
@@ -708,7 +711,7 @@ subroutine ALE_regrid_accelerated(CS, G, GV, US, h, tv, n_itt, u, v, OBC, Reg, d
 
   ! Apply timescale to regridding (for e.g. filtered_grid_motion)
   if (present(dt)) &
-    call ALE_update_regrid_weights(dt, CS)
+      call ALE_update_regrid_weights(dt, CS)
 
   do itt = 1, n_itt
 
@@ -722,12 +725,14 @@ subroutine ALE_regrid_accelerated(CS, G, GV, US, h, tv, n_itt, u, v, OBC, Reg, d
 
     call regridding_main(CS%remapCS, CS%regridCS, G, GV, US, h_loc, tv_local, h, dzInterface)
     if (CS%remap_uv_using_old_alg) &
-      dzIntTotal(:,:,:) = dzIntTotal(:,:,:) + dzInterface(:,:,:)
+        dzIntTotal(:,:,:) = dzIntTotal(:,:,:) + dzInterface(:,:,:)
 
     ! remap from original grid onto new grid
     do j = G%jsc-1,G%jec+1 ; do i = G%isc-1,G%iec+1
-      call remapping_core_h(CS%remapCS, nz, h_orig(i,j,:), tv%S(i,j,:), nz, h(i,j,:), tv_local%S(i,j,:))
-      call remapping_core_h(CS%remapCS, nz, h_orig(i,j,:), tv%T(i,j,:), nz, h(i,j,:), tv_local%T(i,j,:))
+      call remapping_core_h(CS%remapCS, nz, h_orig(i,j,:), tv%S(i,j,:), nz, h(i,j,:), &
+                            tv_local%S(i,j,:))
+      call remapping_core_h(CS%remapCS, nz, h_orig(i,j,:), tv%T(i,j,:), nz, h(i,j,:), &
+                            tv_local%T(i,j,:))
     enddo ; enddo
 
     ! starting grid for next iteration
@@ -1146,7 +1151,7 @@ subroutine ALE_remap_velocities(CS, G, GV, h_old_u, h_old_v, h_new_u, h_new_v, u
   if (CS%id_remap_delta_integ_v2>0) dv2h_tot(:,:) = 0.
 
   if (((CS%id_remap_delta_integ_u2>0) .or. (CS%id_remap_delta_integ_v2>0)) .and. .not.present(dt))&
-    call MOM_error(FATAL, "ALE KE diagnostics requires passing dt into ALE_remap_velocities")
+      call MOM_error(FATAL, "ALE KE diagnostics requires passing dt into ALE_remap_velocities")
 
   nz = GV%ke
 
@@ -1212,7 +1217,7 @@ subroutine ALE_remap_velocities(CS, G, GV, h_old_u, h_old_v, h_new_u, h_new_v, u
     endif
 
     if ((CS%BBL_h_vel_mask > 0.0) .and. (CS%h_vel_mask > 0.0)) &
-      call mask_near_bottom_vel(u_tgt, h2, CS%BBL_h_vel_mask, CS%h_vel_mask, nz)
+        call mask_near_bottom_vel(u_tgt, h2, CS%BBL_h_vel_mask, CS%h_vel_mask, nz)
 
     ! Copy the column of new velocities back to the 3-d array
     do k=1,nz
@@ -1361,13 +1366,14 @@ subroutine ALE_remap_vertex_vals(CS, G, GV, h_old, h_new, vert_val)
 
   do J=G%JscB,G%JecB ; do I=G%IscB,G%IecB
     if ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)) > 0.0 ) then
-      I_mask_sum = 1.0 / ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)))
+      I_mask_sum = 1.0 / ((G%mask2dT(i,j) + G%mask2dT(i+1,j+1)) + &
+                          (G%mask2dT(i+1,j) + G%mask2dT(i,j+1)))
 
     do k=1,nz
       h_src(k) = ((G%mask2dT(i,j) * h_old(i,j,k) + G%mask2dT(i+1,j+1) * h_old(i+1,j+1,k)) + &
-                  (G%mask2dT(i+1,j) * h_old(i+1,j,k) + G%mask2dT(i,j+1) * h_old(i,j+1,k)) ) * I_mask_sum
+          (G%mask2dT(i+1,j) * h_old(i+1,j,k) + G%mask2dT(i,j+1) * h_old(i,j+1,k)) ) * I_mask_sum
       h_tgt(k) = ((G%mask2dT(i,j) * h_new(i,j,k) + G%mask2dT(i+1,j+1) * h_new(i+1,j+1,k)) + &
-                  (G%mask2dT(i+1,j) * h_new(i+1,j,k) + G%mask2dT(i,j+1) * h_new(i,j+1,k)) ) * I_mask_sum
+          (G%mask2dT(i+1,j) * h_new(i+1,j,k) + G%mask2dT(i,j+1) * h_new(i,j+1,k)) ) * I_mask_sum
     enddo
 
     do K=1,nz+1
@@ -1549,7 +1555,8 @@ subroutine ALE_PLM_edge_values( CS, G, GV, h, Q, bdry_extrap, Q_t, Q_b )
   do j = G%jsc-1,G%jec+1 ; do i = G%isc-1,G%iec+1
     slp(1) = 0.
     do k = 2, GV%ke-1
-      slp(k) = PLM_slope_wa(h(i,j,k-1), h(i,j,k), h(i,j,k+1), h_neglect, Q(i,j,k-1), Q(i,j,k), Q(i,j,k+1))
+      slp(k) = PLM_slope_wa(h(i,j,k-1), h(i,j,k), h(i,j,k+1), h_neglect, &
+                            Q(i,j,k-1), Q(i,j,k), Q(i,j,k+1))
     enddo
     slp(GV%ke) = 0.
 
@@ -1562,7 +1569,8 @@ subroutine ALE_PLM_edge_values( CS, G, GV, h, Q, bdry_extrap, Q_t, Q_b )
       mslp = - PLM_extrapolate_slope(h(i,j,2), h(i,j,1), h_neglect, Q(i,j,2), Q(i,j,1))
       Q_t(i,j,1) = Q(i,j,1) - 0.5 * mslp
       Q_b(i,j,1) = Q(i,j,1) + 0.5 * mslp
-      mslp = PLM_extrapolate_slope(h(i,j,GV%ke-1), h(i,j,GV%ke), h_neglect, Q(i,j,GV%ke-1), Q(i,j,GV%ke))
+      mslp = PLM_extrapolate_slope(h(i,j,GV%ke-1), h(i,j,GV%ke), h_neglect, &
+                                   Q(i,j,GV%ke-1), Q(i,j,GV%ke))
       Q_t(i,j,GV%ke) = Q(i,j,GV%ke) - 0.5 * mslp
       Q_b(i,j,GV%ke) = Q(i,j,GV%ke) + 0.5 * mslp
     else
@@ -1630,7 +1638,7 @@ subroutine TS_PPM_edge_values( CS, S_t, S_b, T_t, T_b, G, GV, tv, h, bdry_extrap
     call PPM_reconstruction( GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect, &
                                   answer_date=CS%answer_date )
     if (bdry_extrap) &
-      call PPM_boundary_extrapolation( GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect )
+        call PPM_boundary_extrapolation( GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect )
 
     do k = 1,GV%ke
       S_t(i,j,k) = ppol_E(k,1)
@@ -1651,7 +1659,7 @@ subroutine TS_PPM_edge_values( CS, S_t, S_b, T_t, T_b, G, GV, tv, h, bdry_extrap
     call PPM_reconstruction( GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect, &
                                   answer_date=CS%answer_date )
     if (bdry_extrap) &
-      call PPM_boundary_extrapolation(GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect )
+        call PPM_boundary_extrapolation(GV%ke, hTmp, tmp, ppol_E, ppol_coefs, h_neglect )
 
     do k = 1,GV%ke
       T_t(i,j,k) = ppol_E(k,1)
@@ -1664,7 +1672,8 @@ end subroutine TS_PPM_edge_values
 
 
 !> Initializes regridding for the main ALE algorithm
-subroutine ALE_initRegridding(GV, US, max_depth, param_file, mdl, regridCS)
+subroutine ALE_initRegridding(G, GV, US, max_depth, param_file, mdl, regridCS)
+  type(ocean_grid_type),   intent(in)  :: G          !< Grid structure
   type(verticalGrid_type), intent(in)  :: GV         !< Ocean vertical grid structure
   type(unit_scale_type),   intent(in)  :: US         !< A dimensional unit scaling type
   real,                    intent(in)  :: max_depth  !< The maximum depth of the ocean [Z ~> m].
@@ -1680,7 +1689,7 @@ subroutine ALE_initRegridding(GV, US, max_depth, param_file, mdl, regridCS)
                  trim(regriddingCoordinateModeDoc), &
                  default=DEFAULT_COORDINATE_MODE, fail_if_missing=.true.)
 
-  call initialize_regridding(regridCS, GV, US, max_depth, param_file, mdl, coord_mode, '', '')
+  call initialize_regridding(regridCS, G, GV, US, max_depth, param_file, mdl, coord_mode, '', '')
 
 end subroutine ALE_initRegridding
 

src/ALE/MOM_hybgen_regrid.F90

@@ -988,12 +988,12 @@ subroutine hybgen_column_regrid(CS, nk, thkbot, Rcv_tgt, &
     ! Verify that everything is consistent.
     do k=1,nk
       if (abs((h_col(k) - h_in(k)) + (dp_int(K) - dp_int(K+1))) > 1.0e-13*max(p_int(nk+1), CS%onem)) then
-        write(mesg, '("k ",i4," h ",es13.4," h_in ",es13.4, " dp ",2es13.4," err ",es13.4)') &
+        write(mesg, '("k ",I0," h ",es13.4," h_in ",es13.4, " dp ",2es13.4," err ",es13.4)') &
               k, h_col(k), h_in(k), dp_int(K), dp_int(K+1), (h_col(k) - h_in(k)) + (dp_int(K) - dp_int(K+1))
         call MOM_error(FATAL, "Mismatched thickness changes in hybgen_regrid: "//trim(mesg))
       endif
       if (h_col(k) < 0.0) then ! Could instead do: -1.0e-15*max(p_int(nk+1), CS%onem)) then
-        write(mesg, '("k ",i4," h ",es13.4," h_in ",es13.4, " dp ",2es13.4, " fixlay ",i4)') &
+        write(mesg, '("k ",I0," h ",es13.4," h_in ",es13.4, " dp ",2es13.4, " fixlay ",I0)') &
               k, h_col(k), h_in(k), dp_int(K), dp_int(K+1), fixlay
         call MOM_error(FATAL, "Significantly negative final thickness in hybgen_regrid: "//trim(mesg))
       endif