Fill halos of conformal cubed sphere grid coordinates and metrics in a single pass

[siddharthabishnu]

This PR fills the halos of the grid metrics of the conformal cubed sphere panels at face-face locations using a dedicated function instead of hard-coded values, similar to the existing approach for center-center, face-center, and center-face locations.

cc @navidcy

[navidcy]

@siddharthabishnu, how do we know that now the metrics are filled correctly? you were comparing with a grid from MITgcm?

[siddharthabishnu]

@siddharthabishnu, how do we know that now the metrics are filled correctly? you were comparing with a grid from MITgcm?

Yes. Consider the following sources for the grid metrics:

1. the cs32 grid with one halo layer, used by:
   • Ali for the Rossby-Haurwitz test case in Oceananigans v0.82.0; and
   • yourself to check the interior coordinates and grid metrics of the Oceananigans cc32 grid;
2. the cs32 grid with 4 halo layers created by @jm-c using MITgcm;
3. the cc32 grid created by Oceananigans.

In the validation scripts for solid body rotation and the Rossby-Haurwitz wave within the ncc-glw/cubed-sphere-dynamics branch associated with PR #3306, I compared the metrics of grid (3) against both (1) and (2), and plotted their absolute and relative differences. With the latest modifications, these differences have been minimized though not entirely eliminated.

[navidcy]

@siddharthabishnu, how do we know that now the metrics are filled correctly? you were comparing with a grid from MITgcm?

Yes. Consider the following sources for the grid metrics:

1. the cs32 grid with one halo layer, used by:

   • Ali for the Rossby-Haurwitz test case in Oceananigans v0.82.0; and
   • yourself to check the interior coordinates and grid metrics of the Oceananigans cc32 grid;

2. the cs32 grid with 4 halo layers created by @jm-c using MITgcm;

3. the cc32 grid created by Oceananigans.

In the validation scripts for solid body rotation and the Rossby-Haurwitz wave within the ncc-glw/cubed-sphere-dynamics branch associated with PR #3306, I compared the metrics of grid (3) against both (1) and (2), and plotted their absolute and relative differences. With the latest modifications, these differences have been minimized though not entirely eliminated.

OK, so the benchmark is the cs32 grid by MITgcm. Can we do the comparison in this PR? I'd like to see a test ideally because otherwise how do we assess that the changes we are suggesting here are correct. I can also do that, just give me a code snippet that loads the two grids?

[navidcy]

I cleaned the PR a bit. Didn't change anything.

[siddharthabishnu]

@siddharthabishnu, CUDA.@allowscalar introduced by 3cdd470 is detrimental for performance. Like it induces O(10-100x) slowdown I think....
Is this a temporary solution?
cc @glwagner, @simone-silvestri

I think it is. We can try to see if this works on one GPU. If it does we can keep the allowscalar for the moment otherwise we can remove them. In the end all this will have to live in a kernel. Note that this will still not work on multiple GPUs as you cannot explicitly access one region from another one on a different GPU without switching to the device that holds the data.

@simone-silvestri and @navidcy, I totally agree. I only introduced CUDA.@allowscalar under the impression it was necessary for certain GPU tests to pass. However, now understanding that isn’t the case, I've removed it in commit 7f54c3c.

[navidcy]

is it clean/good idea to extend fill_halo_regions! for (::Tuple{Fields})? I guess we need something like that when dealing with vector-like quantities

cc @glwagner, @siddharthabishnu, @simone-silvestri

[simone-silvestri]

there is already a method for filling the halos of tuples (and namedtuples) of fields. In that case, we launch kernels that fill all the fields together.
See

Oceananigans.jl/src/Fields/field_tuples.jl

Line 56 in f7042b7

function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; kwargs...)

and

Oceananigans.jl/src/BoundaryConditions/fill_halo_regions.jl

Lines 45 to 48 in f7042b7

	const MaybeTupledData = Union{OffsetArray, NTuple{<:Any, OffsetArray}}
	"Fill halo regions in ``x``, ``y``, and ``z`` for a given field's data."
	function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, loc, grid, args...; kwargs...)

[simone-silvestri]

In case of a MultiRegionField (of which a CubedSphereField` is a subset) this behavior is bypassed

Oceananigans.jl/src/MultiRegion/multi_region_boundary_conditions.jl

Lines 34 to 36 in f7042b7

	# This can be implemented once we have a buffer for field_tuples
	@inline function tupled_fill_halo_regions!(full_fields, grid::MultiRegionGrids, args...; kwargs...)
	for field in full_fields

[navidcy]

there are few outstanding questions still so I took back the approval -- sorry