numpy, cupy, numba 速度比較まとめ

1000要素のベクトルのドット積を求めるプログラムで、
numpy, cupy, numba の実行速度を比較してみました。

numba 64bit: 0:00:00.001167
numba 32bit: 0:00:00.000692
cupy 64bit: 0:00:00.000074
cupy 32bit: 0:00:00.000060
numpy 64bit: 0:00:00.000006
numpy 32bit: 0:00:00.000003

となり、numpy 圧勝。cupy が若干遅くなっているのはメイン
メモリとグラフィックメモリの相互移動が発生しているためと
推察します。
numba が著しく遅いのは和の計算に atomic add を使っている
ためでしょう。
(バイナリツリーを用いると劇的に速くなりますが、
面倒なので実装せず。)

最後にテストに用いたプログラムを載せておきます。

# -*- coding: utf-8 -*-
from numba import cuda
import numpy as np
import cupy as cp
import datetime

@cuda.jit
def dotProduct(c, a, b):
        i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
        cuda.atomic.add(c, 0, a[i] * b[i])

N = 1000
a = np.random.rand(N).astype(np.float64)
b = np.random.rand(N).astype(np.float64)
d = np.zeros(1, dtype=np.float64)

a32 = np.random.rand(N).astype(np.float32)
b32 = np.random.rand(N).astype(np.float32)
d32 = np.zeros(1, dtype=np.float32)

ac = cp.asarray(a)
bc = cp.asarray(b)
dc = cp.asarray(d)

ac32 = cp.asarray(a32)
bc32 = cp.asarray(b32)
dc32 = cp.asarray(d32)

threadN = 256
blockN = (N + threadN - 1) // threadN
# burn-in
dotProduct[blockN, threadN](d, a, b)
dotProduct[blockN, threadN](d32, a32, b32)
dc = cp.dot(ac, bc)
dc32 = cp.dot(ac32, bc32)
d_gold = np.dot(a, b)
d32 = np.dot(a32, b32)

d = np.zeros(1, dtype=np.float64)
st = datetime.datetime.now()
dotProduct[blockN, threadN](d, a, b)
ed = datetime.datetime.now()
print("numba 64bit:", ed - st)

d32 = np.zeros(1, dtype=np.float32)
st = datetime.datetime.now()
dotProduct[blockN, threadN](d32, a32, b32)
ed = datetime.datetime.now()
print("numba 32bit:", ed - st)

st = datetime.datetime.now()
dc = cp.dot(ac, bc)
ed = datetime.datetime.now()
print("cupy 64bit:", ed - st)

st = datetime.datetime.now()
dc32 = cp.dot(ac32, bc32)
ed = datetime.datetime.now()
print("cupy 32bit:", ed - st)

st = datetime.datetime.now()
d_gold = np.dot(a, b)
ed = datetime.datetime.now()
print("numpy 64bit:", ed - st)

st = datetime.datetime.now()
d_gold32 = np.dot(a32, b32)
ed = datetime.datetime.now()
print("numpy 32bit:", ed - st)

 

コメントを残す