3. 配列計算とNumba

3.1 NumPyとNumbaを用いたベクトル和とベクトル内積の計算

NumPyとNumbaを用いてベクトル和とベクトル内積を計算し、配列計算の基本的な特性を調べます。
リスト3-1にベクトル和とベクトル内積のソースコードを示します。
計算方法として表3-1の3ケースを考えます。

表3-1 計算方法
計算方法No.	計算方法の詳細
0	単純なfor文
1	NumPyのndarrayを用いた配列演算（以下、配列演算）
2	Numbaを用いたJITコンパイラー（以下、Numba）

リスト3-1 ベクトル和とベクトル内積のソースコード (vector.py)

"""
vector.py
test program of vector operations (NumPy and Numba)
"""

import time
import numpy as np
from numba import jit, prange, set_num_threads

# (vadd-0) sum of two vectors : for version
def vadd_for(a, b, c):
    n = len(a)
    for i in range(n):
        c[i] = a[i] + b[i]

# (vadd-1) sum of two vectors : array version
def vadd_array(a, b, c):
    c[:] = a + b

# (vadd-2) sum of two vectors : NumPy version
def vadd_numpy(a, b, c):
    c[:] = np.add(a, b)

# (vadd-3) sum of two vectors : Numba version
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def vadd_numba(a, b, c):
    n = len(a)
    for i in prange(n):
        c[i] = a[i] + b[i]

# (sdot-0) scalar product of two vectors : for version
def sdot_for(a, b):
    n = len(a)
    s = 0
    for i in range(n):
        s += a[i] * b[i]
    return s

# (sdot-1) scalar product of two vectors : array version
def sdot_array(a, b):
    return sum(a * b)

# (sdot-2) scalar product of two vectors : NumPy version
def sdot_numpy(a, b):
    return np.dot(a, b)

# (sdot-3) scalar product of two vectors : Numba version
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def sdot_numba(a, b):
    n = len(a)
    s = 0
    for i in prange(n):
        s += a[i] * b[i]
    return s

# parameters
N = 10000000
L = 1000
thread = 8
dtype = 'f4'   # 'f4' or 'f8'
fn = 'vadd-2'
#fn = 'sdot-2'

# Numba threads
if (fn == 'vadd-3') or (fn == 'sdot-3'):
    set_num_threads(thread)

t0 = time.time()

# setup
a = np.arange(N).astype(dtype)
b = np.arange(N).astype(dtype)
c = np.zeros(N, dtype)

t1 = time.time()

# calculation
for _ in range(L):
    if   fn == 'vadd-0':
        vadd_for(a, b, c)
    elif fn == 'vadd-1':
        vadd_array(a, b, c)
    elif fn == 'vadd-2':
        vadd_numpy(a, b, c)
    elif fn == 'vadd-3':
        vadd_numba(a, b, c)
    elif fn == 'sdot-0':
        s = sdot_for(a, b)
    elif fn == 'sdot-1':
        s = sdot_array(a, b)
    elif fn == 'sdot-2':
        s = sdot_numpy(a, b)
    elif fn == 'sdot-3':
        s = sdot_numba(a, b)

t2 = time.time()

# check
if fn.startswith('vadd'):
    r1 = np.sum(c)
    r2 = N * (N - 1)
else:
    r1 = s
    r2 = N * (N - 1) * (2 * N - 1) / 6

# output
print('(%s) N=%d, L=%d' % (fn, N, L))
print('%.2f+%.2f[sec], %e, %e' % (t1 - t0, t2 - t1, r1, r2))

# free
a = None
b = None
c = None

3.2 配列演算について

Pythonでは配列の演算はリスト3-1の関数vadd-1,sdot-1 のようにスカラーのように扱うことができ高速化することができます。このようなプログラミングを配列指向プログラミングと呼びます[14][15]。
なお、Python自体はGIL(Global Interpreter Lock)によりマルチスレッドに対応していません。

3.3 NumPyについて

Pythonでは配列の演算はリスト3-1の関数vadd-2,sdot-2のように、 NumPyの関数を使用して計算することができます。
NumPyの関数は高度にチューニングされており、また自動的にマルチコアを使用して並列処理を行います。

3.4 Numbaについて

Pythonでは前項のような配列演算ができないとき、 for文を用いてそのまま記述すると極めて計算が遅くなります。
そのようなときはNumbaを使用します。
Numbaではリスト3-1の関数vadd-3,sdot3のように関数の前に"@jit"行を挿入するだけで、 JIT(Just In Time)コンパイラーが起動してC言語に変換してコンパイルされ、 C言語と同等の速さになります。
@jit行には"(cache=True, nopython=True, nogil=True, parallel=True)" というオプションを指定することができます。それぞれの意味は以下の通りです。

cache=True : コンパイルした結果が__pycache__フォルダに保存され、次回同じ関数を使用するときにそれを使用するので高速に関数を起動することができます。ただし、関数の引数が変わったときはwarningを出して再コンパイルします。
nopython=True : コンパイルに失敗したときはエラーメッセージを出して中止します。
nogil=True : GILを無効にしてマルチスレッドを有効にします。
parallel=True : マルチスレッドにより並列計算します。

Numbaを使用する注意点は以下の通りです。 (将来のバージョンでは変わる可能性があります)

関数の引数はNumPyのサポートする型（'i4','f4','f8','c8',c16'等）に限定される。配列を含む構造体など使えない型がある。
Numbaを使用する関数内で呼び出される関数も@jit宣言されていなければならない。
for文を並列処理するにはrangeをprangeに変更する。（import numba.prangeが必要）
並列計算のスレッド数を指定するにはnumba.set_num_threads関数を使用する。スレッド数の既定値は論理コア数。

3.5 計算時間

リスト3-1において、ベクトルの大きさNと繰り返し回数Lを変え、計算方法にNumPyとNumbaをとり、Numbaについてはスレッド数を変えたときの計算時間を、表3-2（ベクトル和）と表3-3（ベクトル内積）に示します。
計算時間は計算の本体部であり、配列を用意する部分は除いています。
計算時間の誤差を小さくするために繰り返し回数分だけ計算しています。
ベクトルの大きさNと繰り返し回数Lの積はNo.1～No.4で一定10¹⁰です。従ってNo.1～No.4の演算量は同じです。
最適化の効果が見られるケースに色をつけています。

表3-2 ベクトル和の計算時間（実数単精度）
No.	ベクトルの大きさN	繰り返し回数L	vadd-1 (配列演算)	vadd-2 (NumPy)	vadd-3 (Numba)
No.	ベクトルの大きさN	繰り返し回数L	vadd-1 (配列演算)	vadd-2 (NumPy)	1スレッド	2スレッド	4スレッド	8スレッド	16スレッド
1	10,000,000	1,000	13.00秒	12.99秒	5.56秒	4.62秒	4.72秒	4.75秒	5.24秒
2	1,000,000	10,000	14.03秒	13.72秒	5.75秒	5.01秒	4.03秒	3.62秒	4.66秒
3	100,000	100,000	2.09秒	1.98秒	1.40秒	1.87秒	2.37秒	2.26秒	2.47秒
4	10,000	1,000,000	3.23秒	3.24秒	3.30秒	46.25秒	50.29秒	50.98秒	49.60秒

表3-3 ベクトル内積の計算時間（実数単精度）
No.	ベクトルの大きさN	繰り返し回数L	sdot-1 (配列演算)(注1)	sdot-2 (NumPy)	sdot-3 (Numba)
No.	ベクトルの大きさN	繰り返し回数L	sdot-1 (配列演算)(注1)	sdot-2 (NumPy)	1スレッド	2スレッド	4スレッド	8スレッド	16スレッド
1	10,000,000	1,000	1300秒	1.83秒	7.39秒	4.21秒	2.15秒	2.17秒	2.09秒
2	1,000,000	10,000	1300秒	0.19秒	7.74秒	4.44秒	2.42秒	2.10秒	1.77秒
3	100,000	100,000	1300秒	0.39秒	7.48秒	5.07秒	4.43秒	3.80秒	4.21秒
4	10,000	1,000,000	1300秒	2.30秒	9.78秒	24.06秒	45.79秒	47.07秒	27.52秒

(注1)1/100の繰り返し回数時の計算時間を100倍した推定値

表3-2～表3-3から以下のことがわかります。

使用メモリーの少ないNo.2～4ではキャッシュメモリーが有効利用され計算時間が大幅に短縮されることがある。
ベクトル和では配列演算(vadd-1)とNumPy(vadd-2)の計算時間はほぼ同じである。
ベクトル和(No.1～2)ではNumba(vadd-3)が配列演算(vadd-1)とNumPy(vadd-2)の約2倍速い。
ベクトル内積ではNumPy(sdot-2)がNumba(sdot-3)より大幅に速い。（NumPy関数dotの効果）
ベクトル和ではNumba並列の効果はない。(vadd-3)
ベクトル内積ではNumba並列の効果がある。(sdot-3)
ベクトル内積の配列演算(sdot-1)は非常に遅い。
ベクトル内積の並列化はreduction（還元演算）となり並列化は自明ではないが、 NumPyはもちろんNumbaもreduction機能を持っている。
NumbaのNo.4はベクトル和とベクトル内積ともに2スレッド以上で非常に計算時間がかかる。 2スレッド以上でCPUを100%使用しているので特殊（で不適切）な処理をしていると思われる。
Numba並列は速くなるときと変わらないときと大幅に遅くなるときがあるので、注意が必要である。
NumPyが使えるときはNumPyを使うことが概ね最良である。

以上から、アルゴリズムと問題のサイズによって最適な方法が変わります。複数の実装を行って比較することが望ましいです。
一般的には、NumPyが使えるときはNumPyを使い、 NumPyが使えないときはNumbaを使用するのがよいと思われます。

3.6 for文の計算時間

Pythonではfor文を用いてそのまま記述すると極めて計算が遅くなります。
表3-4にfor文のベクトル和(vadd-0)とベクトル内積(sdot-0)の計算時間を示します。
表3-2、表3-3と比べて繰り返し回数Lが1/1000となっていることに注意してください。

表3-4 for文のベクトル和とベクトル内積の計算時間（実数単精度）
No.	ベクトルの大きさN	繰り返し回数L	ベクトル和 vadd-0	ベクトル内積 sdot-0
1	10,000,000	1	3.10秒	3.57秒
2	1,000,000	10	3.09秒	3.53秒
3	100,000	100	3.13秒	3.55秒
4	10,000	1,000	3.10秒	3.65秒

表3-4から以下のことがわかります。

for文では最適化が行われないためにすべての組み合わせの計算時間がほぼ同じである。
表3-4の計算時間は表3-3とほぼ同じオーダーなので、 for文の計算速度はNumPyまたはNumbaの1/1000のオーダーである。

Pythonでは計算の主要部のfor文は本節の方法を用いて高速化することが必須です。