计算推力:: device_vector上的梯度

Question

我目前正在使用thrust::device_vector时将代码从本地C ++传输到CUDA。 现在有一个计算梯度的函数，我不仅需要访问当前元素，还需要访问周围的元素。 在原始代码中，我编写了以下代码：

void calc_grad(Complex *grad, const Complex *data, const int size)
{
    for (int i = 1; i < size - 1; i++) {
        grad[i] = 0.5 * (data[i + 1] - data[i - 1]);
    }
    grad[0] = data[1] - data[0];
    grad[size - 1] = data[size - 1] - data[size - 2];
}

是否可以创建用于推力的functor ，以便可以在thrust::transform()调用它？ 直到现在，我一次只能访问一个元素，而没有获取周围的元素。 还是由于更改之前和之后的元素而无法实现？

该代码的公式来自matlab函数gradient ，如下所示： http : //se.mathworks.com/help/matlab/ref/gradient.html?requestedDomain=www.mathworks.com

Answer 1

一个Thrust::transform()可以完成大部分工作。 您需要做的就是稍微移动数据，使grad[i] ， data[i-1]和data[i+1]对齐。

thrust::transform(data.begin()+2,
                  data.end(),
                  data.begin(),
                  grad.begin()+1,
                  (_1 - _2) * 0.5);

然后，您可以处理边界情况。

编辑

您还可以将边界条件包含在一个转换中。 通过以下形式的转换，您的函子Grad应该能够通过第一个函子参数知道他正在处理的数据的索引。 然后，根据索引，他可以从第二个函子参数（即一个元组）的3个元素中选择2个，以进行正确的计算。

这里的所有内容都未经测试。 我不确定data.begin()-1有效。 您可能还需要注意Complex类型。

thrust::transform(
  thrust::make_counting_iterator(int(0)),
  thrust::make_counting_iterator(int(0)) + size,
  thrust::make_zip_iterator(
      thrust::make_tuple(
          data.begin()-1,
          data.begin(),
          data.begin()+1)),
  grad.begin(),
  Grad(size)
);

函子是这样的。

struct Grad {
  int size_;
  Grad(int s) :
      size_(s) {
  }
  template<typename T, typename Tuple>
  __host__ __device__
  inline T operator()(const int& idx, const Tuple& d) {
    if (idx == 0) {
      return d.get<2>() - d.get<1>();
    } else if (idx == size_ - 1) {
      return d.get<1>() - d.get<0>();
    } else {
      return 0.5 * (d.get<2>() - d.get<0>());
    }
  }
};