CUDA Thrust : reduce_by_key on only some values in an array, 基于 "key"数组中的值

Question

假设我有两个 device_vector 数组， d_keys 和 d_data。

如果 d_data 例如是一个扁平的 2D 3x5 数组(例如 { 1, 2, 3, 4, 5, 6, 7, 8, 9, 8, 7, 6, 5 , 4, 3 } )和 d_keys 是大小为 5 的一维数组(例如 { 1, 0, 0, 1, 1 } )，我怎样才能做一个减少这样我最终如果相应的 d_keys 值为 1(例如，结果为 { 10, 23, 14 } )，则仅在每行基础上添加值？

sum_rows.cu示例允许我在 d_data 中添加每个值，但这并不完全正确。

或者，我可以逐行使用 zip_iterator 并一次将 d_keys 与一行 d_data 组合，并执行 transform_reduce，仅当键值为 1 时才添加，但随后我必须遍历 d_data 数组。

我真正需要的是某种不是内置的 transform_reduce_by_key 功能，但肯定有办法实现它!

Answer 1

基于额外的注释，而不是 3 行，有数千行，我们可以编写一个转换仿函数，对整行求和。基于有数千行的事实，这应该会让机器非常忙碌:

#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/transform.h>
#include <thrust/sequence.h>
#include <thrust/fill.h>

#define ROW   20
#define COL   10

__device__ int *vals;
__device__ int *keys;

struct test_functor
{
  const int a;

  test_functor(int _a) : a(_a) {}

  __device__
  int operator()(int& x, int& y ) {
    int temp = 0;
    for (int i = 0; i<a; i++)
      temp += vals[i + (y*a)] * keys[i];
    return temp;
    }
};

int main(){
  int *s_vals, *s_keys;
  thrust::host_vector<int> h_vals(ROW*COL);
  thrust::host_vector<int> h_keys(COL);
  thrust::sequence(h_vals.begin(), h_vals.end());
  thrust::fill(h_keys.begin(), h_keys.end(), 1);
  h_keys[0] = 0;
  thrust::device_vector<int> d_vals = h_vals;
  thrust::device_vector<int> d_keys = h_keys;
  thrust::device_vector<int> d_sums(ROW);
  thrust::fill(d_sums.begin(), d_sums.end(), 0);
  s_vals = thrust::raw_pointer_cast(&d_vals[0]);
  s_keys = thrust::raw_pointer_cast(&d_keys[0]);
  cudaMemcpyToSymbol(vals, &s_vals, sizeof(int *));
  cudaMemcpyToSymbol(keys, &s_keys, sizeof(int *));
  thrust::device_vector<int> d_idx(ROW);
  thrust::sequence(d_idx.begin(), d_idx.end());
  thrust::transform(d_sums.begin(), d_sums.end(), d_idx.begin(),  d_sums.begin(), test_functor(COL));
  thrust::host_vector<int> h_sums = d_sums;
  std::cout << "Results :" << std::endl;
  for (unsigned i = 0; i<ROW; i++)
    std::cout<<"h_sums["<<i<<"] = " << h_sums[i] << std::endl;
  return 0;
}

这种方法的缺点是通常不会合并对 vals 数组的访问。然而，对于几千行，缓存可能会提供显着的缓解。我们可以通过重新排序以列优先格式存储在展平数组中的数据来解决此问题，并将仿函数中循环中的索引方法更改为如下所示:

for (int i=0; i<a; i++)
  temp += vals[(i*ROW)+y]*keys[i];

如果愿意，您可以将 ROW 作为附加参数传递给仿函数。