使用mpi4py处理不同形状数组的并行收集策略

import numpy as np
from mpi4py import MPI

# 初始化MPI通信器
comm = MPI.COMM_WORLD
size = comm.Get_size() # 获取进程总数
rank = comm.Get_rank() # 获取当前进程的rank

# 根据进程rank创建不同形状的数组
# 示例中，rank 1 创建 (2, 3) 数组，其他 rank 创建 (5, 3) 数组
a = np.zeros((2 if rank == 1 else 5, 3), dtype=float) + rank
print(f"进程 {rank}: 发送数组形状 {a.shape}")

# 使用 comm.gather 收集所有进程的数组
# root=0 表示根进程是 rank 0
b_list = comm.gather(a, root=0)

# 根进程 (rank 0) 对收集到的数组进行拼接
if rank == 0:
    # b_list 现在是一个包含所有进程发送数组的列表
    # 例如：[array_from_rank0, array_from_rank1, ...]
    b_concatenated = np.concatenate(b_list)
    print(f"进程 {rank}: 收集并拼接后的数组形状 {b_concatenated.shape}")
    print(f"进程 {rank}: 拼接后的数组内容:\n{b_concatenated}")
else:
    # 非根进程的 b_list 为 None
    pass

# 所有进程都可以打印自己的结果，但只有根进程有拼接后的数组
# print(f"进程 {rank}: 接收到的数据 (非根进程为None): {b_list}")

import numpy as np
from mpi4py import MPI

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

# 示例中，我们假设只有两个进程 (rank 0 和 rank 1)
# 实际应用中需要动态计算 counts 和 displacements
assert size >= 2, "此示例至少需要两个进程"

# 根据进程rank创建不同形状的数组
if rank == 0:
    a = np.zeros((5, 3), dtype=float) + rank
else: # rank == 1
    a = np.zeros((2, 3), dtype=float) + rank

print(f"进程 {rank}: 发送数组形状 {a.shape}")

# 计算总行数，用于根进程分配接收缓冲区
# 在实际应用中，所有进程可能需要通过 Allgather 等方式共享各自的形状信息
# 这里为了简化，我们硬编码了示例的形状
total_rows = 0
if rank == 0:
    # 假设已知所有进程的行数
    # 实际中，可以先 Allgather 各自的 shape[0]
    all_shapes = comm.gather(a.shape, root=0)
    total_rows = sum([s[0] for s in all_shapes])
else:
    comm.gather(a.shape, root=0) # 非根进程发送 shape

# 根进程分配接收缓冲区
b_recvbuf = None
if rank == 0:
    # 根据所有进程的行数和列数（假设列数固定）分配接收缓冲区
    b_recvbuf = np.zeros((total_rows, a.shape[1]), dtype=float)

# 准备 Gatherv 的参数
# 注意：counts 和 displacements 都是元素总数，不是字节数
# 对于 (N, M) 的数组，元素总数为 N * M

# 示例中只有两个进程，手动指定 counts 和 displacements
# 更通用的方法是先收集所有进程的数组形状，然后计算
if rank == 0:
    # counts: 从每个进程接收的元素总数
    # rank 0 发送 (5, 3) -> 15 元素
    # rank 1 发送 (2, 3) -> 6 元素
    counts = [5 * 3, 2 * 3]

    # displacements: 每个进程的数据在 b_recvbuf 中的起始偏移量
    # rank 0 数据从 b_recvbuf 的第 0 元素开始
    # rank 1 数据从 b_recvbuf 的第 15 元素开始 (即 rank 0 数据之后)
    displacements = [0, 5 * 3]

    # recvbuf_tuple 格式: (接收缓冲区, counts, displacements, 数据类型)
    recvbuf_tuple = (b_recvbuf, counts, displacements, MPI.DOUBLE)
else:
    # 非根进程的 recvbuf_tuple 为 None
    recvbuf_tuple = None

# 执行 Gatherv 操作
# 发送缓冲区 (sendbuf) 是当前进程的数组 a
# 接收缓冲区 (recvbuf_tuple) 仅在根进程上有效
comm.Gatherv(a, recvbuf_tuple, root=0)

# 根进程打印结果
if rank == 0:
    print(f"进程 {rank}: Gatherv 收集后的数组形状 {b_recvbuf.shape}")
    print(f"进程 {rank}: Gatherv 收集后的数组内容:\n{b_recvbuf}")