(0.88.0) MPI communication and computation overlap in the HydrostaticFreeSurfaceModel and NonhydrostaticModel

[simone-silvestri]

Overlapping MPI communication and computation in the HydrostaticFreeSurfaceModel and the NonhydrostaticModel.

In particular, this PR introduces two keyword arguments for the fill_halo_regions! function, active only in case of distributed halo passing boundary conditions:

• theasync::Bool keyword argument that allows launching MPI operations without waiting for the communication to complete.
• the only_local_halos::Bool keyword argument, which fills only the halo in case of a local (i.e., Flux, Value, Gradient, Periodic, and, temporarily, MultiRegionCommunication) boundary condition. This is required for having explicit boundary conditions (like Value or Flux) for turbulent diffusivities (we directly calculate diffusivities in the halos in the case of distributed BC).

This PR allows hiding MPI passing of barotropic auxiliary variables behind the implicit vertical solver and prognostic variables behind the tendency calculations. The latter is done by splitting the tendency kernels into an interior kernel that calculates tendencies between, e.g., i = Hx and i = Nx - Hx, and a boundary kernel, executed once communication is complete, that calculates tendencies adjacent to boundaries.
Before computing tendencies near the boundaries, boundary-adjacent auxiliary diagnostic variables are recalculated (hydrostatic pressure, vertical velocity and diffusivities for the hydrostatic model, hydrostatic pressure and diffusivities for the non-hydrostatic model)

Also, this PR introduces

• 8-way halo passing for distributed grids (west, east, north, south, southwest, southeast, northwest, northeast), something that we will want to implement also for MultiRegion)
• Non uniform partitioning for distributed grids
• the KernelParameters(size, offset) to be passed to the launch! function to start a kernel of size size::Tuple offset by offset::Tuple

todo:

• Implement Offsets in KernelAbstractions (at the moment implemented a fix in src/Utils/kernel_launching.jl to remove when offsets will be implemented)
• Adapt that implementation in this PR
• Generalize the current implementation of kernel offsets for Diffusivities, Pressures and W-velocity
• Find a general way to tag MPI requests to allow large number of cores (currently Nranks < 100)
• Remove views halo passing

API changes

• When calling RectilinearGrid(arch::DistributedArch, size = (Nx, Ny, Nz), ....), (Nx, Ny, Nz) are the per-rank local sizes, not the global size to be divided (easy way to specify non-uniform partitioning, see validation/distributed/mpi_geostrophic_adjustment.jl)
• added the enable_overlapped_communication keyword to DistributedArch (defaults to true)
• removed the use_buffers keyword to DistributedArch (always use buffers, as views did not give significant speedup to justify maintaining two implementations)
• added the keyword argument active_cells_map::Bool = false to ImmersedBoundaryGrid (ex: ImmersedBoundaryGrid(grid, ib, active_cells_map = true)
• added a required_halo_size::Int keyword argument to ScalarDiffusivity (defaults to 1) and ScalarBiharmonicDiffusivity (defaults to 2) to be specified by the user which sets the required halo for the specific ν or κ function (closures have now an explicitly required number of halos)

Major internals change

1. The tendencies are calculated at the end of a time step. Therefore at the end of a simulation model.timestepper will hold tendencies for the last time step completed

2. Removed fill_halo_regions! for hydrostatic pressure in both the non-hydrostatic and the hydrostatic model and for w-velocity in the hydrostatic model. The halos are filled by enlarging the size of the kernels in update_hydrostatic_pressure! and compute_w_from_continuity! to incorporate the needed ghost points.

3. Removed fill_halo_regions! for diffusivities only for halo-passing BC; the halo calculation is now performed by launching the calculate_diffusivity! kernel inside the ghost nodes before recomputing the tendencies. This requires knowing how many halos each closure requires.

4. Added a required_halo parameter to AbstractTurbulenceClosure. This means that each parameterization will have to specify explicitly the number of halos required to calculate the diffusivity:
   e.g

abstract type AbstractTurbulenceClosure{TimeDiscretization, BoundaryBuffer} end
abstract type AbstractScalarDiffusivity{TD, F, N} <: AbstractTurbulenceClosure{TD, N} end
struct TwoDimensionalLeith{FT, CR, GM, M} <: AbstractScalarDiffusivity{ExplicitTimeDiscretization, ThreeDimensionalFormulation, 2}

Where Leith closure requires 2 halos (one for the vorticity calculation and an additional one for the vorticity derivative)

Minor internals change

removed the general calculate_nonlinear_viscosity! and calculate_nonlinear_diffusivity! kernels (to each turbulence closure their own kernel)

Requires JuliaGPU/KernelAbstractions.jl#399 # On hold at the moment
Closes #615
Closes #1882
Closes #3067
Closes #3068
Supersedes #2953

[navidcy]

definitely bump a release (minor or patch; whatever more appropriate) before this is merged :)

[navidcy]

Perhaps we wait for JuliaGPU/KernelAbstractions.jl#399 to be merged and tag a patch release for KA before we merge this? @vchuravy?

[simone-silvestri]

Sure! Anyways, this PR should be in its final form, so ready for review :) (will still wait for KA changes before merging)

[simone-silvestri]

So, it looks like docs take more to build in this PR when compared to main (4:00 hrs vs 3:40 hrs). Probably good to benchmark a bit. I ll try some benchmarking on GPU and CPU

[navidcy]

benchmarking sounds good!

but I thought this PR only made changes to the distributed grids, no?

[simone-silvestri]

It also rearranges how time-stepping works to allow

• minimal fill_halo_regions
• overlapped communication and computation

and removes fill_halo_regions! for computation of w velocity, hydrostatic pressure and diffusivities by enlarging the kernels to include the halo points

[navidcy]

just a few beautification remarks

[navidcy]

oh dear, my PR created merge conflicts......

[navidcy]

@simone-silvestri I think I resolve the conflicts alright but it wouldn't hurt if you have a look at b51e681...

[glwagner]

do we use arch = Distributed() for distributed grids now? Might be good to summarize in the top description

[tomchor]

Requires JuliaGPU/KernelAbstractions.jl#399 # On hold at the moment

At the top comment it says this PR depends on another KA PR that isn't merged. However, this PR is merged (and tagged). Does this not depend on JuliaGPU/KernelAbstractions.jl#399 anymore?

[navidcy]

Requires JuliaGPU/KernelAbstractions.jl#399 # On hold at the moment

At the top comment it says this PR depends on another KA PR that isn't merged. However, this PR is merged (and tagged). Does this not depend on JuliaGPU/KernelAbstractions.jl#399 anymore?

Yeap, it doesn't. I edited the first post ;)