使用 CUDA/Thrust 对多个数组进行排序

Sorting multiple arrays using CUDA/Thrust

我有一个大数组需要在 GPU 上排序。数组本身是多个较小的子数组的串联,满足给定 i < j 的条件,子数组 i 的元素小于子数组 j 的元素。这种数组的一个例子是{5 3 4 2 1 6 9 8 7 10 11}, 其中 5 个元素的第一个子数组的元素小于 6 个元素的第二个子数组的元素。我需要的数组是 {1, 2, 3, 4, 5, 6, 7, 10, 11}。我知道每个子数组在大数组中的起始位置。

我知道我可以简单地对整个数组使用 thrust::sort,但我想知道是否可以启动多个并发排序,每个子数组一个。我希望通过这样做来提高性能。我的假设是,对多个较小的数组进行排序比对包含所有元素的一个大数组进行排序要快。

如果有人能给我一种方法或纠正我的假设以防错误,我将不胜感激。

一种在 thrust 中进行多个并发排序("vectorized" 排序)的方法是通过子数组的标记,并提供一个自定义仿函数,它是一个普通的 thrust 排序仿函数,它也对子数组进行排序通过他们的钥匙。

另一种可能的方法是使用背靠背 thrust::stable_sort_by_key,如 所述。

正如您所指出的,您的情况下的另一种方法只是进行普通排序,因为那最终是您的 objective。

但是我认为任何一种推力排序方法都不太可能比纯排序有显着的加速,尽管您可以尝试一下。 Thrust 有一个快速路径基数排序,它将在某些情况下使用,纯排序方法可能会在您的情况下使用。 (在其他情况下,例如,当您提供自定义仿函数时,thrust 通常会使用较慢的合并排序方法。)

如果子数组的大小在特定范围内,我认为使用 cub 中的块基数排序可能会获得更好的结果(性能方面),每个子数组一个块。

这是一个使用特定大小的示例(因为您没有给出大小范围和其他细节的指示),将推力 "pure sort" 与带函子的推力分段排序与 cub 块排序进行比较方法。对于这种特殊情况,幼崽排序最快:

$ cat t1.cu
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/sort.h>
#include <thrust/scan.h>
#include <thrust/equal.h>
#include <cstdlib>
#include <iostream>


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

const int num_blocks = 2048;
const int items_per = 4;
const int nTPB = 512;
const int block_size = items_per*nTPB; // must be a whole-number multiple of nTPB;
typedef float mt;

unsigned long long dtime_usec(unsigned long long start){

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

struct my_sort_functor
{
        template <typename T, typename T2>
        __host__ __device__
        bool operator()(T t1, T2 t2){
                if (thrust::get<1>(t1) < thrust::get<1>(t2)) return true;
                if (thrust::get<1>(t1) > thrust::get<1>(t2)) return false;
                if (thrust::get<0>(t1) > thrust::get<0>(t2)) return false;
                return true;}
};

// from: https://nvlabs.github.io/cub/example_block_radix_sort_8cu-example.html#_a0
#define CUB_STDERR
#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <cub/block/block_load.cuh>
#include <cub/block/block_store.cuh>
#include <cub/block/block_radix_sort.cuh>
using namespace cub;
//---------------------------------------------------------------------
// Globals, constants and typedefs
//---------------------------------------------------------------------
bool g_verbose = false;
bool g_uniform_keys;
//---------------------------------------------------------------------
// Kernels
//---------------------------------------------------------------------
template <
    typename    Key,
    int         BLOCK_THREADS,
    int         ITEMS_PER_THREAD>
__launch_bounds__ (BLOCK_THREADS)
__global__ void BlockSortKernel(
    Key         *d_in,          // Tile of input
    Key         *d_out)         // Tile of output
{
    enum { TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD };
    // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared memory to a blocked arrangement)
    typedef BlockLoad<Key, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_LOAD_WARP_TRANSPOSE> BlockLoadT;
    // Specialize BlockRadixSort type for our thread block
    typedef BlockRadixSort<Key, BLOCK_THREADS, ITEMS_PER_THREAD> BlockRadixSortT;
    // Shared memory
    __shared__ union TempStorage
    {
        typename BlockLoadT::TempStorage        load;
        typename BlockRadixSortT::TempStorage   sort;
    } temp_storage;
    // Per-thread tile items
    Key items[ITEMS_PER_THREAD];
    // Our current block's offset
    int block_offset = blockIdx.x * TILE_SIZE;
    // Load items into a blocked arrangement
    BlockLoadT(temp_storage.load).Load(d_in + block_offset, items);
    // Barrier for smem reuse
    __syncthreads();
    // Sort keys
    BlockRadixSortT(temp_storage.sort).SortBlockedToStriped(items);
    // Store output in striped fashion
    StoreDirectStriped<BLOCK_THREADS>(threadIdx.x, d_out + block_offset, items);
}

int main(){
        const int ds = num_blocks*block_size;
        thrust::host_vector<mt>      data(ds);
        thrust::host_vector<int>     keys(ds);
        for (int i = block_size; i < ds; i+=block_size) keys[i] = 1; // mark beginning of blocks
        thrust::device_vector<int> d_keys = keys;
        for (int i = 0; i < ds; i++) data[i] = (rand()%block_size) + (i/block_size)*block_size;  // populate data
        thrust::device_vector<mt>  d_data = data;
        thrust::inclusive_scan(d_keys.begin(), d_keys.end(), d_keys.begin());  // fill out keys array  000111222...
        thrust::device_vector<mt> d1 = d_data;  // make a copy of unsorted data
        cudaDeviceSynchronize();
        unsigned long long os = dtime_usec(0);
        thrust::sort(d1.begin(), d1.end());  // ordinary sort
        cudaDeviceSynchronize();
        os = dtime_usec(os);
        thrust::device_vector<mt> d2 = d_data;  // make a copy of unsorted data
        cudaDeviceSynchronize();
        unsigned long long ss = dtime_usec(0);
        thrust::sort(thrust::make_zip_iterator(thrust::make_tuple(d2.begin(), d_keys.begin())), thrust::make_zip_iterator(thrust::make_tuple(d2.end(), d_keys.end())), my_sort_functor());
        cudaDeviceSynchronize();
        ss = dtime_usec(ss);
        if (!thrust::equal(d1.begin(), d1.end(), d2.begin())) {std::cout << "oops1" << std::endl; return 0;}
        std::cout << "ordinary thrust sort: " << os/(float)USECPSEC << "s " << "segmented sort: " << ss/(float)USECPSEC << "s" << std::endl;
        thrust::device_vector<mt> d3(ds);
        cudaDeviceSynchronize();
        unsigned long long cs = dtime_usec(0);
        BlockSortKernel<mt, nTPB, items_per><<<num_blocks, nTPB>>>(thrust::raw_pointer_cast(d_data.data()),  thrust::raw_pointer_cast(d3.data()));
        cudaDeviceSynchronize();
        cs = dtime_usec(cs);
        if (!thrust::equal(d1.begin(), d1.end(), d3.begin())) {std::cout << "oops2" << std::endl; return 0;}
        std::cout << "cub sort: " << cs/(float)USECPSEC << "s" << std::endl;
}
$ nvcc -o t1 t1.cu
$ ./t1
ordinary thrust sort: 0.001652s segmented sort: 0.00263s
cub sort: 0.000265s
$

(CUDA 10.2.89, Tesla V100, Ubuntu 18.04)

我毫不怀疑你的大小和数组维度与我的不符。这里的目的是说明一些可能的方法,而不是适用于您的特定情况的黑盒解决方案。您可能应该自己进行基准比较。我还承认 cub 的块基数排序方法需要大小相等的子数组,而您可能没有。它可能不是适合您的方法,或者您可能希望探索某种padding arrangement。没有必要问我这个问题;根据你问题的信息,我无法回答。

我不声明此代码或我 post 的任何其他代码的正确性。任何使用我 post 的任何代码的人都需要自行承担风险。我只是声称我已尝试解决原始 posting 中的问题,并提供一些解释。我并不是说我的代码没有缺陷,或者它适用于任何特定目的。使用(或不使用)风险自负。