mpi4py Allgatherv for matrices

Question

I have a matrix x of size (n0,N) and (n1,N), for two different MPI processes, repsectively. I am trying to gather them to a single matrix, defined on both processes.

Here's my attempt:

"""
run with: mpiexec -n 2 python test.py

"""

from mpi4py import MPI
import numpy as np


comm=MPI.COMM_WORLD

n0=3
n1=4
N=2


if comm.Get_rank()==0:
    x=np.ones((n0,N),dtype=np.float64)
else:
    x=2.0*np.ones((n1,N),dtype=np.float64)


x_gathered = np.zeros((n0+n1,N), x.dtype  )


comm.Allgatherv([x,   MPI.DOUBLE], [x_gathered,  MPI.DOUBLE], )

print(x_gathered)

I am getting an error:

MPIR_Localcopy(46)..........................: Message truncated; 64 bytes received but buffer size is 56

I noticed that when I set the two sizes n0=n1, then the code runs properly. Can someone explain to me why I can't use Allgatherv in the case of different n0 and n1, and how this Allgatherv-operation can be eventually achieved using mpi4py?

Answer 1

I think I figured it out: Allgatherv requires extra arguments for sendcounts and displacements, see here.

"""
mpiexec -n 2 python test.py

"""

from mpi4py import MPI
import numpy as np
from sys import getsizeof


comm=MPI.COMM_WORLD

n0=4
n1=3
N=2


if comm.Get_rank()==0:
    x=np.ones((n0,N,),dtype=np.float64)
else:
    x=2.0*np.ones((n1,N,),dtype=np.float64)


x_gathered = np.zeros(((n0+n1)*N,),dtype=np.float64  )

sendcountes=(n0*N,n1*N)
displacements=(0,n0*N)

comm.Allgatherv([x, MPI.DOUBLE], [x_gathered, sendcountes, displacements, MPI.DOUBLE], )

print(x_gathered)