如何理解MPI_Alltoallv中的rdispls（recv_displacement）参数？

我认为rdispls（接收位移）应该是recv_count的累积量——对第i个进程来说，接收的数据会被写入recv_data + rdispls(i)*extent(recv_type)的位置。为验证这一点，我将一个Alltoallv示例改写为Python代码，但因recv_disp及发送参数中的位移设置错误，得到了错误结果。

错误代码示例 #

from mpi4py import MPI
import numpy as np

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

assert size == 3

# Send data
if rank == 0:
    send_data = np.array([0, 100, 200],   dtype=np.int32)
    send_rank = np.array([0, 1,   1],     dtype=np.int32)
elif rank == 1:
    send_data = np.array([300, 400, 500], dtype=np.int32)
    send_rank = np.array([2,   2,   2],   dtype=np.int32)
else:
    send_data = np.array([600],           dtype=np.int32)
    send_rank = np.array([0],             dtype=np.int32)

send_length = len(send_data)
send_counts = np.zeros(size, dtype=np.int32)
for procs_id in send_rank:
    send_counts[procs_id] += 1

send_disp     = np.zeros(size, dtype=np.int32)
send_disp[1:] = np.cumsum(send_counts)[:-1]

# Receive data
recv_counts = np.zeros(size, dtype=np.int32)
comm.Alltoall(send_counts, recv_counts)
recv_length = np.sum(recv_counts)
recv_data = np.zeros(recv_length, dtype=np.int32)
recv_disp = np.zeros(size, dtype=np.int32)
recv_disp[1:] = np.cumsum(recv_counts)[:-1]         # Comment this and get correct answer

# Alltoallv
comm.Alltoallv([send_data, send_counts, recv_disp, MPI.INT],
               [recv_data, recv_counts, recv_disp, MPI.INT])

print(f"rank:\n{rank}\n"
      f"send_counts:\n{send_counts}\n"
      f"recv_counts:\n{recv_counts}\n"
      f"send_disp:\n{send_disp}\n"
      f"recv_disp:\n{recv_disp}\n"
      f"send_data:\n{send_data}\n"
      f"recv_data:\n{recv_data}\n"
      f"recv_length:\n{recv_length}\n"
      )

错误点分析 #

1. 发送位移参数传错：调用Alltoallv时，发送参数的位移应该用send_disp，但代码里错误地传入了recv_disp，导致发送数据的起始位置计算错误。
2. 注：recv_disp的计算逻辑是正确的——它确实是recv_counts的累积前缀和（不包含当前元素），用来标记从哪个位置开始写入对应进程发来的数据。

修正后的代码 #

from mpi4py import MPI
import numpy as np

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

assert size == 3

# Send data
if rank == 0:
    send_data = np.array([0, 100, 200],   dtype=np.int32)
    send_rank = np.array([0, 1,   1],     dtype=np.int32)
elif rank == 1:
    send_data = np.array([300, 400, 500], dtype=np.int32)
    send_rank = np.array([2,   2,   2],   dtype=np.int32)
else:
    send_data = np.array([600],           dtype=np.int32)
    send_rank = np.array([0],             dtype=np.int32)

send_length = len(send_data)
send_counts = np.zeros(size, dtype=np.int32)
for procs_id in send_rank:
    send_counts[procs_id] += 1

send_disp     = np.zeros(size, dtype=np.int32)
send_disp[1:] = np.cumsum(send_counts)[:-1]

# Receive data
recv_counts = np.zeros(size, dtype=np.int32)
comm.Alltoall(send_counts, recv_counts)
recv_length = np.sum(recv_counts)
recv_data = np.zeros(recv_length, dtype=np.int32)
recv_disp = np.zeros(size, dtype=np.int32)
recv_disp[1:] = np.cumsum(recv_counts)[:-1]

# Alltoallv：修正发送位移参数为send_disp
comm.Alltoallv([send_data, send_counts, send_disp, MPI.INT],
               [recv_data, recv_counts, recv_disp, MPI.INT])

print(f"rank:\n{rank}\n"
      f"send_counts:\n{send_counts}\n"
      f"recv_counts:\n{recv_counts}\n"
      f"send_disp:\n{send_disp}\n"
      f"recv_disp:\n{recv_disp}\n"
      f"send_data:\n{send_data}\n"
      f"recv_data:\n{recv_data}\n"
      f"recv_length:\n{recv_length}\n"
      )

修正后结果验证 #

• 进程0的recv_data会得到[0, 600]（自身发送的0，加上进程2发送的600）
• 进程1的recv_data会得到[100, 200]（进程0发送的两个数据）
• 进程2的recv_data会得到[300, 400, 500]（进程1发送的三个数据）

内容的提问来源于stack exchange，提问作者Yongxin