mpi4py 进阶之 mpiutil

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

在上一篇中我们介绍了 mpi4py 的若干使用技巧，并且简要介绍了 caput 及其 mpiutil 模块，下面我们将介绍 mpiutil 中提供的若干方便和易用的函数，这些函数可以使我们更加方便地进行 Python 并行编程，并且使我们的程序很容易地做到兼容非 MPI 编程环境。

函数接口

以下介绍的所有函数都可以兼容非 MPI 环境（此时 _comm 为 None），当 mpi4py 可用时 _comm 是 MPI.COMM_WORLD，也可以传递一个其它的通信子，此时将在该通信子上执行相应的操作。

mpilist(full_list, method='con', comm=_comm)

将一个序列 full_list 按照方法 method 划分成 n (n 为 comm 的 size，当 comm 为 None 时， 其 size = 1) 份并分配给每个进程，每个进程分得的元素数目相等（如果可以均分）或相差 1 （如果不能均分，rank 较小的进程会多 1 个），返回每个进程得到的子序列。当 fullist 的元素数目少于进程数目时，后面的进程会因为分不到元素而返回一个空序列。method 的取值可以为：

• ‘con’：连续划分，即 rank 较小的进程得到 full_list 中前面的元素，为默认值；
• ‘alt’：交替划分，即 rank = 0 的进程得到位置为 [0, n, 2n, …] 的元素， rank == 1 的进程得到位置为 [1, n+1, 2n+1, …] 的元素， 等；
• ‘rand’：随机划分，效果相当于对原系列随机重排后再按确定的方式分配给每个进程。

comm 可以是 None 或一个通信子对象，其默认值 _comm 为 None （mpi4py 不可用时） 或者 MPI.COMM_WORLD （mpi4py 可用时），也可以传递一个其它的通信子对象，此时将 full_list 划分给该通信子对象所包含的每个进程。

mpirange(*args, **kargs)

MPI 版本的 range 函数，参数同 range 函数，另外加可选的参数 method （默认值 con） 和 comm （默认值为 _comm），这两个参数的意思同上面介绍的 mpilist。这个函数的执行效果是 mpilist(range(*args), method, comm)，即将 range 函数生成的序列作为 full_list 调用 mpilist 函数，返回每个进程得到的子序列。

barrier(comm=_comm)

栅障同步，参数 comm 的意义同函数 mpilist。

bcast(data, root=0, comm=_comm)

广播操作，数据从 root 进程广播到所有其它进程，参数 comm 的意义同函数 mpilist。

reduce(sendobj, root=0, op=None, comm=_comm)

规约操作，将数据按照方法 op （默认值 None 会执行 MPI.SUM） 规约到 root 进程，参数 comm 的意义同函数 mpilist。

allreduce(sendobj, op=None, comm=_comm)

全规约操作，将数据按照方法 op 作全规约，参数 comm 的意义同函数 mpilist。

gather_list(lst, root=None, comm=_comm)

将各个进程的列表 lst 收集到 root 进程中并合并成一个新的列表，如果 root 为一个整数，则只有 rank 为该整数的进程会收集到数据并返回合并的列表，其它进程返回 None，如果 root 为 None，则所有进程都会收集（全收集操作），每个进程都会返回合并的列表。参数 comm 的意义同函数 mpilist。

parallel_map(func, glist, root=None, method='con', comm=_comm)

将序列 glist 按照方法 method 划分给每个进程，然后将函数 func 作用到每个进程所得的子序列的每个元素上，函数的返回值会被收集到 root 中经合并后返回。如果 root 为一个整数，则只有 rank 为该整数的进程会收集到数据并返回合并的列表，其它进程返回 None，如果 root 为 None，则所有进程都会收集（全收集操作），每个进程都会返回合并的列表。参数 method 和 comm 的意义同函数 mpilist。

split_all(n, comm=_comm)

将一个长度为 n 的序列顺序连续地划分给每个进程，返回一个三元组 (num, start, end)，其中 num，start，end 都为长度为 comm 的 size