Repository: CliMA/Oceananigans.jl
Pull request #4497: Add distributed reductions
Author: login: navidcy
Labels: name: distributed 🕸️
description: Our plan for total cluster domination
Created: 2025-05-12T10:07:31Z | Merged: 2025-05-22T03:40:58Z

This PR adds reductions for distributed fields.

When merged, should allow https://github.com/CliMA/Oceananigans.jl/pull/4470 to go through


Diff:
diff --git a/src/DistributedComputations/distributed_fields.jl b/src/DistributedComputations/distributed_fields.jl
index 8b32a6ac7f4..30d26944d7f 100644
--- a/src/DistributedComputations/distributed_fields.jl
+++ b/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!
@@ -94,3 +103,97 @@ function reconstruct_global_field(field::DistributedField)
 
     return global_field
 end
+
+"""
+    partition_dimensions(arch::Distributed)
+    partition_dimensions(f::DistributedField)
+
+Return the partitioned dimensions of a distributed field or architecture.
+"""
+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
+
+partition_dimensions(f::DistributedField) = partition_dimensions(architecture(f))
+
+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, parent(f), architecture(f))
+    end
+
+    return f
+end
+
+# Allocating and in-place reductions
+for (reduction, all_reduce_op) in zip((:sum, :maximum, :minimum, :all, :any, :prod),
+                                      (:+,   :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!)),
+                                    kwargs...)
+
+            operand = condition_operand(f, a, condition, mask)
+
+            Base.$(reduction!)(identity,
+                               interior(r),
+                               operand;
+                               kwargs...)
+
+            return maybe_all_reduce!($(all_reduce_op), r)
+        end
+
+        function Base.$(reduction!)(r::ReducedAbstractField,
+                                    a::DistributedField;
+                                    condition = nothing,
+                                    mask = get_neutral_mask(Base.$(reduction!)),
+                                    kwargs...)
+
+            Base.$(reduction!)(identity,
+                               interior(r),
+                               condition_operand(a, condition, mask);
+                               kwargs...)
+
+            return maybe_all_reduce!($(all_reduce_op), r)
+        end
+
+        # Allocating
+        function Base.$(reduction)(f::Function,
+                                   c::DistributedField;
+                                   condition = nothing,
+                                   mask = get_neutral_mask(Base.$(reduction!)),
+                                   dims = :)
+
+            conditioned_c = condition_operand(f, c, condition, mask)
+            T = filltype(Base.$(reduction!), c)
+            loc = reduced_location(location(c); dims)
+            r = Field(loc, c.grid, T; indices=indices(c))
+            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)
+
+            if dims isa Colon
+                return @allowscalar first(r)
+            else
+                return r
+            end
+        end
+    end
+end
diff --git a/src/Fields/field.jl b/src/Fields/field.jl
index 5918134c935..80c313609c8 100644
--- a/src/Fields/field.jl
+++ b/src/Fields/field.jl
@@ -727,7 +727,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
     end
 end
 
-# Improve me! We can should both the extrama in one single reduction instead of two
+# Improve me! We can should both the extrema in one single reduction instead of two
 Base.extrema(c::AbstractField; kwargs...) = (minimum(c; kwargs...), maximum(c; kwargs...))
 Base.extrema(f, c::AbstractField; kwargs...) = (minimum(f, c; kwargs...), maximum(f, c; kwargs...))
 
diff --git a/test/test_distributed_models.jl b/test/test_distributed_models.jl
index 56e129a7d69..8e0f87373bd 100644
--- a/test/test_distributed_models.jl
+++ b/test/test_distributed_models.jl
@@ -74,7 +74,6 @@ include_corners ? view(f.data, :, :, right_halo_indices(instantiate(LZ), instant
                                interior_indices(instantiate(LY), instantiate(topology(f, 2)), f.grid.Ny),
                                right_halo_indices(instantiate(LZ), instantiate(topology(f, 3)), f.grid.Nz, f.grid.Hz))
 
-
 function southwest_halo(f::AbstractField)
     Nx, Ny, _ = size(f.grid)
     Hx, Hy, _ = halo_size(f.grid)
@@ -409,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()
@@ -464,6 +462,39 @@ end
         end
     end
 
+    @testset "Distributed reductions" begin
+        child_arch = get(ENV, "TEST_ARCHITECTURE", "CPU") == "GPU" ? GPU() : CPU()
+
+        for partition in [Partition(1, 4), Partition(2, 2), Partition(4, 1)]
+            @info "Time-stepping a distributed NonhydrostaticModel with partition $partition..."
+            arch = Distributed(child_arch; partition)
+            grid = RectilinearGrid(arch, topology=(Periodic, Periodic, Periodic), size=(8, 8, 1), extent=(1, 2, 3))
+            c = CenterField(grid)
+            set!(c, arch.local_rank+1)
+
+            c_reduced = Field{Nothing, Nothing, Nothing}(grid)
+
+            N = grid.Nx * grid.Ny # local rank grid size
+            @test sum(c) == 1*N + 2*N + 3*N + 4*N
+
+            sum!(c_reduced, c)
+            @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


Reviews:
user: login: simone-silvestri
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T10:28:58Z
user: login: simone-silvestri
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T10:29:47Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T13:11:12Z
user: login: simone-silvestri
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T17:04:20Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-19T13:22:03Z

Review comments:
diff_hunk: @@ -94,3 +103,90 @@ function reconstruct_global_field(field::DistributedField)
 
     return global_field
 end
+
+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
path: src/DistributedComputations/distributed_fields.jl
user: login: simone-silvestri
body: maybe we should find a more elegant way to compute this?
created_at: 2025-05-12T10:29:47Z
author_association: COLLABORATOR
diff_hunk: @@ -94,3 +103,90 @@ function reconstruct_global_field(field::DistributedField)
 
     return global_field
 end
+
+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
path: src/DistributedComputations/distributed_fields.jl
user: login: glwagner
body: also docstring? what is it doing?
created_at: 2025-05-12T13:11:12Z
author_association: MEMBER
diff_hunk: @@ -94,3 +103,90 @@ function reconstruct_global_field(field::DistributedField)
 
     return global_field
 end
+
+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
path: src/DistributedComputations/distributed_fields.jl
user: login: simone-silvestri
body: checking which are the ``partitioned'' dimensions
created_at: 2025-05-13T17:04:20Z
author_association: COLLABORATOR
diff_hunk: @@ -116,6 +121,9 @@ function partition_dimensions(arch::Distributed)
     return tuple(dims...)
 end
 
+partition_dimensions(f::DistributedField) = partition_dimensions(architecture(f))
+
path: src/DistributedComputations/distributed_fields.jl
user: login: glwagner
body: ```suggestion
```
created_at: 2025-05-19T13:22:03Z
author_association: MEMBER

PR comments:
user: login: navidcy
created_at: 2025-05-12T11:46:05Z
body: @simone-silvestri could you help with this error?

https://buildkite.com/clima/oceananigans-distributed/builds/7867#0196c42f-389a-4c42-93f3-ff90536a60bd/225-1583

author_association: MEMBER
user: login: glwagner
created_at: 2025-05-12T13:12:00Z
body: How do you reduce across a dimension that is partitioned? You need to gather at the end (and then scatter to all devices)? 
author_association: MEMBER
user: login: simone-silvestri
created_at: 2025-05-13T17:03:21Z
body: we compute the reduction locally and then allreduce through all devices
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-05-13T19:11:04Z
body: > we compute the reduction locally and then allreduce through all devices

can you point to those lines?
author_association: MEMBER
user: login: simone-silvestri
created_at: 2025-05-13T19:41:16Z
body: the reduction is here
https://github.com/CliMA/Oceananigans.jl/blob/f582f255d6014e4188f07343f2725d363272d687/src/DistributedComputations/distributed_fields.jl#L152
which is this function over here
https://github.com/CliMA/Oceananigans.jl/blob/f582f255d6014e4188f07343f2725d363272d687/src/DistributedComputations/distributed_fields.jl#L119-L128
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
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-05-13T20:10:50Z
body: ok nice, for some reason I didn't see that
author_association: MEMBER
user: login: navidcy
created_at: 2025-05-16T04:30:19Z
body: @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}
```
author_association: MEMBER
user: login: simone-silvestri
created_at: 2025-05-20T09:17:29Z
body: > @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.
author_association: COLLABORATOR