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。