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

Question

我正在使用c++和cuda/thrust在GPU上执行计算，这对我来说是一个新领域。不幸的是，我的代码（下面的 MCVE）效率不高，所以我想知道如何优化它。该代码执行以下操作：

有两个键向量和两个值向量。关键向量基本上包含上三角矩阵的 i 和 j（在本例中：大小为 4x4）。

key1 {0, 0, 0, 1, 1, 2} value1: {0.5, 0.5, 0.5, -1.0, -1.0, 2.0}
key2 {1, 2, 3, 2, 3, 3} value2: {-1, 2.0, -3.5, 2.0, -3.5, -3.5}

任务是对具有相同键的所有值求和。为此，我使用 sort_by_key 对第二个值向量进行了排序。结果是：

key1 {0, 0, 0, 1, 1, 2} value1: {0.5, 0.5, 0.5, -1.0, -1.0, 2.0}
key2 {1, 2, 2, 3, 3, 3} value2: {-1.0, 2.0, 2.0, -3.5, -3.5, -3.5}

之后，我使用 merge_by_key 合并了两个值向量，这导致了一个新的键和值向量，其大小是以前的两倍。

key_merge {0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3}
value_merge {0.5, 0.5, 0.5, -1.0, -1.0, -1.0, 2.0, 2.0, 2.0, -3.5, -3.5, -3.5}

最后一步是使用 reduce_by_key 对具有相同键的所有值求和。结果是：

key {0, 1, 2, 3} value: {1.5, -3.0, 6.0, -10.5}

下面执行此操作的代码很慢，恐怕较大尺寸的性能会很差。如何优化？ sort_by_key、merge_by_key 和 reduce_by_key 是否可以融合？由于我事先知道sort_by_key的结果键向量，是否可以转换值向量"from an old to a new key"？在减少它们之前合并两个向量是否有意义，或者对每对 value/key 向量分别使用 reduce_by_key 是否更快？是否可以利用以下事实来加快 reduce_by_key 计算速度，即不同键值的数量已知且相等键的数量始终相同？

#include <stdio.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/reduce.h>
#include <thrust/merge.h>

int main(){
   int key_1[6] = {0, 0, 0, 1, 1, 2};
   int key_2[6] = {1, 2, 3, 2, 3, 3};
   thrust::device_vector<double> k1(key_1,key_1+6);
   thrust::device_vector<double> k2(key_2,key_2+6);

   double value_1[6] = {0.5, 0.5, 0.5, -1.0, -1.0, 2.0};
   double value_2[6] = {-1, 2.0, -3.5, 2.0, -3.5, -3.5};
   thrust::device_vector<double> v1(value_1,value_1+6);
   thrust::device_vector<double> v2(value_2,value_2+6);

   thrust::device_vector<double> mk(12);
   thrust::device_vector<double> mv(12);
   thrust::device_vector<double> rk(4);
   thrust::device_vector<double> rv(4);

   thrust::sort_by_key(k2.begin(), k2.end(), v2.begin());
   thrust::merge_by_key(k1.begin(), k1.end(), k2.begin(), k2.end(),v1.begin(), v2.begin(), mk.begin(), mv.begin());
   thrust::reduce_by_key(mk.begin(), mk.end(), mv.begin(), rk.begin(), rv.begin());

   for (unsigned i=0; i<4; i++) {
     double tmp1 = rk[i];
     double tmp2 = rv[i];
     printf("key value %f is related to %f\n", tmp1, tmp2);
   }
   return 0;
}

结果：

key value 0.000000 is related to 1.500000
key value 1.000000 is related to -3.000000
key value 2.000000 is related to 6.000000
key value 3.000000 is related to -10.500000

Answer 1

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

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

串联向量将包含相等数量的每个键。（我相信你可能已经在你的问题中指出了这一点。）所以在你的情况下，向量将包含三个 0 键、三个 1 键、三个 2 键和三个 3键。我相信这种模式应该适用于任何独立于矩阵维数的上三角模式。因此，维数为 N 的上三角矩阵在连接的索引向量中将有 N 组键，每组由 N-1 个类似的元素组成。
在串联向量中，我们可以 discover/establish 键的一致排序（基于矩阵维数 N），这允许我们以类似键分组的顺序重新排序向量，无需诉诸传统的排序操作。

如果我们把上面2个思路结合起来，那么我们大概可以用一些散点操作来代替sort/mergeactivity，然后再进行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 操作取代，这应该更有效。

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

Cuda Thrust - How to optimize a code using sort_by_key, merge_by_key and reduce_by_key

c++

performance

cuda

thrust