Thrust で 生ポを扱うには

Windows/Visual Studio版ではおなじみの「配列の足し算」: c[i] = a[i] + b[i] (i = 0, 1, ...) を Thrust使って書いてみました。

/*
  DO NOT FORGET: --expt-extended-lambda option to nvcc
 */

#include <cstdio>
#include <cassert>

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/transform.h>

void addWithThrust(thrust::host_vector<int>& c, 
                   const thrust::host_vector<int>& a, 
                   const thrust::host_vector<int>& b);

int main() {
  using namespace std;

  const int arraySize = 5;
  const int a_data[arraySize] = {  1,  2,  3,  4,  5 };
  const int b_data[arraySize] = { 10, 20, 30, 40, 50 };

  thrust::host_vector<int> a(a_data, a_data + arraySize);
  thrust::host_vector<int> b(b_data, b_data + arraySize);
  thrust::host_vector<int> c(arraySize);

  // Add vectors in parallel.
  addWithThrust(c, a, b);

  printf("{1,2,3,4,5} + {10,20,30,40,50} = {%d,%d,%d,%d,%d}\n",
        c[0], c[1], c[2], c[3], c[4]);

  cudaDeviceReset();
}

void addWithThrust(      thrust::host_vector<int>& c, 
                   const thrust::host_vector<int>& a, 
                   const thrust::host_vector<int>& b) {
  using namespace std;

  assert( a.size() == b.size() );
  assert( a.size() == c.size() );

  thrust::device_vector<int> dev_a = a; // copy Host to Device
  thrust::device_vector<int> dev_b = b; // copy Host to Device
  thrust::device_vector<int> dev_c(c.size());

  thrust::transform(begin(dev_a), end(dev_a), // dev_a for input
                    begin(dev_b),             // dev_b for input
                    begin(dev_c),             // dev_c for output
                    [] __device__ (int x, int y) -> int { return x + y; }); // z = x + y

  c = dev_c; // copy Device to Host
}

楽ちんですねー、thrust::host_vector/device_vector にメモリ管理を任せてしまえば cudaMallocはコンストラクタ、cudaMemcpy はoperator=()がやってくれるのでコードが実に涼しげです。

Thrustのコードを追いかけてみたところ、thrust:host_vectorのメモリ確保/解放は new/delete、thrust::device_vectorはcudaMalloc/cudaFreeが使われてるようです。

ここまで簡単になるんだからHost/Deviceで共用できるManaged Memory使えばもっと楽できんじゃないかと試してみたですよ。

#include <cstdio>
#include <cassert>

#include <algorithm>

#include <cuda_runtime.h>

#include <thrust/transform.h>

void addWithThrust(int* c,
                   const int* a,
                   const int* b,
                   int size);

int main() {
  using namespace std;

  const int arraySize = 5;
  const int a_data[arraySize] = {  1,  2,  3,  4,  5 };
  const int b_data[arraySize] = { 10, 20, 30, 40, 50 };

  int* a; cudaMallocManaged(&a, arraySize*sizeof(int)); std::copy( a_data, a_data + arraySize, a);
  int* b; cudaMallocManaged(&b, arraySize*sizeof(int)); std::copy( b_data, b_data + arraySize, b);
  int* c; cudaMallocManaged(&c, arraySize*sizeof(int)); std::fill( c, c+arraySize, 0);

  // Add vectors in parallel.
  addWithThrust(c, a, b, arraySize);

  printf("{1,2,3,4,5} + {10,20,30,40,50} = {%d,%d,%d,%d,%d}\n",
        c[0], c[1], c[2], c[3], c[4]);

  cudaFree(a);
  cudaFree(b);
  cudaFree(c);

  cudaDeviceReset();
}

void addWithThrust(int* c,
                   const int* a,
                   const int* b,
                   int size) {

  thrust::transform(a, a+size, // a for input
                    b,         // b for input
                    c,         // c for output
                    [] __device__ (int x, int y) -> int { return x + y; }); // z = x + y
  cudaDeviceSynchronize();

}

あらら、ぜーんぜん動いてません。原因はココ:

  thrust::transform(a, a+size, // a for input
                    b,         // b for input
                    c,         // c for output
                    [] __device__ (int x, int y) -> int { return x + y; }); // z = x + y

Thrustが提供する多くの関数は、引き渡されるイテレータの型に応じてHostで行うかDeviceで行うかを静的に判断しています。このコードではイテレータとして生のポインタを渡してるんですけど、Thrustではイテレータが(生の)ポインタであったならHostで実行すべきものと判断してるっポいです。

解決策はふたつ。ひとつは第一引数に実行ポリシーを明示的に与えること。

#include <thrust/execution_policy.h>
...
void addWithThrust(int* c,
                   const int* a,
                   const int* b,
                   int size) {

  thrust::transform(thrust::device,
                    a, a+size, // a for input
                    b,         // b for input
                    c,         // c for output
                    [] __host__ __device__ (int x, int y) -> int { return x + y; }); // z = x + y
  cudaDeviceSynchronize();
}

もうひとつは、生のポインタを thrust::device_ptr に変換すること。

#include <thrust/device_ptr.h>
...
void addWithThrust(int* c,
                   const int* a,
                   const int* b,
                   int size) {

  thrust::device_ptr<int> pa(const_cast<int*>(a));
  thrust::device_ptr<int> pb(const_cast<int*>(b));
  thrust::device_ptr<int> pc(c);
  thrust::transform(pa, pa+size, // a for input
                    pb,          // b for input
                    pc,          // c for output
                    [] __host__ __device__ (int x, int y) -> int { return x + y; }); // z = x + y
  cudaDeviceSynchronize();

}