Implement automatically serializing gather (#756)

* Implement automatically serializing `gather` * docs: add missing reference to gather * bump patch --------- Co-authored-by: Simon Byrne <[email protected]>
JuliaParallel · Jul 29, 2023 · 3882775 · 3882775
1 parent fd2c626
commit 3882775
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "MPI"
 uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 authors = []
-version = "0.20.12"
+version = "0.20.13"
 
 [deps]
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"

diff --git a/docs/src/reference/collective.md b/docs/src/reference/collective.md
@@ -22,6 +22,7 @@ MPI.bcast
 ```@docs
 MPI.Gather!
 MPI.Gather
+MPI.gather
 MPI.Gatherv!
 MPI.Allgather!
 MPI.Allgather

diff --git a/src/collective.jl b/src/collective.jl
@@ -294,6 +294,40 @@ Gather(sendbuf::AbstractArray, root::Integer, comm::Comm) =
 Gather(object::T, root::Integer, comm::Comm) where {T} =
     Gather!(Ref(object), Comm_rank(comm) == root ? Array{T}(undef, Comm_size(comm)) : nothing, root, comm)
 
+"""
+    gather(obj, comm::Comm; root::Integer=0)
+
+Gather the objects `obj` from all ranks on `comm` to rank `root`. This is able to to handle arbitrary data. On `root`, it returns a vector of the objects, and `nothing` otherwise.
+
+# See also
+
+- [`Gather!`](@ref)
+"""
+function gather(obj, comm::Comm; root::Integer=0)
+    isroot = Comm_rank(comm) == root
+
+    sendbuf = MPI.serialize(obj)
+    count = length(sendbuf)
+
+    counts = Gather(count, comm; root = root)
+
+    if isroot
+        data = Array{UInt8}(undef, sum(counts))
+        recvbuf = VBuffer(data, counts)
+
+        Gatherv!(sendbuf, recvbuf, comm; root = root)
+
+        objs = [
+            MPI.deserialize(view(recvbuf.data, displ+1:displ+count)) for (displ, count) in zip(recvbuf.displs, recvbuf.counts)
+        ]
+        return objs
+    else
+        Gatherv!(sendbuf, nothing, comm; root = root)
+        return nothing
+    end
+end
+
+
 """
     Gatherv!(sendbuf, recvbuf, comm::Comm; root::Integer=0)
 

diff --git a/test/test_gather.jl b/test/test_gather.jl
@@ -76,6 +76,25 @@ for T in MPITestTypes
         @test A == ArrayType{T}(repeat(1:sz, inner=4))
     end
 
+    # serializing version
+    C = MPI.gather(T(rank+1), comm; root=root)
+    if isroot
+        @test C isa Vector{T}
+        @test C == Vector{T}(1:sz)
+    else
+        @test C === nothing
+    end
+end
+
+
+objs = ["test", 1, Array{Int}, [1,"test"]]
+obj = objs[mod(rank, length(objs))+1]
+
+C = MPI.gather(obj, comm; root=root)
+if isroot
+    @test C == [objs[mod1(i, length(objs))] for i = 1:sz]
+else
+    @test C === nothing
 end
 
 MPI.Finalize()