Anaconda の Numba を CUDA 9.0 で使う

Volta 対応の CUDA 9.0 が出ましたが、Anaconda の cudatoolkit は 8.0 までしか対応していないため、そのままでは Python  の Numba パッケージが使えません。

というわけで、悪戦苦闘した結果を書いておきます。

環境変数の設定

Numba を利用するためには以下の環境変数を設定する必要があります。

• NUMBAPRO_NVVM
• NUMBAPRO_LIBDEVICE

CUDA 9.0 をインストールすると /usr/local/cuda/nvvm/lib64/libnvvm.so.* というライブラリがが作成されますので、このパスを NUMBAPRO_NVVM に設定します。（ファイル名まで含む。）

また、/usr/local/cuda/nvvm/libdevice の下に libdevice.10.bc というファイルがありますので、シンボリックリンクで
libdevice.compute_XX.10.bc というリンクを貼っておきます。（XX は使用している GPU の Compute Capability を指定します。例えば GTX 1080 なら compute_61 となります。）NUMBAPRO_LIBDEVICE には /usr/local/cuda/nvvm/libdevice を指定します。

compute_XX の追加

Anaconda 5.0.0 の Numba には Compute Capability が 5.0 までしか登録されていないため、GPU によっては実行エラーが出てしまいます。ちょっと汚いやり方ですが、Numba パッケージに新しい Compute Capability を追加する方法で回避しました。

GTX 1080 用なら、具体的には ~/anaconda3/lib/python3.6/site-packages/numba/cuda/cudadrv/nvvm.py に以下の2点を追加します。

変数やクラス	修正内容
SUPPORTED_CC	リストの最後に (6, 1) を追加
LibDevice クラスの _known_arch 変数	“compute_61” をリストの最後に追加

nVidia のデモプログラムへの修正

NVidia のサイトで説明しているサンプルは numbapro を使ったものなので、numba 用に修正が必要となります。

import numpy as np
from numpy.random import rand
from numba import vectorize

@vectorize(['float32(float32, float32)'], target='cuda')
def VectorAdd(a, b):
    return a + b

def main():
    N = 100000000 # 100M
    print("wait for initialization...")
    A = rand(N).astype(np.float32)
    B = rand(N).astype(np.float32)
    C = np.zeros(N, dtype=np.float32)
    print("OK. Let's start!")

    C = VectorAdd(A, B)
    print("C[:5] = " + str(C[:5]))

if __name__ == "__main__":
    main()

ちなみに、Anaconda3 は Intel の MKL を抱えているっぽいので、この程度の演算だと GPU の方が CPU より遅いという結果となりました。PCIe 上のメモリ転送がボトルネックを引き起こしているためと思われます。

デバイスメモリに結果をキャッシュするような例なら速くなるはず。