Cython(13) マルチスレッド


2022年 10月 20日

現在のコンピュータには複数のコアがあり、マルチコアマシンと言われている。
とくに、多数のコアの場合には、メニーコアマシン(manycore processor)と呼ぶ。
コア1つにつき、1または2スレッドの並列実行が可能であり、高速化のためには並列処理を駆使することが多い。

マルチスレッドというと、自分でスレッドを作り、スレッドの管理が面倒という意識があるかも知れないが、PythonのプログラムをCython化した場合には、非常に簡単にマルチスレッド化ができる。
Julia集合で非常に時間がかかるところは、2重ループで、1001×1001画素に対する複素平面上の座標値に対して、発散するまで計算を繰り返している。この2重ループ部分にマルチスレッドを適応すると、高速化が実現できる。

まず、マルチスレッド対応したCythonのソースコードを示す。
マークしている行だけが、マルチスレッド対応のために変更(追加)した部分である。

## 固定画像で、この中で完結している。
## matplotlib.pyplot での出力。

# distutils: extra_compile_args = -fopenmp
# distutils: extra_link_args = -fopenmp

# cython: boundscheck=False
# cython: wraparound=False

from threading import Thread
from cython.parallel import *
import numpy as np
import matplotlib.pyplot as plt
import time

cdef int escape( double complex z, double complex c,
        	double z_max, int n_max) nogil:
	cdef:
    	int i = 0
    	double z_max2 = z_max * z_max
   	 
	while norm2(z) < z_max2 and i < n_max:
    	z = z * z + c
    	i += 1
	return i

cdef inline double norm2(double complex z) nogil:
	return z.real * z.real + z.imag * z.imag

def calc_julia( int resolution, double complex c,
            	double bound=1.2, double z_max=2.0, int n_max=1000 ):
	cdef:
    	double step = 2.0 * bound / resolution
    	int i,j
    	double complex z
    	double real, imag
    	int[:,:] counts
	counts = np.zeros((resolution+1,resolution+1), dtype=np.int32)

	for i in prange(resolution+1,nogil=True,schedule='static',chunksize=1):
    	imag = bound - i * step
    	for j in range(resolution+1):
        	real = -bound + j * step
        	z = real + imag *1j
        	counts[i,j] = escape(z,c,z_max,n_max)

	return np.asarray(counts)

def julia():
	tm_start = time.time()
	jl = calc_julia(1000,(0.322+0.047j))
	tm_end = time.time()

	plt.figure(figsize=(10,10))
	plt.imshow(np.log(jl))
	print("time={}".format(tm_end-tm_start))
	plt.show()
    
if __name__ == '__main__':
	julia()

OpenMP

並列処理の標準的なAPIであるOpenMP (open multiprocessing) をコンパイル及びリンク時に利用するために以下の2行が入っている。

# distutils: extra_compile_args = -fopenmp
# distutils: extra_link_args = -fopenmp


OpenMP(オープンエムピー)は、並列計算機環境において共有メモリ・マルチスレッド型の並列アプリケーションソフトウェア開発をサポートするために標準化されたAPIである ,とWikipediaに説明がある。様々な言語がOpenMPをサポートしており、それをここでは利用する。詳しく書くとキリがないので、これ以上の説明は省略。書籍も出ているので、参照されたい。

スレッド利用

以下の2行は、スレッドを利用するためである。

from threading import Thread
from cython.parallel import *

nogil

Pythonには、GIL(Global Interpreter Lock)という排他ロックがあるため、そのままではマルチスレッドが機能しない。これは、安全のためにやっていることなのだが、高速に実行したい場合には、この排他ロックの鍵を外さないといけない。

関数にnogil属性をつけると、GILが消え、その関数を同時にマルチスレッドで実行できるようになる。
関数のパラメータの直後にnogilを入れることで、nogil属性が付き、その関数はマルチスレッド対応となる。

ただし、マルチスレッドにした場合、異なるスレッドで同じメモリ内容を書き換えるなど不都合が起こらないことを確認することを忘れると、上書きが発生し、意味不明な結果になってしまうことがある。発生確率が低く、見つけにくいバグになりやすいので、注意が必要である。

prange

実際にマルチスレッドになるのは、rangeの代わりにprangeを使ったときである。
rangeに似ているが、マルチスレッドに対する様々なオプション指定ができる。

for i in prange(resolution+1,nogil=True,schedule='static',chunksize=1):
  • nogilオプションをTrueにすることで、マルチスレッドになる
  • scheduleオプションには、static, dynamic, guided, runtime が指定できる。ここでは一番単純なstaticを指定してみた。
  • chanksizeオプションは、スレッドの塊のサイズを指定する。ここでは、最小の1を指定していて、これで1行毎に別スレッドになる。
  • num_threadsオプションによりスレッド数の指定が可能だが、ここでは何も指定しなかったので、OpenMPが使用可能スレッド数を検出し、その数のスレッドを生成しているはずである。

たったこれだけの変更で、シングルスレッドで動いていたものがマルチスレッドで動くようになる。forループをマルチスレッド化する場合は、このように簡単である。

in [1]: import julia_cy5                                                   	 

In [2]: julia_cy5.julia()                                                  	 
time=0.021723270416259766

前回、0.321秒だったのが、マルチスレッドしたことで、0.0217秒になったので、14.8倍速になった。
全体では、25.32/0.0217=1166倍とかなり高速になった。1枚のJulia集合を0.02秒台で作れるのなら、動画さえ可能な気がする。