如果你在CUDA里写过一次跨block归约大概率经历过那种“明明逻辑很简单代码却绕出一身汗”的状态先让每个block算出一部分和再把partial写到全局数组然后要么开第二次kernel去收尾要么在kernel里靠原子操作和自旋等待做协调。代码勉强能跑但每次调参、查显存、处理边界条件都要耗掉半天。做GPU并行算法开发的人几乎都被这类基础原语“重复造轮子”的琐碎折磨过。CUB就是为这件事存在的。它是CUDA生态里最老牌的开源并行原语库覆盖归约、扫描、选择、排序、分段归约等高频算子。而这两年CUB在设备级算法上做了大改版从经典的“先查询临时存储大小再分配再调用”两步叙事走向了更符合工程直觉的“单次调用API”对内尽量用一次kernel启动完成全局归约对外配合自己封装的内存管理一行代码就能把一整套并行原语跑完。这篇文章我想聊聊这套东西的设计思路、底层原理以及我实际封装过程中踩过的坑给正在写CUDA工具库或者被归约、扫描折磨的开发者一点参考。1. 单次调用API的设计思路拆解1.1 经典两步调用为什么让人血压升高如果你用过CUB对下面这套模板应该不陌生#include cub/cub.cuh float *d_in nullptr, *d_out nullptr; void *d_temp nullptr; size_t temp_bytes 0; cudaMalloc(d_in, N * sizeof(float)); cudaMalloc(d_out, sizeof(float)); // ... 初始化 d_in ... // 第一步查询所需临时存储大小 cub::DeviceReduce::Sum(nullptr, temp_bytes, d_in, d_out, N); // 第二步分配临时缓冲区 cudaMalloc(d_temp, temp_bytes); // 第三步真正执行归约 cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, N); // 第四步用完释放 cudaFree(d_temp);这套流程本身没毛病它符合CUB作为底层库的定位不替用户决定显存生命周期。但如果你在一个项目里同时用Reduce、Scan、Select、Sort好几个原语你会发现同样的“查询-分配-执行-释放”在你代码里出现了几十遍。更麻烦的是temp_bytes大小和算法实现、GPU架构、元素个数都有关你今天给N传参时忘了加sizeof(float)明天就可能收获一个cudaErrorIllegalAddress。手写kernel也好不到哪去。以跨block归约为例教科书式做法是两级归约第一级kernel让每个block算出局部和第二级kernel只有很少几个block负责把partial数组归约成最终结果。这要求你维护中间数组的显存分配、处理block数量不是输入整除的情况、还得面对“万一中间结果量比一个block能处理的还大”的边界。这些逻辑写一次不难难的是每次换个场景都要重新写一遍而且每遍都可能引入不同的并发bug。所以我在做内部算法库时第一诉求不是性能极致优化而是把“跨block归约”这种高频原语收敛成一行调用。这个诉求催生了我对“单次调用API”的完整设计对外用户只需要传输入、输出、数据量和流对内临时存储的查询、分配、缓存、复用全部收进封装层。代码可读性、可维护性、正确性一起提升。1.2 “单次调用”其实有两层含义理解单次调用API首先要分清两个层面。第一层是API层面的“单次调用”也就是用户侧只写一行代码。第二层是算法层面的“单次内核启动”也就是CUB内部尽量不用两次kernel来完成全局归约。先看算法层。经典的跨block归约需要两次kernel本质原因是没有跨block的全局同步原语。两次kernel启动天然是一个同步栅栏第一遍算完partials第二遍才能安全地去读这些partials。CUB 2.x重写设备级归约后在多数配置下改用“single-pass”方案线程按固定步长跨整个数据集读取元素每个窗口计算出部分结果后通过atomicAdd合并到全局输出用__threadfence()保证内存可见顺序再配合block粒度的协调逻辑一次kernel启动就能得到完整的全局归约结果。这样省掉了一次kernel启动的调度开销也省掉了partial数组的全局写入、再读回的内存流量。再说API层。CUB本身为了灵活性至今仍保留两步调用模式。但就像前面说的工程封装可以把这两步藏起来。我的封装核心是三层临时存储管理、原语调用模板、流与生命周期控制。后文会给出具体的代码实现。这里想强调的是这两层“单次”是互补的API层解决“好不好写”的问题算法层解决“快不快”的问题。你只做API封装性能收益有限只有CUB这代算法也跟上端到端效果才明显。提示不是所有设备级原语都能做到单次内核完成。比如全数据集排序内部依然可能需要多趟kernel。理解“单次调用API”的价值更多是从工程接口层面出发内核是否单次由算法本质决定不要为了单次而单次。2. 核心原语选型与内部原理2.1 常用设备级原语一览CUB的设备级原语覆盖了相当多的并行模式。我按实际使用频率列了个表方便你对照自己的场景选型原语功能典型场景DeviceReduce::Sum/Min/Max全局归约大数组求和、找极值DeviceSegmentedReduce分段归约按行/按用户ID归约稀疏数据DeviceScan::InclusiveSum/ExclusiveSum前缀和流压缩前的索引计算、直方图DeviceSelect::Flagged按标志位筛选过滤符合条件的元素DeviceSelect::Unique去重消除重复元素DeviceMergeSort归并排序键排序、键值排序DeviceRadixSort基数排序海量整数键排序这些原语和手写kernel的关系有点像标准库容器和手写链表的关系。你当然可以手写但库实现里包含的调优策略、架构特化、边界处理是普通开发者不太可能从头复现的。CUB在每次CUDA版本迭代时会对不同算力架构自动选择不同的block大小和items-per-thread配置。你把排序原语从sm_80平台搬到sm_90平台时手写代码可能还得调参CUB内部已经换了更合适的内核配置。选型上我的经验是凡是能映射到这些原语的操作优先用原语。什么时候手写操作逻辑过于定制比如要在归约过程中维护多个辅助数组、或者需要访问邻域数据时才值得手写。也可以用BlockReduce、BlockScan这些块级原语作为积木在自己的kernel里搭这也是CUB非常擅长的场景。2.2 Reduce的single-pass内核到底做了什么理解single-pass归约我建议直接想成“把全局归约拆成原子累加的大循环”。传统两遍方案是第一遍每个block用树形归约算出局部分量写入数组第二遍再来一个block把这个数组读回来。single-pass方案则是所有block协作扫描整个输入每处理完一段数据就把这段结果累加到全局输出上。这里要解决的是竞争问题——多个block可能同时到达谁先谁后必须一致而且不能漏。CUB的内部实现细节非常讲究。它会给线程配一个“全局窗口”每个线程按照blockIdx.x * blockDim.x threadIdx.x的步长去遍历数据一次处理多个元素后先做寄存器级归约再经过block内共享内存归约最后用一个全局原子操作把本block的结果合并给输出。__threadfence()的作用是保证在原子累加之前本线程对全局内存的写入对其它block可见。这个机制在没有全局同步原语的硬件上用“内存栅栏原子操作”模拟出了一个逻辑上的全局同步点令人佩服。那它是不是永远更快不一定。原子操作的竞争在元素量极大、block数量极多时可能成为瓶颈所以CUB内部通常会有策略开关根据数据量和架构选择single-pass还是经典两趟实现。我在实测中也发现数据量几千万级往下时单次启动版本的端到端时间明显占优因为省掉的kernel启动开销是纯赚数据量到几亿级别时两者差距会缩小甚至两趟方案的带宽利用率更高。这个结论不必你自己去验证CUB会自动选策略你只要理解这一点就不会再看到CUB“明明两次kernel更快为什么还选择单次”而感到困惑。2.3 封装层设计临时存储管理的核心矛盾把CUB封装成单次调用API最核心的难点是临时存储管理而不是调用语法本身。CUB在调用时需要一个临时缓冲区这个缓冲区可能只有几十KB也可能是几百MB取决于算法和输入规模。最差的做法是每次调用都cudaMalloc因为cudaMalloc是设备端同步操作成本不低调用频繁时会把性能优势全吃回去。我的封装方案在“每次调用都分配”和“永远不释放”之间取了折中一个按线程局部的临时存储池配合流信息做复用。考虑到多流并发场景我没有用简单的全局静态缓存而是用thread_local存储池同一线程对同一流可以复用缓冲区不同流之间由于可能并行执行缓冲区复用要格外小心我在后文第四章会详细讲这个坑。封装后的最终效果大致是这样namespace cask { template typename T void Sum(const T *d_in, T *d_out, std::size_t num_items, cudaStream_t stream 0) { auto pool detail::TempPool::instance(); std::size_t temp_bytes 0; cub::DeviceReduce::Sum(nullptr, temp_bytes, d_in, d_out, num_items, stream); void *temp pool.acquire(temp_bytes, stream); cub::DeviceReduce::Sum(temp, temp_bytes, d_in, d_out, num_items, stream); pool.release(temp, stream); } } // namespace cask看起来很简单对吧但这个简单背后藏着大量细节。temp_bytes为0时CUB允许传nullptr我的acquire要能处理这种情况多流时池子要判断这个缓冲区当前是否被某条流占用单线程内部调用可重入但跨线程不能共享同一临时存储。这些细节我在第三章会逐个展开。3. 实操过程把CUB封装成自己的单次调用算法库3.1 环境准备与版本选择建议直接使用CUDA 12.x及以上版本自带的CUB也就是/include/cub目录下的头文件版本。CUB是header-only库意思是不存在.so或.lib链接问题只要编译器能找到头文件就能编译使用。这对于希望把算法库分发给团队其他人使用、而不想让人重新编译CUDA工具链的场景特别友好。我建议在编译选项里明确指定目标架构而不是依赖__CUDA_ARCH_LIST__的默认值。比如nvcc -archsm_80 -O3 -stdc17 -I /path/to/cub your_file.cu如果机器是消费级显卡用sm_86或sm_89数据中心卡用sm_80或sm_90。指定架构能让CUB的调优策略在编译期就针对目标架构生效这是性能的基本保障。头文件包含推荐使用新风格#include cub/device/device_reduce.cuh #include cub/device/device_segmented_reduce.cuh #include cub/device/device_select.cuh #include cub/device/device_scan.cuh老风格的#include cub/cub.cuh也能用会把整库头文件全部拉进来增加编译时间。做工程封装时尽量按需包含。3.2 实现一个线程本地的临时存储池临时存储池的实现直接决定了封装的可靠性。我先做一个简化但可用的版本#include cuda_runtime.h #include map #include memory #include mutex namespace cask { namespace detail { struct TempBlock { void *ptr nullptr; size_t capacity 0; }; class TempPool { public: static TempPool instance() { static thread_local TempPool pool; return pool; } ~TempPool() { for (auto entry : blocks_) { if (entry.second.ptr) cudaFree(entry.second.ptr); } } void *acquire(size_t bytes, cudaStream_t stream) { if (bytes 0) return nullptr; auto it blocks_.find(bytes); if (it ! blocks_.end() !it-second.in_use) { it-second.in_use true; return it-second.ptr; } void *ptr nullptr; if (cudaMalloc(ptr, bytes) ! cudaSuccess) { return nullptr; } blocks_[bytes] {ptr, bytes, true}; return ptr; } void release(void *ptr, cudaStream_t stream) { // 简化实现实际应关联 stream 做回收判断 for (auto entry : blocks_) { if (entry.second.ptr ptr) { entry.second.in_use false; return; } } } private: struct TempBlock { void *ptr nullptr; size_t capacity 0; bool in_use false; }; std::mapsize_t, TempBlock blocks_; }; } // namespace detail } // namespace cask这个池子有几个问题需要说明。第一它按bytes精确匹配如果某个算法这次要1024字节、下次要1025字节池子里会同时存在两个缓冲区浪费显存。工程上可以改成按2的幂次对齐分配比如总是分配不小于bytes的最小的2的幂大小提高复用率。第二release没有真正处理流同步如果上游未同步而新的调用又复用了这个缓冲区CUB内部会对缓冲区写入旧计算若还在读取就会读到脏数据。严格的做法是记录最后一次使用的流并在复用前对旧流做cudaStreamQuery或cudaEventQuery确认旧计算完成后再复用。实际项目里我一般建议在封装层要求所有调用都走同一个流或者对外文档明确“同一缓冲区不能跨流并发复用”。临时存储池的价值在频繁调用原语的小数据量场景下特别明显。cudaMalloc一次可能花费几十微秒而封装池复用后这部分成本变成了一条if判断。3.3 用封装API重写三个典型算法有了池子我来演示三个高频场景的完整封装和调用。第一个大数组归约求和。// 文件: cask_algorithm.cuh #pragma once #include cub/device/device_reduce.cuh #include cuda_runtime.h #include temp_pool.hpp namespace cask { template typename T cudaError_t Sum(const T *d_in, T *d_out, std::size_t num_items, cudaStream_t stream 0) { auto pool detail::TempPool::instance(); std::size_t temp_bytes 0; cudaError_t err cub::DeviceReduce::Sum( nullptr, temp_bytes, d_in, d_out, num_items, stream); if (err ! cudaSuccess) return err; void *temp pool.acquire(temp_bytes, stream); if (temp_bytes 0 temp nullptr) return cudaErrorMemoryAllocation; err cub::DeviceReduce::Sum( temp, temp_bytes, d_in, d_out, num_items, stream); pool.release(temp, stream); return err; } } // namespace cask调用侧float *d_data nullptr, *d_result nullptr; cudaMalloc(d_data, N * sizeof(float)); cudaMalloc(d_result, sizeof(float)); cudaMemcpy(d_data, h_data.data(), N * sizeof(float), cudaMemcpyHostToDevice); cudaError_t err cask::Sum(d_data, d_result, N); if (err ! cudaSuccess) { /* 错误处理 */ } float host_result 0.0f; cudaMemcpy(host_result, d_result, sizeof(float), cudaMemcpyDeviceToHost);第二个按行归约。这个场景在矩阵操作里很常见对M x K的矩阵求每一行的和。这时DeviceSegmentedReduce比先转置再归约高效得多。分段信息用d_offsets表示它有M1个元素其中d_offsets[i]和d_offsets[i1]定义第i段的起始和结束位置template typename T cudaError_t SegmentSum(const T *d_in, T *d_out, const int *d_offsets, int num_segments, cudaStream_t stream 0) { auto pool detail::TempPool::instance(); std::size_t temp_bytes 0; cudaError_t err cub::DeviceSegmentedReduce::Sum( nullptr, temp_bytes, d_in, d_out, num_segments, d_offsets, d_offsets 1, stream); if (err ! cudaSuccess) return err; void *temp pool.acquire(temp_bytes, stream); if (temp_bytes 0 temp nullptr) return cudaErrorMemoryAllocation; err cub::DeviceSegmentedReduce::Sum( temp, temp_bytes, d_in, d_out, num_segments, d_offsets, d_offsets 1, stream); pool.release(temp, stream); return err; }注意这里d_offsets必须放在设备端别把主机端的std::vectorint直接传进去那是新手最容易犯的错误。分段大小不一致时DeviceSegmentedReduce会针对不同段长自动分派不同内核不需要你手工做均衡。第三个条件筛选。比如你需要从一个数据流中筛出所有大于阈值的元素同时还要知道筛出了多少个。CUB的DeviceSelect::Flagged会把符合条件的元素紧凑写入d_out并把数量写入d_num_selectedtemplate typename T cudaError_t FlaggedSelect(const T *d_in, const uint8_t *d_flags, T *d_out, int *d_num_selected, std::size_t num_items, cudaStream_t stream 0) { auto pool detail::TempPool::instance(); std::size_t temp_bytes 0; cudaError_t err cub::DeviceSelect::Flagged( nullptr, temp_bytes, d_in, d_flags, d_out, d_num_selected, num_items, stream); if (err ! cudaSuccess) return err; void *temp pool.acquire(temp_bytes, stream); if (temp_bytes 0 temp nullptr) return cudaErrorMemoryAllocation; err cub::DeviceSelect::Flagged( temp, temp_bytes, d_in, d_flags, d_out, d_num_selected, num_items, stream); pool.release(temp, stream); return err; }这三个例子覆盖了无分段、有分段、带额外输出参数三种常见形态。你会发现封装后的调用侧出奇一致传数据、传输出、传流完事。团队里新来的同学看到这种API根本不需要翻手册类型对了就能跑。3.4 临时存储大小会变吗一个重要经验同一个原语在不同输入规模下查询出来的temp_bytes大小会变化。尤其DeviceRadixSort这类算法不仅仅与元素个数有关还可能和bits_per_key有关。因此封装函数里每次调用时都重新走一遍查询步骤是必要的不要试图缓存temp_bytes硬编码到某个固定值。好在一次查询本身非常便宜。它只是根据输入参数估算需要的共享内存和全局内存并不会真正启动内核。把查询放在封装内部、每次调用都执行性能影响可以忽略。4. 踩坑记录与排查心得4.1 性能类问题坑一临时存储池忘了加“按2的幂对齐”导致显存膨胀。我在一个密集排序场景里一次循环要排几十批数据每批数量都不一样。如果严格按字节匹配池子里会积累大量只差几KB的缓冲区块。后来改成总是分配不小于请求大小的2的幂次容量池子里的缓冲区块数量立刻降了一个数量级。坑二低估了cudaMalloc对端到端耗时的拖累。裸CUB调用封装前我在小数据量场景测过一次cudaMalloc加查询加执行三件事耗时接近一比一比一。用上池子之后查询和执行成为主导端到端时间大概缩短了三成。这不是CUB变快了是省去的分配成本被看到了。坑三多流复用同一临时缓冲区导致偶发性能回退。早期封装版本里池子只判断“当前是否空闲”不判断“上一个使用它的流是否已完成”。两条流交替调用同一个算法时第二条流可能在第一条流的内核还在读取旧数据时就开始写缓冲区。行为上表现为两种一种是计算sanitizer报数据竞争另一种是更隐蔽的——性能突然下降因为底层驱动检测到同一个地址上有跨流依赖会强制串行化。我后续在release时记录了旧流复用前对旧流做cudaStreamQuery成功规避了这类问题。4.2 正确性类问题汇总一些实测中高频出现的问题我整理成速查表症状可能原因排查方向cudaErrorIllegalAddress临时存储未分配或分配过小检查第一次查询是否传了nullptr作为temp_storage检查temp_bytes是否被意外改写输出结果一直为0d_num_selected传了空指针或非法指针确认输出数量指针已经在设备端分配结果在Host端拷贝后不对忘记同步流在cudaMemcpy前调用cudaStreamSynchronize或等待事件数据量大时偶尔崩溃临时存储池跨流复用为缓冲区块增加流信息并做同步判断性能远低于手写kernelitems-per-thread配置不佳优先检查是否在目标架构上编译再检查输入规模是否太小小于几万元素时直接CPU跑可能更快还有一类隐蔽问题来自CUB的限定条件。DeviceReduce::Sum支持的类型是有限集合比如double、float、int、unsigned int等标量类型。如果传一个自定义结构体进去编译期就会报很长的模板错误。排查时看看CUB文档中该原语支持的类型列表不要和C模板“什么都能传”的直觉对抗。4.3 版本兼容与动态并行CUB作为头文件库最大的兼容性优势是“随你的工程编译走”。你今天用CUDA 12.8的CUB写代码生成的二进制可以拷贝到另一台只有CUDA 12.2 runtime但硬件相同的机器上运行因为CUB的内核代码已经在你的二进制里了。这一点比依赖动态链接的库省心不少。但也因为它是header-only编译器版本、CUDA版本、标准库实现都会影响具体生成代码。我在升级编译器时遇到过CUB内部静态断言失败的情况通常不是代码问题而是编译器标准支持差异。解决办法是把-stdc17明确打开并确保CUB版本和CUDA toolkit版本匹配。如果你要用到动态并行CDP也就是在设备端kernel里再调用CUB设备级原语需要额外配置。CUB从2.x开始支持CUB_CDP宏开启后允许在设备端调用这些原语。但这会显著增加编译时间和代码体积而且设备端再递归调用kernel的调试难度很高。我的经验是非必要不开启如果一定要用先把主机端封装验证到完全稳定再加一个去掉CUB_CDP的编译分支作为回退方案。注意CUB 2.x后新引入的调优方案会参考__CUDA_ARCH__做内核特化。如果你的编译目标架构里既有sm_80又有sm_86CUB会为不同架构生成不同内核这是正常的不要因为二进制体积增大而移除低版本架构目标。最后分享一个小技巧在实际封装过程中我发现最有价值的不是代码本身而是给每个封装原语都配一个“正确性回归测试”用随机数据调用CUB原语再和CPU上的朴素实现比对。听起来简单但它能自动拦截八成以上的问题。比如临时存储池引入跨流复用bug时测试会在数据量变化后偶然失败比如误传Host端std::vector时测试会稳定崩溃。把这些测试放进持续集成里每次升级CUDA工具包或修改池子实现时跑一遍比任何代码评审都让人安心。另外如果你经常做多流并发建议在封装函数的文档里明确一条约束同一个封装对象或同一个临时存储池不要在同一时刻从多个host线程调用并进入不同流。这不是CUB的限制而是复用显存缓冲区的天然约束。遵守它你会少掉很多半夜被唤起来看崩溃日志的机会。