由于推力库的一些性能问题(有关更多详细信息,请参阅此页面(,我计划重构CUDA应用程序以使用CUB而不是推力。 具体来说,替换推力::sort_by_key和推力::inclusive_scan调用(。 在我的应用程序中的特定点,我需要按键对 3 个数组进行排序。 这就是我用推力做到这一点的方式:
thrust::sort_by_key(key_iter, key_iter + numKeys, indices);
thrust::gather_wrapper(indices, indices + numKeys,
thrust::make_zip_iterator(thrust::make_tuple(values1Ptr, values2Ptr, values3Ptr)),
thrust::make_zip_iterator(thrust::make_tuple(valuesOut1Ptr, valuesOut2Ptr, valuesOut3Ptr))
);
哪里
-
key iter
是一个 thrust::d evice_ptr 指向我要排序的键 -
indices
指向设备内存中的序列(从 0 到 numKeys-1( -
values{1,2,3}Ptr
device_ptrs到我要排序的值 -
values{1,2,3}OutPtr
device_ptrs排序值
使用 CUB 排序对函数,我可以对单个值缓冲区进行排序,但不能一次性对全部 3 个缓冲区进行排序。 问题是我没有看到任何 CUB"类似收集"的实用程序。 建议?
编辑:
我想我可以实现我自己的收集内核,但是除了:
template <typename Index, typename Value>
__global__ void gather_kernel(const unsigned int N, const Index * map,
const Value * src, Value * dst)
{
unsigned int i = blockDim.x * blockIdx.x + threadIdx.x;
if (i < N)
{
dst[i] = src[map[i]];
}
}
非合并的负载和存储让我感到震惊,但如果没有已知的结构,这可能是不可避免的map
.
似乎你想要实现的目标取决于thrust::zip_iterator
。你可以
- 仅将
thrust::sort_by_key
替换为cub::DeviceRadixSort::SortPairs
并保留thrust::gather
,或 - 在使用
cub::DeviceRadixSort::SortPairs
之前将values{1,2,3}
压缩到结构数组中
更新
看完thrust::gather
的实现,
$CUDA_HOME/include/thrust/system/detail/generic/gather.inl
你可以看到它只是一个天真的内核,比如
__global__ gather(int* index, float* in, float* out, int len) {
int i=...;
if (i<len) { out[i] = in[index[i]]; }
}
然后我认为您上面的代码可以毫不费力地被单个内核替换。
在这个内核中,您可以首先使用 CUB 块 wize 原语cub::BlockRadixSort<...>::SortBlockedToStriped
来获取存储在寄存器中的排序索引,然后执行朴素的重新排序复制作为thrust::gather
来填充values{1,2,3}Out
。
在复制values
时,使用 SortBlockedToStriped
而不是 Sort
可以进行合并写入(虽然不是为了阅读(。