Default partition in Distributed (#3496)

* correct partition

* inothing -> isnothing

* does not work on julia 1.9?

* hopefully last one?

* this was the last

* other ones slipped

* other ones remaining

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Author: simone-silvestri

src/AbstractOperations/conditional_operations.jl

@@ -1,6 +1,6 @@
 using Oceananigans.Fields: OneField
 using Oceananigans.Grids: architecture
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 import Oceananigans.Fields: condition_operand, conditional_length, set!, compute_at!, indices
 
 # For conditional reductions such as mean(u * v, condition = u .> 0))

src/DistributedComputations/distributed_architectures.jl

@@ -204,16 +204,19 @@ Keyword arguments
                   if not for testing or developing. Change at your own risk!
 """
 function Distributed(child_architecture = CPU(); 
-                     communicator = MPI.COMM_WORLD,
+                     communicator = nothing,
                      devices = nothing, 
                      synchronized_communication = false,
-                     partition = Partition(MPI.Comm_size(communicator)))
+                     partition = nothing)
 
     if !(MPI.Initialized())
         @info "MPI has not been initialized, so we are calling MPI.Init()."
         MPI.Init()
     end
 
+    communicator = isnothing(communicator) ? MPI.COMM_WORLD : communicator
+    partition    = isnothing(partition) ? Partition(MPI.Comm_size(communicator)) : partition
+
     ranks = size(partition)
     Rx, Ry, Rz = ranks
     total_ranks = Rx * Ry * Rz

src/DistributedComputations/distributed_on_architecture.jl

@@ -1,5 +1,6 @@
 using CUDA: CuArray
 using OffsetArrays
+import Oceananigans.Architectures: on_architecture
 
 # We do not support switching from distributed and serial through `on_architecture`.
 # We only support moving a type from CPU to GPU and the other way around

src/DistributedComputations/partition_assemble.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 
 all_reduce(op, val, arch::Distributed) = 
     MPI.Allreduce(val, op, arch.communicator)

src/ImmersedBoundaries/grid_fitted_bottom.jl

@@ -4,7 +4,6 @@ using OffsetArrays: OffsetArray
 using Oceananigans.Utils: getnamewrapper
 using Oceananigans.Grids: total_size
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Oceananigans.BoundaryConditions: FBC
 using Printf
 

src/ImmersedBoundaries/mask_immersed_field.jl

@@ -1,6 +1,5 @@
 using KernelAbstractions: @kernel, @index
 using Statistics
-using Oceananigans.Architectures: architecture
 using Oceananigans.Fields: location, ZReducedField, Field
 
 instantiate(T::Type) = T()

src/ImmersedBoundaries/partial_cell_bottom.jl

@@ -1,6 +1,5 @@
 using Oceananigans.Utils: prettysummary
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Printf
 
 #####

src/TurbulenceClosures/implicit_explicit_time_discretization.jl

@@ -1,4 +1,3 @@
-using Oceananigans.Utils: on_architecture
 using Oceananigans.Grids: AbstractGrid
 
 abstract type AbstractTimeDiscretization end

src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators: ℑzᵃᵃᶜ

src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators
 using Oceananigans.Grids: inactive_node