前言

归约（Reduce）是GPU并行计算的基础算子。它通常应用在求最大值、平均值、求和等场景。其核心思想是：并行地将一组数据“坍缩”为一个值。

本文将对reduce算子各个版本进行瓶颈分析，并提出相应的解决方案，逐步优化。
为方便理解，示例统一为求和操作。

1. CPU实现

下面的代码为CPU串行实现，只有一个线程迭代计算，效率低下。

int sum = 0;
for(int i=0; i<n; i++) {
    sum += nums[i];
}

举个通俗的例子：校长要统计全校多少人，CPU的做法是把所有学生叫到操场上一个一个数；GPU的做法是每个班长向校长汇报自己班级里有多少个学生，等所有班长汇报完毕，校长再把所有班级的人数相加。下面将介绍GPU是如何加速计算的。

2. reduce_v1

归约操作的基本流程：
总共有N个数据
第 1 轮​： N个数并行计算–>N/2个数
第 2 轮​：N/2个数并行计算–>N/4个数
以此类推：最后变为1个数，即最终结果
总共约$log2N$轮
下面是V1简单的6个数相加的规约示例:

相同的6个数相加，GPU需要迭代3轮，而CPU需要迭代5轮。
下面是V1实现代码，为了合并访问全局显存、加快数据存取速度，使用共享显存sdata暂存一个线程块需要处理的数据。

__global__ void reduce_v1(float *res, float *in, int N){
    __shared__ float sdata[THREAD_PER_BLOCK];
    unsigned int tid = threadIdx.x;
    unsigned int g_id = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = in[g_id];
    __syncthreads();

    for(int stride = 1; stride < blockDim.x; stride *= 2){
        if(tid % (2*stride) == 0){
            sdata[tid] += sdata[tid+stride];
        }
        __syncthreads();
    }
    if(tid==0){ // block内0号线程对应的位置为block内总和
        res[blockIdx.x] = sdata[tid];
    }
}

tid是线程的块内索引，用于共享内存sdata下标；g_id是线程的全局索引，用于全局显存in下标。sdata[0]就是一个线程块内的规约结果。注意：如果只启用一个线程块，即每个线程块包含的线程数THREAD_PER_BLOCK==N，一次规约就能得到全局结果，只需调用一次reduce；如果启用多个线程块，那么至少调用2次reduce，继续规约每个线程块的局部结果。

存在的问题

a.线程浪费

随着轮数和stride的增加, 真正在运行的线程越来越少.
例如: 一个线程块启用128个线程, 首先, 每个线程会从全局显存搬运一个数据到共享显存中; 第一轮, stride=1, 64个线程把相邻索引的数据加到自己身上; 第二轮, stride=2, 32个线程继续执行相同操作; 以此类推.我们发现128个线程里, 真正参与到计算的只有64个线程, 剩下的64个线程只参与数据搬运.这至少浪费了一半的线程.

b.线程束分化(Warp Divergence)

if(tid % (2*stride) == 0)
  // 分支1
else
  // 分支2

CUDA采用单指令多线程(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. 当stride>=32时, 被激活的线程索引为(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中，当stride=32时，sdata[0]、sdata[32]、sdata[64]……这些元素均会映射到Bank0中，且由不同的线程访问，造成Bank Conflict。

3. reduce_v2

a.解决线程浪费问题

前文已经指出，线程利用率随着迭代的进行逐步下降，对GPU并行计算资源造成严重浪费，我们要想办法用更少的线程完成规约操作。
我们不难发现数据搬运和第一轮计算这个过程是线程浪费最严重的阶段, 一半的线程只搬运了数据, 没有参与计算, 于是我们想到让这一半线程也参与一次计算再闲置.
由此引出了预规约操作, 我们让线程在进入第一轮迭代之前, 让相邻两个线程块对应位置的数据提前相加, 即在块内规约之前提前进行一次块间规约. 与块内规约的第一轮类似, 相邻两个线程块对应位置的数据相加, 结果保留在前一个块中.
只需要在v1基础上改动以下两行:

unsigned int g_id = 2 * blockIdx.x * blockDim.x + threadIdx.x;
sdata[tid] = in[g_id] + in[g_id + blockDim.x];

此操作将空闲线程全部激活，在搬运数据的同时，在寄存器完成一次预规约，把读取和计算并入同一条流水线。相较于V1所需线程块数量减少一半，且一个线程从显存加载2个数据，增大了访存粒度。

b.解决线程束分化问题

与V1版本相邻两个数相加的策略不同, 我们改用间隔stride的两个数相加.
每轮迭代中,从V1版本的每隔stride的线程参与计算变为0到stride-1号线程参与计算, 这些线程是连续的, 更容易"填满"一个warp.
当stride>32时, 一个warp的线程全部执行相同操作, 完全消除线程束分化.当stride<32时, 才会出现分支. 故THREAD_PER_BLOCK设置较大时, 效率提升明显. 同时用整除取代了取模操作, 计算速度更快.
只需要把循环改为以下形式即可:

for(int stride = THREAD_PER_BLOCK / 2; stride >= 1; stride /= 2){
        if(tid < stride){
            sdata[tid] += sdata[tid+stride];
        }
        __syncthreads();
    }

c.解决存储体冲突问题

在解决线程数分化的方法实现中，被累加的元素在物理布局上是连续的，不会出现一个线程访问一个Bank的不同位置，因此避免了 bank conflict 的发生.

以下代码是reduce_v2的完整实现:

__global__ void reduce_v2_prereduce_warpConflict(float *res, float *in, int N){
    __shared__ float sdata[THREAD_PER_BLOCK];
    unsigned int tid = threadIdx.x;
    unsigned int g_id = 2 * blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = in[g_id] + in[g_id + blockDim.x];
    __syncthreads();

    for(int stride = THREAD_PER_BLOCK / 2; stride >= 1; stride /= 2){
        if(tid < stride){
            sdata[tid] += sdata[tid+stride];
        }
        __syncthreads();
    }
    if(tid==0){ // block内0号线程对应的位置为block内总和
        res[blockIdx.x] = sdata[tid];
    }
}

4.reduce_v3

在上文中我们提到v2的代码实现, 当stride<32时, 仍然会出现Warp divergence. 为了解决这个问题, 我们进一步对reduce的后半程进行优化.

Warp Shuffle

Shuffle指令是针对warp的指令. Shuffle指令可以让同一个warp内的线程直接在寄存器间交换数据, 不经过任何层级的内存(包括Shared Memory). 具有不访存, 无Bank Conflict, 无Warp Divergence, 自动同步的优点, 速度极快. 若不使用shuffle指令, 线程之间的数据交换只能通过Shared Memory访问彼此的寄存器.
我们这里要用到__shfl_down_sync函数, 数据流动如下图所示:
可以看到此操作与我们之前的sdata[tid] += sdata[tid+stride]是类似的, 只需要把 “=” 换成 “+=” . 故我们可以把最后的32个线程规约操作变为Warp Shuffle控制的Warp内规约操作. 我们依旧采取前文讲到的预规约思想, 可以少写一次Shared Memory, 代码如下:

__device__ void warpReduce(float *cache, unsigned int tid){
    // 少一次写回shared mem和一次同步
    int v = cache[tid] + cache[tid + 32]; 
    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(int stride = THREAD_PER_BLOCK / 2; stride >= 64; stride /= 2){
        if(tid < stride){
            sdata[tid] += sdata[tid+stride];
        }
        __syncthreads();
    }
    if(tid < 32){
        warpReduce(sdata, tid);
    }
    if(tid == 0){ // block内0号线程对应的位置为block内总和
        res[blockIdx.x] = sdata[tid];
    }

继续优化

上文所述, Warp Shuffle只能在一个warp内进行, 即只能规约32个数据. 我们引入二级warpReduce以扩大规约范围.
第一步: 每个warp内部进行规约, 得到一个warp内的计算结果.
第二步: 再使用一个warp规约所有第一步的局部计算结果.
这样我们得出的结果即为一个block内的最终计算结果.
同时我们可以让一个线程处理多个数据, 增加计算强度, 完整代码如下:

constexpr int warpSize = 32;
__device__ float block_reduce(float val){
    const int tid = threadIdx.x;
    int lane = tid % warpSize;  // warp内线程编号
    int warpId = tid / 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__ float warpRes[32];   // 硬编码, 一个线程块内最多32个warp
    if(lane == 0){
        warpRes[warpId] = val;
    }
    __syncthreads();

    if(warpId == 0){    // 用第一个warp继续规约局部和
        val = (tid < blockDim.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);
    }
    return val;
}

template <unsigned int THREAD_PER_BLOCK>
__global__ void reduce_v3(float *res, float *in, int N){
    unsigned int tid = threadIdx.x;
    unsigned int g_id = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    // 一个线程处理多个数据, 把数组的后面的数据加到与启用线程一一对应的前面(预规约)
    for(int i = g_id; i < N; i += gridDim.x * blockDim.x){  
        sum += in[i];
    }
    sum = block_reduce(sum);

    if(tid == 0){ // block内0号线程对应的位置为block内总和
        res[blockIdx.x] = sum;
    }
}

总结

本文以求和操作为示例，围绕 GPU 归约（Reduce）算子展开，逐步介绍了三个版本的 GPU 归约优化方案。
reduce_v1 搭建了基础并行框架，但存在线程浪费、线程束分化和存储体冲突三大问题；reduce_v2 通过预规约和改进循环逻辑，有效解决了上述核心问题，提升了线程利用率和计算效率；
reduce_v3 引入 Warp Shuffle 指令和二级归约，借助寄存器通信最大化并行性能。