c++ - CUB (CUDA UnBound) 相当于 thrust::gather

Question

由于 Thrust 库存在一些性能问题(有关详细信息，请参阅 this page)，我计划重构一个 CUDA 应用程序以使用 CUB 而不是 Thrust。具体来说，就是替换 thrust::sort_by_key 和 thrust::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::device_ptr，它指向我要排序的键
• indices 指向设备内存中的一个序列(从 0 到 numKeys-1)
• values{1,2,3}Ptr 是我要排序的值的 device_ptr
• values{1,2,3}OutPtr 是排序值的 device_ptr

随着CUB SortPairs功能我可以对单个值缓冲区进行排序，但不能一次性对所有 3 个进行排序。问题是我没有看到任何 CUB“类似收集”的实用程序。有什么建议吗？

编辑:

我想我可以实现自己的 gather 内核，但是除了以下方法还有其他更好的方法吗:

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 上没有已知结构，这可能是不可避免的。

Answer 1

看来你想要实现的目标取决于thrust::zip_iterator .你也可以

1. 只替换thrust::sort_by_key通过 cub::DeviceRadixSort::SortPairs并保持 thrust::gather , 或
2. zip 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 block 大小原语 cub::BlockRadixSort<...>::SortBlockedToStriped 获取存储在寄存器中的排序索引，然后执行天真的重新排序拷贝作为 thrust::gather填写values{1,2,3}Out .

使用 SortBlockedToStriped而不是 Sort 复制 values 时可以合并写入(虽然不是为了阅读) .