thrust::for_eachによるヒストグラム

Thrust、個人的にはかなり気に入ってるライブラリなんだけど、なんだか影が薄いんですよね。 cuFFTやcuDNNなどなど、僕らにはとても書けそうにないガチにヘビィな連中とは異なり、Thrustが提供するのは copy/search/sortといった基本的なアルゴリズム群だからでしょうか。

CUDAが、というかnvccが device lambda をサポートしてくれたおかげで、Thrustの使い心地がぐっと良くなりました。その一例。

ヒストグラム(度数分布表)を作ります。たくさんの数値データがあって、それらのとる値の範囲を等間隔に刻んだスロット列を用意します。各数値をスロットに振り分けたとき、各スロットにいくつずつ納まるかってやつ。

アルゴリズムはこんなカンジ:

  float score[N]; // 数値データ列(input)
  int   hist[S]; // ヒストグラム(output) 初期値はall-0
  for ( score 内の各数値 val について　) {
    if ( valがスロットiの範囲にあるなら } {
      ++hist[i};
    }
  }

これをkernelで実装するのはちっとも難しくありません。けどもたったこれだけのためにわざわざkernel書くのが面倒なわけで、Thrustでさらっと実装してみました。

#include <thrust/device_vector.h>
#include <thrust/for_each.h>

#include <curand.h>

#include <iterator>
#include <iostream>
#include <ctime>

int main() {
  using namespace std;

  const int N = 100000;
  thrust::device_vector<float> dscore(N);
  thrust::device_vector<int>   dhist(101, 0);

  { // 平均 50.0, 標準偏差 15.0 の正規分布乱数で dscoreを埋める 
    curandGenerator_t gen;
    curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);
    curandSetPseudoRandomGeneratorSeed(gen, static_cast<unsigned long long>(time(nullptr)));
    curandGenerateNormal(gen, dscore.data().get(), N, 50.30f, 15.0f);
    curandDestroyGenerator(gen);
  }

  // dscoreの各値を四捨五入し、0以上/100以下ならばヒストグラムに加える
  int* histPtr = dhist.data().get(); // for_each内のlambdaがコレをキャプチャする
  thrust::for_each(begin(dscore), end(dscore),
                   [=] __device__ (float val) -> void {
                     int i = static_cast<int>(val+0.5f); // 四捨五入して
                     if ( i >= 0 && i <= 100 ) { // 0以上100以下なら、当該スロットを+1する
                       atomicAdd(histPtr+i, 1); // ++histPtr[i] だと data-race起こすよきっと!
                     }
                   });
//*/
  for ( int item : dhist ) {
    cout << item << endl;
  }
//*/
}

thrust::for_each(こっから, ここまで, なんかする) の"なんかする"を devie lambda で与えています。ヒストグラムを求める領域:dhistのポインタをlambda式の値キャプチャで食わせています。

※ device lambdaでキャプチャできるのは値キャプチャだけ。参照キャプチャするとコンパイル・エラーとなります。Hostにある変数の参照(つまりポインタ)をDevice側に持ち込んでも意味ないもんね。

結果をCSVにリダイレクトし、excelに食わせたものがこちらになります。