c++ - CUB (CUDA UnBound) 相当于推力::聚集

由于 Thrust 库的一些性能问题(有关更多详细信息，请参阅 this page)，我计划重构 CUDA 应用程序以使用 CUB 而不是 Thrust。具体来说，要替换推力::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 是一个推力::device_ptr，它指向我想按

排序的键

indices 指向设备内存中的一个序列(从 0 到 numKeys-1)

values{1,2,3}Ptr 是我要排序的值的 device_ptrs

values{1,2,3}OutPtr 是 device_ptrs 到排序值

使用 CUB SortPairs 函数，我可以对单个值缓冲区进行排序，但不能一次性对所有 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 或

在使用 values{1,2,3}

之前将 cub::DeviceRadixSort::SortPairs 压缩到结构数组中

更新

阅读 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 块级原语 cub::BlockRadixSort<...>::SortBlockedToStriped 来获取存储在寄存器中的排序索引，然后执行一个朴素的重新排序拷贝作为 thrust::gather 来填充 values{1,2,3}Out 。

在复制 SortBlockedToStriped 时，使用 Sort 而不是 values 可以进行合并写入(尽管不是为了读取)。

关于c++ - CUB (CUDA UnBound) 相当于推力::聚集，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/19210652/