8. 連立一次方程式

8.1 NumPyとCuPyのソースコード

リスト8-1にNumPyとCuPyを用いて連立一次方程式を解くソースコードを示します。
NumPyはCPU用、CuPyはGPU用です。 どちらもlinalg.solve関数を使用しています。
実数と複素数、単精度と倍精度を選択することができます。
行列は乱数で作成した非対称密行列です。
CuPyではデバイスメモリーを確実に解放するためにmemory poolを使用しています。

リスト8-1 NumPyとCuPyを用いて連立一次方程式を解くソースコード (leq.py)

"""
leq.py
solve linear equations (NumPy and CuPy)
"""

import numpy as np
import cupy as cp
import time

# parameters
N = 5000  # matrix size
L = 1  # repeat
dkind = 1      # 1=real, 2=complex
precision = 1  # 1=single, 2=double
gpu = 0        # 0=CPU(NumPy), 1=GPU(CuPy)

# dtype
f_dtype = 'f4' if precision == 1 else 'f8'
c_dtype = 'c8' if precision == 1 else 'c16'

# timer
t0 = time.time()

# real
if dkind == 1:
    a = np.random.rand(N, N).astype(f_dtype)
    b = np.random.rand(N).astype(f_dtype)
    x = np.zeros(N, f_dtype)
# complex
elif dkind == 2:
    a = np.random.rand(N, N).astype(f_dtype) \
      + np.random.rand(N, N).astype(f_dtype) * 1j
    b = np.random.rand(N).astype(f_dtype) \
      + np.random.rand(N).astype(f_dtype) * 1j
    x = np.zeros(N, c_dtype)

# alloc device memory
if gpu == 1:
    mempool = cp.get_default_memory_pool()
    # copy to device
    d_a = cp.array(a)
    d_b = cp.array(b)
    d_x = cp.array(x)

# timer
t1 = time.time()

# solve
for _ in range(L):
    if not gpu:
        # CPU (NumPy)
        x = np.linalg.solve(a, b)
    else:
        # GPU (CuPy)
        d_x = cp.linalg.solve(d_a, d_b)

# timer
t2 = time.time()

# free device memory
if gpu == 1:
    # copy to host
    x = cp.asnumpy(d_x)
    # free
    d_a = None
    d_b = None
    d_x = None
    mempool.free_all_blocks()

# check
res = np.dot(a, x) - b
ans = np.linalg.norm(res)

# output
device = 'CPU' if gpu == 0 else 'GPU'
print('(%s) N=%s, %.2f+%.2f[sec], err=%.3e'
    % (device, N, t1 - t0, (t2 - t1) / L, ans))

# free
a = None
b = None
x = None

8.2 実数の計算時間

表8-1に実数の連立一次方程式の計算時間を示します。
計算時間を正確に測定するためにN=5000で50回、N=10000で10回計算して回数で割っています。

表8-1から以下のことがわかります。

• NumPyでは単精度と倍精度の計算時間はほぼ同じである。
• CuPyでは単精度は倍精度の約10倍速い。倍精度はNumPyとほぼ同じである。 ハードウェアに倍精度の制限がかかっている影響が考えられる。
• 計算時間は行列の大きさNの3乗にほぼ比例する。
• CuPyの単精度以外の計算時間はMATLAB[13]とほぼ同じである。
  

以上から連立一次方程式を解くには、計算時間の観点からはCuPyの単精度が優れていますが、 残差を確認するとCuPyの単精度はNumPyの単精度より2桁大きくなっており、 CuPyの単精度を使用するときは計算精度に注意が必要です。

8.3 複素数の計算時間

表8-2に複素数の連立一次方程式の計算時間を示します。
表8-1と同じ傾向が見られます。 複素数の計算時間は実数の4～10倍になります。