Skip to content

Commit

Permalink
Update Distributed (#49)
Browse files Browse the repository at this point in the history
This PR removes the built-in `gather` functionality in favour of the new generic implementation from ImplicitGlobalGrid.jl
  • Loading branch information
omlins authored Jan 23, 2024
1 parent 5fc5896 commit 5c27226
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 24 deletions.
4 changes: 3 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
name = "FastIce"
uuid = "e0de9f13-a007-490e-b696-b07d031015ca"
authors = ["Ludovic Raess <[email protected]>, Ivan Utkin <[email protected]> and contributors"]
authors = ["Ludovic Raess <[email protected]>, Ivan Utkin <[email protected]>, Samuel Omlin <[email protected]> and contributors"]
version = "0.2.0"

[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0"
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
LightXML = "9c8b4983-aa76-5018-a973-4c85ecc9e179"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Expand All @@ -32,6 +33,7 @@ CUDA = "5"
ElasticArrays = "1"
GeometryBasics = "0.4"
HDF5 = "0.17"
ImplicitGlobalGrid = "0.15"
KernelAbstractions = "0.9"
LightXML = "0.9"
MPI = "0.20"
Expand Down
1 change: 1 addition & 0 deletions src/Distributed/Distributed.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import FastIce.BoundaryConditions: apply_boundary_conditions!

using MPI
using KernelAbstractions
import ImplicitGlobalGrid

"Trait structure used as a type parameter to indicate that the Architecture is a distributed MPI Architecture."
struct DistributedMPI end
Expand Down
24 changes: 1 addition & 23 deletions src/Distributed/gather.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,29 +7,7 @@ The array will be gathered on the process with id `root` (`root=0` by default).
Note that the memory for a global array should be allocated only on the process with id `root`, on other processes `dst` can be set to `nothing`.
"""
function gather!(dst::Union{AbstractArray{T,N},Nothing}, src::AbstractArray{T,N}, comm::MPI.Comm; root=0) where {T,N}
dims, _, _ = MPI.Cart_get(comm)
dims = Tuple(dims)
if MPI.Comm_rank(comm) == root
# make subtype for gather
offset = Tuple(0 for _ in 1:N)
subtype = MPI.Types.create_subarray(size(dst), size(src), offset, MPI.Datatype(eltype(dst)))
subtype = MPI.Types.create_resized(subtype, 0, size(src, 1) * Base.elsize(dst))
MPI.Types.commit!(subtype)
# make VBuffer for collective communication
counts = fill(Cint(1), reverse(dims)) # gather one subarray from each MPI rank
displs = zeros(Cint, reverse(dims)) # reverse dims since MPI Cart comm is row-major
csizes = cumprod(size(src)[2:end] .* dims[1:end-1])
strides = (1, csizes...)
for I in CartesianIndices(displs)
offset = reverse(Tuple(I - oneunit(I)))
displs[I] = sum(offset .* strides)
end
recvbuf = MPI.VBuffer(dst, vec(counts), vec(displs), subtype)
MPI.Gatherv!(src, recvbuf, comm; root)
else
MPI.Gatherv!(src, nothing, comm; root)
end
return
ImplicitGlobalGrid.gather!(src, dst, comm; root=root)
end

"""
Expand Down

0 comments on commit 5c27226

Please sign in to comment.