Question

我正在使用thrust::device_vector s将代码从本地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>()); } } };

计算推力:: device_vector上的梯度

1 个答案: