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)

应用场景

1. 数据重排：当需要按照特定顺序重新排列数据时
2. 稀疏矩阵操作：在稀疏矩阵计算中高效地访问非零元素
3. 数据库操作：实现类似SQL中的选择和投影操作
4. 图像处理：像素重映射操作

变体函数

Thrust还提供了一些变体函数：

• gather_if：带条件的gather操作
• scatter_if：带条件的scatter操作
• stable_scatter：保持相对顺序的scatter操作

这些操作在GPU上高度优化，能够充分利用并行计算能力，比在CPU上实现的类似操作要快得多。