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)