tkm2261's blog

研究員(OR屋) → データ分析官 → MLエンジニア → ニート → 米国CS PhDが諸々書いてます

Python高速化　Numba入門　その4

Numba Python

今回はNumbaのGPUコンピューティングについて読んでいきます。
最終回の予定でしたが、エントリが超長くなりそうなので今回はGPUの使用方法、次回に計算速度の検証をして終わりたいと思います。

Writing CUDA-Python — numba 0.15.1 documentation

exprimental扱いなので、商用で使われる方はNumbaProの方をオススメします。

NumbaのGPUコンピューティングはこれまでのデコレータだけれ使えるNumbaの状況とは異なり、かなりCUDAに関する知識を要求されます。

もし、少ない労力でPythonコードを高速化したい方は、『その３』までの内容で十分と思われます。

CUDAについて

とりあえず、Wikipediaから引用

CUDA（Compute Unified Device Architecture：クーダ）とは、NVIDIAが提供するGPU向けのC言語の統合開発環境であり、コンパイラ (nvcc) やライブラリなどから構成されている。アプリケーションを実行する基盤となるプラットフォーム／アーキテクチャそのものをCUDAと呼ぶこともある。 CUDA - Wikipedia

つまりNVIDIAのGPU（GeForce等）でプログラミングするための環境になります。

詳しいことはここがわかりやすかったので紹介。後の議論を理解するには、このページをひと通り読まれることをオススメします。

第６回 CUDAプログラミングモデル① | G-DEP

簡単に説明すると、GPUは

グリッド　->　ブロック　-> スレッド　

という階層構造になっており、1グリッドに65,535×65,535個のブロックを配置でき、1ブロックで512個のスレッドを実行することが出来ます。

つまり65,535×65,535×512並列の計算ができることになります。

計算すると大体2兆並列できるとのこと。

GPUすげー・・・

Numbaでの表記

基本的な表記は以下となっています。

# coding: utf-8

from numba import cuda

@cuda.jit('void(int32[:], int32[:])')
def foo(aryA, aryB):
    ...

if __name__ == '__main__':
    griddim = 1, 2
    blockdim = 3, 4
    foo[griddim, blockdim](aryA, aryB)

ポイントを挙げると

  1. デコレータの指定　・・・　これまでと同様に型を指定 cuda.jitでなくjitにtarget='gpu'を指定しても良い
  2. griddim, blockdimの指定　・・・　今回使用するブロック数(griddim)と、ブロック内のスレッド数(blockdim)を指定する。

griddim, blockdimの数は、スレッドIDをインデックスに使用する場合にはインデックスより多くのスレッドを用意できる様に指定する。

ただしblockdimは32の倍数にしたほうが良いとのことでした。

また@cuda.autojitデコレータでは型を推測してくれますが、基本的に遅くなるので指定した方が無難です。

デコレータの引数

@cuda.jitデコレータには以下の引数を指定することができます。

  • device=True　・・・　GPU上でのみ使える関数
  • inline=True　・・・　上のデバイス関数をインライン化するかどうか。このインラインのことを指していると思われる。インライン展開 - Wikipedia

スレッドの特定

上の第６回 CUDAプログラミングモデル① | G-DEPを読まれた方は見たと思いますが、 CUDAの計算では、どのスレッド何を計算させるかを指定する必要があります。

基本的には（ほぼ必ず）スレッドの番号を使用します。
配列同士の和は、スレッド番号0に1番目の要素の和を、スレッド番号1には2番目の要素の和といった形になります。
グリッドが2次元まで対応しているので2次元配列まではシンプルに計算することが出来ます。

そうでないと2兆個の命令を人間が出すのはほぼ不可能です。

並列計算全般がそうですが、同期の必要のない大量の単純作業をさせるとGPUの効果が最大限発揮されます。

スレッド番号の取得方法は2種類用意されてます。

1: 普通のCUDAと同じに愚直に計算

グリッドが1次元の場合

tx = cuda.threadIdx.x
bx = cuda.blockIdx.x
bw = cuda.blockDim.x
i = tx + bx * bw
array[i] = something(i)

グリッドが2次元の場合

tx = cuda.threadIdx.x
ty = cuda.threadIdx.y
bx = cuda.blockIdx.x
by = cuda.blockIdx.y
bw = cuda.blockDim.x
bh = cuda.blockDim.y
x = tx + bx * bw
y = ty + by * bh
array[x, y] = something(x, y)

行列計算を自分で書いたことある方には馴染み易い方法と思います。

2: Numbaの便利関数を使う

グリッドが1次元の場合

i = cuda.grid(1)
array[i] = something(i)

グリッドが2次元の場合

x, y = cuda.grid(2)
array[x, y] = something(x, y)

好きな方を使えば良いと思います。

メモリ転送

こっからかなりディープな世界に入って来ます。
ここまでの解説で普通の用途では十分な速度向上が見られると思います。

ただし、最大限に性能を引き出すには必須なので、メモリを制する者はGPUを制すといった感じでしょうか

詳細はコチラを見て頂きたいのですが、メモリも色々あります。

第５回 GPUの構造 | G-DEP

ここではメモリの使用を制御する以下の2つを紹介します。

  1. CUDA Streamの使用　　　　　　　　　　　　
  2. Shared Memoryの使用

CUDA Streamの使用

通常のメモリの転送は以下になっています。

  • マシンメモリ→GPUメモリの転送・・・非同期
  • マシンメモリ←GPUメモリの転送・・・同期

CUDA Streamを使用すると『マシンメモリ←GPUメモリの転送』も非同期にすることが出来ます。
使用方法は以下。

stream = cuda.stream()
devary = cuda.to_device(an_array, stream=stream)
a_cuda_kernel[griddim, blockdim, stream](devary)
devary.copy_to_host(an_array, stream=stream)
# data may not be available in an_array
stream.synchronize()
# data available in an_array

with構文で以下のようにもかけます。

stream = cuda.stream()
with stream.auto_synchronize():
    devary = cuda.to_device(an_array, stream=stream)
    a_cuda_kernel[griddim, blockdim, stream](devary)
    devary.copy_to_host(an_array, stream=stream)
# data available in an_array

私自身もCUDA Streamの使用用途を見つけられていませんが、良かったら使って見て下さい。

Shared Memoryの使用

Shared Memoryが同一ブロックのスレッドが参照できるメモリで、かなり高速にアクセスすることが出来ます。Numbaマニュアルの例では行列積をブロック化して計算しています。

有名な例なので日本語でも資料があり、以下が非常にわかりやすい説明でした。

CUDAで行列演算：乗算(シェアードメモリ使用版) - PukiWiki for PBCG Lab

この例でもShared Memoryの使用で5～10倍高速化しており、効果が高いことがわかります。

NumbaマニュアルのShared Memoryを使った行列積コードは以下

bpg = 50
tpb = 32
n = bpg * tpb

@jit(argtypes=[float32[:,:], float32[:,:], float32[:,:]], target='gpu')
def cu_square_matrix_mul(A, B, C):
    sA = cuda.shared.array(shape=(tpb, tpb), dtype=float32)
    sB = cuda.shared.array(shape=(tpb, tpb), dtype=float32)

    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y
    bw = cuda.blockDim.x
    bh = cuda.blockDim.y

    x = tx + bx * bw
    y = ty + by * bh

    acc = 0.
    for i in range(bpg):
        if x < n and y < n:
            sA[ty, tx] = A[y, tx + i * tpb]
            sB[ty, tx] = B[ty + i * tpb, x]

        cuda.syncthreads()

        if x < n and y < n:
            for j in range(tpb):
                acc += sA[ty, j] * sB[j, tx]

        cuda.syncthreads()

    if x < n and y < n:
        C[y, x] = acc

アトミック操作

ここで説明すると長くなるので詳細は以下を参照。簡単に言うとメモリアクセスのブロックが必要な操作のための関数です。

CUDAアトミック関数 - PukiWiki for PBCG Lab

Numbaでの実行方法は以下

@cuda.jit('void(uint32[:], uint32[:])')
def fill_bins(indices, bin_counts):
    offset = cuda.grid(1)
    stride = cuda.gridDim.x * cuda.blockDim.x

    for i in range(offset, indices.shape[0], stride):
        cuda.atomic.add(bin_counts, indices[i], 1) # (対象配列, インデックス, 足す値)

『実際にやってみた』は次回

長々話しましたが、話を見るだけでは正直わからないと思います。
早めに実際のコードも上げるのでそちらを見ながら、当記事を振り返って頂ければと思います。

ライトな用途ではメモリ転送以下の話は無視したほうが良いかもしれません。