thrustで3引数以上のtransformを実行する方法 #CUDA

GPGPU Advent Calendar6日目です！
あっちに書いた通り私自身は全く情報系の人間ではないですしプログラマーじゃないので、期待はしないでください（という前書き。

で、あんまり抽象的なこと書けないので、タイトルみたいな実務的なことを書きます。
thrustのサンプルがあるので、それを見てもらって分かる方はそれを読んでください。

1. やりたいこと

計算力学やってると、（相対）位置ベクトルやら速度ベクトルやらの大きさが欲しい時があります。つまり、

Length2.cpp

#include<iostream>

// 2次元ベクトルCPU
void Length2()
{
    // 要素数
    const int N = 5;

    // x, y方向成分
    double x[N] = {0, 1, 2, 3, 4};
    double y[N] = {1, 2, 3, 4, 5};

    // 各ベクトルの大きさ
    double length[N];

    // 大きさを計算
    for(int i = 0; i < N; i++)
    {
        // √(x^2 + y^2 + z^2)
        length[i] = std::sqrt(x[i]*x[i] + y[i]*y[i]);
    }

    // 結果を表示
    /*
    * 1.00, 2.24, 3.61, 5.00, 6.40ぐらい
    */
    std::cout << "2次元ベクトルCPU" << std::endl;
    for(int i = 0; i < N; i++)
    {
        std::cout << i << ": " << length[i] << std::endl;
    }
}

ってな処理です。

2. じゃあ

これをthrustによるGPGPU by CUDAで並列化するには

Length2Thrust.cpp

#include<iostream>

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

// 2次元ベクトルの大きさを取得する
struct GetLength2 : public thrust::binary_function<const double, const double, double>
{
    __host__ __device__
    double operator()(const double& x, const double& y) const
    {
        return std::sqrt(x*x + y*y);
    }
} getLength2;

// 2次元ベクトルThrust
void Length2Thrust()
{
    // 要素数
    const int N = 5;

    // x, y方向成分
    double x[N] = {0, 1, 2, 3, 4};
    double y[N] = {1, 2, 3, 4, 5};

    // 各ベクトルの大きさ
    double length[N];

    // デバイスの配列を生成
    thrust::device_vector<double> xVector(x, x + N);
    thrust::device_vector<double> yVector(y, y + N);
    thrust::device_vector<double> lengthVector(N);

    // 大きさを計算
    thrust::transform(xVector.begin(), xVector.begin() + N,
        yVector.begin(),
        lengthVector.begin(),
        getLength2);
    thrust::copy_n(lengthVector.begin(), N, length);

    // 結果を表示
    /*
    * 1.00, 2.24, 3.61, 5.00, 6.40ぐらい
    */
    std::cout << "2次元ベクトルThrust" << std::endl;
    for(int i = 0; i < N; i++)
    {
        std::cout << i << ": " << length[i] << std::endl;
    }
}

と、thrust::binary_function構造体を継承した関数オブジェクトを使って演算子を定義したのち、thrust::transform関数を使えばできます。

3. これもしたい

ただ、このthrust::transform関数なんですが、2変数以外にも1変数を引数に取ることができるんですが、3変数以上がないんですよね。
となると、ベクトルが3次元になった場合、

Length3

#include<iostream>

// 3次元ベクトルCPU
void Length3()
{
    // 要素数
    const int N = 5;

    // x, y, z方向成分
    double x[N] = {0, 1, 2, 3, 4};
    double y[N] = {1, 2, 3, 4, 5};
    double z[N] = {2, 3, 4, 5, 6};

    // 各ベクトルの大きさ
    double length[N];

    // 大きさを計算
    for(int i = 0; i < N; i++)
    {
        // √(x^2 + y^2 + z^2)
        length[i] = std::sqrt(x[i]*x[i] + y[i]*y[i] + z[i]*z[i]);
    }

    // 結果を表示
    /*
    * 2.23, 3.74, 5.39, 7.07, 7.77ぐらい
    */
    std::cout << "3次元ベクトルCPU" << std::endl;
    for(int i = 0; i < N; i++)
    {
        std::cout << i << ": " << length[i] << std::endl;
    }
}

に使えません。

4. じゃあどうするの

ということで、これが本題なんですが、3変数以上を使う場合は、タプルを使って複数の引数をまとめる必要があります。

Length3Thrust.cpp

#include<iostream>

#include<thrust/device_vector.h>
#include<thrust/tuple.h>
#include<thrust/transform.h>
#include<thrust/iterator/zip_iterator.h>

// 3次元ベクトル（のタプル）
typedef thrust::tuple<double, double, double> Double3;

// 2次元ベクトルの大きさを取得する
struct GetLength3 : public thrust::unary_function<const Double3, double>
{
    __host__ __device__
    double operator()(const Double3& v) const
    {
        double x = v.get<0>();
        double y = v.get<1>();
        double z = v.get<2>();

        return std::sqrt(x*x + y*y + z*z);
    }
} getLength3;

// 3次元ベクトルThrust
void Length3Thrust()
{
    // 要素数
    const int N = 5;

    // x, y, z方向成分
    double x[N] = {0, 1, 2, 3, 4};
    double y[N] = {1, 2, 3, 4, 5};
    double z[N] = {2, 3, 4, 5, 6};

    // 各ベクトルの大きさ
    double length[N];

    // デバイスの配列を生成
    thrust::device_vector<double> xVector(x, x + N);
    thrust::device_vector<double> yVector(y, y + N);
    thrust::device_vector<double> zVector(z, z + N);
    thrust::device_vector<double> lengthVector(N);

    // タプルを作って、そのタプルのイテレーターを作成
    auto double3Tuple = thrust::make_tuple(xVector.begin(), yVector.begin(), zVector.begin());
    auto double3Iterator = thrust::make_zip_iterator(double3Tuple);

    // 大きさを計算
    thrust::transform(double3Iterator, double3Iterator + N,
        lengthVector.begin(),
        getLength3);
    thrust::copy_n(lengthVector.begin(), N, length);

    // 結果を表示
    /*
    * 2.23, 3.74, 5.39, 7.07, 7.77ぐらい
    */
    std::cout << "3次元ベクトルCPU" << std::endl;
    for(int i = 0; i < N; i++)
    {
        std::cout << i << ": " << length[i] << std::endl;
    }
}

thrustにおいてタプルはthrust::tupleクラスで表現できます。後々面倒なので冒頭（9行目）で型定義しておきます。テンプレートの型は「第1引数の型、第2引数の型、第3引数の型・・・」です（10コまで用意されているっぽい）。
実際に引数を取り出すには、thrust::get関数を（テンプレート引数に引数番号を指定して）呼び出す必要がありますが、関数にしなくてもメソッド形式でも呼び出せるので、素直にdouble first = tuple.get<0>();とかしましょう。

このタプルを作るには、コンストラクタを呼び出すより、thrust::make_tupleを呼び出したほうがいいです。
タプルを作ったら、タプルのイテレーター（厳密には違うんですが）をthrust::make_zip_iterator関数を使ってthrust::zip_iteratorとして作ります。

あとはそのzip_iteratorをtransformの入力引数として入れてやれば完成です。

5. 応用

複数の値を返すこともできます。あんまり使い所ない気もしますが。

6. まとめ

という感じで、thrust便利すぎてOpenCLに戻れない私がいます。

アドベントカレンダーの方まだ人数全然埋まってないところを見ると、あと2-3回書かなきゃいけなそうなので、次回もあったらまたthrust書きます（と言いつつ違うこと書きそう。

明日、2012年12月7日の分は@tmasdaさんが書いてくれます！

A. レポジトリ

全部使える版をgithubのレポジトリにも置いておきますね。

B. 余談

この情報、日本語でなくて英語で調べようとしたんですが、2つがbinary、1つがunaryは知ってたんですが、3つとか知らなくて迷いました。trinaryというらしいです。4つはquanaryらしいです。この情報の方が役に立ちそうです！