并行归约问题

先解释一下什么是 归约（Reduction），归约是将某个计算问题变换为另一个问题的过程。在 CUDA 运算中，在向量中执行满足交换律和结合律的运算，被称为归约问题。每次迭代计算方式都是相同的（归），从一组多个数据最后得到一个数（约）^{[1]}。比如当给定 N 个数值，求其 SUM/MAX/MIN/MEAN 值的操作，也可被称为归约。

如果在 CPU 上，我们只需要一个循环遍历就可以实现以上操作，但使用 GPU 呢？

  1. 将输入划分到更小的数据块中
  2. 用一个 Thread 计算一个数据块的部分和
  3. 对每个数据块的部分和，再求和，得出最终结果

上述方法通常使用迭代成对实现，具体过程见下图：

并行规约

可以看出，用了三轮迭代就计算出了数组总和。

并行归约中的分化

上图是使用相邻配对方法实现的，这里有两个全局内存数组，一个大数组存放整个数组并进行归约；另一个小数组用于存放每个 Block 的部分和。由于归约是就地完成的，意味着在每一步全局内存中都有部分值被替换。使用 __syncthreads() 可以保证 Block 中的任一线程在进入下次迭代之前，当前迭代中的归约结果已写回全局内存。

// 这个 g_idata 存放的是整个数组；g_odata 是全局内存中小数组，存放每个线程块的部分和
__global__ void reduceNeighbored(int* g_idata, int* g_odata, unsigned int n) {
    // 首先获取 thread ID
    unsigned int tid = threadIdx.x;
    // 获取 idx 全局编号
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // 将全局的数据指针转化为这个块的局部数据指针，idata 指向某个块的开始
    int* idata = g_idata + blockIdx.x * blockDim.x;

    // 数组越界检查
    if (idx >= n) {
        return;
    }

    // 就地对全局数据进行规约
    for (int stride = 1; stride < blockDim.x; stride *= 2) {
        if ((tid % (2 * stride)) == 0) {
            idata[tid] += idata[tid + stride];
        }
        // 它能保证线程块内的同步：任一线程在进入下一步迭代之前，
        // 当前迭代中的归约结果已写回全局内存
        __syncthreads();
    }
    // Block 之间不能同步，因此要将 Block 产生的部分和返回 Host
    // 将块的局部结果写到全局内存，blockIdx 的范围 [0, grid.x)
    if (tid == 0)
        g_odata[blockIdx.x] = idata[0];
}

以上示例在全局内存中就地完成归约，便于理解配对过程；实际工程中更常见的做法是先加载到共享内存再归约，以减少全局内存访问开销。

其中 (tid % (2 * stride)) == 0 这个判断可能不太好理解，解释一下：假设和上图一样，8 个数字相加。tid 编号为 0-7，最开始 stride = 1，乘 2 后，只有 0、2、4、6 四个 tid 满足条件，因此这 4 个位置的内存执行操作 idata[tid] += idata[tid + stride]；随后，stride 乘以 2 后等于 2，只有 0、2 满足条件，同样执行加法操作；最后，只有 0 满足条件。

让我们细化到第一轮：

  • tid 为 0 的线程：idata[0] += idata[1]（有效，因为 0 % 2 == 0）
  • tid 为 1 的线程：不执行求和操作（因为 1 % 2 != 0）
  • tid 为 2 的线程：idata[2] += idata[3]（有效，因为 2 % 2 == 0）
  • tid 为 3 的线程：不执行求和操作（因为 3 % 2 != 0）
  • tid 为 4 的线程：idata[4] += idata[5]（有效，因为 4 % 2 == 0）
  • tid 为 5 的线程：不执行求和操作（因为 5 % 2 != 0）
  • tid 为 6 的线程：idata[6] += idata[7]（有效，因为 6 % 2 == 0）
  • tid 为 7 的线程：不执行求和操作（因为 7 % 2 != 0）

改善分化

由于 if ((tid % (2 * stride)) == 0) 只对偶数 ID 线程为 true，即只有偶数 ID 线程会被调度，其他线程闲置，因此会导致线程束分化。可以将访问索引修改为 int index = 2 * stride * tid。注意这里仍是相邻配对，而非交错配对。

当跨度乘 2 后，index 可以后移到有数据的内存，而非使用 tid 对应内存地址，减少了线程空闲。首先我们保证在一个块中前几个执行的线程束是在接近满跑的，而后半部分线程束基本是不需要执行的。当一个线程束内存在分支，而分支都不需要执行的时候，硬件会停止他们调用别人，这样就节省了资源。所以高效体现在这：如果还是所有分支不满足的也要执行，即便整个线程束都不需要执行的时候，那这种方案就无效了^{[1]}。

让我们细化到第一轮：

  • tid 为 0 的线程：idata[0] += idata[1]（有效，因为 0 < 8）
  • tid 为 1 的线程：idata[2] += idata[3]（有效，因为 2 < 8）
  • tid 为 2 的线程：idata[4] += idata[5]（有效，因为 4 < 8）
  • tid 为 3 的线程：idata[6] += idata[7]（有效，因为 6 < 8）
  • tid 为 4 的线程：尝试 idata[8] += idata[9]（无效，因为 8 >= 8）
  • tid 为 5 的线程：尝试 idata[10] += idata[11]（无效，因为 10 >= 8）
  • tid 为 6 的线程：尝试 idata[12] += idata[13]（无效，因为 12 >= 8）
  • tid 为 7 的线程：尝试 idata[14] += idata[15]（无效，因为 14 >= 8）

通过对第一轮操作的细化，应该能够理解，为什么新的索引映射 index = 2 * stride * tid 可以让 Block 的前半部分线程执行求和操作，而后半部分线程自然空闲，从而避免 Warp 分化。

对应的 Kernel 实现如下：

__global__ void reduceNeighbored2(int* g_idata, int* g_odata, unsigned int n) {
    unsigned int tid = threadIdx.x;
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int* idata = g_idata + blockIdx.x * blockDim.x;

    if (idx >= n) return;

    for (int stride = 1; stride < blockDim.x; stride *= 2) {
        int index = 2 * stride * tid;
        if (index < blockDim.x) {
            idata[index] += idata[index + stride];
        }
        __syncthreads();
    }
    if (tid == 0)
        g_odata[blockIdx.x] = idata[0];
}

交错配对归约

交错配对归约

交错配对的方法，初始 stride 是 Block 大小的一半，然后每次迭代减少一半。

核心部分代码如下：

for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
        idata[tid] += idata[tid + stride];
    }
    __syncthreads();
}

其中 stride >>= 1 将 stride 的二进制表示向右移动 1 位。移出的位被丢弃，假设初始值为 8，一次循环后则变为 4。

展开循环

封闭的循环，可以将其迭代次数减少甚至完全消除。循环体的复制次数为 循环展开因子，迭代次数变为 \textstyle \text{原迭代次数} / \text{循环展开因子}。

一个比较简单的例子：

for (int i = 0; i < 100; i++){
    a[i] = b[i] + c[i];
}

展开后：

for (int i = 0; i < 50; i += 2){
    a[i] = b[i] + c[i];
    a[i + 1] = b[i + 1] + c[i + 1];
}

展开的归约

在之前的归约函数中，每个 Block 只能处理一部分数据，但是如果能用一个 Block 处理更多数据，即在数据求和前，使用每个 Thread 进行一次加法，再从别的 Block 取数据，相当于先做一次向量加法，再进行归约。

__global__ void reduceUnroll2(int *g_idata, int *g_odata, unsigned int n){
    // set thread ID
    unsigned int tid = threadIdx.x;
    unsigned int idx = blockDim.x * blockIdx.x * 2 + threadIdx.x;
    // boundary check
    if (idx >= n) return;
    // convert global data pointer to the
    int *idata = g_idata + blockIdx.x * blockDim.x * 2;
    if(idx + blockDim.x < n){
        g_idata[idx] += g_idata[idx + blockDim.x];
    }
    __syncthreads();
    // in-place reduction in global memory
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1){
        if (tid < stride){
            idata[tid] += idata[tid + stride];
        }
        // synchronize within block
        __syncthreads();
    }
    // write result for this block to global mem
    if (tid == 0)  g_odata[blockIdx.x] = idata[0];
}

g_idata[idx] += g_idata[idx + blockDim.x]; 这行代码的作用是让每个 Thread 都与相邻数据块同位置的元素相加。

代码中 unsigned int idx = blockDim.x * blockIdx.x * 2 + threadIdx.x; 以及 int *idata = g_idata + blockIdx.x * blockDim.x * 2; 在确定 Block 对应数据位置时有个乘 2 的偏移量。

展开的归约

根据上图，因为只需要一半的 Block 来处理数据，因此全局数据索引也被调整。

展开线程的归约

当只剩下 32 个或更少线程时，由于 Warp 执行是 SIMT，每条指令后有隐式的 Warp 内同步过程。归约循环的最后 6 个迭代可以用下面语句：

if(tid < 32){
    volatile int *vmem = idata;
    vmem[tid] += vmem[tid + 32];
    vmem[tid] += vmem[tid + 16];
    vmem[tid] += vmem[tid + 8];
    vmem[tid] += vmem[tid + 4];
    vmem[tid] += vmem[tid + 2];
    vmem[tid] += vmem[tid + 1];
}

这个 Warp 展开避免了循环控制和 __syncthreads() 调用。注意：同一 Warp 内的 32 个线程以 SIMT 方式 lockstep 执行，每条语句后都有隐式的 Warp 内同步，因此最后几步归约无需再调用 __syncthreads()。

代码中的 volatile 修饰 vmem 指针，是为了防止编译器将中间结果缓存在寄存器中，强制每次读写都经过内存。由于下一步计算马上要用到上一步的结果，若只留在寄存器或缓存中，后续读取可能拿到旧值。此处 idata 指向全局内存，但即便数据在共享内存中，最后几步 Warp 级归约同样需要 volatile（或 __syncwarp() 等）来保证可见性。

当 Block 内剩余 64 个元素需要合并为 1 个结果时，32 个线程依次执行 6 步加法：

  1. vmem[tid] += vmem[tid + 32]：64 个元素两两相加，得到 32 个部分和
  2. vmem[tid] += vmem[tid + 16]：32 → 16
  3. vmem[tid] += vmem[tid + 8]：16 → 8
  4. vmem[tid] += vmem[tid + 4]：8 → 4
  5. vmem[tid] += vmem[tid + 2]：4 → 2
  6. vmem[tid] += vmem[tid + 1]：2 → 1

最终每个 Block 的归约结果保存在 idata[0] 中^{[2]}。

完全展开的归约

如果已知迭代次数，那么可以把循环完全展开。

__global__ void reduceUnrollWarp8(int *g_idata, int *g_odata, unsigned int n){
    // set thread ID
    unsigned int tid = threadIdx.x;
    unsigned int idx = blockDim.x * blockIdx.x * 8 + threadIdx.x;
    // boundary check
    if (idx >= n) return;
    // convert global data pointer to the
    int *idata = g_idata + blockIdx.x * blockDim.x * 8;
    // unrolling 8
    if(idx + 7 * blockDim.x < n){
        int a1 = g_idata[idx];
        int a2 = g_idata[idx + blockDim.x];
        int a3 = g_idata[idx + 2 * blockDim.x];
        int a4 = g_idata[idx + 3 * blockDim.x];
        int a5 = g_idata[idx + 4 * blockDim.x];
        int a6 = g_idata[idx + 5 * blockDim.x];
        int a7 = g_idata[idx + 6 * blockDim.x];
        int a8 = g_idata[idx + 7 * blockDim.x];
        g_idata[idx] = a1 + a2 + a3 + a4 + a5 + a6 + a7 + a8;
    }
    __syncthreads();
    // in-place reduction in global memory
    for (int stride = blockDim.x / 2; stride > 32; stride >>= 1){
        if (tid < stride){
            idata[tid] += idata[tid + stride];
        }
        // synchronize within block
        __syncthreads();
    }
    // write result for this block to global mem
    if(tid < 32){
        volatile int *vmem = idata;
        vmem[tid] += vmem[tid + 32];
        vmem[tid] += vmem[tid + 16];
        vmem[tid] += vmem[tid + 8];
        vmem[tid] += vmem[tid + 4];
        vmem[tid] += vmem[tid + 2];
        vmem[tid] += vmem[tid + 1];
    }

    if (tid == 0)  g_odata[blockIdx.x] = idata[0];
}

模板函数的归约

既然展开了最后 64 个数据，那也可以展开更多数据。在 __device__ 函数中，CUDA 支持模板参数，我们可以指定展开的数量作为参数。

template <unsigned int iBlockSize>
__device__ void warpReduce(volatile int* idata, int tid){
    if(iBlockSize >= 1024 && tid < 512)
        idata[tid] += idata[tid + 512];
    __syncthreads();
    if(iBlockSize >= 512 && tid < 256)
        idata[tid] += idata[tid + 256];
    __syncthreads();
    if(iBlockSize >= 256 && tid < 128)
        idata[tid] += idata[tid + 128];
    __syncthreads();
    if(iBlockSize >= 128 && tid < 64)
        idata[tid] += idata[tid + 64];
    __syncthreads();
}

对于这部分代码，编译器会去检查 blockDim.x 是否固定。如果固定，则会移除不可能实现的部分。而且当 blockDim.x 确定后就不能更改了。完整实现中，该函数之后通常还需配合前面提到的 Warp 级展开，处理最后 32 个线程的归约。

参考文章

[1] 【CUDA 基础】3.4 避免分支分化

[2] CUDA 编程权威指针（13）- 改善并行规约的分化