Thrust库中的Gather和Scatter操作

Thrust库中的Gather和Scatter操作

Thrust是CUDA提供的一个类似于C++ STL的并行算法库,其中包含两个重要的数据操作:gather(聚集)和scatter(散开)。

Gather操作

Gather操作从一个源数组中按照指定的索引收集元素到目标数组中。

函数原型:

cpp 复制代码
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

示例:

cpp 复制代码
#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操作将源数组的元素按照指定的索引分散到目标数组中。

函数原型:

cpp 复制代码
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

示例:

cpp 复制代码
#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上实现的类似操作要快得多。

相关推荐
weiwei2284418 天前
VS2019+CUDA12.5入门
gpu·vs2019·cuda
扫地的小何尚20 天前
全新NVIDIA Llama Nemotron Nano视觉语言模型在OCR基准测试中准确率夺冠
c++·人工智能·语言模型·机器人·ocr·llama·gpu
nuczzz25 天前
GPU虚拟化
docker·kubernetes·k8s·gpu·nvidia
扫地的小何尚1 个月前
NVIDIA Dynamo:数据中心规模的分布式推理服务框架深度解析
人工智能·分布式·microsoft·链表·语言模型·gpu
探索云原生1 个月前
开源 vGPU 方案:HAMi,实现细粒度 GPU 切分
ai·云原生·kubernetes·gpu
r0ysue_1 个月前
02.上帝之心算法用GPU计算提速50倍
算法·gpu
爱串门的小马驹1 个月前
VScode编译调试debug,gpu的cuda程序,Nsight
vscode·gpu·cuda
摩尔线程1 个月前
推测解码算法在 MTT GPU 的应用实践
算法·语言模型·大模型·gpu算力·gpu·摩尔线程