为Thrust重载“+”运算符,有什么想法吗?

时间:2013-06-03 21:38:01

标签: c++ cuda operator-keyword nvidia nvcc

我正在与CUDA和Thrust合作。我发现键入thrust::transform [plus/minus/divide]单调乏味,所以我只想重载一些简单的运算符。

如果我能做的话会很棒:

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

以下是+的示例摘要:

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;
}

但是,thrust::plus<?>未正确重载,或者我没有正确地执行此操作...其中一个或另一个。 (如果为此重载简单的操作符是个坏主意,请解释原因)。最初,我认为我可以使用?这样的typename T::iterator占位符重载,但这不起作用。

我不确定如何使用向量的类型和向量迭代器的类型来重载+运算符。这有意义吗?

感谢您的帮助!

1 个答案:

答案 0 :(得分:2)

这似乎有效,其他人可能有更好的想法:

#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;
}

请注意,在任何一种情况下,我都可以为结果指定主机或设备向量,因为推力会看到差异并自动生成必要的复制操作。因此,模板中的结果向量类型(主机,设备)并不重要。

另请注意,模板定义中的thrust::transform函数参数不太正确。