Overloading "+" operator for Thrust, any ideas?

Question

I'm working with CUDA and Thrust. I find typing the thrust::transform [plus/minus/divide] tedious so I just want to overload some simple operators.

It would be awesome if I could do :

thrust::[host/device]_vector<float> host;
thrust::[host/device]_vector<float> otherHost;
thrust::[host/device]_vector<float> result = host + otherHost;

Here is an example snippet for +:

template <typename T>
__host__ __device__ T& operator+(T &lhs, const T &rhs) {
    thrust::transform(rhs.begin(), rhs.end(),
                      lhs.begin(), lhs.end(), thrust::plus<?>());
    return lhs;
}

However, the thrust::plus<?> is not correctly overloaded, or I'm not doing it correctly... one or the other. (If overloading simple operators for this is a bad idea, please explain why). Initially, I thought that I could overload the ? placeholder with something like typename T::iterator, but that didn't work.

I'm not sure how to overload the + operator with both the type of the vector and the type of the vector iterator. Does this make sense?

Thanks for your help!

Answer 1

This seems to work, others may have better ideas:

#include <ostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/transform.h>
#include <thrust/functional.h>
#include <thrust/copy.h>
#include <thrust/fill.h>

#define DSIZE 10


template <typename T>
thrust::device_vector<T>  operator+(thrust::device_vector<T> &lhs, const thrust::device_vector<T> &rhs) {
    thrust::transform(rhs.begin(), rhs.end(),
                      lhs.begin(), lhs.begin(), thrust::plus<T>());
    return lhs;
}

template <typename T>
thrust::host_vector<T>  operator+(thrust::host_vector<T> &lhs, const thrust::host_vector<T> &rhs) {
    thrust::transform(rhs.begin(), rhs.end(),
                      lhs.begin(), lhs.begin(), thrust::plus<T>());
    return lhs;
}
int main() {


  thrust::device_vector<float> dvec(DSIZE);
  thrust::device_vector<float> otherdvec(DSIZE);
  thrust::fill(dvec.begin(), dvec.end(), 1.0f);
  thrust::fill(otherdvec.begin(), otherdvec.end(), 2.0f);
  thrust::host_vector<float> hresult1 = dvec + otherdvec;

  std::cout << "result 1: ";
  thrust::copy(hresult1.begin(), hresult1.end(), std::ostream_iterator<float>(std::cout, " "));  std::cout << std::endl;

  thrust::host_vector<float> hvec(DSIZE);
  thrust::fill(hvec.begin(), hvec.end(), 5.0f);
  thrust::host_vector<float> hresult2 = hvec + hresult1;


  std::cout << "result 2: ";
  thrust::copy(hresult2.begin(), hresult2.end(), std::ostream_iterator<float>(std::cout, " "));  std::cout << std::endl;

  // this line would produce a compile error:
  // thrust::host_vector<float> hresult3 = dvec + hvec;

  return 0;
}

Note that in either case, I can specify a host or device vector for the result, since thrust will see the difference and automatically generate the necessary copy operation. So the result vector type (host, device) in my templates is not critical.

Also note that the thrust::transform function parameters you had in your template definition weren't quite right.