はじめに
Chainerのfunctionsのコードを読んでいるとcuda.elementwise
とかcuda.reduce
の呼び出しを見かけることがあります。
これらは独自の処理をGPU上で実行するためのメソッドです。
Chainer functionの実装には欠かせないメソッドと言えるわけで、Chainer中級者になるために必要そうなので調査しました。
この記事ではcuda.elementwise
を扱い、cuda.reduce
については説明しません。
cuda.elementwise
の説明は以下にあります。
http://docs.chainer.org/en/stable/cupy-reference/kernel.html
またSlideShareにPreferred Networksの奥田氏による解説があります。
http://www.slideshare.net/ryokuta/cupy
確認した環境
- Windows 10
- Python 2.7
- Chainer 1.9.0
- CUDA 7.5
cuda.elementwiseは何をするのか?
cuda.elementwise
は、CUDA kernelを定義します。
CUDA kernelとは、CUDAに対応したGPU上で実行するプログラムです。
cuda.elementwise
を呼ぶと、戻り値としてCUDA kernelを実行するための関数kernel invocation functionが得られ、kernel invocation functionを呼ぶとGPU上でCUDA kernelを実行します。
kernel invocation functionの実体はElementwiseKernelオブジェクトで、このオブジェクトはcallableになっています。
最初のサンプルコード
最初の例として、与えられた配列の全ての要素をインクリメントする処理を行います。
まず以下のように必要なモジュールをインポートし、cuda.cupy
の参照としてxp
を定義します。
これ以降のサンプルコードは以下のコードを実行していることを前提とします。
import numpy as np
import chainer
from chainer import cuda
xp = cuda.cupy
次にcuda.elementwise
を使い、配列要素のインクリメントを行います。
x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
y = cuda.elementwise(
'T x',
'T y',
'y = x + 1;',
'sample1_fwd',
)(x)
print(y)
出力は以下のようになりました。
全ての要素がインクリメントされていることがわかります。
beforeの出力の後afterの出力を行うまでに少し時間がかかりますが、裏でnvccによるコンパイルを行っていると思われます。
[[ 2. 3. 4.]
[ 5. 6. 7.]]
サンプルコードの解説
cuda.elementwise
は以下のように2段階に分けて使用します。
ただ実際のコードではcuda.elementwise(...)(x)
のように2つの処理をまとめて記述することが多いです。
-
cuda.elementwise
を呼び出してkernel invocation functionを生成 - 生成したkernel invocation functionを呼び出して処理を実行する
cuda.elementwise
の引数・戻り値の説明はこのメソッドではなく、cupy.ElementwiseKernel
のドキュメントに記載されています。
cuda.elementwise
は内部でcupy.ElementwiseKernel
を呼び出しており、両者の引数は(cuda.elementwise
でname
が必須であることを除いて)同じです。
cuda.elementwise
には以下の引数が必要です。このほかにOptionalな引数がありますが、まだ十分理解できていないので説明は割愛します。
- in_params(str)
- out_params(str)
- opearation(str)
- name(str)
in_params
入力引数を宣言する文字列を指定します。
引数には型と引数名が必要です。
型文字列が1文字の場合にはtype placeholderになります。
type placeholderが表す型はkernel invocation function実行時に渡した変数の型になります。
同じ型の変数を宣言したいときに便利です。
サンプルコードではin_paramsは以下のようになっていました。
'T x',
この文字列によって以下を表現しています。
- 入力引数としてxを宣言
- xの型はT
- Tはtype placeholder
out_params
出力引数を宣言する文字列を指定します。
in_paramsと同様に引数には型と引数名が必要です。
サンプルコードではout_paramsは以下のようになっていました。
'T y',
この文字列によって以下を表現しています。
- 出力引数としてyを宣言
- yの型はT
- Tはtype placeholderなので、yの型はxの型と一緒
operation
実行する処理を定義する文字列を指定します。
サンプルコードでは、y
にx + 1
を代入する処理になっていました。
'y = x + 1;',
name
処理の名前です。
chainer.functions以下のモジュールの実装を見るとforward処理の場合は"function_name_fwd"、backward処理の場合は"function_name_bwd"としています。
kernel invocation functionの実行
kernel invocation functionを実行すると、定義したCUDA kernelを実行します。
実行時に引数としてin_paramsに対応した変数を渡します。
out_paramsに対応した変数は省略可能ですが、in_paramsに対応した変数の後に指定することで明示的に渡すこともできます。
戻り値はout_paramsで指定した引数です。
out_paramsに複数の引数を指定した場合には戻り値がそれらのtupleになります。
out_paramsにも値を渡す
out_paramsにも値を渡してみましょう。
kernel invocation functionの引数にout_paramsに対応する値を渡すだけです。
x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
y = xp.asarray([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
y = cuda.elementwise(
'T x',
'T y',
'y += x;',
'sample2_fwd',
)(x, y)
print(y)
実行結果は以下のようになり、元のyにxが加算されています。
[[ 2. 3. 4.]
[ 6. 7. 8.]]
Broadcasting
配列のブロードキャストは自動で行ってくれます。
x = xp.asarray([1, 2, 3], dtype=np.float32)
y = xp.asarray([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
y = cuda.elementwise(
'T x',
'T y',
'y += x;',
'sample3_fwd',
)(x, y)
print(y)
実行結果:
[[ 2. 3. 4.]
[ 3. 4. 5.]]
配列のサイズが合わずブロードキャストできない場合はエラーが発生します。
x = xp.asarray([1, 2], dtype=np.float32)
y = xp.asarray([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
y = cuda.elementwise(
'T x',
'T y',
'y += x;',
'sample4_fwd',
)(x, y)
print(y)
実行結果:
Traceback (most recent call last):
File "elementwise_sample.py", line 61, in <module>
)(x, y)
File "cupy\core\elementwise.pxi", line 508, in cupy.core.core.ElementwiseKernel.__call__ (cupy\core\core.cpp:34118)
File "cupy\core\elementwise.pxi", line 334, in cupy.core.core._broadcast (cupy\core\core.cpp:31734)
File "cupy\core\core.pyx", line 1504, in cupy.core.core.broadcast.__init__ (cupy\core\core.cpp:50697)
ValueError: Broadcasting failed
Indexing
配列を操作する時にインデックスを指定したいことはよくあります。
以下のようにすることでインデックスを指定できます。
- インデックスを指定してアクセスしたい変数に
raw
をつける - 特殊変数
i
がインデックスを表す -
_ind.size()
がインデックスの数を表す
配列の要素を逆順にするサンプルを示します。
x = xp.asarray([1, 2, 3, 4], dtype=np.float32)
y = xp.zeros_like(x, dtype=np.float32)
y = cuda.elementwise(
'raw T x',
'T y',
'y = x[_ind.size() - i - 1];',
'sample5_fwd',
)(x, y)
print(y)
実行結果は以下のようになります。
[ 4. 3. 2. 1.]
上記のコードはNumpyを使った以下のコードをGPU上で実行していると考えればよいと思います。
x = np.asarray([1, 2, 3, 4], dtype=np.float32)
y = np.zeros_like(x, dtype=np.float32)
i = np.arange(4)
y = x[4 - i - 1]
注意点として、kernel invocation functionにy
を渡すことが必要であることが挙げられます。
以下のようにy
を渡さないとValueError: Loop size is Undecided
というエラーが発生します。
これはrawな引数だけだとインデックスのサイズを決められないために起こるようです。
x = xp.asarray([1, 2, 3, 4], dtype=np.float32)
y = xp.zeros_like(x, dtype=np.float32)
y = cuda.elementwise(
'raw T x',
'T y',
'y = x[_ind.size() - i - 1];',
'sample6_fwd',
)(x)
print(y)
もう少し複雑なIndexing
x
が2次元の配列、t
が1次元の配列のときに、x[t[i]]
(i=0, 1, 2, ...)を取得することを考えます。
これは以下のようにして書けます。
x = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
t = xp.asarray([0, 2, 1], dtype=np.int32)
y = cuda.elementwise(
'raw T x, S t',
'T y',
'int ind[] = {i, t}; y = x[ind];',
'sample7_fwd',
)(x, t)
print(y)
実行結果:
[ 1. 6. 8.]
int ind[] = {i, t};
は、[(0, t[0]), (1, t[1]), (2, t[2])]を示すインデックスを生成します。
forループ
for
、while
などC(正確にはnvcc?)の構文を使用できます。
例として、xの列ごとの累積値を計算します。
y[i, j]
がx[0, j]
からx[i, j]
までの累積になるようにします。
x = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
y = xp.zeros_like(x)
y = cuda.elementwise(
'raw T x, int32 c',
'raw T y',
'''
int ind[] = {0, i};
y[ind] = x[ind];
for (int j = 1; j < c; j++) {
int ind[] = {j, i};
int prev_ind[] = {j - 1, i};
y[ind] = y[prev_ind] + x[ind];
}
''',
'sample8_fwd',
)(x, x.shape[0], y, size=x.shape[1])
print(y)
実行結果:
[[ 1. 2. 3.]
[ 5. 7. 9.]
[ 12. 15. 18.]]
CUDA関数
CUDAの関数も使えます。
ただしどこまでサポートしているかは不明です。
例としてatomicAdd
を使ってみます。
x = xp.zeros((3, 3), dtype=np.float32)
t = xp.asarray([0, 1, 2], dtype=np.int32)
y = cuda.elementwise(
'S t',
'raw T x',
'int ind[] = {i, t}; atomicAdd(&x[ind], 1);',
'sample9_fwd',
)(t, x)
print(y)
実行結果:
[[ 1. 0. 0.]
[ 0. 1. 0.]
[ 0. 0. 1.]]
最後に
cuda.elementwise
を理解することで、Chainerの理解もより深まると思います。
cuda.elementwise
とcuda.reduce
はChainer内部でよく使われており、もっと詳しく知りたい方はそちらを参照するのがよいと思います。