目次

6. Numba for CUDA

6.1 Numba for CUDAとは

Numba for CUDA GPUs を用いると NVIDIA社のグラフィックスボード(GPU)を用いて高速に計算することができます。
C版のCUDAとほぼ同様のプログラムになります。

6.2 Numba for CUDAによるベクトル和とベクトル内積の計算

リスト6-1にベクトル和とベクトル内積を Numba for CUDA で計算するプログラムを示します。
ベクトル内積についはベクトル同士の積（アダマール積）を計算した後、 reduction（@cuda.reduce を用いる）によって和を計算しています。
計算時間を正確に測定するためにeventを使用しています。

リスト6-1 Numba for CUDAによるベクトル和とベクトル内積のソースコード (vector_cuda.py)

"""
vector_cuda.py
test program of vector operations
Numba for CUDA
"""

import numpy as np
from numba import cuda

# (1) add two vectors: c = a + b
@cuda.jit(cache=True)
def vadd(a, b, c):
    tid = cuda.grid(1) # = threadIdx.x + (blockIdx.x * blockDim.x)
    n = len(c)
    if tid < n:
        c[tid] = a[tid] + b[tid]

# (2) scalar product of two vectors (a . b)
def sdot(a, b, c, nblocks, nthreads):
    vmul[nblocks, nthreads](a, b, c)
    return vsum(c)

# multiplicate two vectors: c = a * b
@cuda.jit(cache=True)
def vmul(a, b, c):
    tid = cuda.grid(1)
    n = len(c)
    if tid < n:
        c[tid] = a[tid] * b[tid]

# sum of a vector (reduction)
@cuda.reduce
def vsum(a, b):
    return a + b

# parameters
N = 100000000
L = 100
dtype = 'f4'  # 'f4' or 'f8'
fn = 'vadd'
#fn = 'sdot'

# timer
start = cuda.event()
end   = cuda.event()
start.record()

# host memory
a = np.arange(N).astype(dtype)
b = np.arange(N).astype(dtype)

# device memory
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_c = cuda.device_array_like(d_a)

# Execution Configuration
nthreads = 256
nblocks = (N + (nthreads - 1)) // nthreads

# timer
end.record()
end.synchronize()
t1 = cuda.event_elapsed_time(start, end)
start.record()

# calculation
for _ in range(L):
    if   fn == 'vadd':
        vadd[nblocks, nthreads](d_a, d_b, d_c)
    elif fn == 'sdot':
        s = sdot(d_a, d_b, d_c, nblocks, nthreads)

# timer
end.record()
end.synchronize()
t2 = cuda.event_elapsed_time(start, end)

# output
if   fn == 'vadd':
    s = np.sum(d_c.copy_to_host())
    exact = N * (N - 1)
elif fn == 'sdot':
    exact = N * (N - 1) * (2 * N - 1) / 6
print('N=%d, L=%d' % (N, L))
print('%.2f+%.2f[sec], %e, %e' % (t1 * 1e-3, t2 * 1e-3, s, exact))

# free
a = None
b = None
d_a = None
d_b = None
d_c = None

Numba for CUDA のC版のCUDAとの主な違いは以下の通りです。

6.3 Numba for CUDA の計算時間

表6-1に Numba for CUDA によるベクトル和とベクトル内積の計算時間を示します。
配列の大きさ(=N)と繰り返し回数(=L)の積は一定(=1010)です。 従って全体の演算量は同じです。

表6-1 Numba for CUDA によるベクトル和とベクトル内積の計算時間
No.配列の大きさN繰り返し回数Lベクトル和ベクトル内積
単精度倍精度単精度倍精度
1100,000,000100 0.57秒 1.15秒 1.29秒 1.98秒
210,000,0001,000 0.52秒 1.02秒 1.29秒 2.02秒
31,000,00010,000 0.61秒 1.02秒 7.30秒 7.08秒
4100,000100,000 6.54秒 6.72秒66.30秒66.63秒
510,0001,000,00067.34秒66.78秒660秒(注1)660秒(注1)

（注1）繰り返し回数Lが1/10のときの計算時間を10倍した推定値です。

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