GPU并行计算 -- 归约（Reduce）算子深度优化：从基础实现到极致性能

前言归约Reduce是GPU并行计算的基础算子。它通常应用在求最大值、平均值、求和等场景。其核心思想是并行地将一组数据“坍缩”为一个值。本文将对reduce算子各个版本进行瓶颈分析并提出相应的解决方案逐步优化。为方便理解示例统一为求和操作。1. CPU实现下面的代码为CPU串行实现只有一个线程迭代计算效率低下。intsum0;for(inti0;in;i){sumnums[i];}举个通俗的例子校长要统计全校多少人CPU的做法是把所有学生叫到操场上一个一个数GPU的做法是每个班长向校长汇报自己班级里有多少个学生等所有班长汇报完毕校长再把所有班级的人数相加。下面将介绍GPU是如何加速计算的。2. reduce_v1归约操作的基本流程总共有N个数据第 1 轮​ N个数并行计算–N/2个数第 2 轮​N/2个数并行计算–N/4个数以此类推最后变为1个数即最终结果总共约log2Nlog2Nlog2N轮下面是V1简单的6个数相加的规约示例:相同的6个数相加GPU需要迭代3轮而CPU需要迭代5轮。下面是V1实现代码为了合并访问全局显存、加快数据存取速度使用共享显存sdata暂存一个线程块需要处理的数据。__global__voidreduce_v1(float*res,float*in,intN){__shared__floatsdata[THREAD_PER_BLOCK];unsignedinttidthreadIdx.x;unsignedintg_idblockIdx.x*blockDim.xthreadIdx.x;sdata[tid]in[g_id];__syncthreads();for(intstride1;strideblockDim.x;stride*2){if(tid%(2*stride)0){sdata[tid]sdata[tidstride];}__syncthreads();}if(tid0){// block内0号线程对应的位置为block内总和res[blockIdx.x]sdata[tid];}}tid是线程的块内索引用于共享内存sdata下标g_id是线程的全局索引用于全局显存in下标。sdata[0]就是一个线程块内的规约结果。注意如果只启用一个线程块即每个线程块包含的线程数THREAD_PER_BLOCKN一次规约就能得到全局结果只需调用一次reduce如果启用多个线程块那么至少调用2次reduce继续规约每个线程块的局部结果。存在的问题a.线程浪费随着轮数和stride的增加, 真正在运行的线程越来越少.例如: 一个线程块启用128个线程, 首先, 每个线程会从全局显存搬运一个数据到共享显存中; 第一轮, stride1, 64个线程把相邻索引的数据加到自己身上; 第二轮, stride2, 32个线程继续执行相同操作; 以此类推.我们发现128个线程里, 真正参与到计算的只有64个线程, 剩下的64个线程只参与数据搬运.这至少浪费了一半的线程.b.线程束分化(Warp Divergence)if(tid%(2*stride)0)// 分支1else// 分支2CUDA采用单指令多线程(SIMT)架构, 硬件以Warp为最小单位下发指令, 即一个 warp 内的所有 thread 都执行同一条指令, 无法同时让一半线程执行A, 另一半线程执行B。如果存在 if-else 等逻辑分支会使同一个Warp内的线程走向不同的指令路径. Warp中的每个thread 会执行所有的分支。分支1:当满足索引条件的线程开始计算, 不满足条件的线程在硬件层面被屏蔽, 等待分支1执行完毕.分支2:分支1计算完毕后, 不满足条件判断的线程执行分支2, 满足条件判断的线程被屏蔽, 等待分支2执行完毕.汇总:两个分支执行完毕后, 线程同步, 继续执行下面的代码.虽然V1的代码看起来是并行计算但本质是轮流执行导致并行的代码变成串行执行执行时间翻倍计算资源利用率瞬间下降。c. 存储体冲突(bank conflict)共享内存(Shared Memory)通常由32个存储体(Bank)组成. 连续4B的数据会映射到一个Bank上, 并采取循环映射. 例如: 现在有132B的数据, 0-4B映射在Bank0, 4-8B映射在Bank1, 124-128B映射在Bank127, 128-132B重新映射回Bank0. 当stride32时, 被激活的线程索引为(0, 32, 64, …)均是32的倍数, 映射到同一个Bank.Bank Conflict的本质是每个Bank在一个时钟周期只能处理一个地址的读写请求.注意: 在同一个Warp中, 当不同线程对同一个Bank的不同位置同时访问, 会造成Bank Conflict. 而以下3种特殊情况不会造成冲突:(1) 广播(Broadcast)如果Warp内多个线程访问同一个Bank的同一个地址, 硬件会触发广播机制, 一次读取直接分发给所有线程, 不冲突且速度极快.(2) 单线程同一个线程访问一个Bank的多个不同位置不会冲突. 硬件会按顺序发射两次存取指令. 虽然多个地址在同一个Bank, 但他们属于不同的时间片, 因此不发生冲突, 而是正常的指令排队.(3) 不同的Warp访问一个Bank如果Warp0和Warp1的线程访问了一个Bank, 也不会发生冲突. 因为硬件调度器在某一时刻只能处理一个Warp的访存请求, 他们在时间上本就是错开的.在V0中当stride32时sdata[0]、sdata[32]、sdata[64]……这些元素均会映射到Bank0中且由不同的线程访问造成Bank Conflict。3. reduce_v2a.解决线程浪费问题前文已经指出线程利用率随着迭代的进行逐步下降对GPU并行计算资源造成严重浪费我们要想办法用更少的线程完成规约操作。我们不难发现数据搬运和第一轮计算这个过程是线程浪费最严重的阶段, 一半的线程只搬运了数据, 没有参与计算, 于是我们想到让这一半线程也参与一次计算再闲置.由此引出了预规约操作, 我们让线程在进入第一轮迭代之前, 让相邻两个线程块对应位置的数据提前相加, 即在块内规约之前提前进行一次块间规约. 与块内规约的第一轮类似, 相邻两个线程块对应位置的数据相加, 结果保留在前一个块中.只需要在v1基础上改动以下两行:unsignedintg_id2*blockIdx.x*blockDim.xthreadIdx.x;sdata[tid]in[g_id]in[g_idblockDim.x];此操作将空闲线程全部激活在搬运数据的同时在寄存器完成一次预规约把读取和计算并入同一条流水线。相较于V1所需线程块数量减少一半且一个线程从显存加载2个数据增大了访存粒度。b.解决线程束分化问题与V1版本相邻两个数相加的策略不同, 我们改用间隔stride的两个数相加.每轮迭代中,从V1版本的每隔stride的线程参与计算变为0到stride-1号线程参与计算, 这些线程是连续的, 更容易填满一个warp.当stride32时, 一个warp的线程全部执行相同操作,完全消除线程束分化.当stride32时, 才会出现分支. 故THREAD_PER_BLOCK设置较大时, 效率提升明显. 同时用整除取代了取模操作, 计算速度更快.只需要把循环改为以下形式即可:for(intstrideTHREAD_PER_BLOCK/2;stride1;stride/2){if(tidstride){sdata[tid]sdata[tidstride];}__syncthreads();}c.解决存储体冲突问题在解决线程数分化的方法实现中被累加的元素在物理布局上是连续的不会出现一个线程访问一个Bank的不同位置因此避免了 bank conflict 的发生.以下代码是reduce_v2的完整实现:__global__voidreduce_v2_prereduce_warpConflict(float*res,float*in,intN){__shared__floatsdata[THREAD_PER_BLOCK];unsignedinttidthreadIdx.x;unsignedintg_id2*blockIdx.x*blockDim.xthreadIdx.x;sdata[tid]in[g_id]in[g_idblockDim.x];__syncthreads();for(intstrideTHREAD_PER_BLOCK/2;stride1;stride/2){if(tidstride){sdata[tid]sdata[tidstride];}__syncthreads();}if(tid0){// block内0号线程对应的位置为block内总和res[blockIdx.x]sdata[tid];}}4.reduce_v3在上文中我们提到v2的代码实现, 当stride32时, 仍然会出现Warp divergence. 为了解决这个问题, 我们进一步对reduce的后半程进行优化.Warp ShuffleShuffle指令是针对warp的指令. Shuffle指令可以让同一个warp内的线程直接在寄存器间交换数据, 不经过任何层级的内存(包括Shared Memory). 具有不访存, 无Bank Conflict, 无Warp Divergence, 自动同步的优点, 速度极快. 若不使用shuffle指令, 线程之间的数据交换只能通过Shared Memory访问彼此的寄存器.我们这里要用到__shfl_down_sync函数, 数据流动如下图所示:可以看到此操作与我们之前的sdata[tid] sdata[tidstride]是类似的, 只需要把 “” 换成 “” . 故我们可以把最后的32个线程规约操作变为Warp Shuffle控制的Warp内规约操作. 我们依旧采取前文讲到的预规约思想, 可以少写一次Shared Memory, 代码如下:__device__voidwarpReduce(float*cache,unsignedinttid){// 少一次写回shared mem和一次同步intvcache[tid]cache[tid32];v__shfl_down_sync(0xffffffff,v,16);v__shfl_down_sync(0xffffffff,v,8);v__shfl_down_sync(0xffffffff,v,4);v__shfl_down_sync(0xffffffff,v,2);v__shfl_down_sync(0xffffffff,v,1);cache[tid]v;}需要把v2的后半部分代码改为:for(intstrideTHREAD_PER_BLOCK/2;stride64;stride/2){if(tidstride){sdata[tid]sdata[tidstride];}__syncthreads();}if(tid32){warpReduce(sdata,tid);}if(tid0){// block内0号线程对应的位置为block内总和res[blockIdx.x]sdata[tid];}继续优化上文所述, Warp Shuffle只能在一个warp内进行, 即只能规约32个数据. 我们引入二级warpReduce以扩大规约范围.第一步: 每个warp内部进行规约, 得到一个warp内的计算结果.第二步: 再使用一个warp规约所有第一步的局部计算结果.这样我们得出的结果即为一个block内的最终计算结果.同时我们可以让一个线程处理多个数据, 增加计算强度, 完整代码如下:constexprintwarpSize32;__device__floatblock_reduce(floatval){constinttidthreadIdx.x;intlanetid%warpSize;// warp内线程编号intwarpIdtid/warpSize;//warp编号val__shfl_down_sync(0xffffffff,val,16);val__shfl_down_sync(0xffffffff,val,8);val__shfl_down_sync(0xffffffff,val,4);val__shfl_down_sync(0xffffffff,val,2);val__shfl_down_sync(0xffffffff,val,1);__shared__floatwarpRes[32];// 硬编码, 一个线程块内最多32个warpif(lane0){warpRes[warpId]val;}__syncthreads();if(warpId0){// 用第一个warp继续规约局部和val(tidblockDim.x/warpSize)?warpRes[tid]:0.0f;val__shfl_down_sync(0xffffffff,val,16);val__shfl_down_sync(0xffffffff,val,8);val__shfl_down_sync(0xffffffff,val,4);val__shfl_down_sync(0xffffffff,val,2);val__shfl_down_sync(0xffffffff,val,1);}returnval;}templateunsignedintTHREAD_PER_BLOCK__global__voidreduce_v3(float*res,float*in,intN){unsignedinttidthreadIdx.x;unsignedintg_idblockIdx.x*blockDim.xthreadIdx.x;floatsum0.0f;// 一个线程处理多个数据, 把数组的后面的数据加到与启用线程一一对应的前面(预规约)for(intig_id;iN;igridDim.x*blockDim.x){sumin[i];}sumblock_reduce(sum);if(tid0){// block内0号线程对应的位置为block内总和res[blockIdx.x]sum;}}总结本文以求和操作为示例围绕 GPU 归约Reduce算子展开逐步介绍了三个版本的 GPU 归约优化方案。reduce_v1 搭建了基础并行框架但存在线程浪费、线程束分化和存储体冲突三大问题reduce_v2 通过预规约和改进循环逻辑有效解决了上述核心问题提升了线程利用率和计算效率reduce_v3 引入 Warp Shuffle 指令和二级归约借助寄存器通信最大化并行性能。