如何使用推力根据索引累加数组?

How to use thrust to accumulate array based on index?

我正在尝试根据索引累积数组。我的输入是两个长度相同的向量。第一个向量是索引。第二个向量是值。我的目标是根据索引积累价值。我在 C++ 中有类似的代码。但我是推力编码的新手。我可以用推力设备代码来实现吗?我可以使用哪个功能?我没有发现像功能一样的“地图”。它比 CPU(host) 代码更有效吗? 我的c++版本迷你示例代码。

int a[10]={1,2,3,4,5,1,1,3,4,4};
vector<int> key(a,a+10);
double b[10]={1,2,3,4,5,1,2,3,4,5};
vector<double> val(b,b+10);

unordered_map<size_t,double> M;
for (size_t i = 0;i< 10 ;i++)
{
    M[key[i]] = M[key[i]]+val[i];
}

如评论中所述,执行此操作的规范方法是重新排序数据(键、值),以便将相似的键组合在一起。您可以使用 sort_by_key 执行此操作。 reduce_by_key然后求解。

可以使用提供给 for_each 的具有原子性的仿函数,以一种稍微不像推力的方式解决问题而无需重新排序。

以下说明两者:

$ cat t27.cu
#include <thrust/reduce.h>
#include <thrust/sort.h>
#include <thrust/device_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/for_each.h>
#include <thrust/copy.h>
#include <iostream>
#include <unordered_map>
#include <vector>

// this functor only needed for the non-reordering case
// requires compilation for a cc6.0 or higher GPU e.g. -arch=sm_60
struct my_func {
  double *r;
  my_func(double *_r) : r(_r) {};
  template <typename T>
  __host__ __device__
  void operator()(T t) {
    atomicAdd(r+thrust::get<0>(t)-1, thrust::get<1>(t));  // assumes consecutive keys starting at 1
  }
};

int main(){

  int a[10]={1,2,3,4,5,1,1,3,4,4};
  std::vector<int> key(a,a+10);
  double b[10]={1,2,3,4,5,1,2,3,4,5};
  std::vector<double> val(b,b+10);

  std::unordered_map<size_t,double> M;
  for (size_t i = 0;i< 10 ;i++)
  {
    M[key[i]] = M[key[i]]+val[i];
  }
  for (int i = 1; i < 6; i++) std::cout << M[i] << " ";
  std::cout << std::endl;
  int size_a = sizeof(a)/sizeof(a[0]);
  thrust::device_vector<int>    d_a(a, a+size_a);
  thrust::device_vector<double> d_b(b, b+size_a);
  thrust::device_vector<double> d_r(5); //assumes only 5 keys, for illustration
  thrust::device_vector<int> d_k(5); // assumes only 5 keys, for illustration
  // method 1, without reordering
  thrust::for_each_n(thrust::make_zip_iterator(thrust::make_tuple(d_a.begin(), d_b.begin())), size_a, my_func(thrust::raw_pointer_cast(d_r.data())));
  thrust::host_vector<double> r = d_r;
  thrust::copy(r.begin(), r.end(), std::ostream_iterator<double>(std::cout, " "));
  std::cout << std::endl;
  thrust::fill(d_r.begin(), d_r.end(), 0.0);
  // method 2, with reordering
  thrust::sort_by_key(d_a.begin(), d_a.end(), d_b.begin());
  thrust::reduce_by_key(d_a.begin(), d_a.end(), d_b.begin(), d_k.begin(), d_r.begin());
  thrust::copy(d_r.begin(), d_r.end(), r.begin());
  thrust::copy(r.begin(), r.end(), std::ostream_iterator<double>(std::cout, " "));
  std::cout << std::endl;
}
$ nvcc -o t27 t27.cu -std=c++14 -arch=sm_70
$ ./t27
4 2 6 13 5
4 2 6 13 5
4 2 6 13 5
$

我对这些方法的相对性能不做任何陈述。这可能取决于实际数据集的大小,可能还取决于所使用的 GPU 和其他因素。