CUDA Thrust: redu_by_key только для некоторых значений в массиве, основываясь на значениях в "ключевом" массиве

Question

Допустим, у меня есть два массива device_vector, d_keys а также d_data,

Если d_data представляет собой, например, уплощенную двумерную матрицу 3х5 (например, { 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 значение равно единице (например, заканчивается результатом { 10, 23, 14 })?

Пример sum_rows.cu позволяет мне добавлять каждое значение в d_data, но это не совсем верно.

В качестве альтернативы, я могу, для каждого ряда, использовать zip_iterator и объединить d_keys с одним рядом d_data в то время, и сделать transform_reduceдобавление только в том случае, если значение ключа равно единице, но тогда мне придется перебирать 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 в качестве дополнительного параметра функтору.

Answer 2

Вот пример кода, который делает что-то похожее на то, что вам нужно, используя подход, который я изложил в своем комментарии под вашим вопросом. На самом деле мы хотим использовать 4 кортежа, чтобы подобрать значение ключа. Воспроизведение соответственно измененного комментария здесь:

Вы можете создать итератор zip, который объединит ваши 3 строки вместе с ключевым "строкой" и передаст 4-кортеж специальному функтору. Затем ваш специальный функтор уменьшит массив из 3-х кортежей (также используя ключ) и вернет результат, который является 4-кортежем. Пример продукта Thrust Dot может дать вам несколько идей.

Это один из возможных подходов:

#include <thrust/host_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/sequence.h>
#include <thrust/fill.h>
#include <thrust/tuple.h>

#define N 30  // make this evenly divisible by 3 for this example

typedef thrust::tuple<int, int, int, int>  tpl4int;
typedef thrust::host_vector<int>::iterator intiter;
typedef thrust::tuple<intiter, intiter, intiter, intiter>  tpl4intiter;
typedef thrust::zip_iterator<tpl4intiter>  int4zip;



struct r3key_unary_op : public thrust::unary_function<tpl4int, tpl4int>
{
  __host__ __device__
  tpl4int operator()(const tpl4int& x) const
  {
    tpl4int result;
    thrust::get<0>(result) = x.get<0>()*x.get<3>();
    thrust::get<1>(result) = x.get<1>()*x.get<3>();
    thrust::get<2>(result) = x.get<2>()*x.get<3>();
    thrust::get<3>(result) = 1;
    return result;
   }
};

struct r3key_binary_op : public thrust::binary_function<tpl4int, tpl4int, tpl4int>
{
  __host__ __device__
  tpl4int operator()(const tpl4int& x, const tpl4int& y) const
  {
    tpl4int result;
    thrust::get<0>(result) = x.get<0>()*x.get<3>() + y.get<0>()*y.get<3>();
    thrust::get<1>(result) = x.get<1>()*x.get<3>() + y.get<1>()*y.get<3>();
    thrust::get<2>(result) = x.get<2>()*x.get<3>() + y.get<2>()*y.get<3>();
    thrust::get<3>(result) = 1;
    return result;
  }
};


int main() {

  thrust::host_vector<int> A(N);  // values, in 3 "rows" flattened
  thrust::sequence(A.begin(), A.end());
  thrust::host_vector<int> K(N/3);   // keys in one row
  thrust::fill(K.begin(), K.end(), 1);  // set some keys to 1
  K[9] = 0;  // set some keys to zero

  int4zip first = thrust::make_zip_iterator(thrust::make_tuple(A.begin(), A.begin() + N/3, A.begin() + 2*N/3, K.begin()));
  int4zip  last = thrust::make_zip_iterator(thrust::make_tuple(A.begin() + N/3, A.begin() + 2*N/3, A.end(), K.end()));
  r3key_unary_op my_unary_op;
  r3key_binary_op my_binary_op;
  tpl4int init = my_unary_op(*first);
  // init = thrust::make_tuple((int) 0, (int) 0, (int) 0, (int) 0);
  tpl4int result = thrust::transform_reduce(first, last, my_unary_op, init, my_binary_op);
  std::cout << "row 0 = " << result.get<0>() << std::endl;
  std::cout << "row 1 = " << result.get<1>() << std::endl;
  std::cout << "row 2 = " << result.get<2>() << std::endl;
  return 0;

}

Заметки:

1. Это просто с помощью host_vector, Расширение его для работы с device_vector или шаблонизировать его для работы с чем-то, кроме int должно быть простым.
2. Для полноты я использую унарный функтор для предоставления значения инициализации, отличного от нуля, для уменьшения суммы каждой строки. Возможно, вы захотите изменить значение инициализации на ноль (4 кортежа нулей).