CUDA：查找稀疏数组的 N 个最大值

Question

我有一个数百万个值的数组存储在 GPU 的全局内存中。他们中的大多数都是零，除了几千。这些值是设备上的计算结果。

我想尽快找到1024 的最大值及其索引。

有人有什么建议吗？

Answer 1

[此答案已根据 Robert Crovella 的评论进行编辑]

为了实现一个简单的方法，我建议按降序使用thrust::sort 或thrust::sort_by_key。两种情况都可以通过thrust::greater<int>()实现降序排序。

最简单的方法是按降序使用thrust::sort，这样您就可以从大到小依次访问已排序的元素。

如果要保留原始数据向量的副本以及排序过程的索引，可以按降序使用thrust::sort_by_key。假设您感兴趣的数组具有N 元素。您可以通过thrust::sequence 创建一系列递增索引。在您的情况下，键是N 元素的数组，而值是thrust::sequence 生成的数组。在按降序使用thrust::sort_by_key 之后，值数组将包含您可以访问第一个最大元素的索引。

请注意，您实际上对数据数组是稀疏的情况感兴趣，因此您可能有兴趣仅对数据数组的非消失值进行排序。如果您有兴趣只存储数组的非消失值，则不需要通过thrust::sequence 创建索引向量d_indices，但存储非消失的索引就足够了数据值。如果您已经有一个包含0 的数组，那么您可以在通过thrust::partition 执行排序操作之前提取非消失值。

下面是一个完整的例子，展示了上述所有方法。

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/reverse.h>
#include <thrust/copy.h>    
#include <thrust/sort.h>
#include <cstdlib>
#include <iostream>

using namespace std;

struct is_not_zero
{
    __host__ __device__ bool operator()(const int &x) { return x != 0; }
};

void main(void)
{

    const int N = 8;

    // --- Generate the data vector of random values on the host 
    thrust::host_vector<int> h_vec(N);
    thrust::generate(h_vec.begin(), h_vec.end(), rand);

    // --- Move the data vector to the device
    thrust::device_vector<int> d_vec=h_vec;

    // --- Make two copies of the data vector
    thrust::device_vector<int> d_vec_copy=d_vec;
    thrust::device_vector<int> d_vec_another_copy=d_vec;

    // --- Push back some zero to test thrust::partition
    d_vec_another_copy.push_back(0);
    d_vec_another_copy.push_back(0);
    d_vec_another_copy.push_back(0);

    // --- Display the result
    for(int i = 0; i<N+3; i++)
        cout << d_vec_another_copy[i] << endl;
    cout << endl;

    // --- Generate the indices vector
    thrust::device_vector<int> d_indices(N);
    thrust::sequence(d_indices.begin(), d_indices.end(), 0, 1);

    // --- Sort the indices vector by using the data vector as key in descending order
    thrust::sort_by_key(d_vec.begin(), d_vec.end(), d_indices.begin(),thrust::greater<int>());

    // --- Display the result
    for(int i = 0; i<N; i++) {
        int index = d_indices[i];
        cout << "Original: " << d_vec_copy[index] << " Sorted: " << d_vec[i] << endl;
    }
    cout << endl;

    // --- Use sort in descending order and forget the initial ordering
    thrust::sort(d_vec_copy.begin(), d_vec_copy.end(), thrust::greater<int>());

    // --- Display the result
    for(int i = 0; i<N; i++)
        cout << d_vec_copy[i] << endl;
    cout << endl;

    // --- Use partition prior to sort to extract the non-vanishing elements in descending order
    thrust::partition(d_vec_another_copy.begin(), d_vec_another_copy.end(), is_not_zero());     
    thrust::sort(d_vec_another_copy.begin(), d_vec_another_copy.end(), thrust::greater<int>());

    // --- Display the result
    for(int i = 0; i<N; i++)
        cout << d_vec_another_copy[i] << endl;
    cout << endl;

    getchar();

}