mpi4py 中的打包解包操作

本文从本人简书博客同步过来

在上一篇中我们介绍了 mpi4py 中的数据类型创建方法，下面我们将介绍数据的打包/解包操作。

MPI 除了可以发送或接收连续的数据之外，还可以处理不连续的数据，其基本方法有两种，一是允许用户自定义新的数据类型，这在上一篇中作了相应的介绍，二是数据的打包与解包，即在发送前将不连续空间的数据打包到连续空间，接收端收到数据后再解包恢复到不连续空间。他们也提供了一些 MPI 以前不能提供的功能，例如，一个消息可以被分成几个部分来接收，接收后面部分的接收操作可能依赖于前一部分的内容。

方法接口

mpi4py 中打包解包相关的方法（MPI.Datatype 类方法）接口为：

Pack(self, inbuf, outbuf, int position, Comm comm)

打包操作。将 inbuf 中的数据打包存放到 outbuf 中，起始存放位置由 position 以字节为单位指定，返回打包完成后 outbuf 中的数据长度（以字节为单位）。

Unpack(self, inbuf, int position, outbuf, Comm comm)

解包操作。将 inbuf 中的指定数据解包到 outbuf 中，起始解包位置（以字节为单位）由 position 指定，解包数据的个数是 outbuf 能存放元素的个数，返回解包完成后下一个可用数据的起始位置。

Pack_size(self, int count, Comm comm)

返回打包 count 个该数据类型所需空间的大小。该调用返回的是一个上界，而不是一个精确值，这是因为包装一个消息所需的精确空间可能依赖于上下文。

注意：Pack 和 Unpack 都需要具有缓冲区接口的数据对象，但 Python 内置的一些数据类型却没有相应的缓冲区接口，所以不能直接打包和解包这些数据。例如，我们不能如下打包一个 Python 的 double 类型数据：

...

comm = MPI.COMM_WORLD

pack_buf = bytearray(100)

a = 1.0 # a double

position = 0
position = MPI.INT.Pack(a, pack_buf, position, comm)

但是 numpy 中的标量数字和数组都具有缓冲区接口，我们可以使用它们来进行打包和解包操作，如我们可以使用下列两种方式打包一个 double 数据，使用一个 double 型的 numpy 标量，或含有一个元素的 double 型数组：

...

comm = MPI.COMM_WORLD

pack_buf = bytearray(100)

a = np.array(1.0, dtype='d') # a double scalar
# or
# a = np.array([1.0], dtype='d') # a single element double array

position = 0
position = MPI.INT.Pack(a, pack_buf, position, comm)

另外需要注意的是，发送打包的数据时，需要用 MPI.PACKED 作为数据类型。但是可以用任意数据类型（包括自定义数据类型）来接收消息，类型匹配的规定对于以MPI.PACKED类型发送的消息不再起作用。以任何类型发送的消息（包括MPI.PACKED类型）都可以用MPI.PACKED类型接收，而且这样的消息可以调用 Unpack 方法来解包。

例程

下面给出打包解包相关方法的使用例程。

# pack.py

"""
Demonstrates the usage of Pack, Unpack, Pack_size.

Run this with 2 processes like:
$ mpiexec -n 2 python pack.py
"""

import numpy as np
from mpi4py import MPI


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

# get pack size of some data type
print 'pack size of 3 MPI.INT:', MPI.INT.Pack_size(3, comm)
print 'pack size of 3 MPI.DOUBLE:', MPI.DOUBLE.Pack_size(3, comm)

# --------------------------------------------------------------------------------
# demonstrate Pack and Unpack
buf_size = 100
pack_buf = bytearray(buf_size)

a = np.array([1, 2], dtype='i') # two int type
b = np.array(0.5, dtype='f8') # a double type
# or
# b = np.array([0.5], dtype='f8') # a double type

# pack two int and a double into pack_buf
position = 0
position = MPI.INT.Pack(a, pack_buf, position, comm) # position = 4 * 2
position = MPI.DOUBLE.Pack(b, pack_buf, position, comm) # position = 4 * 2 + 8

a1 = np.array([0, 0], dtype='i') # two int type with initial value 0
b1 = np.array(0.0, dtype='f8') # a double type with initial value 0.0

# unpack two int and a double from pack_buf to a1, b1
position = 0
position = MPI.INT.Unpack(pack_buf, position, a1, comm) # position = 4 * 2
position = MPI.DOUBLE.Unpack(pack_buf, position, b1, comm) # position = 4 * 2 + 8
print 'a1 = [%d, %d], b1 = %f' % (a1[0], a1[1], b1)

# --------------------------------------------------------------------------------
# demonstrate how to send and receive a packed buffer
new_pack_buf = bytearray(buf_size)
if rank == 0:
    a = np.array(1, dtype='i') # an int type
    b = np.array(0.5, dtype='f8') # a double type
    # pack an int and a double into pack_buf
    packsize = 0
    packsize = MPI.INT.Pack(a, pack_buf, packsize, comm) # packsize = 4
    packsize = MPI.DOUBLE.Pack(b, pack_buf, packsize, comm) # packsize = 4 + 8
    # first send packsize to rank 1, it will use it to receive the packed data
    comm.send(packsize, dest=1, tag=11)
    # then send the packed data to rank 1
    comm.Send([pack_buf, packsize, MPI.PACKED], dest=1, tag=22)
    print 'rank 0 send: a = %d, b = %f' % (a, b)
elif rank == 1:
    a = np.array(0, dtype='i') # an int type with initial value 0
    b = np.array(0.0, dtype='f8') # a double type with initial value 0.0
    # receive packsize from rank 0
    packsize = comm.recv(source=0, tag=11)
    # use packsize to receive the packed data from rank 0
    comm.Recv([pack_buf, packsize, MPI.PACKED], source=0, tag=22)
    # unpack an int and a double from pack_buf to a, b
    position = 0
    position = MPI.INT.Unpack(pack_buf, position, a, comm) # position = 4
    position = MPI.DOUBLE.Unpack(pack_buf, position, b, comm) # position = 4 + 8
    print 'rank 1 receives: a = %d, b = %f' % (a, b)

运行结果如下：

$ mpiexec -n 2 python pack.py
pack size of 3 MPI.INT: 12
pack size of 3 MPI.DOUBLE: 24
a1 = [1, 2], b1 = 0.500000
rank 0 send: a = 1, b = 0.500000
pack size of 3 MPI.INT: 12
pack size of 3 MPI.DOUBLE: 24
a1 = [1, 2], b1 = 0.500000
rank 1 receives: a = 1, b = 0.500000

以上我们介绍了 mpi4py 中的数据打包解包操作，在下一篇中我们将介绍进程拓扑。