mpi4py 中的广播操作

在上一篇中我们对集合通信做了一个非常简要的介绍，后面我们将逐步介绍 mpi4py 中提供的各个集合通信操作方法，下面我们从广播操作开始。

对组内通信子上的广播操作而言， 组内被标记为根的进程向组内所有其它进程发送消息。

对于组间通信子上的广播操作，其作用是以某个组内的某进程为根，向另外组的其它所有进程广播消息。参与广播操作的组间通信子内包含的所有进程都要参与执行 Bcast 函数，根进程所在的组内除根外的所有其它进程的 root 参数应该设置成 MPI.PROC_NULL，根进程的 root 参数使用 MPI.ROOT；另外组的所有进程的 root 参数为根进程在其所处组内的 rank 值。

需要注意的是，对于组间通信子上的广播操作，与根处于同一组的其它进程均不接收广播消息，而远程组的所有进程都可接收到广播消息。与此相对，在组内通信的广播操作中，根进程同样也要参与通信并接收来自其自身的消息。

mpi4py 中的广播操作的方法（MPI.Comm 类的方法）接口为：

bcast(self, obj, int root=0)
Bcast(self, buf, int root=0)

其中以小写字母开头的 bcast 可以广播任意可被 pickle 系列化的 Python 对象 obj，并返回所接收到的 obj；以大写字母开头的 Bcast 只能广播具有单段缓冲区接口的类数组对象，如 numpy 数组，参数 buf 可以是一个长度为2或3的 list 或 tuple，类似于 [data, MPI.DOUBLE]，或者 [data, count, MPI.DOUBLE]，以指明发送/接收数据缓冲区，数据计数以及数据类型。当 count 省略时会利用 data 的字节长度和数据类型计算出对应的 count。对 numpy 数组，其计数和数据类型可以自动推断出来，因此可以直接以 data 作为第一个参数传给 buf。

下面给出广播操作的使用例程。

# intra_comm_bcast.py

import numpy as np
from mpi4py import MPI


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

# broadcast a generic object by using bcast
if rank == 1:
    obj = {'a': 1}
else:
    obj = None

obj = comm.bcast(obj, root=1)
print 'rank %d has %s' % (rank, obj)

# broadcast a numpy array by using Bcast
if rank == 2:
    ary = np.arange(10, dtype='i')
else:
    ary = np.empty(10, dtype='i')

comm.Bcast(ary, root=2)
print 'rank %d has %s' % (rank, ary)

运行结果如下：

$ mpiexec -n 3 python intra_comm_bcast.py
rank 1 has {'a': 1}
rank 2 has {'a': 1}
rank 0 has {'a': 1}
rank 1 has [0 1 2 3 4 5 6 7 8 9]
rank 0 has [0 1 2 3 4 5 6 7 8 9]
rank 2 has [0 1 2 3 4 5 6 7 8 9]

# inter_comm_bcast.py

from mpi4py import MPI


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

if rank == 1:
    obj = {'a': 1}
else:
    obj = None

# split comm into three new communicators according to `color`
color = rank % 3
comm_split = comm.Split(color=color, key=rank)
# group 0 with color = 0 -> rank = {0, 3, 6}
# group 1 with color = 1 -> rank = {1, 4, 7}
# group 2 with color = 2 -> rank = {2, 5, 8}

# broadcast `obj` from rank 0 of group1 (rank 1 of COMM_WORLD) to all process in group 2
if color == 1:
    # create an inter-communicator by using rank 0 of local_comm and rank 2 of MPI.COMM_WORLD
    comm_inter = comm_split.Create_intercomm(0, comm, 2, tag=12)
    if comm_inter.rank == 0:
        root = MPI.ROOT
    else:
        root = MPI.PROC_NULL
    comm_inter.bcast(obj, root=root)
elif color == 2:
    # create an inter-communicator by using rank 0 of local_comm and rank 1 of MPI.COMM_WORLD
    comm_inter = comm_split.Create_intercomm(0, comm, 1, tag=12)
    obj = comm_inter.bcast(obj, root=0)

print 'rank %d has %s...' % (rank, obj)

运行结果如下：

$ mpiexec -n 9 python inter_comm_bcast.py
rank 3 has None
rank 6 has None
rank 0 has None
rank 1 has {'a': 1}
rank 2 has {'a': 1}
rank 4 has None
rank 5 has {'a': 1}
rank 7 has None
rank 8 has {'a': 1}

以上我们介绍了 mpi4py 中的广播操作方法，在下一篇中我们将介绍发散操作。