Search code examples
cudagputhrust

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


I would like to call a device function inside a thrust functor but would not even know how to start. This is an obvious need, since there are cases that functor size is big and therefore, needs to be split to functions.

A minimum example is appreciated.

Thank you


Solution

  • You would do this in a fashion very similar to how it would be done in ordinary CUDA C++ code. Define your __device__ function (perhaps mark it __host__ __device__) exactly as you would in ordinary CUDA device code, and use that in a thrust functor operator, just as you would use it anywhere else in CUDA device code.

    Here is a trivial example that demonstrates calling a __host__ __device__ function from a functor operator, as well as calling a CUDA built-in math library function from the functor operator:

    $ 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,
    $