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 中的数据打包解包操作,在下一篇中我们将介绍进程拓扑。

你可能感兴趣的:(mpi4py 中的打包解包操作)