reenable tests

Author: simone-silvestri

src/DistributedComputations/distributed_fields.jl

@@ -1,6 +1,15 @@
 using Oceananigans.Grids: topology
 using Oceananigans.Fields: validate_field_data, indices, validate_boundary_conditions
 using Oceananigans.Fields: validate_indices, set_to_array!, set_to_field!
+using CUDA: @allowscalar
+
+using Oceananigans.Fields: ReducedAbstractField, 
+                           get_neutral_mask,
+                           condition_operand,
+                           initialize_reduced_field!,
+                           filltype,
+                           reduced_dimensions,
+                           reduced_location
 
 import Oceananigans.Fields: Field, location, set!
 import Oceananigans.BoundaryConditions: fill_halo_regions!
@@ -95,63 +104,66 @@ function reconstruct_global_field(field::DistributedField)
     return global_field
 end
 
-function maybe_all_reduce!(op, f::ReducedDistributedField, dims)
-    reduced_dims = reduced_dimensions(f)
+partition_dimensions(f::DistributedField) = partition_dimensions(architecture(f))
+function partition_dimensions(arch::Distributed) 
+    R = ranks(arch) 
+    dims = []
+    for r in eachindex(R)
+        if R[r] > 1
+            push!(dims, r)
+        end
+    end
+    return tuple(dims...)
+end
+
+function maybe_all_reduce!(op, f::ReducedAbstractField)
+    reduced_dims   = reduced_dimensions(f)
+    partition_dims = partition_dimensions(f)
 
-    if any(reduced_dims .∈ tuple(dims...))
+    if any([dim ∈ partition_dims for dim in reduced_dims])
         all_reduce!(op, interior(f), architecture(f))
     end
 
     return f
 end
 
-function maybe_all_reduce!(op, f::ReducedDistributedField, ::Colon)
-    all_reduce!(op, interior(f), architecture(f))
-    return f
-end
-
 # Allocating and in-place reductions
 for (reduction, all_reduce_op) in zip((:sum, :maximum, :minimum, :all, :any, :prod),
-                                      ( +,    max,      min,      &,    |,    *))
+                                      (:+,   :max,     :min,     :&,   :|,   :*))
 
     reduction! = Symbol(reduction, '!')
 
     @eval begin
-
         # In-place
         function Base.$(reduction!)(f::Function,
                                     r::ReducedAbstractField,
                                     a::DistributedField;
                                     condition = nothing,
                                     mask = get_neutral_mask(Base.$(reduction!)),
-                                    dims = :,
                                     kwargs...)
 
             operand = condition_operand(f, a, condition, mask)
 
             Base.$(reduction!)(identity,
                                interior(r),
                                operand;
-                               dims,
                                kwargs...)
 
-            return maybe_all_reduce!(all_reduce_op, r, dims)
+            return maybe_all_reduce!($(all_reduce_op), r)
         end
 
         function Base.$(reduction!)(r::ReducedAbstractField,
                                     a::DistributedField;
                                     condition = nothing,
                                     mask = get_neutral_mask(Base.$(reduction!)),
-                                    dims = :,
                                     kwargs...)
 
             Base.$(reduction!)(identity,
                                interior(r),
                                condition_operand(a, condition, mask);
-                               dims,
                                kwargs...)
 
-            return maybe_all_reduce!(all_reduce_op, r, dims)
+            return maybe_all_reduce!($(all_reduce_op), r)
         end
 
         # Allocating
@@ -168,10 +180,10 @@ for (reduction, all_reduce_op) in zip((:sum, :maximum, :minimum, :all, :any, :pr
             initialize_reduced_field!(Base.$(reduction!), identity, r, conditioned_c)
             Base.$(reduction!)(identity, interior(r), conditioned_c, init=false)
 
-            maybe_all_reduce!(all_reduce_op, r, dims)
+            maybe_all_reduce!($(all_reduce_op), r)
 
             if dims isa Colon
-                return CUDA.@allowscalar first(r)
+                return @allowscalar first(r)
             else
                 return r
             end

test/test_distributed_models.jl

@@ -408,7 +408,6 @@ end
 
 @testset "Distributed MPI Oceananigans" begin
     @info "Testing distributed MPI Oceananigans..."
-    #=
     @testset "Multi architectures rank connectivity" begin
         @info "  Testing multi architecture rank connectivity..."
         test_triply_periodic_rank_connectivity_with_411_ranks()
@@ -462,7 +461,7 @@ end
             @test child_architecture(architecture(cpuosg)) == CPU()
         end
     end
-    =#
+
     @testset "Distributed reductions" begin
         child_arch = get(ENV, "TEST_ARCHITECTURE", "CPU") == "GPU" ? GPU() : CPU()
 
@@ -479,11 +478,24 @@ end
             @test sum(c) == 1*N + 2*N + 3*N + 4*N
 
             sum!(c_reduced, c)
-            @test CUDA.@allowscalar c_reduced.data[1, 1, 1] == 1*N + 2*N + 3*N + 4*N
+            @test CUDA.@allowscalar c_reduced[1, 1, 1] == 1*N + 2*N + 3*N + 4*N
+
+            cbool = CenterField(grid, Bool)
+            cbool_reduced = Field{Nothing, Nothing, Nothing}(grid, Bool)
+            bool_val = arch.local_rank == 0 ? true : false        
+            set!(cbool, bool_val)            
+
+            @test any(cbool) == true
+            @test all(cbool) == false
+
+            any!(cbool_reduced, cbool)
+            @test CUDA.@allowscalar cbool_reduced[1, 1, 1] == true
+            
+            all!(cbool_reduced, cbool)
+            @test CUDA.@allowscalar cbool_reduced[1, 1, 1] == false
         end
     end
 
-    #=
     # Only test on CPU because we do not have a GPU pressure solver yet
     @testset "Time stepping NonhydrostaticModel" begin
         child_arch = get(ENV, "TEST_ARCHITECTURE", "CPU") == "GPU" ? GPU() : CPU()
@@ -521,6 +533,5 @@ end
         @test model isa ShallowWaterModel
         @test model.clock.time ≈ 2
     end
-    =#
 end