mpi4py 在函数之间广播 numpy 数组

Posted 2023-03-27

【中文标题】mpi4py 在函数之间广播 numpy 数组

【英文标题】：mpi4py broadcasting numpy arrays between functions

【发布时间】：2021-12-17 08:27:09

【问题描述】：

我是 MPI 的初学者，我正在尝试制作一个简单的示例程序，其中 master 将一个 numpy 数组发送给 worker。唯一的问题是这些必须在不同的 master 和 worker 函数之间工作。这是一个练习题，用于弄清楚为什么大型项目中的类似设置当前不起作用。我发现的所有使用 Bcast 的示例都没有包含任何关于如何使其在不同模块或函数调用之间工作的提示。

代码如下：

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.rank

def master():
    data = np.arange(4.0)
    comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

def worker():
    data = np.zeros(4)
    data = comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

if rank == 0:
    master()
else:
    worker()

这是当前的输出：

rank 1 None
rank 2 None
rank 0 [0. 1. 2. 3.]
rank 3 None

我不确定为什么这个版本不起作用。使用列表的几乎相同的版本可以正常工作。

【问题讨论】：

您确定您的阵列具有相同的类型吗？使用查询函数打印出数据类型。

在worker，简称comm.Bcast([data, MPI.DOUBLE], root=0)

在调用 Bcast 之前，numpy.ndarray 中worker 中的数据类型，与master 相同。之后它显然只在workers中变成了NoneType。

@Gilles Gouaillardet 谢谢。这解决了它。

【参考方案1】：

我要感谢评论者如此迅速地做出回应。为了让它更明显，工作函数应该是这样的：

def worker():
    data = np.zeros(4)
    #do not set data equal to Bcast call
    comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

感谢 Gilles Gouaillardet 的帮助。这是一个简单的解决方案，它甚至看起来像我的 Fortran 模块的工作版本，但我在尝试用 Python 编写项目时遇到了重新创建示例。