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

Posted 2021-04-21

我想在推力仿函数中调用一个设备函数，但是甚至不知道如何启动。这是一个显而易见的需求，因为有些情况下仿函数的大小很大，因此需要将其拆分为函数。

最低限的例子是值得赞赏的

谢谢

答案

您将以与在普通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,
$