Используйте толчок redu_by_key на структуру - PullRequest
Новая
       60

Используйте толчок redu_by_key на структуру

0 голосов
/ 29 апреля 2020

Я довольно новичок в CUDA и пытаюсь применить операцию reduce_by_key к структуре. 

struct index_and_loc {
  int index;
  int3 location;
}

То, что я хотел бы сделать, у меня есть вектор index_and_loc, где элементы будут одинаковыми index

thrust::host_vector<index_and_loc> my_vector;
my_vector= { {0, {0, 2, 5}},
             {0, {1, 3, 4}},
             {0, {0, 1, 3}},
             {1, {2, 1, 0}},
             {1, {2, 2, 2}}
           }

И я хочу иметь два выхода один будет суммой всех элементов, имеющих одинаковый индекс, а другой вектор будет хранить сколько экземпляров каждого индекса. Таким образом, мои итоговые результаты будут такими: 

// Sum all elements with index 0 and 1
sum_vector = {{0, {1, 6, 12}},
              {1, {4, 3, 2}}
             }

instances_vector = {3, // Number of elements with index = 0
                    2} // Number of elements with index = 1

Читая методы в документации по тяге, я считаю, что я должен использовать reduce_by_key. Я могу объявить мои input1 be my_vector и input2 вектором 1 с той же длиной, что и input1, и я могу уменьшить использование индекса моей структуры и использовать thrust::plus<int> в качестве моей BinaryFunction для вектор 1-х годов. Однако это не сможет применить сумму элементов int3 в моем векторе input1, так как BinaryFunction применяется к input2.

Есть ли способ сделать sh это? Пожалуйста, дайте мне знать, если мой вопрос не очень ясен.

РЕДАКТИРОВАТЬ:

Я переработал проблему и снизил ее до более простого подхода. Вместо этого я добавил столбец экземпляра в my_vector и установил их в 1. Теперь, используя ответ talonmies, я могу получить оба ответа, которые я искал. Вот мой код 

#include <thrust/extrema.h>
#include <thrust/reduce.h>
#include <thrust/execution_policy.h>
#include <thrust/functional.h>

struct index_and_loc {
  int index;
  int3 location;
  int instance;

  index_and_loc() = default;

  __host__ __device__
  index_and_loc(int index_, int3 location_, int instance_) :
  index(index_), instance(instance_) {
    location.x = location_.x; 
    location.y = location_.y;
    location.z = location_.z;
  };

  __host__ __device__
  index_and_loc& operator=(const index_and_loc& y) {
    index = y.index;
    location.x = y.location.x;
    location.y = y.location.y;
    location.z = y.location.z;
    instance = y.instance;
    return *this;
  };

  __host__ __device__
  bool operator==(const index_and_loc& y) const {
    return index == y.index;
  };

  __host__ __device__
  index_and_loc operator+(const index_and_loc& y) const {
    return index_and_loc(index,
      make_int3(location.x + y.location.x, 
                location.y + y.location.y, 
                location.z + y.location.z),
      instance + y.instance);
  };

};

int main() {

  int num_points = 5;
  thrust::device_vector<index_and_loc> my_vector(num_points);
  my_vector.push_back(index_and_loc(0, {0, 2, 5}, 1));
  my_vector.push_back(index_and_loc(0, {1, 3, 4}, 1));
  my_vector.push_back(index_and_loc(0, {0, 1, 3}, 1));
  my_vector.push_back(index_and_loc(1, {2, 1, 0}, 1));
  my_vector.push_back(index_and_loc(1, {2, 2, 2}, 1));

  thrust::device_vector<index_and_loc> same_idx(num_points);
  thrust::device_vector<index_and_loc> sum_locs(num_points);

  thrust::equal_to<index_and_loc> pred;
  thrust::plus<index_and_loc> op;

  auto res = thrust::reduce_by_key(
    thrust::device,
    my_vector.begin(),
    my_vector.end(),
    my_vector.begin(),
    same_idx.begin(),
    sum_locs.begin(),
    pred,
    op);

  for(int i=0; i<2; i++) {
    index_and_loc y = same_idx[i];
    index_and_loc x = sum_locs[i];
    std::cout << y.index << " {" << x.location.x 
                         << " " << x.location.y 
                         << " " << x.location.z 
                         << "}, " << x.instance << std::endl;
  }

  return 0;
}

Выполнение этого на самом деле даст мне 

0 {1 6 12}, 3
1 {4 3 2}, 2

1 Ответ

1 голос
/ 29 апреля 2020

Существует множество способов сделать то, что вы хотите. Я не чувствую, что будет и не будет эффективным, решать вам.

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

#include <thrust/functional.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <iostream>

struct index_and_loc {
  int index;
  int3 location;

  index_and_loc() = default;

  __host__ __device__
  index_and_loc(int index_, int3 location_) {
      index = index_;
      location.x = location_.x;
      location.y = location_.y;
      location.z = location_.z;
  };

  __host__ __device__
  index_and_loc& operator=(const index_and_loc& y) {
      index = y.index;
      location.x = y.location.x;
      location.y = y.location.y;
      location.z = y.location.z;
      return *this;
  };

  __host__ __device__
  bool operator==(const index_and_loc& y) const {
      return index == y.index;
  };

  __host__ __device__
  index_and_loc operator+(const index_and_loc& y) const {
      return index_and_loc(index, make_int3(location.x + y.location.x, 
                                            location.y + y.location.y, 
                                            location.z + y.location.z));
  };
};

int main()
{
    thrust::host_vector<index_and_loc> my_vector(5); 
    my_vector[0] =  {0, {0, 2, 5}};
    my_vector[1] =  {0, {1, 3, 4}};
    my_vector[2] =  {0, {0, 1, 3}};
    my_vector[3] =  {1, {2, 1, 0}};
    my_vector[4] =  {1, {2, 2, 2}};

    thrust::device_vector<index_and_loc> d_vector = my_vector; 
    thrust::device_vector<index_and_loc> keys_out(5);
    thrust::device_vector<index_and_loc> data_out(5);

    thrust::equal_to<index_and_loc> pred;
    thrust::plus<index_and_loc> op;

    auto res = thrust::reduce_by_key(
            d_vector.begin(),
            d_vector.end(),
            d_vector.begin(),
            keys_out.begin(),
            data_out.begin(),
            pred,
            op);

    for(int i=0; i<2; i++) {
        index_and_loc y = keys_out[i];
        index_and_loc x = data_out[i];
        std::cout << y.index << " {" << x.location.x 
                             << " " << x.location.y 
                             << " " << x.location.z 
                             << "}" << std::endl;
    }

    return 0;
}

Они определяют правильное определение operator== для предиката и operator+ для сокращения. Остальное просто необходимо для назначения и построения копии.

Похоже, это делает то, что вы хотите: 

$ nvcc -arch=sm_52 -std=c++11 -o improbable improbable.cu
$ ./improbable 
0 {1 6 12}
1 {4 3 2}

Обратите внимание, что поскольку Thrust ожидает, что ключ и данные являются отдельными итераторами, вы должны иметь отдельные выходные итераторы для сокращенных ключей и данных. Это означает, что вы фактически получите две копии решения. Решите, важно ли это.

...