From 6d76caecec4888e8d3db98a1f61f144b36abe772 Mon Sep 17 00:00:00 2001 From: Ivan Utkin Date: Fri, 27 Oct 2023 17:33:01 +0200 Subject: [PATCH] Fix gather for 2D --- src/Distributed/gather.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/Distributed/gather.jl b/src/Distributed/gather.jl index d5a31b70..c8f066b9 100644 --- a/src/Distributed/gather.jl +++ b/src/Distributed/gather.jl @@ -3,7 +3,8 @@ function gather!(dst::Union{AbstractArray,Nothing}, src::AbstractArray, comm::MP dims = Tuple(dims) if MPI.Comm_rank(comm) == root # make subtype for gather - subtype = MPI.Types.create_subarray(size(dst), size(src), (0, 0), MPI.Datatype(eltype(dst))) + offset = Tuple(0 for _ in 1:ndims(src)) + 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 @@ -15,7 +16,7 @@ function gather!(dst::Union{AbstractArray,Nothing}, src::AbstractArray, comm::MP displs[i, j] = d d += 1 end - d += (size(src, 2) - 1) * dims[2] + d += (size(src, 2) - 1) * dims[1] end # transpose displs as cartesian communicator is row-major recvbuf = MPI.VBuffer(dst, vec(counts), vec(displs'), subtype)