CUDA Thrustライブラリを用いて、所定の値となる要素の数をカウントする方法

本記事では、CUDAのThrustライブラリを用いて、GPU側にある配列から、特定の値となっている配列の要素をカウントする方法を紹介します。

CUDA Thrustについて

ThrustはCUDA C++のテンプレートライブラリで、CUDAをインストールすると自動的に付いてくるライブラリです。

これはC++のSTL（Standard Template Library）と呼ばれるC++の開発でよく使われる標準ライブラリと似た文法で書けることが特徴で、通常のCUDAのプログラミングと比較すると、CPUからGPU側へのデータの転送などを簡易な文法で実現することができます。

よってThrustを使うことで、C++に慣れている人にとっては比較的少ない労力でCUDAを用いたGPUプログラミングをすることができます。

一方、実装されている処理には限りがあり、何でもできるわけではありません。Thrustに実装されている機能の例としては、配列の総和を求めたり、配列をソートしたりする機能が挙げられます。

CUDA Thrustを用いて特定の値を示す要素をカウントするプログラム

以下に、GPU側で特定の値を示す配列の要素をカウントするプログラムを示します。

#include <iostream>
#include <cuda_runtime.h>
#include <thrust/device_vector.h>
#include <thrust/count.h>

int main()
{
    // Creating vector with random numbers
    const int N = 100;
    int* data_host;
    data_host = (int*)malloc(sizeof(int) * N);
    for (int i = 0; i < N; i++) {
        data_host[i] = -1;
    }

    // Data transfer (CPU -> GPU) 
    int* data_dev;
    cudaMalloc(&data_dev, sizeof(int) * N);
    cudaMemcpy(data_dev, data_host, sizeof(int) * N, cudaMemcpyHostToDevice);

    // Calculate the sum in GPU by thrust
    thrust::device_vector<int> d_vec(data_dev, data_dev + N);
    int count = thrust::count(d_vec.begin(), d_vec.end(), -1);

    // Confirmation
    std::cout << "count = " << count << std::endl;

    return 0;
}

機能を使うためには、「thrust/count.h」をインクルードする必要があります。

8～14行目でint型の100個の要素からなる配列を用意し、全ての要素に「-1」を代入します。その後、16～19行目でデータをCPU側からGPU側に転送します。

今回のポイントとなる総和を求める部分は、23行目です。

23行目「int count = thrust::count(d_vec.begin(), d_vec.end(), -1);」は、d_vecが指しているGPU側の配列の中から、第3引数に指定した「-1」の数を数え、要素の数を返します。

今回、全ての配列要素を全て「-1」にしているので、「-1」は100個見つかります。

実行結果は以下の通りです。正常にカウントができていることがわかります。