推力 CUDA 找到每组（段）的最大值

Posted 2023-02-16

【中文标题】推力 CUDA 找到每组（段）的最大值

【英文标题】：Thrust CUDA find maximum per each group(segment)

【发布时间】：2016-08-12 17:40:20

【问题描述】：

我的数据喜欢

value = [1, 2, 3, 4, 5, 6]
key =   [0, 1, 0, 2, 1, 2]

我现在需要为每个组（键）设置最大值（值和索引）。 所以结果应该是

max = [3, 5, 6]
index = [2, 4, 5]
key = [0, 1, 2]

如何通过 cuda 推力获得它？ 我可以做 sort -> reduce_by_key 但它不是很有效。在我的情况下，向量大小 > 10M 和密钥空间 ~ 1K（从 0 开始，没有间隙）。

【问题讨论】：

你试过了吗？

使用thrust::sort_by_key 将相似的键组合在一起。然后使用thrust::reduce_by_key 以及一个zip_iterator 和一个counting_iterator（用于索引）来查找每个键及其索引中的最大值。

@RobertCrovella 我正在寻找更优雅的解决方案。

我花了一分钟才明白你在问什么。您应该编辑第一部分以使其更清晰。

@sh1ng 为什么不先尝试自己编写它，看看您面临什么问题，而不是寻求完整的解决方案。

【参考方案1】：

由于最初的问题集中在推力上，除了我在cmets中提到的之外，我没有任何建议，

但是，根据 cmets 中的进一步对话，我想我会发布一个涵盖 CUDA 和推力的答案。

thrust 方法使用 sort_by_key 操作将相似的键组合在一起，然后使用 reduce_by_key 操作来找到每个键组的最大 + 索引。

CUDA 方法使用我描述的自定义原子方法 here 来查找 32 位最大值加上 32 位索引（对于每个键组）。

对于这个特定的测试用例，CUDA 方法的速度要快得多（~10 倍）。我在这个测试中使用了 10M 的向量大小和 10K 的密钥大小。

我的测试平台是 CUDA 8RC、RHEL 7 和 Tesla K20X GPU。 K20X 是 Kepler 一代的成员，其全局原子比前几代 GPU 快得多。

这是一个完整的示例，涵盖了这两种情况，并提供了时间比较：

$ cat t1234.cu
#include <iostream>
#include <thrust/copy.h>
#include <thrust/reduce.h>
#include <thrust/sort.h>
#include <thrust/device_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/sequence.h>
#include <thrust/functional.h>
#include <cstdlib>

#include <time.h>
#include <sys/time.h>
#define USECPSEC 1000000ULL

unsigned long long dtime_usec(unsigned long long start)

  timeval tv;
  gettimeofday(&tv, 0);
  return ((tv.tv_sec*USECPSEC)+tv.tv_usec)-start;


const size_t ksize = 10000;
const size_t vsize = 10000000;
const int nTPB = 256;

struct my_max_func


  template <typename T1, typename T2>
  __host__ __device__
  T1 operator()(const T1 t1, const T2 t2)
    T1 res;
    if (thrust::get<0>(t1) > thrust::get<0>(t2))
      thrust::get<0>(res) = thrust::get<0>(t1);
      thrust::get<1>(res) = thrust::get<1>(t1);
    else 
      thrust::get<0>(res) = thrust::get<0>(t2);
      thrust::get<1>(res) = thrust::get<1>(t2);
    return res;
    
;

typedef union  
  float floats[2];                 // floats[0] = maxvalue
  int ints[2];                     // ints[1] = maxindex
  unsigned long long int ulong;    // for atomic update
 my_atomics;


__device__ unsigned long long int my_atomicMax(unsigned long long int* address, float val1, int val2)

    my_atomics loc, loctest;
    loc.floats[0] = val1;
    loc.ints[1] = val2;
    loctest.ulong = *address;
    while (loctest.floats[0] <  val1)
      loctest.ulong = atomicCAS(address, loctest.ulong,  loc.ulong);
    return loctest.ulong;



__global__ void my_max_idx(const float *data, const int *keys,const int ds, my_atomics *res)


    int idx = (blockDim.x * blockIdx.x) + threadIdx.x;
    if (idx < ds)
      my_atomicMax(&(res[keys[idx]].ulong), data[idx],idx);



int main()

  float *h_vals = new float[vsize];
  int   *h_keys = new int[vsize];
  for (int i = 0; i < vsize; i++) h_vals[i] = rand(); h_keys[i] = rand()%ksize;
// thrust method
  thrust::device_vector<float> d_vals(h_vals, h_vals+vsize);
  thrust::device_vector<int> d_keys(h_keys, h_keys+vsize);
  thrust::device_vector<int> d_keys_out(ksize);
  thrust::device_vector<float> d_vals_out(ksize);
  thrust::device_vector<int> d_idxs(vsize);
  thrust::device_vector<int> d_idxs_out(ksize);

  thrust::sequence(d_idxs.begin(), d_idxs.end());
  cudaDeviceSynchronize();
  unsigned long long et = dtime_usec(0);

  thrust::sort_by_key(d_keys.begin(), d_keys.end(), thrust::make_zip_iterator(thrust::make_tuple(d_vals.begin(), d_idxs.begin())));
  thrust::reduce_by_key(d_keys.begin(), d_keys.end(), thrust::make_zip_iterator(thrust::make_tuple(d_vals.begin(),d_idxs.begin())), d_keys_out.begin(), thrust::make_zip_iterator(thrust::make_tuple(d_vals_out.begin(), d_idxs_out.begin())), thrust::equal_to<int>(), my_max_func());
  cudaDeviceSynchronize();
  et = dtime_usec(et);
  std::cout << "Thrust time: " << et/(float)USECPSEC << "s" << std::endl;

// cuda method

  float *vals;
  int *keys;
  my_atomics *results;
  cudaMalloc(&keys, vsize*sizeof(int));
  cudaMalloc(&vals, vsize*sizeof(float));
  cudaMalloc(&results, ksize*sizeof(my_atomics));

  cudaMemset(results, 0, ksize*sizeof(my_atomics)); // works because vals are all positive
  cudaMemcpy(keys, h_keys, vsize*sizeof(int), cudaMemcpyHostToDevice);
  cudaMemcpy(vals, h_vals, vsize*sizeof(float), cudaMemcpyHostToDevice);
  et = dtime_usec(0);

  my_max_idx<<<(vsize+nTPB-1)/nTPB, nTPB>>>(vals, keys, vsize, results);
  cudaDeviceSynchronize();
  et = dtime_usec(et);
  std::cout << "CUDA time: " << et/(float)USECPSEC << "s" << std::endl;

// verification

  my_atomics *h_results = new my_atomics[ksize];
  cudaMemcpy(h_results, results, ksize*sizeof(my_atomics), cudaMemcpyDeviceToHost);
  for (int i = 0; i < ksize; i++)
    if (h_results[i].floats[0] != d_vals_out[i]) std::cout << "value mismatch at index: " << i << " thrust: " << d_vals_out[i] << " CUDA: " << h_results[i].floats[0] << std::endl; return -1;
    if (h_results[i].ints[1] != d_idxs_out[i]) std::cout << "index mismatch at index: " << i << " thrust: " << d_idxs_out[i] << " CUDA: " << h_results[i].ints[1] << std::endl; return -1;
    

  std::cout << "Success!" << std::endl;
  return 0;


$ nvcc -arch=sm_35 -o t1234 t1234.cu
$ ./t1234
Thrust time: 0.026593s
CUDA time: 0.002451s
Success!
$

【讨论】：

对于键的有限整数值范围非常快速的解决方案。但是评论中一个问题的创建者补充说：“值只是一个从 0 到 N 的 float key(s)。”从高级系统的经验来看，在 DBMS (MSSQL/Oracle...) 中对所有类型的值和键进行分组通常只使用两种方法：有序匹配（按键排序 + 按键排序分组）和散列匹配（散列具有最小/最大/总和...操作的表）。两者都可以在 CUDA 上实现。

我认为这意味着“值只是一个浮点数”（句号）“键范围从 0 到 N”。 “值只是一个浮动键”对我来说没有多大意义，因为 key 和 value 是独立的概念。我提出的解决方案适用于从 0 到 N 的 int 键，这似乎正是 OP 所要求的。