Cuda Thrust自定義函數

Question

如何在Thrust中實現此功能？

for (i=0;i<n;i++) 
    if (i==pos) 
     h1[i]=1/h1[i]; 
    else 
     h1[i]=-h1[i]/value; 

在CUDA我也喜歡它：

__global__ void inverse_1(double* h1, double value, int pos, int N) 
{ 
    int i = blockDim.x * blockIdx.x + threadIdx.x; 
    if (i < N){ 
     if (i == pos) 
      h1[i] = 1/h1[i]; 
     else 
      h1[i] = -h1[i]/value; 
    } 
} 

謝謝！

Answer 1

您需要創建一個二元仿函數來應用操作，然後使用計數迭代器作爲第二個輸入。您可以將pos和value傳遞給仿函數的構造函數。它看起來像這樣：

struct inv1_functor 
{ 
    const int pos; 
    const double value; 

    inv1_functor(double _value, int _pos) : value(_value), pos(_pos) {} 

    __host__ __device__ 
    double operator()(const double &x, const int &i) const { 
    if (i == pos) 
     return 1.0/x; 
    else 
     return -x/value; 
    } 
}; 

//... 

thrust::transform(d_vec.begin(), d_vec.end(), thrust::counting_iterator<int>(), d_vec.begin(), inv1_functor(value, pos));