计算 thrust::device_vector 上的梯度

Question

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

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 for thrust out of that，这样我就可以在 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>());
    }
  }
};

计算 thrust::device_vector 上的梯度

Calculating the gradient over a thrust::device_vector

c++

cuda

thrust