使用 CUDA/Thrust 對多個 arrays 進行排序

Question

我有一個大數組，需要在 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}。 我知道 position 每個子陣列都從大陣列開始。

我知道我可以簡單地在整個數組上使用thrust::sort ，但我想知道是否可以啟動多個並發排序，每個子數組一個。 我希望通過這樣做來提高性能。 我的假設是，對多個較小的 arrays 進行排序比對一個包含所有元素的大數組進行排序會更快。

如果有人能給我一種方法來做到這一點或糾正我的假設以防萬一它是錯誤的，我將不勝感激。

Answer 1

一種在推力中進行多個並發排序（“矢量化”排序）的方法是通過子 arrays 的標記，並提供一個自定義函子，它是一個普通的推力排序函子，它也通過它們的鍵對子 arrays 進行排序。

另一種可能的方法是使用背靠背thrust::stable_sort_by_key ，如此處所述。

正如您所指出的，在您的情況下，另一種方法就是進行普通排序，因為這最終是您的目標。

但是，我認為任何推力排序方法都不太可能比純排序顯着加快速度，盡管您可以嘗試一下。 Thrust 具有快速路徑基數排序，它將在某些情況下使用，純排序方法可能會在您的情況下使用。 （在其他情況下，例如當您提供自定義函子時，推力通常會使用較慢的合並排序方法。）

如果子 arrays 的大小在一定范圍內，我認為您可能會在 cub 中使用塊基數排序獲得更好的結果（性能方面），每個子陣列一個塊。

這是一個使用特定大小的示例（因為您沒有給出大小范圍和其他詳細信息的指示），將推力“純排序”與具有函子的推力分段排序以及 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、特斯拉 V100、Ubuntu 18.04）

我毫不懷疑您的大小和數組尺寸與我的不符。 此處的目的是說明一些可能的方法，而不是適用於您的特定情況的黑盒解決方案。 您可能應該自己進行基准比較。 我也承認 cub 的塊基數排序方法需要大小相等的子數組，而您可能沒有。 它可能不適合您，或者您可能希望探索某種填充排列。 沒必要問我這個問題； 根據您問題中的信息，我將無法回答。

我不聲明此代碼或我發布的任何其他代碼的正確性。 使用我發布的任何代碼的任何人都需要自擔風險。 我只是聲稱我試圖解決原始帖子中的問題，並提供一些解釋。 我並不是說我的代碼沒有缺陷，或者它適用於任何特定目的。 使用（或不使用）風險自負。