cuda - 在 CUDA 中按键对 3 个数组进行排序(可能使用 Thrust)

Question

我有 3相同大小的数组(超过 300.000 个元素)。一个浮点数数组和两个索引数组。因此，对于每个数字，我都有 2身份证件。

所有 3数组已经在 GPU 全局内存中。我想相应地用他们的 ID 对所有数字进行排序。

有什么办法可以使用 Thrust 库来完成这项任务吗？有没有比 Thrust 库更好的方法？

当然，我不想将它们复制到主机内存或从主机内存复制几次。顺便说一下，它们是数组而不是向量。

提前感谢您的帮助。

暂定方案 ，但这非常慢。几乎需要4秒，我的数组大小的顺序是 300000

thrust::device_ptr<float> keys(afterSum);
thrust::device_ptr<int> vals0(d_index);
thrust::device_ptr<int> vals1(blockId); 

thrust::device_vector<int> sortedIndex(numElements);
thrust::device_vector<int> sortedBlockId(numElements);

thrust::counting_iterator<int> iter(0);
thrust::device_vector<int> indices(numElements);
thrust::copy(iter, iter + indices.size(), indices.begin()); 

thrust::sort_by_key(keys, keys + numElements , indices.begin());    

thrust::gather(indices.begin(), indices.end(), vals0, sortedIndex.begin());
thrust::gather(indices.begin(), indices.end(), vals1, sortedBlockId.begin());

thrust::host_vector<int> h_sortedIndex=sortedIndex;
thrust::host_vector<int> h_sortedBlockId=sortedBlockId;

Answer 1

当然，您可以使用推力。首先，您需要使用 thrust::device_ptr 包装原始 CUDA 设备指针。 .假设您的浮点值在数组中 pkeys ，并且 ID 位于数组 pvals0 中和 pvals1 ，和 numElements 是数组的长度，这样的事情应该工作:

#include <thrust/device_ptr.h>
#include <thrust/sort.h>
#include <thrust/gather.h>
#include <thrust/iterator/counting_iterator.h>

cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

cudaEventRecord(start);

thrust::device_ptr<float> keys(pkeys);
thrust::device_ptr<int> vals0(pvals0);
thrust::device_ptr<int> vals1(pvals1);

// allocate space for the output
thrust::device_vector<int> sortedVals0(numElements);
thrust::device_vector<int> sortedVals1(numElements);

// initialize indices vector to [0,1,2,..]
thrust::counting_iterator<int> iter(0);
thrust::device_vector<int> indices(numElements);
thrust::copy(iter, iter + indices.size(), indices.begin());

// first sort the keys and indices by the keys
thrust::sort_by_key(keys.begin(), keys.end(), indices.begin());

// Now reorder the ID arrays using the sorted indices
thrust::gather(indices.begin(), indices.end(), vals0.begin(), sortedVals0.begin());
thrust::gather(indices.begin(), indices.end(), vals1.begin(), sortedVals1.begin());

cudaEventRecord(stop);
cudaEventSynchronize(stop);
float milliseconds = 0;
cudaEventElapsedTime(&milliseconds, start, stop);
printf("Took %f milliseconds for %d elements\n", milliseconds, numElements);