Distributed: extend MPI functions for arch::Distributed?

[glwagner]

In all the scripts I found on distributed simulations, we seem to query the MPI state using COMM_WORLD, eg:

Oceananigans.jl/validation/distributed_simulations/mpi_nonhydrostatic_two_dimensional_turbulence.jl

Line 74 in 5dd96cf

rank = MPI.Comm_rank(MPI.COMM_WORLD)

and

Oceananigans.jl/validation/distributed_simulations/mpi_nonhydrostatic_two_dimensional_turbulence.jl

Lines 59 to 60 in 5dd96cf

	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)

But Distributed allows the communicator to be set, so arch.communicator can be something other than COMM_WORLD, eg

Oceananigans.jl/src/DistributedComputations/distributed_architectures.jl

Lines 60 to 66 in 5dd96cf

	function Distributed(child_architecture = CPU();
	topology,
	ranks,
	devices = nothing,
	enable_overlapped_computation = true,
	communicator = MPI.COMM_WORLD)

Presumably we should never write scripts that make assumptions about the communicator (and we should strongly discourage users from doing so). So first of all the validation scripts must be changed.

Secondly, I think something that could encouarge clean and good practices would be to extend functions like Comm_rank to arch::Distributed and possibly to the grid as well, so that we can call:

rank = MPI.Comm_rank(arch)
rank = MPI.Comm_rank(grid)
Nranks = MPI.Comm_size(arch)

We'd have the option of throwing an error when non-distributed, or returning some sensible fallback like Comm_rank(arch) = 0 and Comm_size(arch) = 1.

Alternatively, we can develop our own API, since the MPI-based one is rather irregular and abbreviated, thus difficult for newbies to understand (I've been in this boat...)

For example, perhaps simply Base.size and rank would suffice:

Base.size(arch::Distributed) = MPI.Comm_size(arch.communicator)
Base.size(arch::AbstractArchitecture) = 1

rank(arch::Distributed) = MPI.Comm_rank(arch.communicator)
rank(arch::AbstractArchitecture) = 0