CUDA で関数オブジェクト

CUDA (nvcc) で関数オブジェクトが使えるか試してみた。
結論としては使えた。

operator() の前に __device__ を付ければOK。

数値配列を入力として、ある値と等しいかどうかの 0/1 の配列を出力するプログラム。

#include <stdint.h>
#include <cutil.h>
#include <stdio.h>

// 関数オブジェクト、しかもテンプレート
template<typename T>
class EQ {
public:
    EQ(const T v) :
        value(v) {
    }
    __device__ inline bool operator()(const T comp) const {
        return comp == value;
    }
    T value;
};

// GPUで動くkernel関数
template<typename Operator>
__global__ void func(Operator op, uint32_t* gpuMem, int size, uint8_t* bitmap) {
    const unsigned int threadId = threadIdx.x;

    if(threadId < size) {
        bitmap[threadId] = op(gpuMem[threadId]);
    }

};

int main() {
    int totalSize = 32;
    uint32_t* cpuMem;
    uint32_t* gpuMem;
    uint32_t value = 10;
    uint8_t* cpuBmp;
    uint8_t* gpuBmp;

    // CPU 側の領域確保 & 初期化
    cudaHostAlloc((void**) &cpuMem, sizeof(value) * totalSize, 0);
    cudaHostAlloc((void**) &cpuBmp, sizeof(uint8_t) * totalSize, 0);
    memset(cpuBmp, 0, sizeof(uint8_t) * totalSize);
    for (int i = 0; i < totalSize; ++i) {
        cpuMem[i] = i;
    }

    // GPU 側の領域確保 & 初期化
    cudaMalloc((void**) &gpuMem, sizeof(value) * totalSize);
    cudaMemcpy((void *) gpuMem, (void *) cpuMem, sizeof(value) * totalSize, cudaMemcpyHostToDevice);
    cudaMalloc((void**) &gpuBmp, sizeof(uint8_t) * totalSize);
    cudaMemset((void**) &gpuBmp, 0, sizeof(uint8_t) * totalSize);

    // カーネル実行
    dim3 grid, threads;
    threads.x = 32;
    EQ<uint32_t> op(value);
    func< EQ<uint32_t> ><<<grid, threads>>>(op, gpuMem, totalSize, gpuBmp);

    // 値の表示
    cudaMemcpy((void *) cpuBmp, (void *) gpuBmp, sizeof(uint8_t) * totalSize, cudaMemcpyDeviceToHost);
    int j = 0;
    for (int i = 0; i < totalSize; ++i) {
        printf("%d", cpuBmp[i]);
        ++j;
        if (j == 8) {
            printf(" ");
            j = 0;
        }
    }
    printf("\n");

    // CPU側の領域解放
    cudaFreeHost(cpuMem);
    cudaFreeHost(gpuMem);

    // GPU側の領域解法
    cudaFree(gpuMem);
    cudaFree(gpuBmp);

    return 0;
}

んで、実行結果。
予定通り10番目だけ1になった。

00000000 00100000 00000000 00000000 

めでたし、めでたし。