sorting - 使用 CUDA Thrust 同时对多个数组进行排序

Question

我需要按相同的键对 GPU 上已有的 20+ 数组进行排序，每个数组的长度相同。我不能直接使用 sort_by_key() 因为它也会对键进行排序(使它们无法对下一个数组进行排序)。这是我尝试过的:

thrust::device_vector<int>  indices(N); 
thrust::sequence(indices.begin(),indices.end());
thrust::sort_by_key(keys.begin(),keys.end(),indices.begin());

thrust::gather(indices.begin(),indices.end(),a_01,a_01);
thrust::gather(indices.begin(),indices.end(),a_02,a_02);
...
thrust::gather(indices.begin(),indices.end(),a_20,a_20);

这似乎不起作用，因为 gather() 期望输出的数组与输入的数组不同，即这有效:

thrust::gather(indices.begin(),indices.end(),a_01,o_01);
...

但是，我不想为此任务分配 20+ 额外的数组。我知道有一个使用推力::元组、推力::zip_iterator和推力::sort_by_keys()的解决方案，类似于here 。但是，我最多只能在一个元组中组合 10 数组，s.t.我需要再次复制 key 向量。您将如何完成这项任务？

Answer 1

我认为对多个数组进行排序的经典方法是所谓的背对背方法，该方法使用两次thrust::stable_sort_by_key。您需要创建一个键向量，以便同一数组中的元素具有相同的键。例如:

Elements: 10.5 4.3 -2.3 0. 55. 24. 66.
Keys:      0    0    0  1   1   1   1

在本例中，我们有两个数组，第一个包含 3 元素，第二个包含 4 元素。

您首先需要调用 thrust::stable_sort_by_key 将矩阵值作为键，例如

thrust::stable_sort_by_key(d_matrix.begin(),
                           d_matrix.end(),
                           d_keys.begin(),
                           thrust::less<float>());

之后，你就拥有了

Elements: -2.3 0 4.3 10.5 24. 55. 66.
Keys:       0  1  0    0   1   1   1

这意味着数组元素是有序的，而键不是有序的。然后你需要一秒钟来调用 thrust::stable_sort_by_key

thrust::stable_sort_by_key(d_keys.begin(),
                           d_keys.end(),
                           d_matrix.begin(),
                           thrust::less<int>());

因此根据键执行排序。完成该步骤后，您将拥有

Elements: -2.3 4.3 10.5 0 24. 55. 66.
Keys:       0   0   0   1  1   1   1

这是最终想要的结果。

下面是一个完整的工作示例，它考虑以下问题:单独对矩阵的每一行进行排序。这是一种特殊情况，其中所有数组都具有相同的长度，但该方法适用于可能具有不同长度的数组。

#include <cublas_v2.h>

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/generate.h>
#include <thrust/sort.h>
#include <thrust/functional.h>
#include <thrust/random.h>
#include <thrust/sequence.h>

#include <stdio.h>
#include <iostream>

#include "Utilities.cuh"

/**************************************************************/
/* CONVERT LINEAR INDEX TO ROW INDEX - NEEDED FOR APPROACH #1 */
/**************************************************************/
template <typename T>
struct linear_index_to_row_index : public thrust::unary_function<T,T> {

    T Ncols; // --- Number of columns

    __host__ __device__ linear_index_to_row_index(T Ncols) : Ncols(Ncols) {}

    __host__ __device__ T operator()(T i) { return i / Ncols; }
};

/********/
/* MAIN */
/********/
int main()
{
    const int Nrows = 5;     // --- Number of rows
    const int Ncols = 8;     // --- Number of columns

    // --- Random uniform integer distribution between 10 and 99
    thrust::default_random_engine rng;
    thrust::uniform_int_distribution<int> dist(10, 99);

    // --- Matrix allocation and initialization
    thrust::device_vector<float> d_matrix(Nrows * Ncols);
    for (size_t i = 0; i < d_matrix.size(); i++) d_matrix[i] = (float)dist(rng);

    // --- Print result
    printf("Original matrix\n");
    for(int i = 0; i < Nrows; i++) {
        std::cout << "[ ";
        for(int j = 0; j < Ncols; j++)
            std::cout << d_matrix[i * Ncols + j] << " ";
        std::cout << "]\n";
    }

    /*************************/
    /* BACK-TO-BACK APPROACH */
    /*************************/
    thrust::device_vector<float> d_keys(Nrows * Ncols);

    // --- Generate row indices
    thrust::transform(thrust::make_counting_iterator(0),
                      thrust::make_counting_iterator(Nrows*Ncols),
                      thrust::make_constant_iterator(Ncols),
                      d_keys.begin(),
                      thrust::divides<int>());

    // --- Back-to-back approach
    thrust::stable_sort_by_key(d_matrix.begin(),
                               d_matrix.end(),
                               d_keys.begin(),
                               thrust::less<float>());

    thrust::stable_sort_by_key(d_keys.begin(),
                               d_keys.end(),
                               d_matrix.begin(),
                               thrust::less<int>());

    // --- Print result
    printf("\n\nSorted matrix\n");
    for(int i = 0; i < Nrows; i++) {
        std::cout << "[ ";
        for(int j = 0; j < Ncols; j++)
            std::cout << d_matrix[i * Ncols + j] << " ";
        std::cout << "]\n";
    }

    return 0;
}