CUB 库入门GPU 并行计算的瑞士军刀引言为什么需要 CUB如果说 Thrust 是 GPU 编程的高铁——快速、舒适、不用操心细节那 CUB 就是赛车——需要你亲自调校但能跑出极限速度。CUBCUDA UnBound是 NVIDIA 官方提供的底层 CUDA 并行原语库专为追求极致性能的开发者设计。CUB 是什么CUB 提供了一套可复用的、高性能的 CUDA 并行原语覆盖三个层级┌─────────────────────────────────────────┐ │ Device Level设备级 │ ← 整个 GPU │ 如DeviceReduce、DeviceSort │ ├─────────────────────────────────────────┤ │ Block Level线程块级 │ ← 一个 Block 内 │ 如BlockReduce、BlockScan │ ├─────────────────────────────────────────┤ │ Warp Level线程束级 │ ← 一个 Warp32线程内 │ 如WarpReduce、WarpScan │ └─────────────────────────────────────────┘这三个层级对应 CUDA 的三个并行粒度CUB 在每个层级都提供了经过深度优化的实现。CUB 与 Thrust 的关系你的代码 ↓ Thrust高层接口易用 ↓ CUB底层实现高性能 ← CUB 就在这里 ↓ CUDA PTX 指令Thrust 内部大量调用 CUB。当你调用thrust::sort时背后很可能就是 CUB 的DeviceRadixSort在工作。对比项ThrustCUB接口风格STL 风格极简显式控制详细灵活性低高性能好极致学习成本低中等临时内存自动管理手动管理适合场景快速开发性能调优CUB 的核心模块1. Device Level整个 GPU 的并行操作这是最常用的层级一行代码启动整个 GPU 的并行计算。DeviceReduce归约#includecub/cub.cuhintN1000000;float*d_in;// GPU 输入数组float*d_out;// GPU 输出单个值// 第一步查询需要多少临时空间void*d_tempnullptr;size_t temp_bytes0;cub::DeviceReduce::Sum(d_temp,temp_bytes,d_in,d_out,N);// 第二步分配临时空间cudaMalloc(d_temp,temp_bytes);// 第三步执行归约cub::DeviceReduce::Sum(d_temp,temp_bytes,d_in,d_out,N);为什么要两次调用CUB 的设计哲学让调用者控制内存。第一次调用只是询问需要多少临时空间第二次才真正执行。这样你可以复用内存、精确控制显存使用。DeviceReduce 支持的操作cub::DeviceReduce::Sum(...)// 求和cub::DeviceReduce::Min(...)// 求最小值cub::DeviceReduce::Max(...)// 求最大值cub::DeviceReduce::ArgMin(...)// 求最小值的下标cub::DeviceReduce::ArgMax(...)// 求最大值的下标cub::DeviceReduce::Reduce(...)// 自定义归约操作DeviceSort排序// 基数排序整数类型极快cub::DeviceRadixSort::SortKeys(d_temp,temp_bytes,d_keys_in,d_keys_out,// 输入、输出N);// 按键值对排序cub::DeviceRadixSort::SortPairs(d_temp,temp_bytes,d_keys_in,d_keys_out,d_values_in,d_values_out,N);DeviceScan前缀扫描intinput[]{1,2,3,4,5};// 包含扫描inclusive scan// 结果{1, 3, 6, 10, 15}cub::DeviceScan::InclusiveSum(d_temp,temp_bytes,d_in,d_out,N);// 排除扫描exclusive scan// 结果{0, 1, 3, 6, 10}cub::DeviceScan::ExclusiveSum(d_temp,temp_bytes,d_in,d_out,N);DeviceSelect筛选// 筛选出满足条件的元素autoselect_positive[]__device__(floatx){returnx0;};int*d_num_selected;// 输出筛选出了多少个cub::DeviceSelect::If(d_temp,temp_bytes,d_in,d_out,d_num_selected,N,select_positive);2. Block Level线程块内的协作Block Level 的原语在__global__kernel 内部使用让同一个 Block 内的线程高效协作。BlockReduce块内归约#includecub/cub.cuh__global__voidsum_kernel(float*data,float*result,intN){// 声明 BlockReduce 类型256 个线程的块usingBlockReducecub::BlockReducefloat,256;// 在共享内存中分配临时空间__shared__ BlockReduce::TempStorage temp_storage;intidxblockIdx.x*blockDim.xthreadIdx.x;floatval(idxN)?data[idx]:0.0f;// 块内所有线程协作求和floatblock_sumBlockReduce(temp_storage).Sum(val);// 只有线程 0 写结果if(threadIdx.x0){atomicAdd(result,block_sum);}}关键点TempStorage放在共享内存里这是 Block Level 原语高效的秘密——共享内存比全局内存快几十倍。BlockScan块内前缀扫描__global__voidprefix_sum_kernel(int*data,int*output,intN){usingBlockScancub::BlockScanint,256;__shared__ BlockScan::TempStorage temp_storage;intidxblockIdx.x*blockDim.xthreadIdx.x;intval(idxN)?data[idx]:0;intprefix_sum;BlockScan(temp_storage).InclusiveSum(val,prefix_sum);if(idxN)output[idx]prefix_sum;}BlockLoad / BlockStore协作式内存加载这是 CUB 独有的特色功能——让整个 Block 协作加载数据充分利用内存带宽__global__voidprocess_kernel(float*input,float*output,intN){usingBlockLoadcub::BlockLoadfloat,256,4;// 每线程加载4个元素usingBlockStorecub::BlockStorefloat,256,4;__shared__union{BlockLoad::TempStorage load;BlockStore::TempStorage store;}temp_storage;floatitems[4];// 每个线程持有4个元素// 协作式加载256线程 × 4元素 一次加载1024个元素BlockLoad(temp_storage.load).Load(input,items);// 处理数据for(inti0;i4;i)items[i]*2.0f;// 协作式存储BlockStore(temp_storage.store).Store(output,items);}为什么这比普通加载快CUB 会自动选择最优的内存访问模式合并访问、向量化加载等最大化内存带宽利用率。3. Warp Level32 线程的极速协作Warp 是 GPU 的最小调度单位32 个线程Warp Level 原语利用硬件的 warp shuffle 指令无需共享内存速度极快。__global__voidwarp_reduce_kernel(float*data,float*result){usingWarpReducecub::WarpReducefloat;__shared__ WarpReduce::TempStorage temp_storage[4];// 4个warpintwarp_idthreadIdx.x/32;floatvaldata[threadIdx.x];// Warp 内归约无需 __syncthreads()floatwarp_sumWarpReduce(temp_storage[warp_id]).Sum(val);if(threadIdx.x%320){atomicAdd(result,warp_sum);}}实战用 CUB 实现高性能点积任务计算两个大向量的点积sum(a[i] * b[i])。#includecub/cub.cuh#includecuda_runtime.h// 自定义归约操作先乘后加structDotProductOp{float*a;float*b;__device__floatoperator()(inti)const{returna[i]*b[i];}};floatdot_product(float*d_a,float*d_b,intN){// 使用 DeviceReduce::Reduce 配合转换迭代器cub::CountingInputIteratorintcounting_iter(0);// 转换迭代器访问第 i 个元素时自动计算 a[i]*b[i]autotransform_itercub::TransformInputIteratorfloat,DotProductOp,cub::CountingInputIteratorint(counting_iter,{d_a,d_b});float*d_result;cudaMalloc(d_result,sizeof(float));void*d_tempnullptr;size_t temp_bytes0;cub::DeviceReduce::Sum(d_temp,temp_bytes,transform_iter,d_result,N);cudaMalloc(d_temp,temp_bytes);cub::DeviceReduce::Sum(d_temp,temp_bytes,transform_iter,d_result,N);floatresult;cudaMemcpy(result,d_result,sizeof(float),cudaMemcpyDeviceToHost);cudaFree(d_temp);cudaFree(d_result);returnresult;}CUB 的设计哲学1. 显式临时存储// CUB 的风格你来管内存void*d_tempnullptr;size_t temp_bytes0;cub::DeviceReduce::Sum(d_temp,temp_bytes,...);// 查询cudaMalloc(d_temp,temp_bytes);// 分配cub::DeviceReduce::Sum(d_temp,temp_bytes,...);// 执行好处可以在多次调用间复用临时空间精确控制显存使用避免隐式的cudaMalloc它会触发 GPU 同步很慢2. 策略可配置CUB 的很多操作支持通过模板参数调整内部策略// 自定义 BlockReduce 的算法usingBlockReducecub::BlockReducefloat,// 数据类型256,// 线程数cub::BLOCK_REDUCE_WARP_REDUCTIONS// 算法基于warp归约;3. 迭代器抽象CUB 支持各种迭代器让你无需额外内存就能做数据变换// 常量迭代器所有元素都是 1.0fcub::ConstantInputIteratorfloatones(1.0f);// 计数迭代器0, 1, 2, 3, ...cub::CountingInputIteratorintcounter(0);// 变换迭代器对每个元素应用函数autosquaredcub::TransformInputIteratorfloat,Square,float*(d_data,Square{});性能对比CUB vs 手写 Kernel以归约求和为例在 A100 GPU 上对 1 亿个 float 求和实现方式耗时带宽利用率朴素手写 kernel~8ms~40%优化手写 kernel~2ms~85%CUB DeviceReduce~1.2ms~95%CUB 之所以快是因为它针对每种 GPU 架构Volta、Ampere、Hopper 等都有专门调优的实现这些优化积累了 NVIDIA 工程师多年的经验。什么时候用 CUB✅ 适合用 CUB性能是首要目标需要榨干 GPU 的每一分算力需要精确控制内存避免隐式分配在自定义 kernel 内部需要高效的块级/束级原语Thrust 性能不够需要更底层的控制⚠️ 可以先用 Thrust快速原型开发性能要求不极致代码可读性优先算法逻辑复杂不想被内存管理分心快速上手安装CUB 随 CUDA Toolkit 自带CUDA 11.0也可以单独从 GitHub 获取# 随 CUDA 自带直接包含头文件即可#include cub/cub.cuh编译nvcc-archsm_80 my_program.cu-omy_program一个完整的可运行示例#includecub/cub.cuh#includeiostreamintmain(){constintN10;inth_data[]{5,2,8,1,9,3,7,4,6,0};// 分配 GPU 内存int*d_in,*d_out;cudaMalloc(d_in,N*sizeof(int));cudaMalloc(d_out,sizeof(int));cudaMemcpy(d_in,h_data,N*sizeof(int),cudaMemcpyHostToDevice);// CUB 归约求和void*d_tempnullptr;size_t temp_bytes0;cub::DeviceReduce::Sum(d_temp,temp_bytes,d_in,d_out,N);cudaMalloc(d_temp,temp_bytes);cub::DeviceReduce::Sum(d_temp,temp_bytes,d_in,d_out,N);intresult;cudaMemcpy(result,d_out,sizeof(int),cudaMemcpyDeviceToHost);std::coutSum resultstd::endl;// Sum 45cudaFree(d_in);cudaFree(d_out);cudaFree(d_temp);return0;}总结CUB 是什么CUB 是 NVIDIA 官方的底层 CUDA 并行原语库提供设备级、块级、束级三个层次的高性能并行操作是追求极致 GPU 性能的首选工具。核心要点三个层级Device全 GPU、Block线程块内、Warp32线程内显式内存管理两步调用模式精确控制临时空间极致性能针对每种 GPU 架构深度优化接近硬件理论峰值Thrust 的底层Thrust 的高性能来自于 CUB学习路径建议入门 GPU 编程 ↓ 学习 Thrust高层快速上手 ↓ 遇到性能瓶颈 ↓ 学习 CUB Device Level最常用 ↓ 需要自定义 kernel 内部优化 ↓ 学习 CUB Block/Warp Level一句话如果 Thrust 是让你能跑起来那 CUB 是让你跑得飞快。后记2026年8月15日于上海在claude opus 4.8辅助下完成。