Acceleration of Zanna-Bolton-2020 parameterization and new features required for NW2

[Pperezhogin]

Summary

This PR updates the code of Zanna-Bolton-2020 mesoscale parameterization and seeks three purposes:

• Include features required to run NW2 simulations.
• Accelerate the code for many CPUs. Depending on configuration, the new code is 4-100 times faster.
• Add measurement of runtime.

NW2

Simulations in NW2 at 1/4 resolution are now available with namelist parameters:

SMAG_BI_CONST = 0.06
USE_ZB2020 = True
ZB_SCALING = 2.5
ZB_SCHEME = 0
STRESS_SMOOTH_PASS = 4
ZB_KLOWER_R_DISS = 1.

Code details

The code was significantly restructured. The full list of commits is available in a separate branch PR_ZB_2020_NW2_and_acceleration.

Removed/new parameters

The following namelist parameters are no longer available and set to default values:

VG_SMOOTH_PASS=0
VG_SMOOTH_SEL=1
VG_SHARP_SEL=1
STRESS_SMOOTH_SEL=1
ZB_HYPERVISC=0
HYPVISC_GRID_DAMP

These parameters were controlling complicated filters options which are not currently used.

The following two parameters are introduced which control the attenuation of the subgrid model in NW2 configuration with typical values:

ZB_KLOWER_R_DISS=1.
ZB_KLOWER_SHEAR=1

Acceleration of the code

Overall restructuring of the code with storing precomputed arrays led to acceleration ~ 4 times.

Acceleration of filters is achieved with:

• Marching halo with grouping of non-blocking MPI exchanges (~10-100 times faster on thousands of CPUs).
• Decomposition of 2D filter into sequence of 1D filters (30% faster).

Tests

All tests (grid, layout, nan, rotate, dim, restart, openmp) are passed.

Regression

Regression changed twice:

• I removed mask of outcropped points . The results so far did not change qualitatively. Support of the code for this mask is complicated and (apparently) unnecessary.
• The 2D and 1D filters differ by the round-off errors.

[marshallward]

Thanks for putting in some of these improvements to the solver.

Before reviewing this, there's at least one error that need to be fixed.

  ../../../src/parameterizations/lateral/MOM_Zanna_Bolton.F90:314:29:
  
    314 |   if (CS%Stress_iter > 0 .or. CS%HPF_iter) then
        |                             1
  Error: Operands of logical operator ‘.or.’ at (1) are LOGICAL(4)/INTEGER(4)

In logical tests, everything has to be converted to logical (e.g. CS%HPF_iter > 0).

Also a good idea to start fixing up the documentation: https://github.com/NOAA-GFDL/MOM6/actions/runs/6346687254/job/17240553103?pr=484

All function arguments and derived type parameters must always be documented.

[codecov]

Codecov Report

Merging #484 (124afea) into dev/gfdl (43a4fa9) will increase coverage by 0.00%.
The diff coverage is 1.39%.

❗ Current head 124afea differs from pull request most recent head 8cab751. Consider uploading reports for the commit 8cab751 to get more accurate results

@@            Coverage Diff            @@
##           dev/gfdl     #484   +/-   ##
=========================================
  Coverage     37.82%   37.83%           
=========================================
  Files           270      270           
  Lines         78359    78346   -13     
  Branches      14497    14502    +5     
=========================================
+ Hits          29640    29641    +1     
+ Misses        43314    43297   -17     
- Partials       5405     5408    +3     

Files	Coverage Δ
src/parameterizations/lateral/MOM_hor_visc.F90	50.04% <14.28%> (-0.09%)	⬇️
src/parameterizations/lateral/MOM_Zanna_Bolton.F90	1.57% <1.07%> (+0.66%)	⬆️

... and 1 file with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[Pperezhogin]

Hi @marshallward, I addressed the mentioned problems with doxygen and compilation error.

[marshallward]

Hi @Pperezhogin I'm finally getting around to this. Switching from 2D to 3D halo updates looks like a great improvement, and overall structural changes looks good to me as well.

The others may also want to review these changes, but I have added some comments below. Most are related to style (which, as de facto maintainer of this file, you may have some liberty to tell me if I've gone too far 😄). There were a couple places which looked to me like it may improve performance.

• The diff[uv] copy in Zanna_Bolton_2020 seemed unnecessary, since you are already modifying the field. Can this simply be modified in-place?

• The if (CS%Klower_R_diss > 0) in compute_stress_divergence look like they can be moved outside of the do-loops. You may see some improved performance if you invert the if/do blocks.

Other than that, it looks good to me.

[Pperezhogin]

Hi @marshallward

I addressed all your comments. Regression did not change, tests are passed.

Regarding optimizations, I simplified if statements in loops whenever possible. On my system, it does not affect the runtime. Also, now I update diff[u,v] in-place which makes this module 5% faster.

Regarding functions name: I introduced consistent naming for all public elements of the module.

[marshallward]

Gaea regression: https://gitlab.gfdl.noaa.gov/ogrp/MOM6/-/pipelines/21028 ✔️

Thanks for patience, @Pperezhogin