mpi4py 与 OpenMP 混合编程

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

在上一篇中我们介绍了在 cython 中使用 mpi4py 的方法，下面我们将介绍 mpi4py 与 OpenMP 混合编程。

OpenMP 简介

OpenMP （Open Multi-Processing） 是一个跨平台的多线程实现，它本身不是一种独立的并行语言，而是为多处理器上编写并行程序而设计的、指导共享内存多线程并行的编译制导指令和应用程序编程接口（API），可在 C/C++ 和 Fortran 中使用，一般是在这些编程语言的串行代码中以编译器可识别的注释形式出现。OpenMP 可以在大多数的处理器体系和操作系统中运行，包括 GNU/Linux, Mac OS X 和 Microsoft Windows 等。OpenMP 包括一套编译器指令、库和一些能够影响运行行为的环境变量。在使用 OpenMP 时，主线程（顺序的执行指令）生成一系列的子线程，并将任务划分给这些子线程进行执行。这些子线程并行的运行，由运行时环境将线程分配给不同的处理器。

OpenMP 采用可移植的、可扩展的模型，为程序员提供了一个简单而灵活的开发平台，及从标准桌面电脑到超级计算机的并行应用程序接口。

混合并行编程模型构建的应用程序可以同时使用 OpenMP 和 MPI。

OpenMP 的执行模式

OpenMP 的执行模型采用 fork-join 的形式，其中 fork 创建新线程或者唤醒已有线程；join 即多线程的会合。fork-join 执行模型在刚开始执行的时候，只有一个称为“主线程”的运行线程存在。主线程在运行过程中，当遇到需要进行并行计算的时候，派生出线程来执行并行任务。在并行执行的时候，主线程和派生线程共同工作。在并行代码执行结束后，派生线程退出或者阻塞，不再工作，控制流程回到单独的主线程中。

OpenMP 编程要素

OpenMP 编程模型以线程为基础，通过编译制导指令来显式地指导并行化，OpenMP 提供了三种编程要素来实现对并行化的完善控制，它们是：编译制导、API 函数集和环境变量。

在 C/C++ 和 Fortran 中使用 OpenMP 时需要对 OpenMP 的相关知识足够了解和熟悉，读者可以参考 OpenMP 的相关文档和资料，这里不做进一步的介绍。

在 cython 中使用 OpenMP

不能直接在 Python 中使用 OpenMP，不过幸运的是 cython 支持 OpenMP，而我们可以方便而且容易地用 cython 编写 Python 的扩展模块，或者将 C/C++ 的代码包装成 Python 的扩展模块，因此也就可以间接地在使用 OpenMP 的多线程加速了。

cython 通过其 parallel 模块支持本地并行化，目前该模块支持和使用 OpenMP 的多线程并行方案，今后还可能会支持其它并行机制。要使用 cython 的 parallel 模块以实现并行加速，需要释放 Python 的 GIL （Global Interpreter Lock）。

相关函数

下面是 cython.paralle 中的主要函数：

cython.parallel.prange([start,] stop[, step][, nogil=False][, schedule=None[, chunksize=None]][, num_threads=None])

用于并行的 for 循环，只能用在 Python GIL 被释放的情况下。OpenMP 会自动启动一个线程池并将计算任务按照指定的 schedule 分配给线程池中的线程去完成。在 prange 块下面赋值的变量是最后私有的（lastprivate），即该变量会获得其最后一次迭代的值。如果使用 inplace 算符（如 +=, *= 等）为 prange 块中的变量赋值，则该操作会变成规约运算，即在循环完成后，每个线程本地的该变量副本的值会使用该算符执行规约运算并将运算的结果赋值给此变量。prange 的循环变量总是最后私有的。其参数的意义如下：

• start：循环变量的起始值。
• end：循环变量的结束值。
• step：循环步长，不能为 0.
• nogil：如果为 True，整个 prange 循环会被包裹在一个 nogil 代码段中，否则需要手动将该循环放入 nogil 块下。
• schedule：任务分配给线程的调度方式，会传递给 OpemMP，可用的方式有：
  • static：如果设置了 chunksize，计算任务会被提前按照给定的 chunksize 分配到所有线程上。如果没有设置 chunksize，则任务会被划分成近似相同大小的块并提前分配给各个线程。这种分配方式适合于任务能够划分成差不多大小的块并且这些块的执行时间差不多相同的情况。
  • dynamic：任务会被以默认 chunksize 为 1 动态分配给任何空闲的线程。这种方式适合于不同的任务块执行时间不定且事先无法预知的情况。
  • guided：类似于 dynamic 方式，任务会被动态分配给空闲线程，但是分配的任务块会逐步减小（块的大小正比于还未分配的任务数除线程数）。
  • runtime：分配方式和任务块大小在运行过程中确定，如通过 openmp.omp_set_schedule() 函数设置，或通过 OMP_SCHEDULE 环境变量获得。这种方式无法获得在静态编译时的优化，因此在执行性能上可能会差一些。
• num_threads：使用多少线程。如果没有设置，OpenMP 会决定使用多少线程，一般会使用机器上所有可用的核数。使用的线程数也可以通过调用 omp_set_num_threads() 函数来设置，或者通过 MP_NUM_THREADS 环境变量来设置。
• chunksize：任务在各个线程上分配的块大小，只对 static，dynamic 和 guided 方式有用。

下面这个例子展示使用 inplace 算符以实现规约运算：

from cython.parallel import prange

cdef int i
cdef int n = 30
cdef int sum = 0

for i in prange(n, nogil=True):
    sum += i

print(sum)

cython.parallel.parallel(num_threads=None)

用在 with 语句中以并行执行一个代码段。可以用来建立会被 prange 使用的线程本地缓冲区。包含在此代码段中的 prange 循环任务会被分配给各个进程协同并行执行。

cython.parallel.threadid()

返回线程 id。对 n 个线程，线程 id 从 0 到 n-1 编号。

使用 OpemMP 函数

从 cython.parallel 中 cimport openmp 后可以使用 OpenMP 中的相关函数，举例如下：

# tag: openmp
# You can ignore the previous line.
# It's for internal testing of the Cython documentation.

from cython.parallel cimport parallel
cimport openmp

cdef int num_threads

openmp.omp_set_dynamic(1)
with nogil, parallel():
    num_threads = openmp.omp_get_num_threads()
    # ...

编译方法

要在 cython 中使用 OpenMP，必须告诉 C/C++ 编译器启用 OpenMP。例如对 gcc 及某些编译器，可以使用类似下面的 setup.py 脚本来编译。

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize

ext_modules = [
    Extension(
        "hello",
        ["hello.pyx"],
        extra_compile_args=['-fopenmp'],
        extra_link_args=['-fopenmp'],
    )
]

setup(
    name='hello-parallel-world',
    ext_modules=cythonize(ext_modules),
)

mpi4py 与 OpenMP 混合编程

混合编程是指在程序中同时使用多种编程模型，如同时使用 MPI 模型与多线程模型。在MPI 中多线程的使用中我们已经介绍过一种使用 mpi4py 与 Python threading 模块进行混合编程的方法，下面介绍使用 mpi4py 与 OpenMP （借助于 cython）进行混合编程的方法。

在上一篇中我们介绍了在 cython 中使用 mpi4py 的方法，现在我们又介绍了在 cython 中使用 OpenMP 的方法，将两种结合起来就能实现 mpi4py 与 OpenMP （借助于 cython）的混合编程。不过需要注意的是，对 Python 对象的很多操作都需要 GIL，而 cython.parallel 的一些操作（如 prange）却需要释放 GIL，因此需要在一些地方使用 with gil 或 wigh nogil 语句以获得或释放 GIL。

例程

下面给出简单的使用例程。

首先是使用了 mpi4py 与 OpenMP 的 cython 代码。

# hello.pyx

from cython import parallel
from mpi4py import MPI


# function that uses MPI and OpenMP hybrid programming
def say_hello(comm):
    cdef unsigned int thread_id
    cdef int i

    # use 2 OpenMP threads to execute the following code
    with nogil, parallel.parallel(num_threads=2):
        # allocate 3 tasks to the 2 threads with a chunk size of 2
        # and a static schedule, so thread 0 will have task 0 and 1,
        # thread 1 will have task 2
        for i in parallel.prange(3, schedule="static", chunksize=2):
            # get the thread id
            thread_id = parallel.threadid()
            # acquire the GIL for Python operation like print
            with gil:
                # get the rank of the MPI process
                rank = comm.rank
                # get the processor name
                pname = MPI.Get_processor_name()
                print 'MPI rank %d, OpenMP thread %d in %s says hello' % (rank, thread_id, pname)

使用下面的脚本将其编译成 Python 扩展模块。

# setup.py

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize

ext_modules = [
    Extension(
        "hello",
        ["hello.pyx"],
        extra_compile_args=['-fopenmp'],
        extra_link_args=['-fopenmp'],
    )
]

setup(
    name='hello-parallel-world',
    ext_modules=cythonize(ext_modules),
)

编译以上扩展模块的命令如下：

$ python setup.py build_ext --inplace

下面是 Python 测试程序。

# mpi_openmp.py

"""
Demonstrates how to use mpi4py and OpenMP hybrid programming.

2un this with 2 processes like:
$ mpiexec -n 2 python mpi_openmp.py
"""


from mpi4py import MPI
import hello


comm = MPI.COMM_WORLD

hello.say_hello(comm)

执行的结果如下：

$ mpiexec -n 2 -host node1,node2 python mpi_openmp.py
MPI rank 0, OpenMP thread 0 in node1 says hello
MPI rank 0, OpenMP thread 1 in node1 says hello
MPI rank 0, OpenMP thread 0 in node1 says hello
MPI rank 1, OpenMP thread 0 in node2 says hello
MPI rank 1, OpenMP thread 0 in node2 says hello
MPI rank 1, OpenMP thread 1 in node2 says hello

以上介绍了 mpi4py 与 OpenMP 混合编程，在下一篇中我们将介绍在 IPython 中使用 mpi4py。