CUDA 推力排序或 CUB::DeviceRadixSort

Question

我有一个由 float4 数组表示的粒子池，其中 w 分量是粒子在 [0, 1] 范围内的当前生命周期。

我需要根据粒子的生命周期降序对这个数组进行排序，以便我可以准确计算当前有多少粒子 "active"（生命周期大于 0）。我需要这个计数器，因为当我需要激活更多粒子（随机发生）时，它可以让我索引到阵列中的正确位置。

我的粒子数组存储在设备内存中，似乎我应该能够对数组进行排序而无需将数组传输到主机内存。

我不太幸运地在网上找到了展示我如何使用 Thrust 或 CUB 执行此操作的示例。此外，我对使用 Thrust 犹豫不决，因为我不知道如何防止它退化为合并排序（这比基数排序慢得多），因为我需要基于 w 组件进行排序。至于 CUB，我根本没有找到任何关于如何做到这一点的资源。

我也更愿意将生命周期存储在 w 组件中，因为这使我在代码的其他部分的工作变得更加轻松。

有没有简单的方法可以做到这一点？感谢您的帮助。

Answer 1

在 cub 或 thrust 中，我们可以仅对 .w "keys" 进行排序，进行键值排序，其中值只是线性递增索引：

0, 1, 2, 3, ...

然后我们可以使用索引序列的结果重排一步对原始 float4 数组重新排序（按 .w 有效排序）。这将允许您保持基数排序速度（在立方体或推力中）并且也可能相当有效，因为 float4 数量只需要 moved/rearranged 一次，而不是在排序操作期间连续移动.

这里有一个完整的 thrust 示例，在 32M 元素上，演示了 "ordinary" thrust 排序，使用函子指定对 .w 元素（sort_f4_w）的排序，接下来是上述方法。在这种情况下，在我的特定设置（Fedora 20、CUDA 7、Quadro5000）上，第二种方法似乎快了大约 5 倍：

$ cat t686.cu
#include <iostream>
#include <vector_types.h>
#include <stdlib.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/sequence.h>
#include <thrust/copy.h>
#include <thrust/equal.h>

#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;
}

#define DSIZE (32*1048576)

struct sort_f4_w
{
  __host__ __device__
  bool operator()(const float4 &a, const float4 &b) const {
    return (a.w < b.w);}
};
// functor to extract the .w element from a float4
struct f4_to_fw : public thrust::unary_function<float4, float>
{
  __host__ __device__
  float operator()(const float4 &a) const {
    return a.w;}
};
// functor to extract the .x element from a float4
struct f4_to_fx : public thrust::unary_function<float4, float>
{
  __host__ __device__
  float operator()(const float4 &a) const {
    return a.x;}
};


bool validate(thrust::device_vector<float4> &d1, thrust::device_vector<float4> &d2){
  return thrust::equal(thrust::make_transform_iterator(d1.begin(), f4_to_fx()), thrust::make_transform_iterator(d1.end(), f4_to_fx()), thrust::make_transform_iterator(d2.begin(), f4_to_fx()));
}


int main(){
  unsigned long long t1_time, t2_time;
  float4 *mydata = new float4[DSIZE];
  for (int i = 0; i < DSIZE; i++){
    mydata[i].x = i;
    mydata[i].y = i;
    mydata[i].z = i;
    mydata[i].w = rand()/(float)RAND_MAX;}

  thrust::host_vector<float4>   h_data(mydata, mydata+DSIZE);
  // do once as a warm-up run, then report timings on second run
  for (int i = 0; i < 2; i++){
    thrust::device_vector<float4> d_data1 = h_data;
    thrust::device_vector<float4> d_data2 = h_data;

  // first time sort using typical thrust approach
    t1_time = dtime_usec(0);
    thrust::sort(d_data1.begin(), d_data1.end(), sort_f4_w());
    cudaDeviceSynchronize();
    t1_time = dtime_usec(t1_time);
  // now extract keys and create index values, sort, then rearrange
    t2_time = dtime_usec(0);
    thrust::device_vector<float> keys(DSIZE);
    thrust::device_vector<int> vals(DSIZE);
    thrust::copy(thrust::make_transform_iterator(d_data2.begin(), f4_to_fw()), thrust::make_transform_iterator(d_data2.end(), f4_to_fw()), keys.begin());
    thrust::sequence(vals.begin(), vals.end());
    thrust::sort_by_key(keys.begin(), keys.end(), vals.begin());
    thrust::device_vector<float4> result(DSIZE);
    thrust::copy(thrust::make_permutation_iterator(d_data2.begin(), vals.begin()), thrust::make_permutation_iterator(d_data2.begin(), vals.end()), result.begin());
    cudaDeviceSynchronize();
    t2_time = dtime_usec(t2_time);
    if (!validate(d_data1, result)){
      std::cout << "Validation failure " << std::endl;
      }
    }
  std::cout << "thrust t1 time: " << t1_time/(float)USECPSEC << "s, t2 time: " << t2_time/(float)USECPSEC << std::endl;
}


$ nvcc -o t686 t686.cu
$ ./t686
thrust t1 time: 0.731456s, t2 time: 0.149959
$

CUDA 推力排序或 CUB::DeviceRadixSort

CUDA Thrust sort or CUB::DeviceRadixSort

sorting

cuda

thrust

cub