Add distributed reductions

[navidcy]

This PR adds reductions for distributed fields.

When merged, should allow #4470 to go through

[simone-silvestri]

maybe we should find a more elegant way to compute this?

[glwagner]

also docstring? what is it doing?

[simone-silvestri]

checking which are the ``partitioned'' dimensions

[navidcy]

@simone-silvestri could you help with this error?

https://buildkite.com/clima/oceananigans-distributed/builds/7867#0196c42f-389a-4c42-93f3-ff90536a60bd/225-1583

[glwagner]

How do you reduce across a dimension that is partitioned? You need to gather at the end (and then scatter to all devices)?

[simone-silvestri]

we compute the reduction locally and then allreduce through all devices

[glwagner]

we compute the reduction locally and then allreduce through all devices

can you point to those lines?

[simone-silvestri]

the reduction is here

Oceananigans.jl/src/DistributedComputations/distributed_fields.jl

Line 152 in f582f25

return maybe_all_reduce!($(all_reduce_op), r)

which is this function over here

Oceananigans.jl/src/DistributedComputations/distributed_fields.jl

Lines 119 to 128 in f582f25

	function maybe_all_reduce!(op, f::ReducedAbstractField)
	reduced_dims = reduced_dimensions(f)
	partition_dims = partition_dimensions(f)
	if any([dim ∈ partition_dims for dim in reduced_dims])
	all_reduce!(op, interior(f), architecture(f))
	end
	return f
	end

It has to be done only if (at least one) of the dimensions we are reducing are partitioned in different ranks, otherwise it is not necessary

[glwagner]

ok nice, for some reason I didn't see that

[navidcy]

@simone-silvestri any idea why this comes up?
https://buildkite.com/clima/oceananigans-distributed/builds/7902#0196ca9b-cf01-44ee-bc78-8b4226e78e05/207-1426

  ArgumentError: Illegal conversion of a CUDA.DeviceMemory to a Ptr{Float64}

[simone-silvestri]

@simone-silvestri any idea why this comes up? https://buildkite.com/clima/oceananigans-distributed/builds/7902#0196ca9b-cf01-44ee-bc78-8b4226e78e05/207-1426

  ArgumentError: Illegal conversion of a CUDA.DeviceMemory to a Ptr{Float64}

Looks like using parent(field) rather than interior(field) works in the reduction. On hindsight this is quite clear since CUDA-aware MPI does not support passing non-contiguous data.