はじめに

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内部でよく使われており、もっと詳しく知りたい方はそちらを参照するのがよいと思います。

chainer.cuda.elementwiseを使ってGPUで処理を行う

はじめに

確認した環境

cuda.elementwiseは何をするのか？

最初のサンプルコード

サンプルコードの解説

in_params

out_params

operation

name

kernel invocation functionの実行

out_paramsにも値を渡す

Broadcasting

Indexing

もう少し複雑なIndexing

forループ

CUDA関数

最後に