モンテカルロ法で円周率を求める(改)

さっきのarticle、総計を求めるためだけに thrust::reduce を使いましたけど、Thrustはすっごく便利で使いでがあるんですよ。

たとえば thrust::device_vector<T> は cudaMalloc/cudaFree で取得/解放するデバイス・メモリーが可変長配列であるかのようにふるまってくれるクラス。 デバイス・メモリーがstd::vector<T> みたいに使えちゃう。

「モンテカルロで円周率」を Thrust色濃いメにre-writeしてみました。
ぱっと見これがCUDAとは思えぬコードができましたよ♪

/*
  DO NOT FORGET OPTION: --expt-extended-lambda 
 */
// CUDA Runtime
#include <cuda_runtime.h>
#include <device_launch_parameters.h>

// cuRAND
#include <curand.h>

// Thrust
#include <thrust/device_vector.h>
#include <thrust/device_ptr.h>
#include <thrust/reduce.h>
#include <thrust/transform.h>

#include <iostream>
using namespace std;

int main() {

  const size_t N = 1000 * 1000;

  thrust::device_vector<float> x(N);
  thrust::device_vector<float> y(N);
  thrust::device_vector<float> count(N);

  curandGenerator_t gen;
  // メルセンヌ・ツイスタ
  curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_MTGP32);
  curandSetPseudoRandomGeneratorSeed(gen, 1234ULL); // seedはテキトーな値

  // 0～1の一様乱数列を生成する
  cout << "--- generating uniform random numbers ---" << endl;
  curandGenerateUniform(gen, thrust::raw_pointer_cast(x.data()), N);
  curandGenerateUniform(gen, thrust::raw_pointer_cast(y.data()), N);

  // i番目の点(x[i],y[i])が1/4円の中に入っていれば count[i] = 1
  cout << "--- is the point in the circle? ---" << endl;
  thrust::transform(begin(x), end(x), begin(y), begin(count), 
                    [] __device__ (float x, float y) { 
                      return ( x*x + y*y < 1.0f ) ? 1.0f : 0.0f; 
                    }); 

  // count[N]の総計をNで割って4倍すれば π になるハズ
  cout << "--- accunulating the result(pi) ---" << endl;
  float pi = thrust::reduce(begin(count), end(count)) * 4.0f / N;
  cout << pi << endl;

  curandDestroyGenerator(gen);

}

(original: 2015-04-06)
追記: CUDA 7.5 以降、device-lambda が使えるようになりました。
コンパイル・オプション: --expt-extended-lambda をお忘れなく。