如何使用推力根据索引累积数组？

问题描述

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

int a[10]={1,2,3,4,5,1,4};
vector<int> key(a,a+10);
double b[10]={1,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,4};
  std::vector<int> key(a,a+10);
  double b[10]={1,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_b.begin(),d_k.begin(),d_r.begin());
  thrust::copy(d_r.begin(),r.begin());
  thrust::copy(r.begin()," "));
  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 和其他因素。