Cuda Thrust - 如何使用sort_by_key、merge_by_key和reduce_by_key优化代码

这是一种可能的方法，我认为它可能比你的序列更快。 关键思想是，我们希望避免在提前知道顺序的情况下对数据进行排序。 如果我们可以利用我们拥有的订单知识，而不是对数据进行排序，我们可以简单地将其重新排序为所需的排列。

让我们对数据进行一些观察。 如果你的key1和key2实际上是上三角矩阵的 i，j 索引，那么我们可以对这两个向量的串联进行一些观察：

1. 串联的向量将包含相等数量的每个键。(我相信你可能在你的问题中已经指出了这一点。 因此，在您的情况下，矢量将包含三个0键、三个1键、三个2键和三个3键。 我相信这种模式应该适用于任何独立于矩阵维度的上三角形图案。 因此，上三角形的 N 维矩阵在串联索引向量中将有 N 组键，每组由 N-1 个类似元素组成。

2. 在级联向量中，我们可以发现/建立一致的键顺序(基于矩阵维 N(，这允许我们以相似键分组的顺序对向量重新排序，而无需诉诸传统的排序操作。

如果我们结合上述 2 个想法，那么我们可能可以通过一些分散操作来解决整个问题，以替换排序/合并活动，然后是thrust::reduce_by_key操作。 散点运算可以通过thrust::copy到适当的thrust::permutation_iterator与适当的索引计算函子相结合来完成。 由于我们确切地知道重新排序的串联key向量会是什么样子(在您的维度 4 示例中：{0,0,0,1,1,1,2,2,2,3,3,3}(，我们不需要对其显式执行重新排序。 但是，我们必须使用相同的映射对value向量重新排序。 因此，让我们开发该映射的算术：

dimension (N=)4 example
vector index:          0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11
desired (group) order: 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3
concatenated keys:     0, 0, 0, 1, 1, 2, 1, 2, 3, 2, 3, 3
group start idx:       0, 0, 0, 3, 3, 6, 3, 6, 9, 6, 9, 9
group offset idx:      0, 1, 2, 0, 1, 0, 2, 1, 0, 2, 1, 2
destination idx:       0, 1, 2, 3, 4, 6, 5, 7, 9, 8,10,11

dimension (N=)5 example
vector index:          0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,15,16,17,18,19
desired (group) order: 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4
concatenated keys:     0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 1, 2, 3, 4, 2, 3, 4, 3, 4, 4
group start idx:       0, 0, 0, 0, 4, 4, 4, 8, 8,12, 4, 8,12,16, 8,12,16,12,16,16
group offset idx:      0, 1, 2, 3, 0, 1, 2, 0, 1, 0, 3, 2, 1, 0, 3, 2, 1, 3, 2, 3 
destination idx:       0, 1, 2, 3, 4, 5, 6,10, 7, 8,11,14, 9,12,15,17,13,16,18,19

我们可以观察到，在每种情况下，目标索引(即将所选键或值移动到的位置，对于所需的组顺序(等于组开始索引加上组偏移索引。 组启动索引只是键乘以 (N-1(。 组偏移索引是一种类似于上三角形或下三角形索引模式的模式(对于级联向量的每一半，有 2 个不同的化身(。 串联键只是key1和key2向量的串联(我们将使用 permutation_iterator 虚拟创建此连接(。 所需的群顺序是先验已知的，它只是一个整数群序列，其中 N 个组由 N-1 个元素组成。 它等效于级联键向量的排序版本。 因此，我们可以直接计算函子中的目标索引。

为了创建组偏移指数模式，我们可以减去两个关键向量(并减去额外的 1(：

key2:                  1, 2, 3, 2, 3, 3
key1:                  0, 0, 0, 1, 1, 2
key1+1:                1, 1, 1, 2, 2, 3
p1 = key2-(key1+1):    0, 1, 2, 0, 1, 0
p2 = (N-2)-p1:         2, 1, 0, 2, 1, 2
grp offset idx=p1|p2:  0, 1, 2, 0, 1, 0, 2, 1, 0, 2, 1, 2

下面是一个完整的示例，使用您的示例数据演示了上述概念：

$ cat t1133.cu
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <thrust/copy.h>
#include <thrust/transform.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <iostream>
// "triangular sort" index generator
struct idx_functor
{
  int n;
  idx_functor(int _n): n(_n) {};
  template <typename T>
  __host__ __device__
  int operator()(const T &t){
    int k1 = thrust::get<0>(t);
    int k2 = thrust::get<1>(t);
    int id = thrust::get<2>(t);
    int go,k;
    if (id < (n*(n-1))/2){ // first half
      go = k2-k1-1;
      k  = k1;
      }
    else { // second half
      go = n-k2+k1-1;
      k  = k2;
      }
    return k*(n-1)+go;
  }
};

const int N = 4;
using namespace thrust::placeholders;
int main(){
   // useful dimensions
   int d1 = N*(N-1);
   int d2 = d1/2;
   // iniitialize keys
   int key_1[] = {0, 0, 0, 1, 1, 2};
   int key_2[] = {1, 2, 3, 2, 3, 3};
   thrust::device_vector<int> k1(key_1, key_1+d2);
   thrust::device_vector<int> k2(key_2, key_2+d2);
   // initialize values
   double value_1[] = {0.5, 0.5, 0.5, -1.0, -1.0, 2.0};
   double value_2[] = {-1, 2.0, -3.5, 2.0, -3.5, -3.5};
   thrust::device_vector<double> v(d1);
   thrust::device_vector<double> vg(d1);
   thrust::copy_n(value_1, d2, v.begin());
   thrust::copy_n(value_2, d2, v.begin()+d2);
   // reorder (group) values by key
   thrust::copy(v.begin(), v.end(), thrust::make_permutation_iterator(vg.begin(), thrust::make_transform_iterator(thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(k1.begin(), thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1%d2)), thrust::make_permutation_iterator(k2.begin(), thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1%d2)), thrust::counting_iterator<int>(0))), idx_functor(N))));
   // sum results
   thrust::device_vector<double> rv(N);
   thrust::device_vector<int> rk(N);
   thrust::reduce_by_key(thrust::make_transform_iterator(thrust::counting_iterator<int>(0), _1/(N-1)), thrust::make_transform_iterator(thrust::counting_iterator<int>(d1), _1/(N-1)), vg.begin(), rk.begin(), rv.begin());
   // print results
   std::cout << "Keys:" << std::endl;
   thrust::copy_n(rk.begin(), N, std::ostream_iterator<int>(std::cout, ", "));
   std::cout << std::endl <<  "Sums:" << std::endl;
   thrust::copy_n(rv.begin(), N, std::ostream_iterator<double>(std::cout, ", "));
   std::cout << std::endl;
   return 0;
}
$ nvcc -std=c++11 -o t1133 t1133.cu
$ ./t1133
Keys:
0, 1, 2, 3,
Sums:
1.5, -3, 6, -10.5,
$

净效果是，您的thrust::sort_by_key和thrust::merge_by_key操作已被单个thrust::copy操作所取代，这应该更有效率。