Thrust库中的Gather和Scatter操作
Thrust库中的Gather和Scatter操作
Thrust是CUDA提供的一个类似于C++ STL的并行算法库,其中包含两个重要的数据操作:gather(聚集)和scatter(散开)。
Gather操作
Gather操作从一个源数组中按照指定的索引收集元素到目标数组中。
函数原型:
template<typename InputIterator1, typename InputIterator2, typename OutputIterator>
OutputIterator gather(InputIterator1 map_first, InputIterator1 map_last,InputIterator2 input_first, OutputIterator result);
工作方式:
result[i] = input[map[i]] 对于 map中的每个索引i
示例:
#include <thrust/gather.h>
#include <thrust/device_vector.h>// 源数据
thrust::device_vector<int> input(4);
input[0] = 10; input[1] = 20; input[2] = 30; input[3] = 40;// 索引映射
thrust::device_vector<int> map(3);
map[0] = 3; map[1] = 1; map[2] = 2;// 目标向量
thrust::device_vector<int> result(3);// 执行gather操作
thrust::gather(map.begin(), map.end(), input.begin(), result.begin());
// result现在包含 [40, 20, 30]
Scatter操作
Scatter操作将源数组的元素按照指定的索引分散到目标数组中。
函数原型:
template<typename InputIterator1, typename InputIterator2, typename InputIterator3, typename OutputIterator>
OutputIterator scatter(InputIterator1 first, InputIterator1 last,InputIterator2 map_first, InputIterator3 stencil,OutputIterator result);
工作方式:
result[map[i]] = input[i] 对于 map中的每个索引i
示例:
#include <thrust/scatter.h>
#include <thrust/device_vector.h>// 源数据
thrust::device_vector<int> input(3);
input[0] = 10; input[1] = 20; input[2] = 30;// 索引映射
thrust::device_vector<int> map(3);
map[0] = 3; map[1] = 1; map[2] = 2;// 目标向量(需要足够大)
thrust::device_vector<int> result(4);// 执行scatter操作
thrust::scatter(input.begin(), input.end(), map.begin(), result.begin());
// result现在包含 [0, 20, 30, 10] (初始值为0)
应用场景
- 数据重排:当需要按照特定顺序重新排列数据时
- 稀疏矩阵操作:在稀疏矩阵计算中高效地访问非零元素
- 数据库操作:实现类似SQL中的选择和投影操作
- 图像处理:像素重映射操作
变体函数
Thrust还提供了一些变体函数:
gather_if
:带条件的gather操作scatter_if
:带条件的scatter操作stable_scatter
:保持相对顺序的scatter操作
这些操作在GPU上高度优化,能够充分利用并行计算能力,比在CPU上实现的类似操作要快得多。