前言

MPI C/Fortran P2P通信一般需要指定<数据首地址,数据类型，数据数目>三元组信息，根据三元组信息可以通信一段连续内存数据（MPI派生类型允许内存不连续数据一次通信完成，但底层实现仍然要求发送内存连续数据）。

mpi4py是MPI的Python绑定。如果你对Python标准实现Cpython有了解的话，Python对象都不是一个简单的“裸”数据，Python对象在Cpython实现中一般是一个C结构体，比如Python列表对象如下所示：

typedef struct {
    PyObject_VAR_HEAD
    /* Vector of pointers to list elements.  list[0] is ob_item[0], etc. */
    PyObject **ob_item;

    /* ob_item contains space for 'allocated' elements.  The number
     * currently in use is ob_size.
     * Invariants:
     *     0 <= ob_size <= allocated
     *     len(list) == ob_size
     *     ob_item == NULL implies ob_size == allocated == 0
     * list.sort() temporarily sets allocated to -1 to detect mutations.
     *
     * Items must normally not be NULL, except during construction when
     * the list is not yet visible outside the function that builds it.
     */
    Py_ssize_t allocated;
} PyListObject;

对于MPI的Python绑定，为了支持Python对象的通信，必须对Python对象进行处理满足MPI通信要求。那么mpi4py是如何实现的呢？

实现

不同Python对象类型使用不同的C结构体实现，为了实现通用化处理，mpi4py利用Python pickle模块序列化/反序列化功能先将python对象转换成字节序列，在进行MPI通信。

比如MPI Send/Recv流程为：

• 使用pickle.dumps()将Python对象序列化为字节流
• 调用MPI_Send的C实现发送数据
• 调用MPI_Recv的C实现接收字节流数据
• 使用pickle.load()将字节流反序列化为Python对象，实现Python对象的通信。

![[Drawing 2024-03-16 10.20.02.excalidraw.png]]

通用Python对象通信

mpi4py中使用小写的类方法对通用python对象通信，比如Comm.send、Comm.recv。

比如Comm.send接口：

Comm.recv接口：

具体实现细节

Comm.send类方法定义在Comm.pyx文件中（大概1942行）

    def send(
        self,
        obj: Any,
        int dest: int,
        int tag: int = 0,
    ) -> None:
        """Send in standard mode."""
        cdef MPI_Comm comm = self.ob_mpi
        return PyMPI_send(obj, dest, tag, comm)

实际上调用的是函数PyMPI_send，该函数是在msgpickle.pxi中定义。

cdef object PyMPI_send(object obj, int dest, int tag,
                       MPI_Comm comm):
    cdef Pickle pickle = PyMPI_PICKLE
    #
    cdef void *sbuf = NULL
    cdef MPI_Count scount = 0
    cdef MPI_Datatype stype = MPI_BYTE
    #
    cdef object unuseds = None
    if dest != MPI_PROC_NULL:
        unuseds = pickle_dump(pickle, obj, &sbuf, &scount)
    with nogil: CHKERR( MPI_Send_c(
        sbuf, scount, stype,
        dest, tag, comm) )
    return None

Python对象obj通过函数pickle_dump()，然后调用MPI_Send_c函数进行MPI通信。

cdef object pickle_dump(Pickle pkl, object obj, void **p, MPI_Count *n):
    cdef object buf = cdumps(pkl, obj)
    p[0] = PyBytes_AsString(buf)
    n[0] = PyBytes_Size(buf)
    return buf

其中cdumps实现为

cdef object cdumps(Pickle pkl, object obj):
    if pkl.ob_PROTO is not None:
        return pkl.ob_dumps(obj, pkl.ob_PROTO)
    else:
        return pkl.ob_dumps(obj)

再往上追溯可以看出是调用pickle模块的dumps方法。

from pickle import dumps as PyPickle_dumps 
...
self.ob_dumps = PyPickle_dumps

配对的Comm.recv则是接收数据后反序列化为Python对象，这里就不细讲了。

小结

• 通过序列化功能使得Python对象能够调用MPI C接口进行通信。
• 序列化/反序列化会带来开销（包括时间/内存消耗），对于大数据通信性能损失有点大。
• 因此MPI4PY也支持通过Python缓冲协议（buffer protocol）对某些对象进行直接通信，比如Numpy数组，提升性能。

参考资料

https://mpi4py.readthedocs.io/en/stable/

https://docs.python.org/3/library/pickle.html