CODEWITHSUNDEEP
pythonnumpympimpi4py
ray_lv
ray_lv

Reputation: 105

Not able to correctly use `comm.Gatherv()` to gather uneven sized numpy arrays

I am learning MPI4Py and I wanted to implement a simple program.

Explanation
Here, each rank has a send_array of size rank+1 and values equal to rank+1 respectively.

rank0 = [1]
rank1 = [2 2]
rank2 = [3 3 3]
rank3 = [4 4 4 4]

I want to gather the values to rank=0 to the buffer rbuf. Its size is equal to the total size of all local send_arrays, i.e. 1+2+3+4 = 10.

Program

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

send_array = np.ones(rank+1).astype(int) * (rank + 1)

print(rank, send_array)
if rank == 0:
    gather_size = np.array([sum([i+1 for i in range(size)])])
    print(gather_size)

    rbuf = np.zeros(gather_size[0]).astype(int)
else:
    gather_size = None
    rbuf = None

# comm.Gatherv(sendbuf, recvbuf=(recvbuf, recvcounts, displs, datatype), root=0)
comm.Gatherv(sendbuf=send_array, recvbuf=(rbuf, (1,2,3,4),(0,1,3,6), MPI.INT), root=0)

 
if rank == 0:
    print(rbuf, len(rbuf))

Observation/Errors/Doubts

I manually hard-coded the values to be more explicit as I have doubts on the values of recvcounts and displs. I am getting the following error:

comm.Gatherv(sendbuf=send_array, recvbuf=(rbuf, (1,2,3,4),(0,1,3,6), MPI.INT), root=0)
  File "mpi4py/MPI/Comm.pyx", line 724, in mpi4py.MPI.Comm.Gatherv
mpi4py.MPI.Exception: MPI_ERR_TRUNCATE: message truncated

The result I am expecting is:

[1 2 2 3 3 3 4 4 4 4]

After playing around a little, I also realized that the code works if I keep the size of rbuf equal to [maximum_local_array_size*number_of_ranks], which in this case could be 4*4, given I have 4 ranks.

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# send_array = np.ones(10).astype(int) * (rank + 1)
send_array = np.ones(rank+1).astype(int) * (rank + 1)

print(rank, send_array)
if rank == 0:
    gather_size = np.array([16])
    print(gather_size)

    rbuf = np.zeros(gather_size[0]).astype(int)
else:
    gather_size = None
    rbuf = None

# comm.Gatherv(sendbuf, recvbuf=(recvbuf, recvcounts, displs, datatype), root=0)
comm.Gatherv(sendbuf=send_array, recvbuf=(rbuf, MPI.INT), root=0)
 
if rank == 0:
    print(rbuf, len(rbuf))

The result I get is not what I wanted, but I don't get any errors. Note that this time I also don't specify recvcounts and displs. So I am not sure what exact values are used by default for these arguments, since those seem to work.

[1 0 0 0 2 2 0 0 3 3 3 0 4 4 4 4]

Upvotes: 0

Views: 59

Answers (1)

Gilles Gouaillardet
Gilles Gouaillardet

Reputation: 8380

Here is an example on how to manually specify receive counts and displacements

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

sbuf = np.ones(rank+1).astype(int) * (rank + 1)

def triangular(n):
    return n * (n+1) // 2

if rank == 0:
    gather_size = triangular(size)
    rcounts = np.arange(1, size+1)
    rdispls = np.array([triangular(i) for i in range(size)])
    rbuf = np.zeros(gather_size, dtype=int)
else:
    gather_size = None
    rbuf = None
    rcounts = None
    rdispls = None

comm.Gatherv(sendbuf=sbuf, recvbuf=(rbuf, (rcounts, rdispls)), root=0)
 
if rank == 0:
    print(rbuf, len(rbuf))

Upvotes: 0

Related Questions