More than 1 year has passed since last update.

概要

GPUやCPUで高速に動作するライブラリthrustを用いてhistogram作成を行ってみた。
ソースコードはgistにある。

histogram作成

各binに属している要素を数えたい。すなわち

key = {1, 2, 2, 3, 3, 3, 1, 1, 5, 5, 0, 0}

と与えられたら

bin = {2, 3, 2, 3, 0, 2}

とkeyが0のものの個数、1のものの個数...を配列として返すようにしたい。
これをthrustだけを用いて行ってみることにする。

histogram作成に必要な関数

thrust::reduce_by_key

これは実例を見せたほうが早い。

histogram.cu
// input value
const int key[]   {1, 1, 2, 2, 2, 3, 3, 1, 1};
const int value[] {1, 1, 1, 1, 1, 1, 1, 1, 1};

// copy data from host to device
thrust::device_vector<int> key_in(key, key + in_size);
thrust::device_vector<int> value_in(value, value + in_size);

// create output buffer
thrust::device_vector<int> key_out(in_size, 0);
thrust::device_vector<int> value_out(in_size, 0);

// do reduction operation
auto new_end = thrust::reduce_by_key(key_in.begin(),
                                     key_in.end(),
                                     value_in.begin(),
                                     key_out.begin(),
                                     value_out.begin());

とあった時に、outputは

key_out   = {1, 2, 3, 1};
value_out = {2, 3, 2, 2};

のようになる。つまり、隣接するkeyで同じ値を持っているものについてのみreductionを行う。
valueを全て1に設定することで各keyに属する要素が何個あるのかを計算している。

しかし、ここで注意すべきは、隣接するkeyについてしかreductionが行われないということ。このままではkeyが1となっている要素のカウントを誤ってしまう。

thrust::sort

そこで一旦keyをsortする。thrustにはSTLのようにsortを行うための関数が幾つか用意されている。

histogram.cu
thrust::sort(key_in.begin(), key_in.end());

keyをsortしたあとでreductionを行う。

histogram.cu
thrust::device_vector<int> uni_vect(in_size, 1); // 要素が全て1のベクトル
thrust::device_vector<int> num_in_bins(in_size); // 各keyに所属している要素の数を入れるためのbuffer
new_end = thrust::reduce_by_key(key_in.begin(),  // histogram計算
                                key_in.end(),
                                uni_vect.begin(),
                                key_out.begin(),
                                num_in_bins.begin());

すると結果は

key_out = {1, 2, 3};
num_in_bins = {4, 3, 2};

となり、正しくhistogramが作成できていることがわかる。

各キーの開始アドレスを計算

histogram作成後、各keyの開始アドレスが欲しい時がある。
例をあげると、

key = {1, 1, 1, 1, 2, 2, 2, 3, 3};
       ↑           ↑        ↑
ptr = {0,          4,       7};

↑で示した箇所のアドレスを配列ptrとして計算するということ。

この時にはscanを用いる。scanとは

\begin{align}
&\mathrm{in} \,\,\,\,\,\,\, a_{0} \,\,\,\,\, a_{1} \,\,\,\,\, a_{2} \,\,\,\,\, a_{3}  \\
&\mathrm{out} \,\,\,\, a_{0} \,\,\,\,\, a_{0} + a_{1} \,\,\,\,\, a_{0} + a_{1} + a_{2} \,\,\,\,\, a_{0} + a_{1} + a_{2} +a_{3}
\end{align}

のような処理を指す。 histogramが格納されている配列に対して上記の処理を行うと先頭アドレスが得られることがわかると思う。

上記のような処理はthrustでinclusive_scanとexclusive_scanとして実装されている。
二つの関数の違いは、配列の先頭データの取り扱い方にある。

num = {1, 3, 5, 4, 8};

// after inclusive scan
out_incl = {1, 4, 9, 13, 21};

// after exclusive scan
out_excl = {0, 1, 4, 9 , 13};

ここではinvlusive_scanを用いて各キーの開始アドレスを計算してみることにする。
単純にnum_in_binsに対してthrust::inclusive_scan関数を呼べば良い。

histogram.cu
thrust::device_vector<int> pointer(out_size + 1, 0);
thrust::inclusive_scan(num_in_bins.begin(),
                       num_in_bins.end(),
                       pointer.begin() + 1);

pointerには各キーの開始アドレスが格納されることになる。

補足

以上あげたthrustの関数だが、同じ名前の関数がboostやSTL(C++17から)でも使えるようになっている。
boostやSTL版はもちろんCPU上で動くものだが、今後マルチコア対応のものも出るようだ。

C++17から
std::inclusive_scan
std::exclusive_scan

boostから
boost::compute::reduce_by_key

並列版
parallel stl

まとめ

thrustを用いたhistogram作成の方法についてまとめた。
分子動力学法でVerlet listを作成する方法である"Grid search method"を実装する際にこの手の処理が必要になる。