是否可以在 Thrust 仿函数中调用设备函数?

is it possible to call a device function inside a Thrust functor?

我想在推力仿函数中调用设备函数,但不知道如何开始。这是一个明显的需求,因为在某些情况下仿函数大小很大,因此需要拆分为函数。

一个最小的例子表示赞赏。

谢谢

您执行此操作的方式与在普通 CUDA C++ 代码中的方式非常相似。完全按照在普通 CUDA 设备代码中的方式定义 __device__ 函数(也许将其标记为 __host__ __device__),并在推力函子运算符中使用它,就像在其他任何地方使用它一样在 CUDA 设备代码中。

这是一个简单的示例,演示了从仿函数运算符调用 __host__ __device__ 函数,以及从仿函数运算符调用 CUDA 内置数学库函数:

$ cat t9.cu
#include <iostream>
#include <thrust/transform.h>
#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <vector>
#include <math.h>

template <typename T>
__host__ __device__ T square(T &s){
  return s*s;
}

struct my_math
{
  __host__ __device__
  float operator()(float &r){
   return log2f(square(r));
  }
};

int main(){

  std::vector<float> h_data = {1,2,4,8,16};
  thrust::device_vector<float> d_data = h_data;
  thrust::transform(d_data.begin(), d_data.end(), d_data.begin(), my_math());
  thrust::copy(d_data.begin(), d_data.end(), std::ostream_iterator<float>(std::cout, ","));
  std::cout << std::endl;
}


$ nvcc -arch=sm_35 -std=c++11 -o t9 t9.cu
$ ./t9
0,2,4,6,8,
$