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)

 

CUDA の Tensor Core を使ってみた

折角 TITAN V を購入したので、SM 7.0 以上で利用できる
Tensor Core を使った行列計算テストプログラム
書いてみました。

CUDA Toolkit のサンプルとして Tensor Core を
使ったプログラムがありますが、余計な最適化処理を
していてちょっと分かりづらい。

さらに2のn乗(n >= 4)でないと動作しないコードだった
ので、任意の行列に対応するようにしました。

で、その効果ですが、やや微妙です。。。

start wmma version...
calc_fm elapsed time without TensorCore:7.1383ms
calc_fm elapsed time without TensorCore:4.08064ms
calc_fm elapsed time with TensorCore:3.6567ms
start CPU gold...
calc_fm_gold elapsed time:43.6952m

MNIST を VGG とかで解析する際の第1段めを想定した実験
ですが、効果は10%程度となりました。っていうか、Tensor Core
を使わなくても TITAN V の CUDA Core が速すぎるというか。

WMMA を使うにあたっていくつか TIPS がありました。

  1. 行列の内積を求める際は row_major を使う。
    サンプルでは matrix_a に対して col_major を指定しています
    が、row_major にしないと正しく積が計算されません。
  2. __syncthreads() を呼ばなくてもいい。
    wmma ロジック内部でスレッド同期してから返される
    ようなので、__syncthreads() 等のスレッド同期を外部で
    行う必要はありませんでした。
  3. タイル数をむやみに増やしても効果がない。
    nVidia のサンプルでは全ての部分行列(タイル)を読み込んで
    高速化を図っていますが、行列をキャッシュして使い回す以上の
    意味はなく、それ以外の積和演算のパフォーマンスへの
    インパクトがありません。
    これは wmma のアクセス関数が *_sync と同期タイプしか
    提供されていないためだと思われます。
    (正直、どうやったら 640 コア使えるのか分かりません。)
  4. 部分行列の総和は matrix_c を読み書きする。
    accumulator はload/storeが可能です。
    というか、この機能がないとwmmaの外で総和を計算
    することになり、若干不便です。(私が複数タイル用に
    書いたようなコードを埋め込む必要が出てきます。)
  5. 行列の要素数によっては効果がないことがある。
    一般に、要素数が多いほど効果は上がるのですが、
    n, m, k の関係によっては Tensor Core を使った方が
    遅くなる現象が見られました。発生条件は調査中です。

nbody を cupy で書いてみた

CUDA サンプルにもある nbody (多体問題) を cupy で解いてみた。
うまくはまれば namba を使うよりも少ない行数で記述できる。
実行速度は namba とほぼ変わらなかった。

初速は銀河面上を円運動するように与えているため、シミュレーション
時間を長くしてもあまり面白い絵になっていませんので、
適宜工夫してみて下さい。

ソースコードは以下の通り。

# -*- coding: utf-8 -*-

import datetime
import cupy as cp
import numpy as np
import matplotlib.pyplot as plt

# 質点数
N = 1024
REP = 5000
AREAX = 1e3
AREAY = 1e3
FLIP = 0

# 物理定数
cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
MASS = cp.float32(1.0)
M = cp.ones(N, dtype=cp.float32) * MASS
G = cp.float32(1.0)
EPSILON = cp.float32(1e-6)
DT = cp.float32(0.01)

# 分布、初速を決定する
DAX = cp.zeros(N, dtype=cp.float32)
DAY = cp.zeros(N, dtype=cp.float32)

def advance(x0, y0, xx1, xx2, yy1, yy2, M, EPSILON):
        diffx = x0[xx2] - x0[xx1]
        diffy = y0[yy2] - y0[yy1]
        r = cp.sqrt(diffx * diffx + diffy * diffy)
        r3 = r * r * r + EPSILON
        dax = cp.sum(M * diffx / r3, axis = 1)
        day = cp.sum(M * diffy / r3, axis = 1)
        return dax, day

def update(x0, y0, vx0, vy0, DAX, DAY, G, DT):
        vx1 = vx0 + DAX * G * DT
        vy1 = vy0 + DAY * G * DT
        x1 = x0 + vx1 * DT
        y1 = y0 + vy1 * DT
        return x1, y1, vx1, vy1

def eccentric(x):
        return (cp.exp(2 - 8 * x) - cp.exp(2)) / (cp.exp(2 - 8) - cp.exp(2))

def init_dist():
        global AREAX, AREAY, MASS, G
        RR = eccentric(cp.random.rand(N)) + 1e-7
        THETA = cp.random.rand(N)
        XS = RR * cp.cos(THETA * 2 * cp.pi)
        YS = RR * cp.sin(THETA * 2 * cp.pi)
        VAREA = cp.ones(N, dtype=cp.float32) * cp.sqrt(G/N) / RR
        RR = cp.random.rand(N) / 4
        THETA = cp.random.rand(N)
        VSX = RR * cp.cos(THETA * 2 * cp.pi)
        VSY = RR * cp.sin(THETA * 2 * cp.pi)

        r = cp.sqrt(XS * XS + YS * YS)
        VSX = VAREA * (-YS + VSX * (1 - r))
        VSY = VAREA * (XS + VSY * (1 - r))
        XS *= AREAX
        YS *= AREAY

        # 重心を求める
        x0 = cp.sum(XS) / N
        y0 = cp.sum(YS) / N
        vx0 = cp.sum(VSX) / N
        vy0 = cp.sum(VSY) / N

        # 重心座標系に移行する
        XS -= x0
        YS -= y0
        VSX -= vx0
        VSY -= vy0
        return XS, YS, VSX, VSY

# 分布、初速を決定する
XS, YS, VSX, VSY = init_dist()
X = [XS, cp.zeros(N, dtype=cp.float32)]
Y = [YS, cp.zeros(N, dtype=cp.float32)]
VX = [VSX, cp.zeros(N, dtype=cp.float32)]
VY = [VSY, cp.zeros(N, dtype=cp.float32)]
print("intialized")

# REP 時間繰り返す
x0s = range(N)
y0s = range(N)
xx1, xx2 = np.meshgrid(x0s, x0s)
yy1, yy2 = np.meshgrid(y0s, y0s)
st = datetime.datetime.now()
for t in range(REP):
        DAX, DAY = advance(X[FLIP], Y[FLIP], xx1, xx2, yy1, yy2, M, EPSILON)
        X[1-FLIP], Y[1-FLIP], VX[1-FLIP], VY[1-FLIP] = \
                update(X[FLIP], Y[FLIP], VX[FLIP], VY[FLIP], \
                DAX, DAY, G, DT)
        if t % (REP // 50) == 0:
                print('*', end = '', flush = True)
        # 結果をグラフで表示
        if t == 0:
                plt.plot(cp.asnumpy(X[FLIP]), cp.asnumpy(Y[FLIP]), 'x', color='red')
        FLIP = 1 - FLIP

ed = datetime.datetime.now()
print("\nelapsed time:", ed - st)
x = cp.asnumpy(X[FLIP])
y = cp.asnumpy(Y[FLIP])
plt.plot(x, y, '.', color='green')
plt.show()

 

cupy のメモリアロケータ

cupy の中の人の投稿によると、cupy.cuda.set_allocator() で
メモリプールを使い回すようにすると高速になると書いてある。

実際そのとおりなのだが、投稿にある test() を2回連続で
回すと、アロケータをセットした方が圧倒的に遅くなった。
ソースコードを見ていないが、どこかに実行結果をキャッシュしている?

numpy: 0:00:00.004890 sec
cupy: 0:00:00.000179 sec
cupy(memory pool): 0:00:00.000456 sec

なお、記事では触れていないが、いわいる「バーンイン」を
行わないで GPU 回りのメソッドを呼ぶと初期化設定のため、
最初の1回だけは非常に時間がかかることに注意。

使用したテストコードは以下の通り。

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

def test(xp):
    x = xp.arange(1000000).reshape(1000, -1)
    return x.T * 2

def test_with_time(title, xp):
    t1 = datetime.datetime.now()
    test(xp)
    t2 = datetime.datetime.now()
    print(title, t2 - t1, "sec")

# バーンイン
#cp.arange(1)
test(cp)
 
test_with_time("numpy:", np)
test_with_time("cupy:", cp)
cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
test_with_time("cupy(memory pool):", cp)

 

numba パッケージの記事を追加しました

numba パッケージを使う場合のヒントに関する記事をこちらにアップしました。

ループ数が小さかったり配列数が小さかったりすると CUDA の効果はほぼないです。
手軽に CUDA を使って実験したい場合に重宝しそうです。

なお、一定数ループを回ると実行速度がぐんと上がる現象が見られています。
(python3 の実行エンジンの関係かもしれませんが)

CUDA9.2 を使って JCuda0.9.0d をコンパイルする

JCuda – ダウンロードとコンパイルで解説しているように実行すると、JCufft を作成する時にエラーで止まります。

これは9.1から cufftSetCompatibilityMode() が削除されたためなので、jcufft/JCufftJNI/src/JCufft.cpp の最後の関数の戻り値を JCUFFT_INTERNAL_ERROR 等に変更します。

これでコンパイルが通るようになりますが、今度は jcuda-main に移動して mvn clean するときに失敗します。この原因は pom.xml に cuDNN の設定が入っていることと、デフォルトの Ubuntu 16.04.x に openjdk-8-jdk パッケージがインストールされていないためです。

pom.xml から cuDNN の設定を削除し、openjdk-8-jdk をインストールすると、mvn install が成功し、output ディレクトリに jar ファイルが作成されます。

そのうち HP の方も修正します。修正しました。

久々の投稿: python で cavity 問題を解いてみた

python3 で2次元の cavity 問題を解いたサイトがないか調査していたが、探しきれなかったので、

https://github.com/minosys-jp/Fluid.git

の方に登録してみました。高速化は numpy, numba で実行していますが、numpy のみを使った方が成績が良いようです。

出力をアニメーションしようとすると多大な時間がかかります。matplotlib の実装が悪いからだと思います。

後日談: OpenFOAM があることに気づきました。でもたまには自分の手を動かすのもよいかと。

jcurandの記事を追加しました

cuRAND を JCuda 環境で使用するための API である、jcurand の使い方をアップしました。

お膳立てが若干ありますが、それさえクリアすれば、シンプルな呼び出しで高速に乱数生成ができます。

メモリはデバイス側に確保する必要がありますが、メモリコピーを使えばホスト側に持ってきて使用することが可能です。

JCuda 版ニューラルネット記事追加

JCuda 版ニューラルネットのコードに関する説明記事をアップしました。CPU 版と比較して、以下の変更があります、

  • CPU版では batchsize 方向の並列化はしなかったが、JCuda 版では batchsize 方向も並列計算する。
  • CPU 版ではホストメモリの管理のみだったが、JCuda 版ではホストメモリおよびデバイスメモリの両方を管理する。
  • JCuda 版では計算はデバイスメモリで閉じるようにする。
  • メソッドはほぼ共通して使えるように、デバイスメモリ管理をオブジェクト定義に追加して隠蔽する。

 

SimpleNet.java バグ修正しました

SimpleNet.java の forward() 計算で、bias を使用しないというバグがあったため、JCudaNNKernel.cu とともに修正しました。ただ、bias があってもなくても正答率に目立った変化はないようです。

JCuda 版ニューラルネットのソースコード解説は、現在作成中です。完成にはあと2〜3日かかると思いますので、好ご期待。

AutoEncoder 系のライブラリも作成中ですが、まだ動きません。